fix(ui): reconcile stale active tools after reload (#3396)

* fix: reconcile stale active tools on materialization

* fix: recover stale turns after reload

packages/ui/src/sync/__tests__/materialization.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test"
-import type { Message, Part } from "@opencode-ai/sdk/v2/client"
+import type { Message, Part, ToolPart } from "@opencode-ai/sdk/v2/client"
 import {
   getSessionMaterializationRequestKey,
   getSessionMaterializationStatus,
@@ -16,6 +16,23 @@ function userMessage(id: string, sessionID = "ses_1"): Message {
   return { id, sessionID, role: "user", time: { created: 1 } } as Message
 }
 
+function completedAssistantMessage(id: string, sessionID = "ses_1"): Message {
+  return {
+    id,
+    sessionID,
+    role: "assistant",
+    time: { created: 1, completed: 4000 },
+    parentID: "msg_parent",
+    modelID: "model",
+    providerID: "provider",
+    mode: "mode",
+    agent: "agent",
+    path: { cwd: "/repo", root: "/repo" },
+    cost: 0,
+    tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+  }
+}
+
 function part(id: string, messageID: string, type = "text", text = id): Part {
   return { id, messageID, sessionID: "ses_1", type, text } as Part
 }
@@ -28,6 +45,69 @@ describe("getSessionMaterializationRequestKey", () => {
 })
 
 describe("materializeSessionSnapshots", () => {
+  test("finalizes an active tool under a completed assistant message", () => {
+    const completedMessage = completedAssistantMessage("msg_1")
+    const staleRunningTool = {
+      id: "prt_1",
+      messageID: "msg_1",
+      sessionID: "ses_1",
+      type: "tool",
+      tool: "bash",
+      state: { status: "running", input: { command: "ls" }, time: { start: 1000 } },
+      callID: "call-prt_1",
+    } satisfies ToolPart
+
+    const result = materializeSessionSnapshots(
+      { message: {}, part: {} },
+      "ses_1",
+      [{ info: completedMessage, parts: [staleRunningTool] }],
+    )
+
+    const reconciledPart = result.part.msg_1[0]
+    if (!reconciledPart || reconciledPart.type !== "tool") throw new Error("Expected tool part")
+    if (reconciledPart.state.status !== "error") throw new Error("Expected interrupted tool part")
+    expect(reconciledPart.state.error).toBe("Interrupted")
+    expect(reconciledPart.state.time).toEqual({ start: 1000, end: 4000 })
+    expect(getStaleRunningToolMessageID(result, "ses_1")).toBe(undefined)
+  })
+
+  test("preserves a terminal tool already observed when a completed snapshot is stale", () => {
+    const completedMessage = completedAssistantMessage("msg_1")
+    const terminalTool = {
+      id: "prt_1",
+      messageID: "msg_1",
+      sessionID: "ses_1",
+      type: "tool",
+      tool: "bash",
+      state: {
+        status: "completed",
+        input: { command: "ls" },
+        output: "done",
+        title: "bash",
+        metadata: {},
+        time: { start: 1000, end: 2000 },
+      },
+      callID: "call-prt_1",
+    } satisfies ToolPart
+    const staleRunningTool = {
+      ...terminalTool,
+      state: { status: "running", input: {}, time: { start: 1000 } },
+    } satisfies ToolPart
+    const state = {
+      message: { ses_1: [completedMessage] },
+      part: { msg_1: [terminalTool] },
+    }
+
+    const result = materializeSessionSnapshots(
+      state,
+      "ses_1",
+      [{ info: completedMessage, parts: [staleRunningTool] }],
+    )
+
+    expect(result.part).toBe(state.part)
+    expect(result.part.msg_1[0]).toBe(terminalTool)
+  })
+
   test("marks an empty successful page as materialized", () => {
     const result = materializeSessionSnapshots(
       { message: {}, part: {} },

packages/ui/src/sync/__tests__/message-completion-status-poll.test.ts

@@ -7,7 +7,7 @@
  */
 import { beforeEach, describe, expect, mock, test } from "bun:test"
 import { create, type StoreApi } from "zustand"
-import type { SessionStatus } from "@opencode-ai/sdk/v2/client"
+import type { Message, Part, SessionStatus } from "@opencode-ai/sdk/v2/client"
 import { INITIAL_STATE } from "../types"
 import type { DirectoryStore } from "../child-store"
 
@@ -29,17 +29,43 @@ mock.module("@/lib/runtime-switch", () => ({
   getRuntimeKey: () => "test-runtime",
 }))
 
-import { maybePollStatusAfterMessageCompletion, MESSAGE_COMPLETION_STATUS_POLL_DELAY_MS } from "../sync-context"
+import {
+  maybePollStatusAfterMessageCompletion,
+  MESSAGE_COMPLETION_STATUS_POLL_DELAY_MS,
+  recoverInterruptedTurnAfterMessageLoad,
+} from "../sync-context"
 
-const createStore = (status: SessionStatus): StoreApi<DirectoryStore> => {
+const createStore = (status?: SessionStatus): StoreApi<DirectoryStore> => {
+  const session_status: DirectoryStore["session_status"] = {}
+  if (status) session_status.ses_1 = status
   return create<DirectoryStore>()((set) => ({
     ...INITIAL_STATE,
-    session_status: { ses_1: status },
+    session_status,
     patch: (partial) => set(partial),
     replace: (next) => set(next),
   }))
 }
 
+// SAFETY: The recovery path reads only the identity, role, and completion time
+// fields from this synthetic assistant message.
+const unfinishedAssistant = {
+  id: "msg_1",
+  sessionID: "ses_1",
+  role: "assistant",
+  time: { created: 1 },
+} as Message
+
+// SAFETY: The recovery path reads only the tool discriminator and state fields
+// from this synthetic part.
+const runningTool = {
+  id: "part_1",
+  messageID: "msg_1",
+  sessionID: "ses_1",
+  type: "tool",
+  tool: "bash",
+  state: { status: "running", time: { start: 1 }, input: {} },
+} as Part
+
 const sleep = (ms: number): Promise<void> => new Promise((resolve) => setTimeout(resolve, ms))
 
 /** Past the deferral, plus room for the background-network task chain. */
@@ -138,4 +164,23 @@ describe("maybePollStatusAfterMessageCompletion (issue OPE-193)", () => {
     expect(statusSnapshotCalls).toEqual(["/test/project", "/test/project"])
     expect(store.getState().session_status?.ses_1?.type).toBe("idle")
   })
+
+  test("recovers an unfinished turn after reload when status was initially unknown", async () => {
+    const store = createStore()
+    store.getState().patch({
+      message: { ses_1: [unfinishedAssistant] },
+      part: { msg_1: [runningTool] },
+    })
+
+    await recoverInterruptedTurnAfterMessageLoad("/test/project", store, "ses_1")
+
+    expect(statusSnapshotCalls).toEqual(["/test/project"])
+    expect(store.getState().session_status?.ses_1?.type).toBe("idle")
+    const message = store.getState().message.ses_1[0]
+    expect(message?.role).toBe("assistant")
+    if (message?.role === "assistant") expect(message.time.completed).toBeDefined()
+    const part = store.getState().part.msg_1[0]
+    expect(part?.type).toBe("tool")
+    if (part?.type === "tool") expect(part.state.status).toBe("error")
+  })
 })

packages/ui/src/sync/__tests__/session-status-snapshot.test.ts

@@ -86,6 +86,13 @@ describe("applySessionStatusSnapshot", () => {
       expect(changed).toBe(true)
       expect(store.getState().session_status.ses_a).toEqual({ type: "idle" })
     })
+
+    test("seeds idle for a candidate with no previous status entry", () => {
+      const store = createDirectoryStore({ session_status: {} })
+      const changed = applySessionStatusSnapshot(store, {} as StatusSnapshot, ["ses_a"], "authoritative")
+      expect(changed).toBe(true)
+      expect(store.getState().session_status.ses_a).toEqual({ type: "idle" })
+    })
   })
 })