From ebac703532fcbd47617f000d0d003210acf8f64e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 22 Jan 2026 09:47:02 +0100 Subject: [PATCH 1/2] Fix sometimes missing tool results --- crates/code_assistant/src/agent/runner.rs | 98 ++++++- crates/code_assistant/src/agent/tests.rs | 299 ++++++++++++++++++++++ 2 files changed, 393 insertions(+), 4 deletions(-) diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index 0c241790..a2b7898c 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -1606,8 +1606,14 @@ impl Agent { Ok(()) } - /// Prepare messages for LLM request, dynamically rendering tool outputs - fn render_tool_results_in_messages(&self) -> Vec { + /// Prepare messages for LLM request, dynamically rendering tool outputs. + /// + /// This function also handles cancelled tool executions: if an assistant message + /// contains `ToolUse` blocks but there's no corresponding `ToolResult` in the + /// following user message (or no following user message at all), we generate + /// a synthetic "user cancelled" `ToolResult` to satisfy the API requirement that + /// every `tool_use` must have a corresponding `tool_result`. + pub fn render_tool_results_in_messages(&self) -> Vec { // Start with a clean slate let mut messages = Vec::new(); @@ -1625,11 +1631,95 @@ impl Agent { tool_outputs.insert(tool_use_id.clone(), rendered_output); } - // Now rebuild the message history, replacing tool outputs with our dynamically rendered versions + // Build a set of all tool_use_ids that have corresponding tool_results in the message history + let mut tool_ids_with_results: std::collections::HashSet = + std::collections::HashSet::new(); + for msg in self.active_messages() { + if let MessageContent::Structured(blocks) = &msg.content { + for block in blocks { + if let ContentBlock::ToolResult { tool_use_id, .. } = block { + tool_ids_with_results.insert(tool_use_id.clone()); + } + } + } + } + + // Now rebuild the message history + let active_msgs: Vec<_> = self.active_messages().to_vec(); + for (idx, msg) in active_msgs.iter().enumerate() { match &msg.content { MessageContent::Structured(blocks) => { - // Look for ToolResult blocks + if msg.role == MessageRole::Assistant { + // Check for ToolUse blocks that need synthetic ToolResults + let tool_use_ids: Vec = blocks + .iter() + .filter_map(|block| { + if let ContentBlock::ToolUse { id, .. } = block { + Some(id.clone()) + } else { + None + } + }) + .collect(); + + // Find tool_use_ids without corresponding tool_results + let missing_results: Vec<&String> = tool_use_ids + .iter() + .filter(|id| !tool_ids_with_results.contains(*id)) + .collect(); + + if !missing_results.is_empty() { + // We need to add the assistant message first, then add a synthetic + // user message with cancelled tool results + messages.push(msg.clone()); + + // Generate synthetic ToolResult blocks for cancelled tools + let cancelled_blocks: Vec = missing_results + .iter() + .map(|tool_id| { + debug!( + "Generating synthetic 'cancelled' tool result for tool_use_id: {}", + tool_id + ); + ContentBlock::ToolResult { + tool_use_id: (*tool_id).clone(), + content: "Tool execution was cancelled by user.".to_string(), + is_error: Some(true), + start_time: None, + end_time: None, + } + }) + .collect(); + + // Check if the next message is already a user message with tool results + // In that case, we need to merge the cancelled results + let next_msg = active_msgs.get(idx + 1); + let should_create_new_message = match next_msg { + Some(next) if next.role == MessageRole::User => { + // Check if this user message has tool results + match &next.content { + MessageContent::Structured(next_blocks) => !next_blocks + .iter() + .any(|b| matches!(b, ContentBlock::ToolResult { .. })), + _ => true, + } + } + _ => true, + }; + + if should_create_new_message { + // Insert a new user message with the cancelled tool results + let cancelled_msg = + Message::new_user_content(cancelled_blocks.clone()); + messages.push(cancelled_msg); + } + // If next message already has tool results, we'll handle merging when we process it + continue; + } + } + + // Look for ToolResult blocks and update with rendered output let mut new_blocks = Vec::new(); let mut need_update = false; diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index eafc6756..7d0097de 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -2153,3 +2153,302 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() Ok(()) } + +#[tokio::test] +async fn test_render_tool_results_generates_cancelled_results_for_missing_executions() -> Result<()> +{ + // This test verifies that when an assistant message contains ToolUse blocks + // but there's no corresponding ToolResult in the message history (because the + // user cancelled the tool execution), we generate synthetic "user cancelled" + // ToolResult blocks to satisfy the API requirement. + + let mock_llm = MockLLMProvider::new(vec![]); + let components = AgentComponents { + llm_provider: Box::new(mock_llm), + project_manager: Box::new(MockProjectManager::new()), + command_executor: Box::new(create_command_executor_mock()), + ui: Arc::new(MockUI::default()), + state_persistence: Box::new(MockStatePersistence::new()), + permission_handler: None, + sub_agent_runner: None, + }; + + let session_config = SessionConfig { + tool_syntax: ToolSyntax::Native, + ..SessionConfig::default() + }; + + let mut agent = Agent::new(components, session_config); + agent.disable_naming_reminders(); + + // Simulate a scenario where: + // 1. User asks a question + // 2. Assistant responds with a tool call + // 3. User cancels the tool execution (no ToolResult message added) + // 4. User asks a follow-up question + + // Add user message + agent.append_message(Message::new_user("Please run the tests."))?; + + // Add assistant message with a tool call (simulating what LLM returned) + agent.append_message( + Message::new_assistant_content(vec![ + ContentBlock::new_text("I'll run the tests for you."), + ContentBlock::new_tool_use( + "tool-1-1", + "execute_command", + serde_json::json!({ + "project": "test", + "command_line": "cargo test" + }), + ), + ]) + .with_request_id(1), + )?; + + // Note: We do NOT add a ToolResult message - simulating user cancellation + + // Add another user message (user continues the conversation) + agent.append_message(Message::new_user("Never mind, let's do something else."))?; + + // Now call render_tool_results_in_messages - this is what gets sent to the LLM + let rendered_messages = agent.render_tool_results_in_messages(); + + // We should have: + // 1. Original user message + // 2. Assistant message with tool call + // 3. Synthetic user message with cancelled tool result + // 4. Follow-up user message + assert_eq!( + rendered_messages.len(), + 4, + "Expected 4 messages: user, assistant, cancelled tool result, follow-up user" + ); + + // Verify the synthetic cancelled tool result was inserted + let cancelled_message = &rendered_messages[2]; + assert_eq!(cancelled_message.role, MessageRole::User); + + if let MessageContent::Structured(blocks) = &cancelled_message.content { + assert_eq!(blocks.len(), 1); + if let ContentBlock::ToolResult { + tool_use_id, + content, + is_error, + .. + } = &blocks[0] + { + assert_eq!(tool_use_id, "tool-1-1"); + assert!(content.contains("cancelled")); + assert!(is_error.unwrap_or(false)); + } else { + panic!("Expected ToolResult block"); + } + } else { + panic!("Expected Structured content for cancelled tool result"); + } + + // Verify the follow-up user message is still present + let followup_message = &rendered_messages[3]; + assert_eq!(followup_message.role, MessageRole::User); + if let MessageContent::Text(text) = &followup_message.content { + assert_eq!(text, "Never mind, let's do something else."); + } else { + panic!("Expected Text content for follow-up message"); + } + + Ok(()) +} + +#[tokio::test] +async fn test_render_tool_results_preserves_existing_tool_results() -> Result<()> { + // This test verifies that when tool results already exist, we don't add + // synthetic cancelled results for them. + + let mock_llm = MockLLMProvider::new(vec![]); + let components = AgentComponents { + llm_provider: Box::new(mock_llm), + project_manager: Box::new(MockProjectManager::new()), + command_executor: Box::new(create_command_executor_mock()), + ui: Arc::new(MockUI::default()), + state_persistence: Box::new(MockStatePersistence::new()), + permission_handler: None, + sub_agent_runner: None, + }; + + let session_config = SessionConfig { + tool_syntax: ToolSyntax::Native, + ..SessionConfig::default() + }; + + let mut agent = Agent::new(components, session_config); + agent.disable_naming_reminders(); + + // Simulate a normal scenario where tool execution completed successfully + + // Add user message + agent.append_message(Message::new_user("Read the README."))?; + + // Add assistant message with a tool call + agent.append_message( + Message::new_assistant_content(vec![ + ContentBlock::new_text("I'll read the README file."), + ContentBlock::new_tool_use( + "tool-1-1", + "read_files", + serde_json::json!({ + "project": "test", + "paths": ["README.md"] + }), + ), + ]) + .with_request_id(1), + )?; + + // Add the tool result (normal execution completed) + agent.append_message(Message::new_user_content(vec![ContentBlock::ToolResult { + tool_use_id: "tool-1-1".to_string(), + content: "File contents here".to_string(), + is_error: None, + start_time: None, + end_time: None, + }]))?; + + // Now call render_tool_results_in_messages + let rendered_messages = agent.render_tool_results_in_messages(); + + // We should have exactly 3 messages - no synthetic cancelled results added + assert_eq!( + rendered_messages.len(), + 3, + "Expected 3 messages: user, assistant, tool result" + ); + + // Verify the tool result is the original one (not a cancelled one) + let result_message = &rendered_messages[2]; + if let MessageContent::Structured(blocks) = &result_message.content { + if let ContentBlock::ToolResult { + tool_use_id, + content, + is_error, + .. + } = &blocks[0] + { + assert_eq!(tool_use_id, "tool-1-1"); + // Content should be the original, not "cancelled" + assert!(content.contains("File contents") || content.is_empty()); + // Should not be marked as error + assert!(!is_error.unwrap_or(false)); + } + } + + Ok(()) +} + +#[tokio::test] +async fn test_render_tool_results_handles_multiple_cancelled_tools() -> Result<()> { + // This test verifies that multiple cancelled tool calls are all handled correctly. + + let mock_llm = MockLLMProvider::new(vec![]); + let components = AgentComponents { + llm_provider: Box::new(mock_llm), + project_manager: Box::new(MockProjectManager::new()), + command_executor: Box::new(create_command_executor_mock()), + ui: Arc::new(MockUI::default()), + state_persistence: Box::new(MockStatePersistence::new()), + permission_handler: None, + sub_agent_runner: None, + }; + + let session_config = SessionConfig { + tool_syntax: ToolSyntax::Native, + ..SessionConfig::default() + }; + + let mut agent = Agent::new(components, session_config); + agent.disable_naming_reminders(); + + // Add user message + agent.append_message(Message::new_user("Check the project."))?; + + // Add assistant message with multiple tool calls (all cancelled) + agent.append_message( + Message::new_assistant_content(vec![ + ContentBlock::new_text("I'll check multiple things."), + ContentBlock::new_tool_use( + "tool-1-1", + "read_files", + serde_json::json!({ + "project": "test", + "paths": ["file1.txt"] + }), + ), + ContentBlock::new_tool_use( + "tool-1-2", + "execute_command", + serde_json::json!({ + "project": "test", + "command_line": "ls -la" + }), + ), + ]) + .with_request_id(1), + )?; + + // No tool results added - both cancelled + + // Now call render_tool_results_in_messages + let rendered_messages = agent.render_tool_results_in_messages(); + + // We should have: + // 1. Original user message + // 2. Assistant message with tool calls + // 3. Synthetic user message with both cancelled tool results + assert_eq!( + rendered_messages.len(), + 3, + "Expected 3 messages: user, assistant, cancelled tool results" + ); + + // Verify the synthetic cancelled results + let cancelled_message = &rendered_messages[2]; + assert_eq!(cancelled_message.role, MessageRole::User); + + if let MessageContent::Structured(blocks) = &cancelled_message.content { + assert_eq!(blocks.len(), 2, "Should have 2 cancelled tool results"); + + // Check first cancelled result + if let ContentBlock::ToolResult { + tool_use_id, + content, + is_error, + .. + } = &blocks[0] + { + assert_eq!(tool_use_id, "tool-1-1"); + assert!(content.contains("cancelled")); + assert!(is_error.unwrap_or(false)); + } else { + panic!("Expected ToolResult block for first cancelled tool"); + } + + // Check second cancelled result + if let ContentBlock::ToolResult { + tool_use_id, + content, + is_error, + .. + } = &blocks[1] + { + assert_eq!(tool_use_id, "tool-1-2"); + assert!(content.contains("cancelled")); + assert!(is_error.unwrap_or(false)); + } else { + panic!("Expected ToolResult block for second cancelled tool"); + } + } else { + panic!("Expected Structured content for cancelled tool results"); + } + + Ok(()) +} From a9990299dec1fcff45a656719ee4e0bfee0d531e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Thu, 22 Jan 2026 09:51:10 +0100 Subject: [PATCH 2/2] Restore comment about dynamic tool results --- crates/code_assistant/src/agent/runner.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index a2b7898c..b3eda77e 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -1645,7 +1645,7 @@ impl Agent { } } - // Now rebuild the message history + // Now rebuild the message history, replacing tool outputs with our dynamically rendered versions let active_msgs: Vec<_> = self.active_messages().to_vec(); for (idx, msg) in active_msgs.iter().enumerate() { match &msg.content {