Skip to content

Commit 12498a4

Browse files
ericdalloericdallo
committed
Rollback improvements
1 parent d2be7d5 commit 12498a4

File tree

7 files changed

+182
-53
lines changed

7 files changed

+182
-53
lines changed

‎CHANGELOG.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22

33
## Unreleased
44

5+
- Support rollback file changes done by `write_file`, `edit_file` and `move_file`. #218
6+
- Improve rollback to keep consistent UI before the rollback, fixing tool names and user messages.
7+
58
## 0.80.4
69

710
- Fix binary for macos amd64. #217

‎src/eca/db.clj‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,11 +58,14 @@
5858
:status (or :idle :running :stopping :login)
5959
:created-at :number
6060
:login-provider :string
61-
:messages [::chat-message]
61+
:messages [{:role (or "user" "assistant" "tool_call" "tool_call_output" "reason")
62+
:content (or :string [::any-map]) ;; string for simple text, map/vector for structured content
63+
:content-id :string}]
6264
:tool-calls {"<tool-call-id>"
6365
{:status (or :initial :preparing :check-approval :waiting-approval
6466
:execution-approved :executing :rejected :cleanup
6567
:completed :stopping)
68+
6669
:name :string
6770
:full-name :string
6871
:server :string
@@ -73,7 +76,9 @@
7376
:future-cleanup-complete?* :promise
7477
:start-time :long
7578
:future :future
76-
:resources :map}}}}
79+
:resources ::any-map
80+
:rollback-changes [{:path :string
81+
:content (or :string nil)}]}}}}
7782
:auth {"<provider-name>" {:step (or :login/start :login/waiting-login-method
7883
:login/waiting-provider-code :login/waiting-api-key
7984
:login/waiting-user-confirmation :login/done :login/renew-token)

‎src/eca/features/chat.clj‎

Lines changed: 61 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
(ns eca.features.chat
22
(:require
33
[cheshire.core :as json]
4+
[clojure.java.io :as io]
45
[clojure.set :as set]
56
[clojure.string :as string]
67
[eca.config :as config]
@@ -721,12 +722,14 @@
721722
total-time-ms (- (System/currentTimeMillis) start-time)]
722723
(add-to-history! {:role "tool_call"
723724
:content (assoc tool-call
725+
:name name
724726
:details details
725727
:summary summary
726728
:origin origin
727729
:server server-name)})
728730
(add-to-history! {:role "tool_call_output"
729731
:content (assoc tool-call
732+
:name name
730733
:error (:error result)
731734
:output result
732735
:total-time-ms total-time-ms
@@ -886,49 +889,59 @@
886889
:text error-message})
887890
(prompt-messages! messages chat-ctx))))
888891

889-
(defn ^:private message-content->chat-content [role message-content]
892+
(defn ^:private message-content->chat-content [role message-content content-id]
890893
(case role
891894
("user"
892895
"system"
893-
"assistant") [(reduce
894-
(fn [m content]
895-
(case (:type content)
896-
:text (assoc m
897-
:type :text
898-
:text (str (:text m) "\n" (:text content)))
899-
m))
900-
{}
901-
message-content)]
902-
"tool_call" [{:type :toolCallPrepare
903-
:origin (:origin message-content)
904-
:name (:name message-content)
905-
:server (:server message-content)
906-
:arguments-text ""
907-
:id (:id message-content)}]
908-
"tool_call_output" [{:type :toolCalled
909-
:origin (:origin message-content)
910-
:name (:name message-content)
911-
:server (:server message-content)
912-
:arguments (:arguments message-content)
913-
:total-time-ms (:total-time-ms message-content)
914-
:error (:error message-content)
896+
"assistant") [{:role role
897+
:content (reduce
898+
(fn [m content]
899+
(case (:type content)
900+
:text (assoc m
901+
:type :text
902+
:text (str (:text m) "\n" (:text content)))
903+
m))
904+
(assoc-some {} :content-id content-id)
905+
message-content)}]
906+
"tool_call" [{:role :assistant
907+
:content {:type :toolCallPrepare
908+
:origin (:origin message-content)
909+
:name (:name message-content)
910+
:server (:server message-content)
911+
:summary (:summary message-content)
912+
:details (:details message-content)
913+
:arguments-text ""
914+
:id (:id message-content)}}]
915+
"tool_call_output" [{:role :assistant
916+
:content {:type :toolCalled
917+
:origin (:origin message-content)
918+
:name (:name message-content)
919+
:server (:server message-content)
920+
:arguments (:arguments message-content)
921+
:total-time-ms (:total-time-ms message-content)
922+
:summary (:summary message-content)
923+
:details (:details message-content)
924+
:error (:error message-content)
925+
:id (:id message-content)
926+
:outputs (:contents (:output message-content))}}]
927+
"reason" [{:role :assistant
928+
:content {:type :reasonStarted
929+
:id (:id message-content)}}
930+
{:role :assistant
931+
:content {:type :reasonText
915932
:id (:id message-content)
916-
:outputs (:contents (:output message-content))}]
917-
"reason" [{:type :reasonStarted
918-
:id (:id message-content)}
919-
{:type :reasonText
920-
:id (:id message-content)
921-
:text (:text message-content)}
922-
{:type :reasonFinished
923-
:id (:id message-content)
924-
:total-time-ms (:total-time-ms message-content)}]))
933+
:text (:text message-content)}}
934+
{:role :assistant
935+
:content {:type :reasonFinished
936+
:id (:id message-content)
937+
:total-time-ms (:total-time-ms message-content)}}]))
925938

926939
(defn ^:private send-chat-contents! [messages chat-ctx]
927940
(doseq [message messages]
928-
(doseq [chat-content (message-content->chat-content (:role message) (:content message))]
941+
(doseq [{:keys [role content]} (message-content->chat-content (:role message) (:content message) (:content-id message))]
929942
(send-content! chat-ctx
930-
(:role message)
931-
chat-content))))
943+
role
944+
content))))
932945

933946
(defn ^:private handle-command! [{:keys [command args]} chat-ctx]
934947
(let [{:keys [type on-finished-side-effect] :as result} (f.commands/handle-command! command args chat-ctx)]
@@ -1115,7 +1128,19 @@
11151128
Then notify to clear chat and then the kept messages."
11161129
[{:keys [chat-id content-id]} db* messenger]
11171130
(let [all-messages (get-in @db* [:chats chat-id :messages])
1118-
new-messages (vec (take-while #(not= (:content-id %) content-id) all-messages))]
1131+
tool-calls (get-in @db* [:chats chat-id :tool-calls])
1132+
new-messages (vec (take-while #(not= (:content-id %) content-id) all-messages))
1133+
removed-messages (vec (drop-while #(not= (:content-id %) content-id) all-messages))
1134+
rollback-changes (->> removed-messages
1135+
(filter #(= "tool_call_output" (:role %)))
1136+
(keep #(get-in tool-calls [(:id (:content %)) :rollback-changes]))
1137+
flatten
1138+
reverse)]
1139+
(doseq [{:keys [path content]} rollback-changes]
1140+
(logger/info (format "Rolling back change for '%s' to content: '%s'" path content))
1141+
(if content
1142+
(spit path content)
1143+
(io/delete-file path true)))
11191144
(swap! db* assoc-in [:chats chat-id :messages] new-messages)
11201145
(messenger/chat-cleared
11211146
messenger

‎src/eca/features/tools.clj‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,11 @@
102102
(f.mcp/call-tool! name arguments {:db db}))))]
103103
(logger/debug logger-tag "Tool call result: " result)
104104
(metrics/count-up! "tool-called" {:name name :error (:error result)} metrics)
105-
result)
105+
(if-let [r (:rollback-changes result)]
106+
(do
107+
(swap! db* assoc-in [:chats chat-id :tool-calls tool-call-id :rollback-changes] r)
108+
(dissoc result :rollback-changes))
109+
result))
106110
(catch Exception e
107111
(logger/warn logger-tag (format "Error calling tool %s: %s\n%s" name (.getMessage e) (with-out-str (.printStackTrace e))))
108112
(metrics/count-up! "tool-called" {:name name :error true} metrics)

‎src/eca/features/tools/filesystem.clj‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -96,10 +96,13 @@
9696

9797
(defn ^:private write-file [arguments _]
9898
(let [path (get arguments "path")
99-
content (get arguments "content")]
99+
content (get arguments "content")
100+
old-content (try (slurp path) (catch Exception _ nil))]
100101
(fs/create-dirs (fs/parent (fs/path path)))
101102
(spit path content)
102-
(tools.util/single-text-content (format "Successfully wrote to %s" path))))
103+
(assoc (tools.util/single-text-content (format "Successfully wrote to %s" path))
104+
:rollback-changes [{:path path
105+
:content old-content}])))
103106

104107
(defn ^:private write-file-summary [{:keys [args]}]
105108
(if-let [path (get args "path")]
@@ -239,7 +242,6 @@
239242
(text-match/apply-content-change-to-string file-content original-content new-content all? path)
240243
(smart-edit/apply-smart-edit file-content original-content new-content path)))
241244

242-
243245
(defn ^:private edit-file [arguments {:keys [_db]}]
244246
(or (tools.util/invalid-arguments arguments (concat (path-validations)
245247
[["path" fs/readable? "File $path is not readable"]]))
@@ -251,12 +253,14 @@
251253
result (apply-file-edit-strategy initial-content original-content new-content all? path)
252254
write! (fn [res]
253255
(spit path (:new-full-content res))
254-
(handle-file-change-result res path (format "Successfully replaced content in %s." path)))]
256+
(-> (handle-file-change-result res path (format "Successfully replaced content in %s." path))
257+
(assoc :rollback-changes [{:path path
258+
:content initial-content}])))]
255259
(if (:new-full-content result)
256260
(let [current-content (slurp path)]
257261
(if (= current-content (:original-full-content result))
258262
(write! result)
259-
;; Optimistic retry once against latest content
263+
;; Optimistic retry once against latest content
260264
(let [retry (apply-file-edit-strategy current-content original-content new-content all? path)]
261265
(if (:new-full-content retry)
262266
(write! retry)
@@ -288,9 +292,14 @@
288292
(or (tools.util/invalid-arguments arguments [["source" fs/exists? "$source is not a valid path"]
289293
["destination" (complement fs/exists?) "Path $destination already exists"]])
290294
(let [source (get arguments "source")
291-
destination (get arguments "destination")]
295+
destination (get arguments "destination")
296+
source-content (slurp source)]
292297
(fs/move source destination {:replace-existing false})
293-
(tools.util/single-text-content (format "Successfully moved %s to %s" source destination)))))
298+
(assoc (tools.util/single-text-content (format "Successfully moved %s to %s" source destination))
299+
:rollback-changes [{:path destination
300+
:content nil}
301+
{:path source
302+
:content source-content}]))))
294303

295304
(def definitions
296305
{"directory_tree"

‎test/eca/features/chat_test.clj‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -534,3 +534,71 @@
534534
:prompt "prompt"
535535
:args ["arg1" "arg2"]}
536536
(#'f.chat/message->decision "/server:prompt arg1 arg2" {} {}))))))
537+
538+
(deftest rollback-chat-test
539+
(testing "Rollback chat removes messages after content-id"
540+
(h/reset-components!)
541+
(let [{:keys [chat-id]}
542+
(prompt!
543+
{:message "Count with me: 1"}
544+
{:all-tools-mock (constantly [])
545+
:api-mock
546+
(fn [{:keys [on-first-response-received
547+
on-message-received]}]
548+
(on-first-response-received {:type :text :text "2"})
549+
(on-message-received {:type :text :text "2"})
550+
(on-message-received {:type :finish}))})
551+
first-content-id (get-in (h/db) [:chats chat-id :messages 0 :content-id])
552+
_ (is (some? first-content-id) "first-content-id should exist")]
553+
;; Verify initial state
554+
(is (match?
555+
{chat-id {:id chat-id
556+
:messages [{:role "user" :content [{:type :text :text "Count with me: 1"}] :content-id first-content-id}
557+
{:role "assistant" :content [{:type :text :text "2"}]}]}}
558+
(:chats (h/db))))
559+
560+
;; Add second message
561+
(h/reset-messenger!)
562+
(prompt!
563+
{:message "3"
564+
:chat-id chat-id}
565+
{:all-tools-mock (constantly [])
566+
:api-mock
567+
(fn [{:keys [on-first-response-received
568+
on-message-received]}]
569+
(on-first-response-received {:type :text :text "4"})
570+
(on-message-received {:type :text :text "4"})
571+
(on-message-received {:type :finish}))})
572+
(let [second-content-id (get-in (h/db) [:chats chat-id :messages 2 :content-id])]
573+
574+
;; Verify we now have 4 messages
575+
(is (match?
576+
{chat-id {:id chat-id
577+
:messages [{:role "user" :content [{:type :text :text "Count with me: 1"}] :content-id first-content-id}
578+
{:role "assistant" :content [{:type :text :text "2"}]}
579+
{:role "user" :content [{:type :text :text "3"}] :content-id second-content-id}
580+
{:role "assistant" :content [{:type :text :text "4"}]}]}}
581+
(:chats (h/db))))
582+
583+
;; Rollback to second message (keep first 2 messages, remove last 2)
584+
(h/reset-messenger!)
585+
(is (= {} (f.chat/rollback-chat {:chat-id chat-id :content-id second-content-id} (h/db*) (h/messenger))))
586+
587+
;; Verify messages after content-id are removed (keeps messages before content-id)
588+
(is (match?
589+
{chat-id {:id chat-id
590+
:messages [{:role "user" :content [{:type :text :text "Count with me: 1"}] :content-id first-content-id}
591+
{:role "assistant" :content [{:type :text :text "2"}]}]}}
592+
(:chats (h/db))))
593+
594+
;; Verify messenger received chat-clear and then messages
595+
(is (match?
596+
{:chat-clear [{:chat-id chat-id :messages true}]
597+
:chat-content-received
598+
[{:chat-id chat-id
599+
:content {:type :text :text "\nCount with me: 1" :content-id first-content-id}
600+
:role "user"}
601+
{:chat-id chat-id
602+
:content {:type :text :text "\n2"}
603+
:role "assistant"}]}
604+
(h/messages)))))))

0 commit comments

Comments
 (0)