fix(tui): undo pending session input

packages/tui/src/context/data.tsx

@@ -438,6 +438,24 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
             )
           })
           break
+        case "session.input.withdrawn": {
+          removePending(event.data.sessionID, event.data.inputID)
+          if (store.session.input[event.data.sessionID]?.includes(event.data.inputID))
+            message.update(event.data.sessionID, (draft, index) => {
+              const position = index.get(event.data.inputID)
+              if (position === undefined) return
+              draft.splice(position, 1)
+              index.clear()
+              draft.forEach((message, indexValue) => index.set(message.id, indexValue))
+            })
+          setStore(
+            "session",
+            "input",
+            event.data.sessionID,
+            (store.session.input[event.data.sessionID] ?? []).filter((id) => id !== event.data.inputID),
+          )
+          break
+        }
         case "session.instructions.updated":
           const instructions = event.metadata?.instructions
           if (

packages/tui/src/mini/stream-v2.subagent.ts

@@ -643,6 +643,10 @@ export function createSubagentTracker(input: SubagentTrackerInput): SubagentTrac
       if (event.data.input.type === "user") child.prompts.set(event.data.inputID, event.data.input.data.text)
       return
     }
+    if (event.type === "session.input.withdrawn") {
+      child.prompts.delete(event.data.inputID)
+      return
+    }
     if (event.type === "session.input.promoted") {
       const prompt = child.prompts.get(event.data.inputID)
       if (prompt === undefined) return

packages/tui/src/mini/stream-v2.transport.ts

@@ -911,6 +911,12 @@ export async function createSessionTransport(input: StreamInput): Promise<Sessio
       })
       return
     }
+    if (event.type === "session.input.withdrawn") {
+      state.admitted.delete(event.data.inputID)
+      state.pending.delete(event.data.inputID)
+      syncPending()
+      return
+    }
     if (event.type === "session.input.promoted") {
       const waiting = state.wait?.messageID === event.data.inputID
       if (state.wait) promoteWait(state.wait, true, event.data.inputID)

packages/tui/src/routes/session/dialog-message.tsx

@@ -8,6 +8,7 @@ import { errorMessage } from "../../util/error"
 import { DialogFork } from "./dialog-fork"
 import type { PromptInfo } from "../../prompt/history"
 import { projectedPromptInput } from "../../prompt/codec"
+import { undoMessage } from "./undo"
 
 export function DialogMessage(props: {
   messageID: string
@@ -42,9 +43,11 @@ export function DialogMessage(props: {
                 pasted: [],
               })
             }
-            void client.api.session.revert
-              .stage({ sessionID: props.sessionID, messageID: props.messageID })
-              .catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
+            void undoMessage(client.api, {
+              sessionID: props.sessionID,
+              messageID: props.messageID,
+              pending: data.session.input.has(props.sessionID, props.messageID),
+            }).catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
             dialog.clear()
           },
         },

packages/tui/src/routes/session/index.tsx

@@ -64,6 +64,7 @@ import { useToast } from "../../ui/toast"
 import stripAnsi from "strip-ansi"
 import { usePromptRef } from "../../context/prompt"
 import { projectedPromptInput } from "../../prompt/codec"
+import { undoMessage } from "./undo"
 import { useEpilogue } from "../../context/epilogue"
 import { normalizePath } from "../../util/path"
 import { PermissionPrompt } from "./permission"
@@ -590,9 +591,11 @@ export function Session() {
           dialog.clear()
           return
         }
-        void client.api.session.revert
-          .stage({ sessionID: route.sessionID, messageID: message.id })
-          .catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
+        void undoMessage(client.api, {
+          sessionID: route.sessionID,
+          messageID: message.id,
+          pending: data.session.input.has(route.sessionID, message.id),
+        }).catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
         prompt()?.set({
           ...projectedPromptInput(message),
           pasted: [],

packages/tui/src/routes/session/undo.ts

@@ -0,0 +1,12 @@
+import type { OpenCodeClient } from "@opencode-ai/client"
+
+export function undoMessage(
+  client: OpenCodeClient,
+  input: { readonly sessionID: string; readonly messageID: string; readonly pending: boolean },
+) {
+  const revert = () => client.session.revert.stage(input).then(() => undefined)
+  if (!input.pending) return revert()
+  return client.session.pending
+    .withdraw({ sessionID: input.sessionID, inputID: input.messageID })
+    .then((withdrawn) => (withdrawn ? undefined : revert()))
+}

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

@@ -853,6 +853,65 @@ test("completes exploration when a queued prompt is promoted", async () => {
   }
 })
 
+test("removes optimistic input when it is withdrawn", async () => {
+  const events = createEventStream()
+  const sessionID = "session-withdrawal"
+  const calls = createFetch((url) => {
+    if (url.pathname === `/api/session/${sessionID}/message`) return json({ data: [], cursor: {} })
+  }, events)
+  let data!: ReturnType<typeof useData>
+  let client!: ReturnType<typeof useClient>
+
+  function Probe() {
+    data = useData()
+    client = useClient()
+    return <box />
+  }
+
+  const app = await testRender(() => (
+    <TestTuiContexts>
+      <ClientProvider api={createApi(calls.fetch)}>
+        <ProjectProvider>
+          <DataProvider>
+            <Probe />
+          </DataProvider>
+        </ProjectProvider>
+      </ClientProvider>
+    </TestTuiContexts>
+  ))
+
+  try {
+    await wait(() => client.connection.status() === "connected")
+    emitEvent(events, {
+      id: "evt_prompt_admitted",
+      created: 1,
+      type: "session.input.admitted",
+      durable: durable(sessionID, 1),
+      data: {
+        sessionID,
+        inputID: "message-user",
+        input: { type: "user", data: { text: "Never mind" }, delivery: "steer" },
+      },
+    })
+    await wait(() => data.session.message.get(sessionID, "message-user") !== undefined)
+    expect(data.session.input.has(sessionID, "message-user")).toBe(true)
+
+    emitEvent(events, {
+      id: "evt_prompt_withdrawn",
+      created: 2,
+      type: "session.input.withdrawn",
+      durable: durable(sessionID, 2),
+      data: { sessionID, inputID: "message-user" },
+    })
+
+    await wait(() => data.session.message.get(sessionID, "message-user") === undefined)
+    expect(data.session.input.has(sessionID, "message-user")).toBe(false)
+    expect(data.session.pending.list(sessionID)).toEqual([])
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
 test("classifies live tool rows independently of their call ID", async () => {
   const events = createEventStream()
   const sessionID = "session-tool-call-id"

packages/tui/test/cli/tui/undo.test.ts

@@ -0,0 +1,27 @@
+import { expect, test } from "bun:test"
+import { OpenCode } from "@opencode-ai/client"
+import { undoMessage } from "../../../src/routes/session/undo"
+
+test.each([
+  { pending: true, withdrawn: true, expected: ["withdraw"] },
+  { pending: true, withdrawn: false, expected: ["withdraw", "revert"] },
+  { pending: false, withdrawn: false, expected: ["revert"] },
+])("routes undo for pending=$pending withdrawn=$withdrawn", async ({ pending, withdrawn, expected }) => {
+  const calls: string[] = []
+  const client = OpenCode.make({
+    baseUrl: "http://localhost:3000",
+    fetch: Object.assign(
+      async (input: URL | RequestInfo, init?: BunFetchRequestInit | RequestInit) => {
+        const request = input instanceof Request ? input : new Request(input, init)
+        const operation = request.url.endsWith("/withdraw") ? "withdraw" : "revert"
+        calls.push(operation)
+        return Response.json({ data: operation === "withdraw" ? withdrawn : { messageID: "msg_user" } })
+      },
+      { preconnect: fetch.preconnect },
+    ),
+  })
+
+  await undoMessage(client, { sessionID: "ses_test", messageID: "msg_user", pending })
+
+  expect(calls).toEqual([...expected])
+})