Handle streaming input in title updates

Ben Brandt and Antonio Scandurra created 3 months ago

Co-authored-by: Antonio Scandurra <[e-mail]>

Change summary

crates/agent2/src/tests/mod.rs | 32 +++++++++++++++++++++++++++++---
crates/agent2/src/thread.rs    | 29 ++++++++++++++++++++---------
2 files changed, 49 insertions(+), 12 deletions(-)

Detailed changes

crates/agent2/src/tests/mod.rs

@@ -647,6 +647,19 @@ async fn test_tool_updates_to_completion(cx: &mut TestAppContext) {
     let mut events = thread.update(cx, |thread, cx| thread.send(model.clone(), "Think", cx));
     cx.run_until_parked();
 
+    // Simulate streaming partial input.
+    let input = json!({});
+    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
+        LanguageModelToolUse {
+            id: "1".into(),
+            name: ThinkingTool.name().into(),
+            raw_input: input.to_string(),
+            input,
+            is_input_complete: false,
+        },
+    ));
+
+    // Input streaming completed
     let input = json!({ "content": "Thinking hard!" });
     fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
         LanguageModelToolUse {
@@ -665,12 +678,12 @@ async fn test_tool_updates_to_completion(cx: &mut TestAppContext) {
         tool_call,
         acp::ToolCall {
             id: acp::ToolCallId("1".into()),
-            title: "Thinking".into(),
+            title: "thinking".into(),
             kind: acp::ToolKind::Think,
             status: acp::ToolCallStatus::Pending,
             content: vec![],
             locations: vec![],
-            raw_input: Some(json!({ "content": "Thinking hard!" })),
+            raw_input: Some(json!({})),
             raw_output: None,
         }
     );
@@ -680,7 +693,20 @@ async fn test_tool_updates_to_completion(cx: &mut TestAppContext) {
         acp::ToolCallUpdate {
             id: acp::ToolCallId("1".into()),
             fields: acp::ToolCallUpdateFields {
-                status: Some(acp::ToolCallStatus::InProgress,),
+                title: Some("Thinking".into()),
+                kind: Some(acp::ToolKind::Think),
+                raw_input: Some(json!({ "content": "Thinking hard!" })),
+                ..Default::default()
+            },
+        }
+    );
+    let update = expect_tool_call_update(&mut events).await;
+    assert_eq!(
+        update,
+        acp::ToolCallUpdate {
+            id: acp::ToolCallId("1".into()),
+            fields: acp::ToolCallUpdateFields {
+                status: Some(acp::ToolCallStatus::InProgress),
                 ..Default::default()
             },
         }

crates/agent2/src/thread.rs

@@ -474,8 +474,17 @@ impl Thread {
             }
         });
 
+        let mut title = SharedString::from(&tool_use.name);
+        let mut kind = acp::ToolKind::Other;
+        if let Some(tool) = tool.as_ref() {
+            if let Ok(initial_title) = tool.initial_title(tool_use.input.clone()) {
+                title = initial_title;
+            }
+            kind = tool.kind();
+        }
+
         if push_new_tool_use {
-            event_stream.send_tool_call(tool.as_ref(), &tool_use);
+            event_stream.send_tool_call(&tool_use.id, title, kind, tool_use.input.clone());
             last_message
                 .content
                 .push(MessageContent::ToolUse(tool_use.clone()));
@@ -483,6 +492,8 @@ impl Thread {
             event_stream.send_tool_call_update(
                 &tool_use.id,
                 acp::ToolCallUpdateFields {
+                    title: Some(title.into()),
+                    kind: Some(kind),
                     raw_input: Some(tool_use.input.clone()),
                     ..Default::default()
                 },
@@ -842,17 +853,17 @@ impl AgentResponseEventStream {
 
     fn send_tool_call(
         &self,
-        tool: Option<&Arc<dyn AnyAgentTool>>,
-        tool_use: &LanguageModelToolUse,
+        id: &LanguageModelToolUseId,
+        title: SharedString,
+        kind: acp::ToolKind,
+        input: serde_json::Value,
     ) {
         self.0
             .unbounded_send(Ok(AgentResponseEvent::ToolCall(Self::initial_tool_call(
-                &tool_use.id,
-                tool.and_then(|t| t.initial_title(tool_use.input.clone()).ok())
-                    .map(|i| i.into())
-                    .unwrap_or_else(|| tool_use.name.to_string()),
-                tool.map(|t| t.kind()).unwrap_or(acp::ToolKind::Other),
-                tool_use.input.clone(),
+                id,
+                title.to_string(),
+                kind,
+                input,
             ))))
             .ok();
     }