diff --git a/packages/opencode/src/cli/cmd/tui/context/sync-optimistic.ts b/packages/opencode/src/cli/cmd/tui/context/sync-optimistic.ts new file mode 100644 index 0000000000..cede9bd313 --- /dev/null +++ b/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 + fetched: { info: Message; parts: Part[] }[] + optimisticMessages: ReadonlySet +}) { + const fetchedIDs = new Set(input.fetched.map((message) => message.info.id)) + const messages = input.fetched.map((message) => message.info) + const parts = new Map() + const resolved = new Set() + + 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 } +} diff --git a/packages/opencode/src/cli/cmd/tui/context/sync.tsx b/packages/opencode/src/cli/cmd/tui/context/sync.tsx index cc87e24353..6afd9aeb6f 100644 --- a/packages/opencode/src/cli/cmd/tui/context/sync.tsx +++ b/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 ?? [] }), diff --git a/packages/opencode/test/cli/tui-sync-optimistic.test.ts b/packages/opencode/test/cli/tui-sync-optimistic.test.ts new file mode 100644 index 0000000000..576d9ae7a5 --- /dev/null +++ b/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: "", + }, + ]) + }) +})