From 1362a5872a2b783f308eba5ba74c32a837dea3cd Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Fri, 24 Oct 2025 17:44:19 +0100
Subject: [PATCH] Add revert button to user messages
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Add revert icon (↶) button in message header for all user messages
- Clicking revert button reverts session to that specific message
- Restores the message text back to the prompt input
- Style revert button with hover effects and border highlight
- Revert action updates UI automatically via SSE session.updated event
---
 src/App.tsx                       | 38 +++++++++++++++++++++++++++++++
 src/components/message-item.tsx   | 17 ++++++++++++++
 src/components/message-stream.tsx |  8 ++++++-
 src/index.css                     | 27 ++++++++++++++++++++++
 4 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/src/App.tsx b/src/App.tsx
index dcca4cb2..20f18e2e 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -77,6 +77,43 @@ const SessionView: Component<{
     await updateSessionModel(props.instanceId, props.sessionId, model)
   }
 
+  async function handleRevert(messageId: string) {
+    const instance = instances().get(props.instanceId)
+    if (!instance || !instance.client) return
+
+    try {
+      console.log("Reverting to message:", messageId)
+
+      await instance.client.session.revert({
+        path: { id: props.sessionId },
+        body: { messageID: messageId },
+      })
+
+      // Restore the message to input
+      const currentSession = session()
+      if (currentSession) {
+        const revertedMessage = currentSession.messages.find((m) => m.id === messageId)
+        const revertedInfo = currentSession.messagesInfo.get(messageId)
+
+        if (revertedMessage && revertedInfo?.role === "user") {
+          const textParts = revertedMessage.parts.filter((p: any) => p.type === "text")
+          if (textParts.length > 0) {
+            const textarea = document.querySelector(".prompt-input") as HTMLTextAreaElement
+            if (textarea) {
+              textarea.value = textParts.map((p: any) => p.text).join("\n")
+              textarea.focus()
+            }
+          }
+        }
+      }
+
+      console.log("Reverted to message - UI will update via SSE")
+    } catch (error) {
+      console.error("Failed to revert:", error)
+      alert("Failed to revert to message")
+    }
+  }
+
   return (
     <Show
       when={session()}
@@ -94,6 +131,7 @@ const SessionView: Component<{
             messages={s().messages || []}
             messagesInfo={s().messagesInfo}
             revert={s().revert}
+            onRevert={handleRevert}
           />
           <PromptInput
             instanceId={props.instanceId}
diff --git a/src/components/message-item.tsx b/src/components/message-item.tsx
index 44929d08..7acd15e2 100644
--- a/src/components/message-item.tsx
+++ b/src/components/message-item.tsx
@@ -6,6 +6,7 @@ interface MessageItemProps {
   message: Message
   messageInfo?: any
   isQueued?: boolean
+  onRevert?: (messageId: string) => void
 }
 
 export default function MessageItem(props: MessageItemProps) {
@@ -42,11 +43,27 @@ export default function MessageItem(props: MessageItemProps) {
     return !hasContent() && props.messageInfo?.time?.completed === 0
   }
 
+  const handleRevert = () => {
+    if (props.onRevert && isUser()) {
+      props.onRevert(props.message.id)
+    }
+  }
+
   return (
     <div class={`message-item ${isUser() ? "user" : "assistant"}`}>
       <div class="message-header">
         <span class="message-sender">{isUser() ? "You" : "Assistant"}</span>
         <span class="message-timestamp">{timestamp()}</span>
+        <Show when={isUser() && props.onRevert}>
+          <button
+            class="message-revert-button"
+            onClick={handleRevert}
+            title="Revert to this message"
+            aria-label="Revert to this message"
+          >
+            ↶
+          </button>
+        </Show>
       </div>
 
       <div class="message-content">
diff --git a/src/components/message-stream.tsx b/src/components/message-stream.tsx
index bc958da4..0e334cfc 100644
--- a/src/components/message-stream.tsx
+++ b/src/components/message-stream.tsx
@@ -17,6 +17,7 @@ interface MessageStreamProps {
     diff?: string
   }
   loading?: boolean
+  onRevert?: (messageId: string) => void
 }
 
 interface DisplayItem {
@@ -179,7 +180,12 @@ export default function MessageStream(props: MessageStreamProps) {
                   </div>
                 }
               >
-                <MessageItem message={item.data} messageInfo={item.messageInfo} isQueued={item.data.isQueued} />
+                <MessageItem
+                  message={item.data}
+                  messageInfo={item.messageInfo}
+                  isQueued={item.data.isQueued}
+                  onRevert={props.onRevert}
+                />
               </Show>
             )
           }}
diff --git a/src/index.css b/src/index.css
index 3c49382e..29e5d84d 100644
--- a/src/index.css
+++ b/src/index.css
@@ -188,6 +188,33 @@ body {
   color: var(--text-muted);
 }
 
+.message-revert-button {
+  background: none;
+  border: 1px solid var(--border-color);
+  color: var(--text-muted);
+  cursor: pointer;
+  padding: 2px 8px;
+  border-radius: 4px;
+  font-size: 16px;
+  line-height: 1;
+  transition: all 0.2s;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  min-width: 28px;
+  height: 24px;
+}
+
+.message-revert-button:hover {
+  background-color: var(--hover-bg);
+  border-color: var(--accent-color);
+  color: var(--accent-color);
+}
+
+.message-revert-button:active {
+  transform: scale(0.95);
+}
+
 .message-content {
   padding-top: 6px;
   line-height: 1.6;