fix(tui): render submitted prompts optimistically

bun.lock

@@ -1003,6 +1003,7 @@
         "@opencode-ai/client": "workspace:*",
         "@opencode-ai/core": "workspace:*",
         "@opencode-ai/plugin": "workspace:*",
+        "@opencode-ai/schema": "workspace:*",
         "@opencode-ai/sdk": "workspace:*",
         "@opencode-ai/simulation": "workspace:*",
         "@opencode-ai/ui": "workspace:*",

packages/tui/package.json

@@ -55,6 +55,7 @@
     "@opencode-ai/client": "workspace:*",
     "@opencode-ai/core": "workspace:*",
     "@opencode-ai/plugin": "workspace:*",
+    "@opencode-ai/schema": "workspace:*",
     "@opencode-ai/sdk": "workspace:*",
     "@opencode-ai/simulation": "workspace:*",
     "@opencode-ai/ui": "workspace:*",

packages/tui/src/component/prompt/index.tsx

@@ -40,6 +40,7 @@ import { useRenderer, useTerminalDimensions, type JSX } from "@opentui/solid"
 import type { UserMessage } from "@opencode-ai/sdk/v2"
 import { Locale } from "../../util/locale"
 import { errorMessage } from "../../util/error"
+import { SessionMessage } from "@opencode-ai/schema/session-message"
 import { createColors, createFrames } from "../../ui/spinner"
 import { useDialog } from "../../ui/dialog"
 import { DialogIntegration } from "../dialog-integration"
@@ -1146,9 +1147,18 @@ export function Prompt(props: PromptProps) {
           return false
         }
       }
+      const messageID = SessionMessage.ID.create()
+      data.session.input.optimistic(sessionID, {
+        id: messageID,
+        type: "user",
+        text: inputText,
+        agents: store.prompt.agents,
+        time: { created: Date.now() },
+      })
       const error = await sdk.api.session
         .prompt({
           sessionID,
+          id: messageID,
           text: inputText,
           files: store.prompt.files,
           agents: store.prompt.agents,
@@ -1158,6 +1168,7 @@ export function Prompt(props: PromptProps) {
           (error) => error,
         )
       if (error) {
+        data.session.input.rollback(sessionID, messageID)
         toast.show({ title: "Failed to send prompt", message: errorMessage(error), variant: "error" })
         return false
       }

packages/tui/src/context/data.tsx

@@ -21,6 +21,7 @@ import type {
   SessionMessageAssistantText,
   SessionMessageAssistantTool,
   SessionInfo,
+  SessionMessageUser,
   Shell,
   SkillInfo,
 } from "@opencode-ai/sdk/v2"
@@ -105,6 +106,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
       directory: process.cwd(),
     })
     const messageIndex = new Map<string, Map<string, number>>()
+    const optimisticInput = new Map<string, Set<string>>()
     let bootstrapping: Promise<void> | undefined
     let connected = false
 
@@ -214,6 +216,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
 
     function removeSession(sessionID: string) {
       messageIndex.delete(sessionID)
+      optimisticInput.delete(sessionID)
       setStore(
         "session",
         produce((draft) => {
@@ -327,32 +330,35 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           )
           break
         }
-        case "session.input.admitted":
+        case "session.input.admitted": {
+          const inputs = optimisticInput.get(event.data.sessionID)
+          if (inputs?.delete(event.data.inputID) && inputs.size === 0) optimisticInput.delete(event.data.sessionID)
           if (!store.session.input[event.data.sessionID]?.includes(event.data.inputID))
             setStore("session", "input", event.data.sessionID, [
               ...(store.session.input[event.data.sessionID] ?? []),
               event.data.inputID,
             ])
           message.update(event.data.sessionID, (draft, index) => {
-            message.append(
+            const item =
-              draft,
-              index,
               event.data.input.type === "user"
                 ? {
                     id: event.data.inputID,
-                    type: "user",
+                    type: "user" as const,
                     ...event.data.input.data,
                     time: { created: event.created },
                   }
                 : {
                     id: event.data.inputID,
-                    type: "synthetic",
+                    type: "synthetic" as const,
                     ...event.data.input.data,
                     time: { created: event.created },
-                  },
+                  }
-            )
+            const position = index.get(event.data.inputID)
+            if (position === undefined) return message.append(draft, index, item)
+            draft[position] = item
           })
           break
+        }
         case "session.instructions.updated":
           message.update(event.data.sessionID, (draft, index) => {
             message.append(draft, index, {
@@ -814,6 +820,30 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           return store.session.status[sessionID] ?? "idle"
         },
         input: {
+          optimistic(sessionID: string, item: SessionMessageUser) {
+            const inputs = optimisticInput.get(sessionID) ?? new Set<string>()
+            inputs.add(item.id)
+            optimisticInput.set(sessionID, inputs)
+            if (!store.session.input[sessionID]?.includes(item.id))
+              setStore("session", "input", sessionID, [...(store.session.input[sessionID] ?? []), item.id])
+            message.update(sessionID, (draft, index) => message.append(draft, index, item))
+          },
+          rollback(sessionID: string, inputID: string) {
+            const inputs = optimisticInput.get(sessionID)
+            if (!inputs?.delete(inputID)) return
+            if (inputs.size === 0) optimisticInput.delete(sessionID)
+            setStore(
+              "session",
+              produce((draft) => {
+                draft.input[sessionID] = (draft.input[sessionID] ?? []).filter((id) => id !== inputID)
+                const messages = draft.message[sessionID]
+                const position = messageIndex.get(sessionID)?.get(inputID)
+                if (!messages || position === undefined) return
+                messages.splice(position, 1)
+                messageIndex.set(sessionID, new Map(messages.map((item, index) => [item.id, index])))
+              }),
+            )
+          },
           list(sessionID: string) {
             return store.session.input[sessionID] ?? []
           },

packages/tui/test/cli/tui/data.test.tsx

@@ -1064,9 +1064,7 @@ test("tracks session status from active sessions and execution events", async ()
       const message = data.session.message.get("session-manual", "message-compaction")
       return message?.type === "compaction" && message.status === "running" && message.summary === "Streamed summary"
     })
-    const compactionRow = manualRows.find(
+    const compactionRow = manualRows.find((row) => row.type === "message" && row.messageID === "message-compaction")
-      (row) => row.type === "message" && row.messageID === "message-compaction",
-    )
     emitEvent(events, {
       id: "evt_manual_compaction_ended",
       created: 3,
@@ -1108,9 +1106,7 @@ test("tracks session status from active sessions and execution events", async ()
       const message = data.session.message.get("session-live", "msg_compaction_started")
       return message?.type === "compaction" && message.status === "running" && message.summary === "Live summary"
     })
-    const autoCompactionRow = rows.find(
+    const autoCompactionRow = rows.find((row) => row.type === "message" && row.messageID === "msg_compaction_started")
-      (row) => row.type === "message" && row.messageID === "msg_compaction_started",
-    )
 
     emitEvent(events, {
       id: "evt_compaction_ended",
@@ -2059,6 +2055,14 @@ test("renders admitted prompts immediately and tracks them until promoted", asyn
     await mounted
     const received: string[] = []
     const unsubscribe = sync.listen((event) => received.push(event.name))
+    sync.session.input.optimistic(sessionID, {
+      id: messageID,
+      type: "user",
+      text: "optimistic",
+      time: { created: -1 },
+    })
+    expect(sync.session.message.get(sessionID, messageID)).toMatchObject({ text: "optimistic", time: { created: -1 } })
+    expect(sync.session.input.list(sessionID)).toEqual([messageID])
     emitEvent(events, {
       id: "evt_admitted_1",
       created: 0,
@@ -2070,11 +2074,13 @@ test("renders admitted prompts immediately and tracks them until promoted", asyn
         input: { type: "user", data: { text: "hello" }, delivery: "steer" },
       },
     })
-    await wait(() => sync.session.message.list(sessionID)?.length === 1)
+    await wait(() => sync.session.message.get(sessionID, messageID)?.time.created === 0)
     const admitted = sync.session.message.list(sessionID)?.[0]
     expect(admitted).toMatchObject({ id: messageID, type: "user", text: "hello" })
     expect(admitted?.metadata).toBeUndefined()
     expect(sync.session.input.list(sessionID)).toEqual([messageID])
+    sync.session.input.rollback(sessionID, messageID)
+    expect(sync.session.message.ids(sessionID)).toEqual([messageID])
 
     await sync.session.message.refresh(sessionID)
     expect(sync.session.message.list(sessionID)?.[0]?.metadata).toBeUndefined()
@@ -2109,6 +2115,66 @@ test("renders admitted prompts immediately and tracks them until promoted", asyn
   }
 })
 
+test("rolls back unconfirmed optimistic prompts and accepts later admission", async () => {
+  const events = createEventStream()
+  const sessionID = "session-1"
+  const messageID = "msg_user_rollback"
+  const calls = createFetch(undefined, events)
+  let sync!: ReturnType<typeof useData>
+  let ready!: () => void
+  const mounted = new Promise<void>((resolve) => {
+    ready = resolve
+  })
+
+  function Probe() {
+    sync = useData()
+    onMount(ready)
+    return <box />
+  }
+
+  const app = await testRender(() => (
+    <TestTuiContexts>
+      <SDKProvider client={createClient(calls.fetch)} api={createApi(calls.fetch)}>
+        <ProjectProvider>
+          <DataProvider>
+            <Probe />
+          </DataProvider>
+        </ProjectProvider>
+      </SDKProvider>
+    </TestTuiContexts>
+  ))
+
+  try {
+    await mounted
+    sync.session.input.optimistic(sessionID, {
+      id: messageID,
+      type: "user",
+      text: "optimistic",
+      time: { created: -1 },
+    })
+    sync.session.input.rollback(sessionID, messageID)
+    expect(sync.session.message.ids(sessionID)).toEqual([])
+    expect(sync.session.input.list(sessionID)).toEqual([])
+
+    emitEvent(events, {
+      id: "evt_admitted_after_rollback",
+      created: 1,
+      type: "session.input.admitted",
+      durable: durable(sessionID),
+      data: {
+        sessionID,
+        inputID: messageID,
+        input: { type: "user", data: { text: "accepted" }, delivery: "steer" },
+      },
+    })
+    await wait(() => sync.session.message.ids(sessionID).length === 1)
+    expect(sync.session.message.get(sessionID, messageID)).toMatchObject({ text: "accepted", time: { created: 1 } })
+    expect(sync.session.input.list(sessionID)).toEqual([messageID])
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
 test("projects live instruction updates with their message ID", async () => {
   const events = createEventStream()
   const calls = createFetch(undefined, events)
@@ -2180,8 +2246,7 @@ function sessionInfo(id: string, parentID: string | undefined, cost = 0) {
 async function mountData(parents: Record<string, string>, costs: Record<string, number> = {}) {
   const calls = createFetch((url) => {
     const match = url.pathname.match(/^\/api\/session\/([^/]+)$/)
-    if (match && match[1] !== "active")
+    if (match && match[1] !== "active") return json({ data: sessionInfo(match[1], parents[match[1]], costs[match[1]]) })
-      return json({ data: sessionInfo(match[1], parents[match[1]], costs[match[1]]) })
   })
   let data!: ReturnType<typeof useData>
   let ready!: () => void
@@ -2250,10 +2315,7 @@ test("indexes arbitrarily deep nesting under a single root", async () => {
 })
 
 test("totals family cost for roots and keeps subagent cost scoped", async () => {
-  const { data, app } = await mountData(
+  const { data, app } = await mountData({ grandchild: "child", child: "root" }, { root: 1, child: 2, grandchild: 3 })
-    { grandchild: "child", child: "root" },
-    { root: 1, child: 2, grandchild: 3 },
-  )
   try {
     await data.session.refresh("grandchild")
     await data.session.refresh("child")