fix(ui): use part_text_accum_delta to prevent markdown cutoff during streaming (#26822)

Author: Brendonovich

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)

packages/app/src/context/global-sync/event-reducer.test.ts

@@ -81,6 +81,7 @@ const baseState = (input: Partial<State> = {}) =>
     limit: 10,
     message: {},
     part: {},
+    part_text_accum_delta: {},
     ...input,
   }) as State
 

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,

packages/app/src/context/global-sync/session-cache.test.ts

@@ -39,6 +39,7 @@ describe("app session cache", () => {
       part: Record<string, Part[] | undefined>
       permission: Record<string, PermissionRequest[] | undefined>
       question: Record<string, QuestionRequest[] | undefined>
+      part_text_accum_delta: Record<string, string | undefined>
     } = {
       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<string, Part[] | undefined>
       permission: Record<string, PermissionRequest[] | undefined>
       question: Record<string, QuestionRequest[] | undefined>
+      part_text_accum_delta: Record<string, string | undefined>
     } = {
       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"])

packages/app/src/context/global-sync/session-cache.ts

@@ -18,6 +18,7 @@ type SessionCache = {
   part: Record<string, Part[] | undefined>
   permission: Record<string, PermissionRequest[] | undefined>
   question: Record<string, QuestionRequest[] | undefined>
+  part_text_accum_delta: Record<string, string | undefined>
 }
 
 export function dropSessionCaches(store: SessionCache, sessionIDs: Iterable<string>) {
@@ -27,6 +28,9 @@ export function dropSessionCaches(store: SessionCache, sessionIDs: Iterable<stri
   for (const key of Object.keys(store.part)) {
     const parts = store.part[key]
     if (!parts?.some((part) => stale.has(part?.sessionID ?? ""))) continue
+    for (const part of parts) {
+      delete store.part_text_accum_delta[part.id]
+    }
     delete store.part[key]
   }
 

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 = {

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 (
     <Show when={text()}>

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