diff --git a/packages/ui/src/components/chat/message/parts/ToolPart.tsx b/packages/ui/src/components/chat/message/parts/ToolPart.tsx index b862860b..d1126352 100644 --- a/packages/ui/src/components/chat/message/parts/ToolPart.tsx +++ b/packages/ui/src/components/chat/message/parts/ToolPart.tsx @@ -2055,7 +2055,10 @@ const ToolPartContent: React.FC = ({ 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(undefined); const [localFinalizedAt, setLocalFinalizedAt] = React.useState(undefined); diff --git a/packages/ui/src/sync/__tests__/materialization.test.ts b/packages/ui/src/sync/__tests__/materialization.test.ts index 68d137b4..6ce540dc 100644 --- a/packages/ui/src/sync/__tests__/materialization.test.ts +++ b/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", () => { diff --git a/packages/ui/src/sync/materialization.ts b/packages/ui/src/sync/materialization.ts index cb2e436f..acb06ca9 100644 --- a/packages/ui/src/sync/materialization.ts +++ b/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 + const nextState = (next as Record).state as Record | undefined + const newState = { ...(nextState ?? {}), time: { start: preservedStart, end: preservedEnd } } + mergedRecord.state = newState + } + } + return merged }