Support rollback chats

Fixes #42

Author: ericdallo

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add support to rollback messages via `chat/rollback` and `chat/clear` messages. #42
+
 ## 0.79.1
 
 - Improve system prompt to add project env context.

docs/protocol.md

@@ -555,12 +555,19 @@ type ChatContent =
     | ChatToolCalledContent
     | ChatToolCallRejectedContent
     | ChatMetadataContent;
-
+    
 /**
  * Simple text message from the LLM
  */
-interface ChatTextContent {
+interface ChatTextContent extends BaseChatContent {
     type: 'text';
+    
+    /**
+     * The unique identifier of this content.
+     * Mostly used to rollback messages.
+     */
+    contentId: string;
+
     /**
      * The text content
      */
@@ -1224,6 +1231,62 @@ interface ChatPromptStopParams {
 }
 ```
 
+### Chat rollback (↩️)
+
+A client request to rollback chat messages to before a specific user sent message using `contentId`.
+Clients should show an option close to user sent messages in chat to rollback, calling this method.
+Server will then remove the messages from its memory after that contentId and produce `chat/cleared` followed 
+with `chat/contentReceived` with the kept messages.
+
+_Request:_ 
+
+* method: `chat/rollback`
+* params: `ChatRollbackParams` defined as follows:
+
+```typescript
+interface ChatRollbackParams {
+    /**
+     * The chat session identifier.
+     */
+    chatId: string;
+    
+    /**
+     * The message content id.
+     */
+    contentId: string;
+}
+```
+
+_Response:_
+
+```typescript
+interface ChatRollbackResponse {}
+```
+
+### Chat cleared (⬅️)
+
+A server notification to clear a chat UI, currently supporting removing only messages of the chat.
+
+_Request:_ 
+
+* method: `chat/cleared`
+* params: `ChatClearedParams` defined as follows:
+
+```typescript
+interface ChatClearedParams {
+
+    /**
+     * The chat session identifier.
+     */
+    chatId: string;
+
+    /**
+     * Whether to clear the messages of a chat.
+     */
+    messages: boolean;
+}
+```
+
 ### Chat delete (↩️)
 
 A client request to delete a existing chat, removing all previous messages and used tokens/costs from memory, good for reduce context or start a new clean chat.
@@ -1243,6 +1306,12 @@ interface ChatDeleteParams {
 }
 ```
 
+_Response:_
+
+```typescript
+interface ChatDeleteResponse {}
+```
+
 ### Chat selected behavior changed (➡️)
 
 A client notification for server telling the user selected a different behavior in chat.
@@ -1261,12 +1330,6 @@ interface ChatSelectedBehaviorChanged {
 }
 ```
 
-_Response:_
-
-```typescript
-interface ChatDeleteResponse {}
-```
-
 ### Editor diagnostics (↪️)
 
 A server request to retrieve LSP or any other kind of diagnostics if available from current workspaces.

src/eca/features/chat.clj

@@ -24,6 +24,9 @@
 
 (def ^:private logger-tag "[CHAT]")
 
+(defn ^:private new-content-id []
+  (str (random-uuid)))
+
 (defn default-model [db config]
   (llm-api/default-model db config))
 
@@ -600,7 +603,7 @@
         :on-first-response-received (fn [& _]
                                       (assert-chat-not-stopped! chat-ctx)
                                       (doseq [message user-messages]
-                                        (add-to-history! message))
+                                        (add-to-history! (assoc message :content-id (:user-content-id chat-ctx))))
                                       (send-content! chat-ctx :system {:type :progress
                                                                        :state :running
                                                                        :text "Generating"}))
@@ -622,7 +625,8 @@
 
                                                   (finish-chat-prompt! :idle chat-ctx))
                                  :finish (do
-                                           (add-to-history! {:role "assistant" :content [{:type :text :text @received-msgs*}]})
+                                           (add-to-history! {:role "assistant"
+                                                             :content [{:type :text :text @received-msgs*}]})
                                            (finish-chat-prompt! :idle chat-ctx))))
         :on-prepare-tool-call (fn [{:keys [id full-name arguments-text]}]
                                 (assert-chat-not-stopped! chat-ctx)
@@ -919,24 +923,29 @@
                :id (:id message-content)
                :total-time-ms (:total-time-ms message-content)}]))
 
+(defn ^:private send-chat-contents! [messages chat-ctx]
+  (doseq [message messages]
+    (doseq [chat-content (message-content->chat-content (:role message) (:content message))]
+      (send-content! chat-ctx
+                     (:role message)
+                     chat-content))))
+
 (defn ^:private handle-command! [{:keys [command args]} chat-ctx]
   (let [{:keys [type on-finished-side-effect] :as result} (f.commands/handle-command! command args chat-ctx)]
     (case type
       :chat-messages (do
                        (doseq [[chat-id {:keys [messages title]}] (:chats result)]
-                         (doseq [message messages]
-                           (let [new-chat-ctx (assoc chat-ctx :chat-id chat-id)]
-                             (doseq [chat-content (message-content->chat-content (:role message) (:content message))]
-                               (send-content! new-chat-ctx
-                                              (:role message)
-                                              chat-content))
-                             (when title
-                               (send-content! new-chat-ctx :system (assoc-some
-                                                                    {:type :metadata}
-                                                                    :title title))))))
+                         (let [new-chat-ctx (assoc chat-ctx :chat-id chat-id)]
+                           (send-chat-contents! messages new-chat-ctx)
+                           (when title
+                             (send-content! new-chat-ctx :system (assoc-some
+                                                                  {:type :metadata}
+                                                                  :title title)))))
                        (finish-chat-prompt! :idle chat-ctx))
       :new-chat-status (finish-chat-prompt! (:status result) chat-ctx)
-      :send-prompt (prompt-messages! [{:role "user" :content (:prompt result)}] (assoc chat-ctx :on-finished-side-effect on-finished-side-effect))
+      :send-prompt (prompt-messages! [{:role "user"
+                                       :content (:prompt result)}]
+                                     (assoc chat-ctx :on-finished-side-effect on-finished-side-effect))
       nil)))
 
 (defn prompt
@@ -995,6 +1004,7 @@
                   :db* db*
                   :metrics metrics
                   :config config
+                  :user-content-id (new-content-id)
                   :messenger messenger}
         decision (message->decision message db config)
         hook-outputs* (atom [])
@@ -1014,6 +1024,7 @@
                         user-messages)]
     (swap! db* assoc-in [:chats chat-id :status] :running)
     (send-content! chat-ctx :user {:type :text
+                                   :content-id (:user-content-id chat-ctx)
                                    :text (str message "\n")})
     (case (:type decision)
       :mcp-prompt (send-mcp-prompt! decision chat-ctx)
@@ -1098,3 +1109,20 @@
   [{:keys [chat-id]} db* metrics]
   (swap! db* update :chats dissoc chat-id)
   (db/update-workspaces-cache! @db* metrics))
+
+(defn rollback-chat
+  "Remove messages from chat in db until content-id matches.
+   Then notify to clear chat and then the kept messages."
+  [{:keys [chat-id content-id]} db* messenger]
+  (let [all-messages (get-in @db* [:chats chat-id :messages])
+        new-messages (vec (take-while #(not= (:content-id %) content-id) all-messages))]
+    (swap! db* assoc-in [:chats chat-id :messages] new-messages)
+    (messenger/chat-cleared
+     messenger
+     {:chat-id chat-id
+      :messages true})
+    (send-chat-contents!
+     new-messages
+     {:chat-id chat-id
+      :messenger messenger})
+    {}))

src/eca/features/prompt.clj

@@ -64,7 +64,7 @@
    "<contexts description=\"User-Provided. This content is current and accurate. Treat this as sufficient context for answering the query.\">"
    ""
    (reduce
-    (fn [context-str {:keys [type path position content lines-range uri] :as a}]
+    (fn [context-str {:keys [type path position content lines-range uri]}]
       (str context-str (case type
                          :file (if lines-range
                                  (format "<file line-start=%s line-end=%s path=\"%s\">%s</file>\n\n"

src/eca/handlers.clj

@@ -116,6 +116,10 @@
     (f.chat/delete-chat params db* metrics)
     {}))
 
+(defn chat-rollback [{:keys [db* metrics messenger]} params]
+  (metrics/task metrics :eca/chat-rollback
+    (f.chat/rollback-chat params db* messenger)))
+
 (defn mcp-stop-server [{:keys [db* messenger metrics config]} params]
   (metrics/task metrics :eca/mcp-stop-server
     (f.tools/stop-server! (:name params) db* messenger config metrics)))
@@ -165,4 +169,4 @@
   [{:keys [db* config metrics messenger]} params]
   (metrics/task metrics :eca/rewrite-prompt
     (handle-expected-errors
-      (f.rewrite/prompt params db* config messenger metrics))))
+     (f.rewrite/prompt params db* config messenger metrics))))

src/eca/llm_providers/anthropic.clj

@@ -131,24 +131,26 @@
              :content [{:type "thinking"
                         :signature (:external-id content)
                         :thinking (:text content)}]}
-            (update msg :content (fn [c]
+            (-> msg
+                (dissoc :content-id)
+                (update :content (fn [c]
                                    (if (string? c)
                                      (string/trim c)
                                      (vec
-                                      (keep #(case (name (:type %))
+                                       (keep #(case (name (:type %))
 
-                                               "text"
-                                               (update % :text string/trim)
+                                                "text"
+                                                (update % :text string/trim)
 
-                                               "image"
-                                               (when supports-image?
-                                                 {:type "image"
-                                                  :source {:data (:base64 %)
-                                                           :media_type (:media-type %)
-                                                           :type "base64"}})
+                                                "image"
+                                                (when supports-image?
+                                                  {:type "image"
+                                                   :source {:data (:base64 %)
+                                                            :media_type (:media-type %)
+                                                            :type "base64"}})
 
-                                               %)
-                                            c)))))))
+                                                %)
+                                            c))))))))
         past-messages))
 
 (defn ^:private add-cache-to-last-message [messages]

src/eca/llm_providers/ollama.clj

@@ -95,7 +95,7 @@
                          (assoc :name (:full-name tool)))})
         tools))
 
-(defn ^:private normalize-messages [past-messages]
+(defn ^:private normalize-messages [messages]
   (mapv (fn [{:keys [role content] :as msg}]
           (case role
             "tool_call" {:role "assistant" :tool-calls [{:type "function"
@@ -111,7 +111,7 @@
              ;; TODO add image supprt
              ;; :images []
              }))
-        past-messages))
+        messages))
 
 (defn chat! [{:keys [model user-messages reason? instructions api-url past-messages tools]}
              {:keys [on-message-received on-error on-prepare-tool-call on-tools-called

src/eca/llm_providers/openai.clj

@@ -94,7 +94,9 @@
                                  [{:type "summary_text"
                                    :text (:text content)}])
                       :encrypted_content (:external-id content)}
-            (update msg :content (fn [c]
+            (-> msg
+                (dissoc :content-id)
+                (update :content (fn [c]
                                    (if (string? c)
                                      c
                                      (keep #(case (name (:type %))
@@ -112,7 +114,7 @@
                                                                     (:base64 %))})
 
                                               %)
-                                           c))))))
+                                           c)))))))
         messages))
 
 (defn ^:private ->tools [tools]

src/eca/messenger.clj

@@ -6,6 +6,7 @@
 
 (defprotocol IMessenger
   (chat-content-received [this data])
+  (chat-cleared [this params])
   (rewrite-content-received [this data])
   (tool-server-updated [this params])
   (config-updated [this params])

src/eca/server.clj

@@ -68,6 +68,9 @@
 (defmethod jsonrpc.server/receive-request "chat/delete" [_ components params]
   (handlers/chat-delete (with-config components) params))
 
+(defmethod jsonrpc.server/receive-request "chat/rollback" [_ components params]
+  (handlers/chat-rollback (with-config components) params))
+
 (defmethod jsonrpc.server/receive-notification "mcp/stopServer" [_ components params]
   (handlers/mcp-stop-server (with-config components) params))
 
@@ -112,6 +115,9 @@
   (chat-content-received [_this content]
     (jsonrpc.server/discarding-stdout
      (jsonrpc.server/send-notification server "chat/contentReceived" content)))
+  (chat-cleared [_this params]
+    (jsonrpc.server/discarding-stdout
+     (jsonrpc.server/send-notification server "chat/cleared" params)))
   (rewrite-content-received [_this content]
     (jsonrpc.server/discarding-stdout
      (jsonrpc.server/send-notification server "rewrite/contentReceived" content)))