From e58462b926f08021f36d4cf08ec8b1a932bb9853 Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Fri, 24 Oct 2025 18:21:11 +0100
Subject: [PATCH] Fix revert/undo state synchronization and message cleanup

- Clear session.revert when server sends undefined (fixes inability to send messages after revert)
- Handle message.removed and message.part.removed SSE events to reload messages
- Ensures reverted messages are removed from UI in real-time without requiring app restart
---
 src/lib/sse-manager.ts |  8 ++++++++
 src/stores/sessions.ts | 20 +++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/lib/sse-manager.ts b/src/lib/sse-manager.ts
index 52053d12..f7dc407a 100644
--- a/src/lib/sse-manager.ts
+++ b/src/lib/sse-manager.ts
@@ -89,6 +89,12 @@ class SSEManager {
       case "message.part.updated":
         this.onMessageUpdate?.(instanceId, event)
         break
+      case "message.removed":
+        this.onMessageRemoved?.(instanceId, event)
+        break
+      case "message.part.removed":
+        this.onMessagePartRemoved?.(instanceId, event)
+        break
       case "session.updated":
         this.onSessionUpdate?.(instanceId, event)
         break
@@ -136,6 +142,8 @@ class SSEManager {
   }
 
   onMessageUpdate?: (instanceId: string, event: MessageUpdateEvent) => void
+  onMessageRemoved?: (instanceId: string, event: any) => void
+  onMessagePartRemoved?: (instanceId: string, event: any) => void
   onSessionUpdate?: (instanceId: string, event: SessionUpdateEvent) => void
   onSessionCompacted?: (instanceId: string, event: any) => void
   onSessionError?: (instanceId: string, event: any) => void
diff --git a/src/stores/sessions.ts b/src/stores/sessions.ts
index ca96305a..dcf45607 100644
--- a/src/stores/sessions.ts
+++ b/src/stores/sessions.ts
@@ -666,7 +666,7 @@ function handleSessionUpdate(instanceId: string, event: any): void {
             snapshot: info.revert.snapshot,
             diff: info.revert.diff,
           }
-        : existingSession.revert,
+        : undefined,
     }
 
     setSessions((prev) => {
@@ -875,7 +875,25 @@ function handleSessionError(instanceId: string, event: any): void {
   alert(`Error: ${message}`)
 }
 
+function handleMessageRemoved(instanceId: string, event: any): void {
+  const sessionID = event.properties?.sessionID
+  if (!sessionID) return
+
+  console.log(`[SSE] Message removed from session ${sessionID}, reloading messages`)
+  loadMessages(instanceId, sessionID, true).catch(console.error)
+}
+
+function handleMessagePartRemoved(instanceId: string, event: any): void {
+  const sessionID = event.properties?.sessionID
+  if (!sessionID) return
+
+  console.log(`[SSE] Message part removed from session ${sessionID}, reloading messages`)
+  loadMessages(instanceId, sessionID, true).catch(console.error)
+}
+
 sseManager.onMessageUpdate = handleMessageUpdate
+sseManager.onMessageRemoved = handleMessageRemoved
+sseManager.onMessagePartRemoved = handleMessagePartRemoved
 sseManager.onSessionUpdate = handleSessionUpdate
 sseManager.onSessionCompacted = handleSessionCompacted
 sseManager.onSessionError = handleSessionError