fix(tui): restore clicked reverted prompt (#36567)

Co-authored-by: Aiden Cline <63023139+rekram1-node@users.noreply.github.com>

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

@@ -6,8 +6,13 @@ import { useToast } from "../../ui/toast"
 import { useSDK } from "../../context/sdk"
 import { errorMessage } from "../../util/error"
 import { DialogFork } from "./dialog-fork"
+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 +27,26 @@ 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?.({
+                text: value.text,
+                files: value.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: value.agents?.map((agent) => ({
+                  name: agent.name,
+                  mention: agent.mention ? { ...agent.mention } : undefined,
+                })),
+                pasted: [],
+              })
+            }
+            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

@@ -1468,6 +1468,7 @@ function UserMessage(props: { message: SessionMessageUser }) {
   )
   const dialog = useDialog()
   const renderer = useRenderer()
+  const promptRef = usePromptRef()
 
   return (
     <Show when={props.message.text.trim() || files().length}>
@@ -1486,7 +1487,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}