fix(tui): pass prompt setter to revert dialog

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

@@ -6,8 +6,14 @@ import { useToast } from "../../ui/toast"
 import { useSDK } from "../../context/sdk"
 import { errorMessage } from "../../util/error"
 import { DialogFork } from "./dialog-fork"
+import { revertedPrompt } from "../../util/revert-prompt"
+import type { PromptInfo } from "../../prompt/history"
 
-export function DialogMessage(props: { messageID: string; sessionID: string }) {
+export function DialogMessage(props: {
+  messageID: string
+  sessionID: string
+  setPrompt?: (prompt: PromptInfo) => void
+}) {
   const data = useData()
   const clipboard = useClipboard()
   const toast = useToast()
@@ -22,9 +28,11 @@ export function DialogMessage(props: { messageID: string; sessionID: string }) {
           title: "Revert",
           value: "session.revert",
           description: "undo messages and file changes",
-          onSelect: async (dialog) => {
-            await sdk.api.session
-              .revert.stage({ sessionID: props.sessionID, messageID: props.messageID })
+          onSelect: (dialog) => {
+            const value = message()
+            if (value?.type === "user") props.setPrompt?.(revertedPrompt(value))
+            void sdk.api.session.revert
+              .stage({ sessionID: props.sessionID, messageID: props.messageID })
               .catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
             dialog.clear()
           },

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

@@ -75,6 +75,7 @@ import { usePathFormatter } from "../../context/path-format"
 import { LocationProvider } from "../../context/location"
 import { createSessionRows, resolvePart, type PartRef, type SessionRow } from "./rows"
 import { switchLabel } from "../../util/model"
+import { revertedPrompt } from "../../util/revert-prompt"
 
 addDefaultParsers(parsers.parsers)
 
@@ -424,20 +425,7 @@ export function Session() {
         void sdk.api.session.revert
           .stage({ sessionID: route.sessionID, messageID: message.id })
           .catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
-        prompt?.set({
-          text: message.text,
-          files: message.files?.map((file) => ({
-            uri: file.source.type === "uri" ? file.source.uri : `data:${file.mime};base64,${file.data}`,
-            name: file.name,
-            description: file.description,
-            mention: file.mention ? { ...file.mention } : undefined,
-          })),
-          agents: message.agents?.map((agent) => ({
-            name: agent.name,
-            mention: agent.mention ? { ...agent.mention } : undefined,
-          })),
-          pasted: [],
-        })
+        prompt?.set(revertedPrompt(message))
         dialog.clear()
       },
     },
@@ -1477,6 +1465,7 @@ function UserMessage(props: { message: SessionMessageUser }) {
   )
   const dialog = useDialog()
   const renderer = useRenderer()
+  const promptRef = usePromptRef()
 
   return (
     <Show when={props.message.text.trim() || files().length}>
@@ -1495,7 +1484,13 @@ function UserMessage(props: { message: SessionMessageUser }) {
           }}
           onMouseUp={() => {
             if (renderer.getSelection()?.getSelectedText()) return
-            dialog.replace(() => <DialogMessage messageID={props.message.id} sessionID={ctx.sessionID} />)
+            dialog.replace(() => (
+              <DialogMessage
+                messageID={props.message.id}
+                sessionID={ctx.sessionID}
+                setPrompt={(value) => promptRef.current?.set(value)}
+              />
+            ))
           }}
           paddingTop={1}
           paddingBottom={1}

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

@@ -0,0 +1,19 @@
+import type { SessionMessageUser } from "@opencode-ai/client"
+import type { PromptInfo } from "../prompt/history"
+
+export function revertedPrompt(message: SessionMessageUser): PromptInfo {
+  return {
+    text: message.text,
+    files: message.files?.map((file) => ({
+      uri: file.source.type === "uri" ? file.source.uri : `data:${file.mime};base64,${file.data}`,
+      name: file.name,
+      description: file.description,
+      mention: file.mention ? { ...file.mention } : undefined,
+    })),
+    agents: message.agents?.map((agent) => ({
+      name: agent.name,
+      mention: agent.mention ? { ...agent.mention } : undefined,
+    })),
+    pasted: [],
+  }
+}

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

@@ -0,0 +1,39 @@
+import { describe, expect, test } from "bun:test"
+import type { SessionMessageUser } from "@opencode-ai/client"
+import { revertedPrompt } from "../../src/util/revert-prompt"
+
+describe("reverted prompt", () => {
+  test("restores message text and references", () => {
+    const message = {
+      id: "message-1",
+      type: "user",
+      text: "Fix @test.ts with @review",
+      time: { created: 1 },
+      files: [
+        {
+          mime: "text/typescript",
+          name: "test.ts",
+          description: "test file",
+          data: "",
+          source: { type: "uri", uri: "file:///repo/test.ts" },
+          mention: { start: 4, end: 12, text: "@test.ts" },
+        },
+      ],
+      agents: [{ name: "review", mention: { start: 18, end: 25, text: "@review" } }],
+    } satisfies SessionMessageUser
+
+    expect(revertedPrompt(message)).toEqual({
+      text: "Fix @test.ts with @review",
+      files: [
+        {
+          uri: "file:///repo/test.ts",
+          name: "test.ts",
+          description: "test file",
+          mention: { start: 4, end: 12, text: "@test.ts" },
+        },
+      ],
+      agents: [{ name: "review", mention: { start: 18, end: 25, text: "@review" } }],
+      pasted: [],
+    })
+  })
+})