From 71be35f7a0e8f3ebb3eb008c35abf7d748b2e545 Mon Sep 17 00:00:00 2001
From: Aiden Cline <63023139+rekram1-node@users.noreply.github.com>
Date: Sun, 12 Jul 2026 22:31:08 +0000
Subject: [PATCH] fix(tui): pass prompt setter to revert dialog

---
 packages/tui/src/routes/session/dialog-message.tsx | 11 +++++++----
 packages/tui/src/routes/session/index.tsx          |  9 ++++++++-
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/packages/tui/src/routes/session/dialog-message.tsx b/packages/tui/src/routes/session/dialog-message.tsx
index 1ddb629952..4d43d6c868 100644
--- a/packages/tui/src/routes/session/dialog-message.tsx
+++ b/packages/tui/src/routes/session/dialog-message.tsx
@@ -6,15 +6,18 @@ import { useToast } from "../../ui/toast"
 import { useSDK } from "../../context/sdk"
 import { errorMessage } from "../../util/error"
 import { DialogFork } from "./dialog-fork"
-import { usePromptRef } from "../../context/prompt"
 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()
   const sdk = useSDK()
-  const promptRef = usePromptRef()
   const message = createMemo(() => data.session.message.get(props.sessionID, props.messageID))
 
   return (
@@ -27,7 +30,7 @@ export function DialogMessage(props: { messageID: string; sessionID: string }) {
           description: "undo messages and file changes",
           onSelect: (dialog) => {
             const value = message()
-            if (value?.type === "user") promptRef.current?.set(revertedPrompt(value))
+            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 }))
diff --git a/packages/tui/src/routes/session/index.tsx b/packages/tui/src/routes/session/index.tsx
index 940f72ce3c..84c0115d49 100644
--- a/packages/tui/src/routes/session/index.tsx
+++ b/packages/tui/src/routes/session/index.tsx
@@ -1465,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}>
@@ -1483,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}