fix(tui): restore reverted prompt in v2

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

@@ -72,6 +72,7 @@ import { usePathFormatter } from "../../context/path-format"
 import { LocationProvider } from "../../context/location"
 import { createSessionRows, type PartRef, type SessionRow } from "./rows"
 import { switchLabel } from "../../util/model"
+import { revertedPrompt } from "../../util/revert-prompt"
 
 addDefaultParsers(parsers.parsers)
 
@@ -403,24 +404,28 @@ export function Session() {
         void (async () => {
           const boundary = session()?.revert?.messageID
           const list = messages()
-          let target: string | undefined
+          let target: SessionMessageUser | undefined
           for (let i = list.length - 1; i >= 0; i--) {
             const message = list[i]
             if (message.type !== "user" || !message.text.trim()) continue
             if (boundary && message.id >= boundary) continue
-            target = message.id
+            target = message
             break
           }
           if (!target) {
             toast.show({ message: "Nothing to undo", variant: "error", duration: 3000 })
             dialog.clear()
             return
           }
-          const error = await sdk.api.session.revertStage({ sessionID: route.sessionID, messageID: target }).then(
+          const error = await sdk.api.session.revertStage({ sessionID: route.sessionID, messageID: target.id }).then(
             () => undefined,
             (error) => error,
           )
           if (error) toast.show({ message: errorMessage(error), variant: "error", duration: 5000 })
+          if (!error && prompt) {
+            const restored = revertedPrompt(prompt.current, target)
+            if (restored) prompt.set(restored)
+          }
           dialog.clear()
         })()
       },

packages/tui/src/util/revert-prompt.ts

@@ -0,0 +1,41 @@
+import type { SessionMessageUser } from "@opencode-ai/sdk/v2"
+import type { PromptInfo } from "../prompt/history"
+
+export function revertedPrompt(current: PromptInfo, message: SessionMessageUser): PromptInfo | undefined {
+  const input = current.input.trim()
+  if (current.parts.length || (input && !"/undo".startsWith(input))) return
+
+  return {
+    input: message.text,
+    parts: [
+      ...(message.files ?? []).map((file) => ({
+        type: "file" as const,
+        mime: file.mime,
+        filename: file.name,
+        url: file.uri,
+        source: file.source
+          ? {
+              type: "file" as const,
+              path: file.name ?? file.uri,
+              text: {
+                start: file.source.start,
+                end: file.source.end,
+                value: file.source.text,
+              },
+            }
+          : undefined,
+      })),
+      ...(message.agents ?? []).map((agent) => ({
+        type: "agent" as const,
+        name: agent.name,
+        source: agent.source
+          ? {
+              start: agent.source.start,
+              end: agent.source.end,
+              value: agent.source.text,
+            }
+          : undefined,
+      })),
+    ],
+  }
+}

packages/tui/test/util/revert-prompt.test.ts

@@ -0,0 +1,65 @@
+import { describe, expect, test } from "bun:test"
+import type { SessionMessageUser } from "@opencode-ai/sdk/v2"
+import { revertedPrompt } from "../../src/util/revert-prompt"
+
+const message: SessionMessageUser = {
+  id: "message-1",
+  type: "user",
+  text: "Fix the tests",
+  time: { created: 1 },
+  files: [
+    {
+      uri: "file:///repo/test.ts",
+      mime: "text/typescript",
+      name: "test.ts",
+      source: { start: 0, end: 8, text: "@test.ts" },
+    },
+  ],
+  agents: [{ name: "review", source: { start: 9, end: 16, text: "@review" } }],
+}
+
+describe("reverted prompt", () => {
+  test("restores the reverted user message into an empty prompt", () => {
+    expect(revertedPrompt({ input: "", parts: [] }, message)).toEqual({
+      input: "Fix the tests",
+      parts: [
+        {
+          type: "file",
+          mime: "text/typescript",
+          filename: "test.ts",
+          url: "file:///repo/test.ts",
+          source: {
+            type: "file",
+            path: "test.ts",
+            text: { start: 0, end: 8, value: "@test.ts" },
+          },
+        },
+        {
+          type: "agent",
+          name: "review",
+          source: { start: 9, end: 16, value: "@review" },
+        },
+      ],
+    })
+  })
+
+  test.each(["/", "/u", "/un", "/und", "/undo", " /undo "])("replaces the undo autocomplete query %p", (input) => {
+    expect(revertedPrompt({ input, parts: [] }, message)?.input).toBe("Fix the tests")
+  })
+
+  test("preserves an existing text draft", () => {
+    expect(revertedPrompt({ input: "Keep this", parts: [] }, message)).toBeUndefined()
+  })
+
+  test("preserves an existing attachment draft", () => {
+    expect(
+      revertedPrompt(
+        {
+          input: "",
+          parts: [{ type: "agent", name: "build" }],
+        },
+        message,
+      ),
+    ).toBeUndefined()
+  })
+})