fix(chat): preserve tool duration across session switches (#1712)

* fix(chat): preserve tool duration across session switches

Fix #1636: ToolPart.tsx reset pinnedTime to empty on unmount/remount,
causing LiveDuration to not render on first paint. Now initializes
pinnedTime from server-provided time?.start/time?.end in the useState
initializer, eliminating the one-frame gap.

* fix(sync): preserve tool state.time in materialization merge

---------

Co-authored-by: Leonid Skorobogatyy <bash@opencode.itc.local>

packages/ui/src/components/chat/message/parts/ToolPart.tsx

@@ -2055,7 +2055,10 @@ const ToolPartContent: React.FC<ToolPartProps> = ({
     const input = stateWithData.input;
     const time = stateWithData.time;
 
-    const [pinnedTime, setPinnedTime] = React.useState<{ start?: number; end?: number }>({});
+    const [pinnedTime, setPinnedTime] = React.useState<{ start?: number; end?: number }>(() => ({
+        start: typeof time?.start === 'number' ? time.start : undefined,
+        end: typeof time?.end === 'number' ? time.end : undefined,
+    }));
     const [localStartAt, setLocalStartAt] = React.useState<number | undefined>(undefined);
     const [localFinalizedAt, setLocalFinalizedAt] = React.useState<number | undefined>(undefined);
 

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

@@ -122,6 +122,37 @@ describe("materializeSessionSnapshots", () => {
 
     expect(result.part.msg_1).toEqual([serverPart])
   })
+
+  test("preserves state.time from existing part when snapshot drops it", () => {
+    const livePart = {
+      id: "prt_1",
+      messageID: "msg_1",
+      sessionID: "ses_1",
+      type: "tool",
+      state: { status: "completed", time: { start: 1000, end: 2000 } },
+    } as unknown as Part
+    const snapshotPart = {
+      id: "prt_1",
+      messageID: "msg_1",
+      sessionID: "ses_1",
+      type: "tool",
+      state: { status: "completed" },
+    } as unknown as Part
+    const state = {
+      message: { ses_1: [message("msg_1")] },
+      part: { msg_1: [livePart] },
+    }
+
+    const result = materializeSessionSnapshots(
+      state,
+      "ses_1",
+      [{ info: message("msg_1"), parts: [snapshotPart] }],
+    )
+
+    const mergedPart = result.part.msg_1[0] as { state?: { time?: { start?: number; end?: number } } }
+    expect(mergedPart.state?.time?.start).toBe(1000)
+    expect(mergedPart.state?.time?.end).toBe(2000)
+  })
 })
 
 describe("getSessionMaterializationStatus", () => {

packages/ui/src/sync/materialization.ts

@@ -77,6 +77,15 @@ function hasLiveStreamingField(part: Part): boolean {
   })
 }
 
+function getPartStateTime(part: Part): { start?: number; end?: number } | undefined {
+  const stateTime = (part as { state?: { time?: { start?: unknown; end?: unknown } } }).state?.time
+  if (!stateTime || typeof stateTime !== "object") return undefined
+  const start = typeof stateTime.start === "number" ? stateTime.start : undefined
+  const end = typeof stateTime.end === "number" ? stateTime.end : undefined
+  if (start === undefined && end === undefined) return undefined
+  return { start, end }
+}
+
 function mergeMaterializedPart(existing: Part | undefined, next: Part): Part {
   if (!existing || getPartEndTime(next) !== undefined) return next
 
@@ -94,6 +103,20 @@ function mergeMaterializedPart(existing: Part | undefined, next: Part): Part {
     mergedRecord[field] = existingValue
   }
 
+  const existingTime = getPartStateTime(existing)
+  if (existingTime) {
+    const nextTime = getPartStateTime(next)
+    const preservedStart = nextTime?.start ?? existingTime.start
+    const preservedEnd = nextTime?.end ?? existingTime.end
+    if (preservedStart !== nextTime?.start || preservedEnd !== nextTime?.end) {
+      if (merged === next) merged = { ...next }
+      const mergedRecord = merged as Record<string, unknown>
+      const nextState = (next as Record<string, unknown>).state as Record<string, unknown> | undefined
+      const newState = { ...(nextState ?? {}), time: { start: preservedStart, end: preservedEnd } }
+      mergedRecord.state = newState
+    }
+  }
+
   return merged
 }