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.
2025-12-27 20:36:43 +04:00
parent 5022a23aeb
commit 1e991d9ebd
8 changed files with 235 additions and 20 deletions
--- 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/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
+        if (!session.messageIds.includes(newMessage.id)) {
            session.messageIds.push(newMessage.id)
-        session.updatedAt = now
+        }
+        session.updatedAt = updatedAt

        await this.saveStore(workspaceId)
        return newMessage
--- a/packages/ui/src/components/chat/multix-v2/index.tsx
+++ b/packages/ui/src/components/chat/multix-v2/index.tsx
@@ -77,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 = () => {
@@ -140,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
@@ -149,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 || []);
@@ -251,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();
@@ -304,7 +325,7 @@ export default function MultiXV2(props: MultiXV2Props) {
                syncFromStore();

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

                const s = soloState();
                if (s.isAutonomous) {
@@ -357,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
@@ -683,7 +683,25 @@ Now analyze the project and report your findings.`
  })

  const handleSessionSelect = (sessionId: string) => {
+    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)
--- 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/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)