assistant edit tool: Fix editing files in context (#26751)

When the user attached context in the thread, the editor model request
would fail because its tool use wouldn't be removed properly leading to
an API error.

Also, after an edit, we'd keep the old file snapshot in the context.
This would make the model think that the edits didn't apply and make it
go in a loop.

Release Notes:

- N/A

crates/assistant2/src/thread.rs

@@ -2,7 +2,7 @@ use std::io::Write;
 use std::sync::Arc;
 
 use anyhow::{Context as _, Result};
-use assistant_tool::ToolWorkingSet;
+use assistant_tool::{ActionLog, ToolWorkingSet};
 use chrono::{DateTime, Utc};
 use collections::{BTreeMap, HashMap, HashSet};
 use futures::future::Shared;
@@ -104,6 +104,7 @@ pub struct Thread {
     prompt_builder: Arc<PromptBuilder>,
     tools: Arc<ToolWorkingSet>,
     tool_use: ToolUseState,
+    action_log: Entity<ActionLog>,
     scripting_session: Entity<ScriptingSession>,
     scripting_tool_use: ToolUseState,
     initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
@@ -134,6 +135,7 @@ impl Thread {
             tool_use: ToolUseState::new(),
             scripting_session: cx.new(|cx| ScriptingSession::new(project.clone(), cx)),
             scripting_tool_use: ToolUseState::new(),
+            action_log: cx.new(|_| ActionLog::new()),
             initial_project_snapshot: {
                 let project_snapshot = Self::project_snapshot(project, cx);
                 cx.foreground_executor()
@@ -191,6 +193,7 @@ impl Thread {
             prompt_builder,
             tools,
             tool_use,
+            action_log: cx.new(|_| ActionLog::new()),
             scripting_session,
             scripting_tool_use,
             initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
@@ -750,7 +753,13 @@ impl Thread {
 
         for tool_use in pending_tool_uses {
             if let Some(tool) = self.tools.tool(&tool_use.name, cx) {
-                let task = tool.run(tool_use.input, &request.messages, self.project.clone(), cx);
+                let task = tool.run(
+                    tool_use.input,
+                    &request.messages,
+                    self.project.clone(),
+                    self.action_log.clone(),
+                    cx,
+                );
 
                 self.insert_tool_output(tool_use.id.clone(), task, cx);
             }
@@ -857,8 +866,15 @@ impl Thread {
     pub fn send_tool_results_to_model(
         &mut self,
         model: Arc<dyn LanguageModel>,
+        updated_context: Vec<ContextSnapshot>,
         cx: &mut Context<Self>,
     ) {
+        self.context.extend(
+            updated_context
+                .into_iter()
+                .map(|context| (context.id, context)),
+        );
+
         // Insert a user message to contain the tool results.
         self.insert_user_message(
             // TODO: Sending up a user message without any content results in the model sending back
@@ -1057,6 +1073,10 @@ impl Thread {
         Ok(String::from_utf8_lossy(&markdown).to_string())
     }
 
+    pub fn action_log(&self) -> &Entity<ActionLog> {
+        &self.action_log
+    }
+
     pub fn cumulative_token_usage(&self) -> TokenUsage {
         self.cumulative_token_usage.clone()
     }