test: cover TUI optimistic prompt sync

Author: kitlangton

packages/opencode/src/cli/cmd/tui/context/sync-optimistic.ts

@@ -0,0 +1,50 @@
+import type { AgentPartInput, FilePartInput, Message, Part, SubtaskPartInput, TextPartInput } from "@opencode-ai/sdk/v2"
+import { Binary } from "@opencode-ai/core/util/binary"
+
+export type OptimisticPromptPart = (TextPartInput | FilePartInput | AgentPartInput | SubtaskPartInput) & { id: string }
+
+export function optimisticParts(input: { sessionID: string; messageID: string; parts: OptimisticPromptPart[] }) {
+  return input.parts.map((part): Part => {
+    const withIDs = {
+      ...part,
+      sessionID: input.sessionID,
+      messageID: input.messageID,
+    }
+    if (withIDs.type === "file") return { ...withIDs, url: "" }
+    return withIDs
+  })
+}
+
+export function mergeFetchedMessages(input: {
+  currentMessages: Message[]
+  currentParts: Record<string, Part[] | undefined>
+  fetched: { info: Message; parts: Part[] }[]
+  optimisticMessages: ReadonlySet<string>
+}) {
+  const fetchedIDs = new Set(input.fetched.map((message) => message.info.id))
+  const messages = input.fetched.map((message) => message.info)
+  const parts = new Map<string, Part[]>()
+  const resolved = new Set<string>()
+
+  for (const message of input.currentMessages) {
+    if (input.optimisticMessages.has(message.id) && !fetchedIDs.has(message.id)) {
+      Binary.insert(messages, message, (item) => item.id)
+    }
+  }
+
+  for (const message of input.fetched) {
+    if (message.parts.length > 0) {
+      resolved.add(message.info.id)
+      parts.set(message.info.id, message.parts)
+      continue
+    }
+    if (input.optimisticMessages.has(message.info.id)) {
+      const current = input.currentParts[message.info.id]
+      if (current) parts.set(message.info.id, current)
+      continue
+    }
+    parts.set(message.info.id, message.parts)
+  }
+
+  return { messages, parts, resolved }
+}

packages/opencode/src/cli/cmd/tui/context/sync.tsx

@@ -4,10 +4,6 @@ import type {
   Provider,
   Session,
   Part,
-  TextPartInput,
-  FilePartInput,
-  AgentPartInput,
-  SubtaskPartInput,
   Config,
   Todo,
   Command,
@@ -36,8 +32,7 @@ import * as Log from "@opencode-ai/core/util/log"
 import { emptyConsoleState, type ConsoleState } from "@/config/console-state"
 import path from "path"
 import { useKV } from "./kv"
-
-type OptimisticPromptPart = (TextPartInput | FilePartInput | AgentPartInput | SubtaskPartInput) & { id: string }
+import { mergeFetchedMessages, optimisticParts, type OptimisticPromptPart } from "./sync-optimistic"
 
 export const { use: useSync, provider: SyncProvider } = createSimpleContext({
   name: "Sync",
@@ -226,6 +221,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
           break
 
         case "session.deleted": {
+          for (const message of store.message[event.properties.info.id] ?? []) optimisticMessages.delete(message.id)
           const result = Binary.search(store.session, event.properties.info.id, (s) => s.id)
           if (result.found) {
             setStore(
@@ -258,7 +254,6 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
         }
 
         case "message.updated": {
-          optimisticMessages.delete(event.properties.info.id)
           const messages = store.message[event.properties.info.sessionID]
           if (!messages) {
             setStore("message", event.properties.info.sessionID, [event.properties.info])
@@ -313,6 +308,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
           break
         }
         case "message.part.updated": {
+          optimisticMessages.delete(event.properties.part.messageID)
           const parts = store.part[event.properties.part.messageID]
           if (!parts) {
             setStore("part", event.properties.part.messageID, [event.properties.part])
@@ -386,6 +382,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
       const workspace = project.workspace.current()
       if (workspace !== syncedWorkspace) {
         fullSyncedSessions.clear()
+        optimisticMessages.clear()
         syncedWorkspace = workspace
       }
       const projectPromise = project.sync()
@@ -544,21 +541,6 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
               ...(input.variant ? { variant: input.variant } : {}),
             },
           }
-          const parts = input.parts.map((part): Part => {
-            const withIDs = {
-              ...part,
-              sessionID: input.sessionID,
-              messageID: input.messageID,
-            }
-            if (withIDs.type === "file") {
-              return {
-                ...withIDs,
-                url: "",
-              }
-            }
-            return withIDs
-          })
-
           batch(() => {
             if (!messages) {
               setStore("message", input.sessionID, [info])
@@ -571,7 +553,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
                 }),
               )
             }
-            setStore("part", input.messageID, reconcile(parts))
+            setStore("part", input.messageID, reconcile(optimisticParts(input)))
           })
         },
         removeOptimisticPrompt(sessionID: string, messageID: string) {
@@ -607,21 +589,21 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
           setStore(
             produce((draft) => {
               const match = Binary.search(draft.session, sessionID, (s) => s.id)
-              const fetched = messages.data!
-              const fetchedIDs = new Set(fetched.map((message) => message.info.id))
-              const optimistic = (draft.message[sessionID] ?? []).filter(
-                (message) => optimisticMessages.has(message.id) && !fetchedIDs.has(message.id),
-              )
+              const merged = mergeFetchedMessages({
+                currentMessages: draft.message[sessionID] ?? [],
+                currentParts: draft.part,
+                fetched: messages.data!,
+                optimisticMessages,
+              })
               if (match.found) draft.session[match.index] = session.data!
               if (!match.found) draft.session.splice(match.index, 0, session.data!)
               draft.todo[sessionID] = todo.data ?? []
-              draft.message[sessionID] = fetched.map((x) => x.info)
-              for (const message of optimistic) {
-                Binary.insert(draft.message[sessionID], message, (item) => item.id)
+              draft.message[sessionID] = merged.messages
+              for (const messageID of merged.resolved) {
+                optimisticMessages.delete(messageID)
               }
-              for (const message of fetched) {
-                optimisticMessages.delete(message.info.id)
-                draft.part[message.info.id] = message.parts
+              for (const [messageID, parts] of merged.parts) {
+                draft.part[messageID] = parts
               }
               draft.session_diff[sessionID] = diff.data ?? []
             }),

packages/opencode/test/cli/tui-sync-optimistic.test.ts

@@ -0,0 +1,92 @@
+import { describe, expect, test } from "bun:test"
+import type { Message, Part } from "@opencode-ai/sdk/v2"
+import { mergeFetchedMessages, optimisticParts } from "@/cli/cmd/tui/context/sync-optimistic"
+
+function user(id: string): Message {
+  return {
+    id,
+    sessionID: "ses_test",
+    role: "user",
+    time: { created: 1 },
+    agent: "build",
+    model: { providerID: "test", modelID: "model" },
+  }
+}
+
+function text(messageID: string, text: string): Part {
+  return {
+    id: `part_${messageID}`,
+    sessionID: "ses_test",
+    messageID,
+    type: "text",
+    text,
+  }
+}
+
+describe("TUI optimistic prompt sync", () => {
+  test("keeps an optimistic message while session sync has not fetched it yet", () => {
+    const merged = mergeFetchedMessages({
+      currentMessages: [user("msg_2")],
+      currentParts: { msg_2: [text("msg_2", "optimistic")] },
+      fetched: [{ info: user("msg_1"), parts: [text("msg_1", "persisted")] }],
+      optimisticMessages: new Set(["msg_2"]),
+    })
+
+    expect(merged.messages.map((message) => message.id)).toEqual(["msg_1", "msg_2"])
+    expect(merged.parts.get("msg_1")?.map((part) => (part.type === "text" ? part.text : ""))).toEqual(["persisted"])
+    expect(merged.resolved.has("msg_2")).toBe(false)
+  })
+
+  test("preserves optimistic parts when sync fetches the message before its parts", () => {
+    const merged = mergeFetchedMessages({
+      currentMessages: [user("msg_1")],
+      currentParts: { msg_1: [text("msg_1", "optimistic")] },
+      fetched: [{ info: user("msg_1"), parts: [] }],
+      optimisticMessages: new Set(["msg_1"]),
+    })
+
+    expect(merged.messages.map((message) => message.id)).toEqual(["msg_1"])
+    expect(merged.parts.get("msg_1")?.map((part) => (part.type === "text" ? part.text : ""))).toEqual(["optimistic"])
+    expect(merged.resolved.has("msg_1")).toBe(false)
+  })
+
+  test("replaces optimistic parts once real fetched parts arrive", () => {
+    const merged = mergeFetchedMessages({
+      currentMessages: [user("msg_1")],
+      currentParts: { msg_1: [text("msg_1", "optimistic")] },
+      fetched: [{ info: user("msg_1"), parts: [text("msg_1", "persisted")] }],
+      optimisticMessages: new Set(["msg_1"]),
+    })
+
+    expect(merged.parts.get("msg_1")?.map((part) => (part.type === "text" ? part.text : ""))).toEqual(["persisted"])
+    expect(merged.resolved.has("msg_1")).toBe(true)
+  })
+
+  test("strips file URLs from optimistic render parts", () => {
+    const parts = optimisticParts({
+      sessionID: "ses_test",
+      messageID: "msg_1",
+      parts: [
+        {
+          id: "part_file",
+          type: "file",
+          mime: "image/png",
+          filename: "image.png",
+          url: "data:image/png;base64,large",
+        },
+      ],
+    })
+
+    expect(parts).toEqual([
+      {
+        id: "part_file",
+        sessionID: "ses_test",
+        messageID: "msg_1",
+        type: "file",
+        mime: "image/png",
+        filename: "image.png",
+        url: "",
+      },
+    ])
+  })
+})