fix: complete session persistence overhaul (Codex 5.2)

1. Implemented auto-selection of tasks in MultiXV2 to prevent empty initial state. 2. Added force-loading logic for task session messages with debouncing. 3. Updated session-actions to return full assistant text and immediately persist native messages. 4. Fixed caching logic in instance-shell2 to retain active task sessions in memory.
fix: flush session persistence on chat unmount to prevent data loss on navigation
2025-12-27 20:36:43 +04:00 · 2025-12-27 13:10:39 +04:00 · 2025-12-27 13:02:35 +04:00 · 2025-12-27 12:55:57 +04:00 · 2025-12-27 12:46:06 +04:00 · 2025-12-27 12:41:23 +04:00
14 changed files with 299 additions and 45 deletions
--- a/packages/server/src/api-types.ts
+++ b/packages/server/src/api-types.ts
@@ -16,6 +16,8 @@ export interface Task {
  status: TaskStatus
  timestamp: number
  messageIds?: string[] // IDs of messages associated with this task
+  taskSessionId?: string
+  archived?: boolean
 }

 export interface SessionTasks {
--- a/packages/server/src/server/routes/native-sessions.ts
+++ b/packages/server/src/server/routes/native-sessions.ts
@@ -209,6 +209,51 @@ export function registerNativeSessionsRoutes(app: FastifyInstance, deps: NativeS
        }
    })

+    // Append messages to a session (client-side persistence)
+    app.post<{
+        Params: { workspaceId: string; sessionId: string }
+        Body: {
+            messages: Array<{
+                id?: string
+                role: "user" | "assistant" | "system" | "tool"
+                content?: string
+                createdAt?: number
+                updatedAt?: number
+                status?: "pending" | "streaming" | "completed" | "error"
+            }>
+        }
+    }>("/api/native/workspaces/:workspaceId/sessions/:sessionId/messages", async (request, reply) => {
+        const { workspaceId, sessionId } = request.params
+        const payload = request.body?.messages
+        if (!Array.isArray(payload)) {
+            reply.code(400)
+            return { error: "messages array is required" }
+        }
+
+        try {
+            const results: SessionMessage[] = []
+            for (const entry of payload) {
+                if (!entry || typeof entry.role !== "string") {
+                    continue
+                }
+                const saved = await sessionManager.addMessage(workspaceId, sessionId, {
+                    id: entry.id,
+                    role: entry.role,
+                    content: entry.content,
+                    createdAt: entry.createdAt,
+                    updatedAt: entry.updatedAt,
+                    status: entry.status,
+                })
+                results.push(saved)
+            }
+            return { messages: results }
+        } catch (error) {
+            logger.error({ error }, "Failed to append messages")
+            reply.code(500)
+            return { error: "Failed to append messages" }
+        }
+    })
+
    // Add a message (user prompt) and get streaming response
    app.post<{
        Params: { workspaceId: string; sessionId: string }
--- a/packages/server/src/server/routes/storage.ts
+++ b/packages/server/src/server/routes/storage.ts
@@ -3,7 +3,7 @@ import { z } from "zod"
 import { InstanceStore } from "../../storage/instance-store"
 import { EventBus } from "../../events/bus"
 import { ModelPreferenceSchema } from "../../config/schema"
-import type { InstanceData, Task, SessionTasks } from "../../api-types"
+import type { InstanceData } from "../../api-types"
 import { WorkspaceManager } from "../../workspaces/manager"

 interface RouteDeps {
@@ -18,6 +18,8 @@ const TaskSchema = z.object({
  status: z.enum(["completed", "interrupted", "in-progress", "pending"]),
  timestamp: z.number(),
  messageIds: z.array(z.string()).optional(),
+  taskSessionId: z.string().optional(),
+  archived: z.boolean().optional(),
 })

 const InstanceDataSchema = z.object({
--- a/packages/server/src/storage/session-store.ts
+++ b/packages/server/src/storage/session-store.ts
@@ -27,6 +27,12 @@ export interface SessionMessage {
    status?: "pending" | "streaming" | "completed" | "error"
 }

+type IncomingSessionMessage = Omit<SessionMessage, "id" | "sessionId" | "createdAt" | "updatedAt"> & {
+    id?: string
+    createdAt?: number
+    updatedAt?: number
+}
+
 export interface MessagePart {
    type: "text" | "tool_call" | "tool_result" | "thinking" | "code"
    content?: string
@@ -260,23 +266,29 @@ export class NativeSessionManager {
            .filter((msg): msg is SessionMessage => msg !== undefined)
    }

-    async addMessage(workspaceId: string, sessionId: string, message: Omit<SessionMessage, "id" | "sessionId" | "createdAt" | "updatedAt">): Promise<SessionMessage> {
+    async addMessage(workspaceId: string, sessionId: string, message: IncomingSessionMessage): Promise<SessionMessage> {
        const store = await this.loadStore(workspaceId)
        const session = store.sessions[sessionId]
        if (!session) throw new Error(`Session not found: ${sessionId}`)

        const now = Date.now()
+        const messageId = message.id ?? ulid()
+        const createdAt = typeof message.createdAt === "number" ? message.createdAt : now
+        const updatedAt = typeof message.updatedAt === "number" ? message.updatedAt : createdAt
+
        const newMessage: SessionMessage = {
            ...message,
-            id: ulid(),
+            id: messageId,
            sessionId,
-            createdAt: now,
-            updatedAt: now,
+            createdAt,
+            updatedAt,
        }

        store.messages[newMessage.id] = newMessage
-        session.messageIds.push(newMessage.id)
-        session.updatedAt = now
+        if (!session.messageIds.includes(newMessage.id)) {
+            session.messageIds.push(newMessage.id)
+        }
+        session.updatedAt = updatedAt

        await this.saveStore(workspaceId)
        return newMessage
--- a/packages/ui/src/App.tsx
+++ b/packages/ui/src/App.tsx
@@ -45,6 +45,7 @@ import {
  clearActiveParentSession,
  createSession,
  fetchSessions,
+  flushSessionPersistence,
  updateSessionAgent,
  updateSessionModel,
 } from "./stores/sessions"
@@ -217,6 +218,7 @@ const App: Component = () => {

    if (!confirmed) return

+    clearActiveParentSession(instanceId)
    await stopInstance(instanceId)
  }

@@ -244,6 +246,12 @@ const App: Component = () => {
      return
    }

+    try {
+      await flushSessionPersistence(instanceId)
+    } catch (error) {
+      log.error("Failed to flush session persistence before closing", error)
+    }
+
    clearActiveParentSession(instanceId)

    try {
@@ -303,7 +311,7 @@ const App: Component = () => {
      const tauriBridge = (window as { __TAURI__?: { event?: { listen: (event: string, handler: (event: { payload: unknown }) => void) => Promise<() => void> } } }).__TAURI__
      if (tauriBridge?.event) {
        let unlistenMenu: (() => void) | null = null
-        
+
        tauriBridge.event.listen("menu:newInstance", () => {
          handleNewInstanceRequest()
        }).then((unlisten) => {
@@ -321,7 +329,7 @@ const App: Component = () => {

  // Check if this is OAuth callback
  const isOAuthCallback = window.location.pathname === '/auth/qwen/callback'
-  
+
  if (isOAuthCallback) {
    return <QwenOAuthCallback />
  }
@@ -391,29 +399,29 @@ const App: Component = () => {
                onNew={handleNewInstanceRequest}
                onOpenRemoteAccess={() => setRemoteAccessOpen(true)}
              />
- 
+
              <For each={Array.from(instances().values())}>
                {(instance) => {
                  const isActiveInstance = () => activeInstanceId() === instance.id
                  const isVisible = () => isActiveInstance() && !showFolderSelection()
-                    return (
-                      <div class="flex-1 min-h-0 overflow-hidden" style={{ display: isVisible() ? "flex" : "none" }}>
-                        <InstanceMetadataProvider instance={instance}>
-                          <InstanceShell
-                            instance={instance}
-                            escapeInDebounce={escapeInDebounce()}
-                            paletteCommands={paletteCommands}
-                            onCloseSession={(sessionId) => handleCloseSession(instance.id, sessionId)}
-                            onNewSession={() => handleNewSession(instance.id)}
-                            handleSidebarAgentChange={(sessionId, agent) => handleSidebarAgentChange(instance.id, sessionId, agent)}
-                            handleSidebarModelChange={(sessionId, model) => handleSidebarModelChange(instance.id, sessionId, model)}
-                            onExecuteCommand={executeCommand}
-                            tabBarOffset={instanceTabBarHeight()}
-                          />
-                        </InstanceMetadataProvider>
+                  return (
+                    <div class="flex-1 min-h-0 overflow-hidden" style={{ display: isVisible() ? "flex" : "none" }}>
+                      <InstanceMetadataProvider instance={instance}>
+                        <InstanceShell
+                          instance={instance}
+                          escapeInDebounce={escapeInDebounce()}
+                          paletteCommands={paletteCommands}
+                          onCloseSession={(sessionId) => handleCloseSession(instance.id, sessionId)}
+                          onNewSession={() => handleNewSession(instance.id)}
+                          handleSidebarAgentChange={(sessionId, agent) => handleSidebarAgentChange(instance.id, sessionId, agent)}
+                          handleSidebarModelChange={(sessionId, model) => handleSidebarModelChange(instance.id, sessionId, model)}
+                          onExecuteCommand={executeCommand}
+                          tabBarOffset={instanceTabBarHeight()}
+                        />
+                      </InstanceMetadataProvider>

-                      </div>
-                    )
+                    </div>
+                  )

                }}
              </For>
@@ -458,9 +466,9 @@ const App: Component = () => {
            </div>
          </div>
        </Show>
- 
+
        <RemoteAccessOverlay open={remoteAccessOpen()} onClose={() => setRemoteAccessOpen(false)} />
- 
+
        <AlertDialog />

        <Toaster
--- a/packages/ui/src/components/chat/multix-v2/index.tsx
+++ b/packages/ui/src/components/chat/multix-v2/index.tsx
@@ -10,6 +10,7 @@
 import { createSignal, Show, onMount, For, onCleanup, batch } from "solid-js";
 import toast from "solid-toast";
 import { sessions, activeSessionId, setActiveSession } from "@/stores/session-state";
+import { loadMessages, fetchSessions, flushSessionPersistence } from "@/stores/sessions";
 import { sendMessage, compactSession, updateSessionAgent, updateSessionModelForSession, forceReset, abortSession } from "@/stores/session-actions";
 import { addTask, setActiveTask, archiveTask } from "@/stores/task-actions";
 import { messageStoreBus } from "@/stores/message-v2/bus";
@@ -76,6 +77,8 @@ export default function MultiXV2(props: MultiXV2Props) {
    const [soloState, setSoloState] = createSignal({ isApex: false, isAutonomous: false, autoApproval: false, activeTaskId: null as string | null });
    const [lastAssistantIndex, setLastAssistantIndex] = createSignal(-1);
    const [bottomSentinel, setBottomSentinel] = createSignal<HTMLDivElement | null>(null);
+    const [hasUserSelection, setHasUserSelection] = createSignal(false);
+    const forcedLoadTimestamps = new Map<string, number>();

    // Helper to check if CURRENT task is sending
    const isSending = () => {
@@ -139,6 +142,10 @@ export default function MultiXV2(props: MultiXV2Props) {
                setVisibleTasks(allTasks.filter(t => !t.archived));
                // NOTE: Don't overwrite selectedTaskId from store - local state is authoritative
                // This prevents the reactive cascade when the store updates
+                if (!selectedTaskId() && !hasUserSelection() && allTasks.length > 0) {
+                    const preferredId = session.activeTaskId || allTasks[0].id;
+                    setSelectedTaskIdLocal(preferredId);
+                }
            }

            // Get message IDs for currently selected task
@@ -148,6 +155,20 @@ export default function MultiXV2(props: MultiXV2Props) {
                if (task) {
                    const store = getMessageStore();
                    if (task.taskSessionId) {
+                        const cachedIds = store.getSessionMessageIds(task.taskSessionId);
+                        if (cachedIds.length === 0) {
+                            const lastForced = forcedLoadTimestamps.get(task.taskSessionId) ?? 0;
+                            if (Date.now() - lastForced > 1000) {
+                                forcedLoadTimestamps.set(task.taskSessionId, Date.now());
+                                loadMessages(props.instanceId, task.taskSessionId, true).catch((error) =>
+                                    log.error("Failed to load task session messages", error)
+                                );
+                            }
+                        } else {
+                            loadMessages(props.instanceId, task.taskSessionId).catch((error) =>
+                                log.error("Failed to load task session messages", error)
+                            );
+                        }
                        setMessageIds(store.getSessionMessageIds(task.taskSessionId));
                    } else {
                        setMessageIds(task.messageIds || []);
@@ -216,6 +237,9 @@ export default function MultiXV2(props: MultiXV2Props) {
            setSendingTasks(new Set<string>());
        });

+        // Initialize
+        loadMessages(props.instanceId, props.sessionId);
+        fetchSessions(props.instanceId);
        syncFromStore();

        const interval = setInterval(syncFromStore, 150);
@@ -230,6 +254,8 @@ export default function MultiXV2(props: MultiXV2Props) {
        onCleanup(() => {
            clearInterval(interval);
            scrollContainer?.removeEventListener('scroll', handleScroll);
+            // Ensure any pending task updates are saved immediately before we potentially reload them
+            flushSessionPersistence(props.instanceId);
        });
    });

@@ -245,6 +271,7 @@ export default function MultiXV2(props: MultiXV2Props) {
    const setSelectedTaskId = (id: string | null) => {
        // Update local state immediately (fast)
        setSelectedTaskIdLocal(id);
+        setHasUserSelection(true);

        // Immediately sync to load the new task's agent/model
        syncFromStore();
@@ -298,7 +325,7 @@ export default function MultiXV2(props: MultiXV2Props) {
                syncFromStore();

                // Set the selected task
-                setSelectedTaskIdLocal(taskId);
+                setSelectedTaskId(taskId);

                const s = soloState();
                if (s.isAutonomous) {
@@ -351,7 +378,7 @@ export default function MultiXV2(props: MultiXV2Props) {
        setTimeout(async () => {
            try {
                const result = await addTask(props.instanceId, props.sessionId, title);
-                setSelectedTaskIdLocal(result.id);
+                setSelectedTaskId(result.id);
                setTimeout(() => syncFromStore(), 50);
            } catch (error) {
                log.error("handleCreateTask failed", error);
--- a/packages/ui/src/components/instance/instance-shell2.tsx
+++ b/packages/ui/src/components/instance/instance-shell2.tsx
@@ -66,7 +66,7 @@ import SessionView from "../session/session-view"
 import { Sidebar, type FileNode } from "./sidebar"
 import { Editor } from "./editor"
 import { serverApi } from "../../lib/api-client"
-import { Sparkles, Layout as LayoutIcon, Terminal as TerminalIcon, Search, Loader2, Zap, Shield, Settings, FileArchive } from "lucide-solid"
+import { Sparkles, Layout as LayoutIcon, Terminal as TerminalIcon, Search, Loader2, Zap, Shield, Settings, FileArchive, ArrowLeft } from "lucide-solid"
 import { formatTokenTotal } from "../../lib/formatters"
 import { sseManager } from "../../lib/sse-manager"
 import { getLogger } from "../../lib/logger"
@@ -683,7 +683,25 @@ Now analyze the project and report your findings.`
  })

  const handleSessionSelect = (sessionId: string) => {
-    setActiveSession(props.instance.id, sessionId)
+    if (sessionId === "info") {
+      setActiveSession(props.instance.id, sessionId)
+      return
+    }
+
+    const instanceSessions = sessions().get(props.instance.id)
+    const session = instanceSessions?.get(sessionId)
+
+    if (session?.parentId) {
+      setActiveParentSession(props.instance.id, session.parentId)
+      const parentSession = instanceSessions?.get(session.parentId)
+      const matchingTask = parentSession?.tasks?.find((task) => task.taskSessionId === sessionId)
+      if (matchingTask) {
+        setActiveTask(props.instance.id, session.parentId, matchingTask.id)
+      }
+      return
+    }
+
+    setActiveParentSession(props.instance.id, sessionId)
  }


@@ -731,6 +749,7 @@ Now analyze the project and report your findings.`
    const sessionsMap = activeSessions()
    const parentId = parentSessionIdForInstance()
    const activeId = activeSessionIdForInstance()
+    const instanceSessions = sessions().get(props.instance.id)
    setCachedSessionIds((current) => {
      const next: string[] = []
      const append = (id: string | null) => {
@@ -743,6 +762,16 @@ Now analyze the project and report your findings.`
      append(parentId)
      append(activeId)

+      const parentSessionId = parentId || activeId
+      const parentSession = parentSessionId ? instanceSessions?.get(parentSessionId) : undefined
+      const activeTaskId = parentSession?.activeTaskId
+      if (activeTaskId && parentSession?.tasks?.length) {
+        const activeTask = parentSession.tasks.find((task) => task.id === activeTaskId)
+        if (activeTask?.taskSessionId) {
+          append(activeTask.taskSessionId)
+        }
+      }
+
      const limit = parentId ? SESSION_CACHE_LIMIT + 1 : SESSION_CACHE_LIMIT
      const trimmed = next.length > limit ? next.slice(0, limit) : next
      const trimmedSet = new Set(trimmed)
@@ -1342,6 +1371,14 @@ Now analyze the project and report your findings.`
          <div class="flex items-center space-x-4">
            <Show when={activeSessionIdForInstance() && activeSessionIdForInstance() !== "info"}>
              <div class="flex items-center space-x-2">
+                <button
+                  onClick={() => props.onCloseSession(activeSessionIdForInstance()!)}
+                  class="flex items-center gap-1.5 px-2.5 py-1 text-[11px] font-semibold text-zinc-400 hover:text-white hover:bg-white/10 border border-transparent hover:border-white/10 transition-all rounded-full"
+                  title="Back to Sessions"
+                >
+                  <ArrowLeft size={14} strokeWidth={2} />
+                  <span>Back</span>
+                </button>
                {/* Compact Button */}
                <button
                  onClick={handleCompact}
--- a/packages/ui/src/components/settings/AntigravitySettings.tsx
+++ b/packages/ui/src/components/settings/AntigravitySettings.tsx
@@ -126,7 +126,8 @@ const AntigravitySettings: Component = () => {
            })

            if (!response.ok) {
-                throw new Error('Failed to start authentication')
+                const errorData = await response.json().catch(() => ({}))
+                throw new Error(errorData.error || errorData.details || 'Failed to start authentication')
            }

            const data = await response.json() as {
--- a/packages/ui/src/lib/lite-mode.ts
+++ b/packages/ui/src/lib/lite-mode.ts
@@ -165,6 +165,28 @@ export const nativeSessionApi = {
        return data.messages
    },

+    async appendMessages(
+        workspaceId: string,
+        sessionId: string,
+        messages: Array<{
+            id?: string
+            role: "user" | "assistant" | "system" | "tool"
+            content?: string
+            createdAt?: number
+            updatedAt?: number
+            status?: "pending" | "streaming" | "completed" | "error"
+        }>
+    ): Promise<NativeMessage[]> {
+        const response = await fetch(`${CODENOMAD_API_BASE}/api/native/workspaces/${encodeURIComponent(workspaceId)}/sessions/${encodeURIComponent(sessionId)}/messages`, {
+            method: "POST",
+            headers: { "Content-Type": "application/json" },
+            body: JSON.stringify({ messages })
+        })
+        if (!response.ok) throw new Error("Failed to append messages")
+        const data = await response.json()
+        return data.messages
+    },
+
    /**
     * Import sessions from SDK mode to Native mode
     */
--- a/packages/ui/src/stores/session-actions.ts
+++ b/packages/ui/src/stores/session-actions.ts
@@ -555,7 +555,7 @@ async function streamOllamaChat(
  messageId: string,
  assistantMessageId: string,
  assistantPartId: string,
-): Promise<void> {
+): Promise<string> {
  const controller = new AbortController()
  const timeoutId = setTimeout(() => controller.abort(), STREAM_TIMEOUT_MS)

@@ -682,6 +682,8 @@ async function streamOllamaChat(
      isEphemeral: false,
    })
  })
+
+  return fullText
 }

 async function streamQwenChat(
@@ -695,7 +697,7 @@ async function streamQwenChat(
  messageId: string,
  assistantMessageId: string,
  assistantPartId: string,
-): Promise<void> {
+): Promise<string> {
  const controller = new AbortController()
  const timeoutId = setTimeout(() => controller.abort(), STREAM_TIMEOUT_MS)

@@ -831,6 +833,8 @@ async function streamQwenChat(
    updatedAt: Date.now(),
    isEphemeral: false,
  })
+
+  return fullText
 }

 async function streamOpenCodeZenChat(
@@ -842,7 +846,7 @@ async function streamOpenCodeZenChat(
  messageId: string,
  assistantMessageId: string,
  assistantPartId: string,
-): Promise<void> {
+): Promise<string> {
  const controller = new AbortController()
  const timeoutId = setTimeout(() => controller.abort(), STREAM_TIMEOUT_MS)

@@ -978,6 +982,8 @@ async function streamOpenCodeZenChat(
    updatedAt: Date.now(),
    isEphemeral: false,
  })
+
+  return fullText
 }

 async function streamZAIChat(
@@ -989,7 +995,7 @@ async function streamZAIChat(
  messageId: string,
  assistantMessageId: string,
  assistantPartId: string,
-): Promise<void> {
+): Promise<string> {
  const controller = new AbortController()
  const timeoutId = setTimeout(() => controller.abort(), STREAM_TIMEOUT_MS)

@@ -1116,6 +1122,8 @@ async function streamZAIChat(
    updatedAt: Date.now(),
    isEphemeral: false,
  })
+
+  return fullText
 }

 async function streamAntigravityChat(
@@ -1127,7 +1135,7 @@ async function streamAntigravityChat(
  messageId: string,
  assistantMessageId: string,
  assistantPartId: string,
-): Promise<void> {
+): Promise<string> {
  const controller = new AbortController()
  const timeoutId = setTimeout(() => controller.abort(), STREAM_TIMEOUT_MS)

@@ -1254,6 +1262,26 @@ async function streamAntigravityChat(
    updatedAt: Date.now(),
    isEphemeral: false,
  })
+
+  return fullText
+}
+
+async function persistNativeMessages(
+  instanceId: string,
+  sessionId: string,
+  messages: Array<{
+    id: string
+    role: "user" | "assistant" | "system" | "tool"
+    content: string
+    createdAt: number
+    updatedAt: number
+  }>,
+): Promise<void> {
+  try {
+    await nativeSessionApi.appendMessages(instanceId, sessionId, messages)
+  } catch (error) {
+    log.warn("Failed to persist native messages", { instanceId, sessionId, error })
+  }
 }

 async function sendMessage(
@@ -1267,6 +1295,7 @@ async function sendMessage(
  if (!instance || !instance.client) {
    throw new Error("Instance not ready")
  }
+  const isNative = instance.binaryPath === "__nomadarch_native__"

  const instanceSessions = sessions().get(instanceId)
  const session = instanceSessions?.get(sessionId)
@@ -1450,10 +1479,11 @@ async function sendMessage(
      })
    })

+    let assistantText = ""
    try {
      if (providerId === "ollama-cloud") {
        const tStream1 = performance.now()
-        await streamOllamaChat(
+        assistantText = await streamOllamaChat(
          instanceId,
          sessionId,
          providerId,
@@ -1466,7 +1496,7 @@ async function sendMessage(
        const tStream2 = performance.now()
        addDebugLog(`Stream Complete: ${Math.round(tStream2 - tStream1)}ms`, "info")
      } else if (providerId === "opencode-zen") {
-        await streamOpenCodeZenChat(
+        assistantText = await streamOpenCodeZenChat(
          instanceId,
          sessionId,
          providerId,
@@ -1477,7 +1507,7 @@ async function sendMessage(
          assistantPartId,
        )
      } else if (providerId === "zai") {
-        await streamZAIChat(
+        assistantText = await streamZAIChat(
          instanceId,
          sessionId,
          providerId,
@@ -1488,7 +1518,7 @@ async function sendMessage(
          assistantPartId,
        )
      } else if (providerId === "antigravity") {
-        await streamAntigravityChat(
+        assistantText = await streamAntigravityChat(
          instanceId,
          sessionId,
          providerId,
@@ -1526,7 +1556,7 @@ async function sendMessage(
          return messageId
        }

-        await streamQwenChat(
+        assistantText = await streamQwenChat(
          instanceId,
          sessionId,
          providerId,
@@ -1539,6 +1569,25 @@ async function sendMessage(
          assistantPartId,
        )
      }
+      if (isNative) {
+        const completedAt = Date.now()
+        await persistNativeMessages(instanceId, sessionId, [
+          {
+            id: messageId,
+            role: "user",
+            content: resolvedPrompt,
+            createdAt: now,
+            updatedAt: now,
+          },
+          {
+            id: assistantMessageId,
+            role: "assistant",
+            content: assistantText,
+            createdAt: now,
+            updatedAt: completedAt,
+          },
+        ])
+      }
      return messageId
    } catch (error: any) {
      if (providerId === "opencode-zen") {
--- a/packages/ui/src/stores/session-api.ts
+++ b/packages/ui/src/stores/session-api.ts
@@ -486,6 +486,28 @@ async function fetchSessions(instanceId: string): Promise<void> {
      })
    }

+    if (isNative) {
+      const updates: Promise<unknown>[] = []
+      for (const [parentId, tasks] of Object.entries(sessionTasks)) {
+        if (!Array.isArray(tasks)) continue
+        for (const task of tasks as Array<{ taskSessionId?: string }>) {
+          const childId = task?.taskSessionId
+          if (!childId) continue
+          const childSession = sessionMap.get(childId)
+          if (!childSession) continue
+          if (childSession.parentId === parentId) continue
+          sessionMap.set(childId, {
+            ...childSession,
+            parentId,
+          })
+          updates.push(nativeSessionApi.updateSession(instanceId, childId, { parentId }).catch(() => undefined))
+        }
+      }
+      if (updates.length > 0) {
+        await Promise.allSettled(updates)
+      }
+    }
+
    const validSessionIds = new Set(sessionMap.keys())

    setSessions((prev) => {
@@ -1055,7 +1077,10 @@ async function loadMessages(instanceId: string, sessionId: string, force = false

      messagesInfo.set(messageId, info)

-      const parts: any[] = (apiMessage.parts || []).map((part: any) => normalizeMessagePart(part))
+      let parts: any[] = (apiMessage.parts || []).map((part: any) => normalizeMessagePart(part))
+      if (parts.length === 0 && typeof apiMessage.content === "string" && apiMessage.content.trim().length > 0) {
+        parts = [normalizeMessagePart({ id: `part-${messageId}`, type: "text", text: apiMessage.content })]
+      }

      const message: Message = {
        id: messageId,
--- a/packages/ui/src/stores/session-state.ts
+++ b/packages/ui/src/stores/session-state.ts
@@ -171,6 +171,15 @@ function schedulePersist(instanceId: string) {
  persistTimers.set(instanceId, timer)
 }

+async function flushSessionPersistence(instanceId: string) {
+  const existing = persistTimers.get(instanceId)
+  if (existing) {
+    clearTimeout(existing)
+    persistTimers.delete(instanceId)
+  }
+  await persistSessionTasks(instanceId)
+}
+
 async function persistSessionTasks(instanceId: string) {
  try {
    const instanceSessions = sessions().get(instanceId)
@@ -439,6 +448,7 @@ export {
  pruneDraftPrompts,
  withSession,
  persistSessionTasks,
+  flushSessionPersistence,
  setSessionCompactionState,
  setSessionPendingPermission,
  setActiveSession,
--- a/packages/ui/src/stores/sessions.ts
+++ b/packages/ui/src/stores/sessions.ts
@@ -26,6 +26,7 @@ import {
  setActiveParentSession,
  setActiveSession,
  setSessionDraftPrompt,
+  flushSessionPersistence,
 } from "./session-state"

 import { getDefaultModel } from "./session-models"
@@ -113,5 +114,6 @@ export {
  setSessionDraftPrompt,
  updateSessionAgent,
  updateSessionModel,
+  flushSessionPersistence,
 }
 export type { SessionInfo }
--- a/packages/ui/src/stores/task-actions.ts
+++ b/packages/ui/src/stores/task-actions.ts
@@ -3,6 +3,8 @@ import { Task, TaskStatus } from "../types/session"
 import { nanoid } from "nanoid"
 import { createSession } from "./session-api"
 import { showToastNotification } from "../lib/notifications"
+import { instances } from "./instances"
+import { nativeSessionApi } from "../lib/lite-mode"

 export function setActiveTask(instanceId: string, sessionId: string, taskId: string | undefined): void {
  withSession(instanceId, sessionId, (session) => {
@@ -35,6 +37,16 @@ export async function addTask(
        taskSession.model = { ...parentModel }
      }
    })
+    const instance = instances().get(instanceId)
+    if (instance?.binaryPath === "__nomadarch_native__") {
+      try {
+        await nativeSessionApi.updateSession(instanceId, taskSessionId, {
+          parentId: sessionId,
+        })
+      } catch (error) {
+        console.warn("[task-actions] Failed to persist parent session", error)
+      }
+    }
    // console.log("[task-actions] task session created", { taskSessionId });
  } catch (error) {
    console.error("[task-actions] Failed to create session for task", error)