diff --git a/packages/app/src/context/global-sync/child-store.ts b/packages/app/src/context/global-sync/child-store.ts index 0138310cdc..737c6bedc9 100644 --- a/packages/app/src/context/global-sync/child-store.ts +++ b/packages/app/src/context/global-sync/child-store.ts @@ -231,6 +231,7 @@ export function createChildStoreManager(input: { limit: 5, message: {}, part: {}, + part_text_accum_delta: {}, }) children[key] = child disposers.set(key, dispose) diff --git a/packages/app/src/context/global-sync/event-reducer.test.ts b/packages/app/src/context/global-sync/event-reducer.test.ts index 892129788e..f02ac5c7be 100644 --- a/packages/app/src/context/global-sync/event-reducer.test.ts +++ b/packages/app/src/context/global-sync/event-reducer.test.ts @@ -81,6 +81,7 @@ const baseState = (input: Partial = {}) => limit: 10, message: {}, part: {}, + part_text_accum_delta: {}, ...input, }) as State diff --git a/packages/app/src/context/global-sync/event-reducer.ts b/packages/app/src/context/global-sync/event-reducer.ts index 5f43c341bc..13d34ef6c5 100644 --- a/packages/app/src/context/global-sync/event-reducer.ts +++ b/packages/app/src/context/global-sync/event-reducer.ts @@ -211,6 +211,12 @@ export function applyDirectoryEvent(input: { const result = Binary.search(messages, props.messageID, (m) => m.id) if (result.found) messages.splice(result.index, 1) } + const parts = draft.part[props.messageID] + if (parts) { + for (const part of parts) { + delete draft.part_text_accum_delta[part.id] + } + } delete draft.part[props.messageID] }), ) @@ -219,6 +225,11 @@ export function applyDirectoryEvent(input: { case "message.part.updated": { const part = (event.properties as { part: Part }).part if (SKIP_PARTS.has(part.type)) break + input.setStore( + produce((draft) => { + delete draft.part_text_accum_delta[part.id] + }), + ) const parts = input.store.part[part.messageID] if (!parts) { input.setStore("part", part.messageID, [part]) @@ -240,6 +251,11 @@ export function applyDirectoryEvent(input: { } case "message.part.removed": { const props = event.properties as { messageID: string; partID: string } + input.setStore( + produce((draft) => { + delete draft.part_text_accum_delta[props.partID] + }), + ) const parts = input.store.part[props.messageID] if (!parts) break const result = Binary.search(parts, props.partID, (p) => p.id) @@ -263,6 +279,7 @@ export function applyDirectoryEvent(input: { if (!parts) break const result = Binary.search(parts, props.partID, (p) => p.id) if (!result.found) break + input.setStore("part_text_accum_delta", props.partID, (existing) => (existing ?? "") + props.delta) input.setStore( "part", props.messageID, diff --git a/packages/app/src/context/global-sync/session-cache.test.ts b/packages/app/src/context/global-sync/session-cache.test.ts index 472ac219e9..4b2be505ea 100644 --- a/packages/app/src/context/global-sync/session-cache.test.ts +++ b/packages/app/src/context/global-sync/session-cache.test.ts @@ -39,6 +39,7 @@ describe("app session cache", () => { part: Record permission: Record question: Record + part_text_accum_delta: Record } = { session_status: { ses_1: { type: "busy" } as SessionStatus }, session_diff: { ses_1: [] }, @@ -47,12 +48,14 @@ describe("app session cache", () => { part: { msg_1: [part("prt_1", "ses_1", "msg_1")] }, permission: { ses_1: [] as PermissionRequest[] }, question: { ses_1: [] as QuestionRequest[] }, + part_text_accum_delta: { prt_1: "streamed text" }, } dropSessionCaches(store, ["ses_1"]) expect(store.message.ses_1).toBeUndefined() expect(store.part.msg_1).toBeUndefined() + expect(store.part_text_accum_delta.prt_1).toBeUndefined() expect(store.todo.ses_1).toBeUndefined() expect(store.session_diff.ses_1).toBeUndefined() expect(store.session_status.ses_1).toBeUndefined() @@ -70,6 +73,7 @@ describe("app session cache", () => { part: Record permission: Record question: Record + part_text_accum_delta: Record } = { session_status: {}, session_diff: {}, @@ -78,6 +82,7 @@ describe("app session cache", () => { part: { [m.id]: [part("prt_1", "ses_1", m.id)] }, permission: {}, question: {}, + part_text_accum_delta: {}, } dropSessionCaches(store, ["ses_1"]) diff --git a/packages/app/src/context/global-sync/session-cache.ts b/packages/app/src/context/global-sync/session-cache.ts index 6f4d81062b..05cdc84643 100644 --- a/packages/app/src/context/global-sync/session-cache.ts +++ b/packages/app/src/context/global-sync/session-cache.ts @@ -18,6 +18,7 @@ type SessionCache = { part: Record permission: Record question: Record + part_text_accum_delta: Record } export function dropSessionCaches(store: SessionCache, sessionIDs: Iterable) { @@ -27,6 +28,9 @@ export function dropSessionCaches(store: SessionCache, sessionIDs: Iterable stale.has(part?.sessionID ?? ""))) continue + for (const part of parts) { + delete store.part_text_accum_delta[part.id] + } delete store.part[key] } diff --git a/packages/app/src/context/global-sync/types.ts b/packages/app/src/context/global-sync/types.ts index e3ec83c5ee..6bf42a0737 100644 --- a/packages/app/src/context/global-sync/types.ts +++ b/packages/app/src/context/global-sync/types.ts @@ -72,6 +72,9 @@ export type State = { part: { [messageID: string]: Part[] } + part_text_accum_delta: { + [partID: string]: string + } } export type VcsCache = { diff --git a/packages/ui/src/components/message-part.tsx b/packages/ui/src/components/message-part.tsx index 137f689756..7a7d5b15fa 100644 --- a/packages/ui/src/components/message-part.tsx +++ b/packages/ui/src/components/message-part.tsx @@ -1461,7 +1461,7 @@ PART_MAPPING["text"] = function TextPartDisplay(props) { const streaming = createMemo( () => props.message.role === "assistant" && typeof (props.message as AssistantMessage).time.completed !== "number", ) - const text = () => (part().text ?? "").trim() + const text = () => (data.store.part_text_accum_delta?.[part().id] ?? part().text ?? "").trim() const isLastTextPart = createMemo(() => { const last = (data.store.part?.[props.message.id] ?? []) .filter((item): item is TextPart => item?.type === "text" && !!item.text?.trim()) @@ -1521,11 +1521,12 @@ PART_MAPPING["text"] = function TextPartDisplay(props) { } PART_MAPPING["reasoning"] = function ReasoningPartDisplay(props) { + const data = useData() const part = () => props.part as ReasoningPart const streaming = createMemo( () => props.message.role === "assistant" && typeof (props.message as AssistantMessage).time.completed !== "number", ) - const text = () => part().text.trim() + const text = () => (data.store.part_text_accum_delta?.[part().id] ?? part().text).trim() return ( diff --git a/packages/ui/src/context/data.tsx b/packages/ui/src/context/data.tsx index 632bed0cfa..3d015257f3 100644 --- a/packages/ui/src/context/data.tsx +++ b/packages/ui/src/context/data.tsx @@ -24,6 +24,9 @@ type Data = { part: { [messageID: string]: Part[] } + part_text_accum_delta?: { + [partID: string]: string + } } export type NavigateToSessionFn = (sessionID: string) => void