diff --git a/packages/tui/src/routes/session/index.tsx b/packages/tui/src/routes/session/index.tsx
index d809061f5f..b96bf2d2eb 100644
--- a/packages/tui/src/routes/session/index.tsx
+++ b/packages/tui/src/routes/session/index.tsx
@@ -430,29 +430,34 @@ export function Session() {
       category: "Session",
       slash: { name: "undo" },
       run: () => {
-        void (async () => {
-          const boundary = session()?.revert?.messageID
-          const list = messages()
-          let target: string | 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
-            break
-          }
-          if (!target) {
-            toast.show({ message: "Nothing to undo", variant: "error", duration: 3000 })
-            dialog.clear()
-            return
-          }
-          const error = await sdk.api.session.revert.stage({ sessionID: route.sessionID, messageID: target }).then(
-            () => undefined,
-            (error) => error,
-          )
-          if (error) toast.show({ message: errorMessage(error), variant: "error", duration: 5000 })
+        const boundary = session()?.revert?.messageID
+        const message = messages().findLast(
+          (message): message is SessionMessageUser =>
+            message.type === "user" && !!message.text.trim() && (!boundary || message.id < boundary),
+        )
+        if (!message) {
+          toast.show({ message: "Nothing to undo", variant: "error", duration: 3000 })
           dialog.clear()
-        })()
+          return
+        }
+        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: [],
+        })
+        dialog.clear()
       },
     },
     {