diff --git a/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx b/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx index c6bcb89924a6..f84ad1dee8da 100644 --- a/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx +++ b/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx @@ -1170,25 +1170,27 @@ export function Prompt(props: PromptProps) { })), }) } else { + const parts = [ + ...editorParts, + { + id: PartID.ascending(), + type: "text" as const, + text: inputText, + }, + ...nonTextParts.map(assign), + ] + const request = { + sessionID, + messageID, + agent: agent.name, + model: selectedModel, + variant, + parts, + } + sync.session.addOptimisticPrompt(request) sdk.client.session - .prompt({ - sessionID, - ...selectedModel, - messageID, - agent: agent.name, - model: selectedModel, - variant, - parts: [ - ...editorParts, - { - id: PartID.ascending(), - type: "text", - text: inputText, - }, - ...nonTextParts.map(assign), - ], - }) - .catch(() => {}) + .prompt(request) + .catch(() => sync.session.removeOptimisticPrompt(request.sessionID, request.messageID)) if (editorParts.length > 0) editor.markSelectionSent() } history.append({ 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 000000000000..cede9bd3131a --- /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 24609dd81e4f..6afd9aeb6ffc 100644 --- a/packages/opencode/src/cli/cmd/tui/context/sync.tsx +++ b/packages/opencode/src/cli/cmd/tui/context/sync.tsx @@ -32,6 +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" +import { mergeFetchedMessages, optimisticParts, type OptimisticPromptPart } from "./sync-optimistic" export const { use: useSync, provider: SyncProvider } = createSimpleContext({ name: "Sync", @@ -112,6 +113,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({ const kv = useKV() const fullSyncedSessions = new Set() + const optimisticMessages = new Set() let syncedWorkspace = project.workspace.current() function sessionListQuery(): { scope?: "project"; path?: string } { @@ -219,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( @@ -290,6 +293,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({ break } case "message.removed": { + optimisticMessages.delete(event.properties.messageID) const messages = store.message[event.properties.sessionID] const result = Binary.search(messages, event.properties.messageID, (m) => m.id) if (result.found) { @@ -304,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]) @@ -377,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() @@ -512,6 +518,66 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({ if (last.role === "user") return "working" return last.time.completed ? "idle" : "working" }, + addOptimisticPrompt(input: { + sessionID: string + messageID: string + agent: string + model: { providerID: string; modelID: string } + variant?: string + parts: OptimisticPromptPart[] + }) { + optimisticMessages.add(input.messageID) + const messages = store.message[input.sessionID] + const match = messages ? Binary.search(messages, input.messageID, (m) => m.id) : undefined + const info: Message = { + id: input.messageID, + sessionID: input.sessionID, + role: "user", + time: { created: Date.now() }, + agent: input.agent, + model: { + providerID: input.model.providerID, + modelID: input.model.modelID, + ...(input.variant ? { variant: input.variant } : {}), + }, + } + batch(() => { + if (!messages) { + setStore("message", input.sessionID, [info]) + } else if (!match?.found) { + setStore( + "message", + input.sessionID, + produce((draft) => { + Binary.insert(draft, info, (message) => message.id) + }), + ) + } + setStore("part", input.messageID, reconcile(optimisticParts(input))) + }) + }, + removeOptimisticPrompt(sessionID: string, messageID: string) { + if (!optimisticMessages.delete(messageID)) return + const messages = store.message[sessionID] + const match = messages ? Binary.search(messages, messageID, (m) => m.id) : undefined + batch(() => { + if (match?.found) { + setStore( + "message", + sessionID, + produce((draft) => { + draft.splice(match.index, 1) + }), + ) + } + setStore( + "part", + produce((draft) => { + delete draft[messageID] + }), + ) + }) + }, async sync(sessionID: string) { if (fullSyncedSessions.has(sessionID)) return const [session, messages, todo, diff] = await Promise.all([ @@ -523,12 +589,21 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({ setStore( produce((draft) => { const match = Binary.search(draft.session, sessionID, (s) => s.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] = messages.data!.map((x) => x.info) - for (const message of messages.data!) { - draft.part[message.info.id] = message.parts + draft.message[sessionID] = merged.messages + for (const messageID of merged.resolved) { + optimisticMessages.delete(messageID) + } + 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 000000000000..576d9ae7a511 --- /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: "", + }, + ]) + }) +})