agent: Fix "tool cancelled" status being overapplied to failed tool calls (#30021)

Cole Miller created 7 months ago

Release Notes:

- Agent Beta: Fixed a bug that caused past failed tool calls to
incorrectly display as cancelled by the user.

Change summary

crates/agent/src/tool_use.rs | 35 +++++++++++++++++++++--------------
1 file changed, 21 insertions(+), 14 deletions(-)

Detailed changes

crates/agent/src/tool_use.rs

@@ -110,20 +110,27 @@ impl ToolUseState {
     }
 
     pub fn cancel_pending(&mut self) -> Vec<PendingToolUse> {
-        let mut pending_tools = Vec::new();
-        for (tool_use_id, tool_use) in self.pending_tool_uses_by_id.drain() {
-            self.tool_results.insert(
-                tool_use_id.clone(),
-                LanguageModelToolResult {
-                    tool_use_id,
-                    tool_name: tool_use.name.clone(),
-                    content: "Tool canceled by user".into(),
-                    is_error: true,
-                },
-            );
-            pending_tools.push(tool_use.clone());
-        }
-        pending_tools
+        let mut cancelled_tool_uses = Vec::new();
+        self.pending_tool_uses_by_id
+            .retain(|tool_use_id, tool_use| {
+                if matches!(tool_use.status, PendingToolUseStatus::Error { .. }) {
+                    return true;
+                }
+
+                let content = "Tool canceled by user".into();
+                self.tool_results.insert(
+                    tool_use_id.clone(),
+                    LanguageModelToolResult {
+                        tool_use_id: tool_use_id.clone(),
+                        tool_name: tool_use.name.clone(),
+                        content,
+                        is_error: true,
+                    },
+                );
+                cancelled_tool_uses.push(tool_use.clone());
+                false
+            });
+        cancelled_tool_uses
     }
 
     pub fn pending_tool_uses(&self) -> Vec<&PendingToolUse> {