assistant2: Decouple scripting tool from the Tool trait (#26382)

This PR decouples the scripting tool from the `Tool` trait while still
allowing it to be used as a tool from the model's perspective.

This will allow us to evolve the scripting tool as more of a first-class
citizen while still retaining the ability to have the model call it as a
regular tool.

Release Notes:

- N/A

crates/assistant2/src/active_thread.rs

@@ -457,9 +457,13 @@ impl ActiveThread {
 
         let context = thread.context_for_message(message_id);
         let tool_uses = thread.tool_uses_for_message(message_id);
+        let scripting_tool_uses = thread.scripting_tool_uses_for_message(message_id);
 
         // Don't render user messages that are just there for returning tool results.
-        if message.role == Role::User && thread.message_has_tool_results(message_id) {
+        if message.role == Role::User
+            && (thread.message_has_tool_results(message_id)
+                || thread.message_has_scripting_tool_results(message_id))
+        {
             return Empty.into_any();
         }
 
@@ -609,16 +613,22 @@ impl ActiveThread {
                 .id(("message-container", ix))
                 .child(message_content)
                 .map(|parent| {
-                    if tool_uses.is_empty() {
+                    if tool_uses.is_empty() && scripting_tool_uses.is_empty() {
                         return parent;
                     }
 
                     parent.child(
-                        v_flex().children(
-                            tool_uses
-                                .into_iter()
-                                .map(|tool_use| self.render_tool_use(tool_use, cx)),
-                        ),
+                        v_flex()
+                            .children(
+                                tool_uses
+                                    .into_iter()
+                                    .map(|tool_use| self.render_tool_use(tool_use, cx)),
+                            )
+                            .children(
+                                scripting_tool_uses
+                                    .into_iter()
+                                    .map(|tool_use| self.render_scripting_tool_use(tool_use, cx)),
+                            ),
                     )
                 }),
             Role::System => div().id(("message-container", ix)).py_1().px_2().child(
@@ -727,6 +737,15 @@ impl ActiveThread {
                 }),
         )
     }
+
+    fn render_scripting_tool_use(
+        &self,
+        tool_use: ToolUse,
+        cx: &mut Context<Self>,
+    ) -> impl IntoElement {
+        // TODO: Add custom rendering for scripting tool uses.
+        self.render_tool_use(tool_use, cx)
+    }
 }
 
 impl Render for ActiveThread {

crates/assistant2/src/thread.rs

@@ -13,13 +13,14 @@ use language_model::{
     Role, StopReason,
 };
 use project::Project;
+use scripting_tool::ScriptingTool;
 use serde::{Deserialize, Serialize};
 use util::{post_inc, TryFutureExt as _};
 use uuid::Uuid;
 
 use crate::context::{attach_context_to_message, ContextId, ContextSnapshot};
 use crate::thread_store::SavedThread;
-use crate::tool_use::{PendingToolUse, ToolUse, ToolUseState};
+use crate::tool_use::{ToolUse, ToolUseState};
 
 #[derive(Debug, Clone, Copy)]
 pub enum RequestKind {
@@ -75,6 +76,7 @@ pub struct Thread {
     project: Entity<Project>,
     tools: Arc<ToolWorkingSet>,
     tool_use: ToolUseState,
+    scripting_tool_use: ToolUseState,
 }
 
 impl Thread {
@@ -97,6 +99,7 @@ impl Thread {
             project,
             tools,
             tool_use: ToolUseState::new(),
+            scripting_tool_use: ToolUseState::new(),
         }
     }
 
@@ -115,6 +118,7 @@ impl Thread {
                 .unwrap_or(0),
         );
         let tool_use = ToolUseState::from_saved_messages(&saved.messages);
+        let scripting_tool_use = ToolUseState::new();
 
         Self {
             id,
@@ -138,6 +142,7 @@ impl Thread {
             project,
             tools,
             tool_use,
+            scripting_tool_use,
         }
     }
 
@@ -198,31 +203,46 @@ impl Thread {
         )
     }
 
-    pub fn pending_tool_uses(&self) -> Vec<&PendingToolUse> {
-        self.tool_use.pending_tool_uses()
-    }
-
     /// Returns whether all of the tool uses have finished running.
     pub fn all_tools_finished(&self) -> bool {
+        let mut all_pending_tool_uses = self
+            .tool_use
+            .pending_tool_uses()
+            .into_iter()
+            .chain(self.scripting_tool_use.pending_tool_uses());
+
         // If the only pending tool uses left are the ones with errors, then that means that we've finished running all
         // of the pending tools.
-        self.pending_tool_uses()
-            .into_iter()
-            .all(|tool_use| tool_use.status.is_error())
+        all_pending_tool_uses.all(|tool_use| tool_use.status.is_error())
     }
 
     pub fn tool_uses_for_message(&self, id: MessageId) -> Vec<ToolUse> {
         self.tool_use.tool_uses_for_message(id)
     }
 
+    pub fn scripting_tool_uses_for_message(&self, id: MessageId) -> Vec<ToolUse> {
+        self.scripting_tool_use.tool_uses_for_message(id)
+    }
+
     pub fn tool_results_for_message(&self, id: MessageId) -> Vec<&LanguageModelToolResult> {
         self.tool_use.tool_results_for_message(id)
     }
 
+    pub fn scripting_tool_results_for_message(
+        &self,
+        id: MessageId,
+    ) -> Vec<&LanguageModelToolResult> {
+        self.scripting_tool_use.tool_results_for_message(id)
+    }
+
     pub fn message_has_tool_results(&self, message_id: MessageId) -> bool {
         self.tool_use.message_has_tool_results(message_id)
     }
 
+    pub fn message_has_scripting_tool_results(&self, message_id: MessageId) -> bool {
+        self.scripting_tool_use.message_has_tool_results(message_id)
+    }
+
     pub fn insert_user_message(
         &mut self,
         text: impl Into<String>,
@@ -313,16 +333,25 @@ impl Thread {
         let mut request = self.to_completion_request(request_kind, cx);
 
         if use_tools {
-            request.tools = self
-                .tools()
-                .tools(cx)
-                .into_iter()
-                .map(|tool| LanguageModelRequestTool {
-                    name: tool.name(),
-                    description: tool.description(),
-                    input_schema: tool.input_schema(),
-                })
-                .collect();
+            let mut tools = Vec::new();
+            tools.push(LanguageModelRequestTool {
+                name: ScriptingTool::NAME.into(),
+                description: ScriptingTool::DESCRIPTION.into(),
+                input_schema: ScriptingTool::input_schema(),
+            });
+
+            tools.extend(
+                self.tools()
+                    .tools(cx)
+                    .into_iter()
+                    .map(|tool| LanguageModelRequestTool {
+                        name: tool.name(),
+                        description: tool.description(),
+                        input_schema: tool.input_schema(),
+                    }),
+            );
+
+            request.tools = tools;
         }
 
         self.stream_completion(request, model, cx);
@@ -357,6 +386,8 @@ impl Thread {
                 RequestKind::Chat => {
                     self.tool_use
                         .attach_tool_results(message.id, &mut request_message);
+                    self.scripting_tool_use
+                        .attach_tool_results(message.id, &mut request_message);
                 }
                 RequestKind::Summarize => {
                     // We don't care about tool use during summarization.
@@ -373,6 +404,8 @@ impl Thread {
                 RequestKind::Chat => {
                     self.tool_use
                         .attach_tool_uses(message.id, &mut request_message);
+                    self.scripting_tool_use
+                        .attach_tool_uses(message.id, &mut request_message);
                 }
                 RequestKind::Summarize => {
                     // We don't care about tool use during summarization.
@@ -450,9 +483,15 @@ impl Thread {
                                     .iter()
                                     .rfind(|message| message.role == Role::Assistant)
                                 {
-                                    thread
-                                        .tool_use
-                                        .request_tool_use(last_assistant_message.id, tool_use);
+                                    if tool_use.name.as_ref() == ScriptingTool::NAME {
+                                        thread
+                                            .scripting_tool_use
+                                            .request_tool_use(last_assistant_message.id, tool_use);
+                                    } else {
+                                        thread
+                                            .tool_use
+                                            .request_tool_use(last_assistant_message.id, tool_use);
+                                    }
                                 }
                             }
                         }
@@ -572,6 +611,7 @@ impl Thread {
 
     pub fn use_pending_tools(&mut self, cx: &mut Context<Self>) {
         let pending_tool_uses = self
+            .tool_use
             .pending_tool_uses()
             .into_iter()
             .filter(|tool_use| tool_use.status.is_idle())
@@ -585,6 +625,20 @@ impl Thread {
                 self.insert_tool_output(tool_use.id.clone(), task, cx);
             }
         }
+
+        let pending_scripting_tool_uses = self
+            .scripting_tool_use
+            .pending_tool_uses()
+            .into_iter()
+            .filter(|tool_use| tool_use.status.is_idle())
+            .cloned()
+            .collect::<Vec<_>>();
+
+        for scripting_tool_use in pending_scripting_tool_uses {
+            let task = ScriptingTool.run(scripting_tool_use.input, self.project.clone(), cx);
+
+            self.insert_scripting_tool_output(scripting_tool_use.id.clone(), task, cx);
+        }
     }
 
     pub fn insert_tool_output(
@@ -613,6 +667,32 @@ impl Thread {
             .run_pending_tool(tool_use_id, insert_output_task);
     }
 
+    pub fn insert_scripting_tool_output(
+        &mut self,
+        tool_use_id: LanguageModelToolUseId,
+        output: Task<Result<String>>,
+        cx: &mut Context<Self>,
+    ) {
+        let insert_output_task = cx.spawn(|thread, mut cx| {
+            let tool_use_id = tool_use_id.clone();
+            async move {
+                let output = output.await;
+                thread
+                    .update(&mut cx, |thread, cx| {
+                        thread
+                            .scripting_tool_use
+                            .insert_tool_output(tool_use_id.clone(), output);
+
+                        cx.emit(ThreadEvent::ToolFinished { tool_use_id });
+                    })
+                    .ok();
+            }
+        });
+
+        self.scripting_tool_use
+            .run_pending_tool(tool_use_id, insert_output_task);
+    }
+
     pub fn send_tool_results_to_model(
         &mut self,
         model: Arc<dyn LanguageModel>,

crates/assistant2/src/tool_use.rs

@@ -267,6 +267,7 @@ impl ToolUseState {
 pub struct PendingToolUse {
     pub id: LanguageModelToolUseId,
     /// The ID of the Assistant message in which the tool use was requested.
+    #[allow(unused)]
     pub assistant_message_id: MessageId,
     pub name: Arc<str>,
     pub input: serde_json::Value,