From ed4d92932ac5d2d9123a8448aac4627bb8bb2d7c Mon Sep 17 00:00:00 2001 From: rajulbhatnagar Date: Thu, 26 Mar 2026 00:02:41 -0700 Subject: [PATCH] fix(agent): discard truncated tool calls when finish_reason == Length (#1631) (#1632) --- src/agent/agentic_loop.rs | 126 +++++++++++++++++++++++++++++++++++++- src/agent/dispatcher.rs | 2 + src/llm/mod.rs | 3 +- src/llm/reasoning.rs | 110 ++++++++++++++++++++++++++++++++- src/worker/job.rs | 2 + 5 files changed, 239 insertions(+), 4 deletions(-) diff --git a/src/agent/agentic_loop.rs b/src/agent/agentic_loop.rs index e61856dc..27c2ab72 100644 --- a/src/agent/agentic_loop.rs +++ b/src/agent/agentic_loop.rs @@ -10,7 +10,7 @@ use std::borrow::Cow; use crate::agent::session::PendingApproval; use crate::error::Error; -use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult}; +use crate::llm::{ChatMessage, FinishReason, Reasoning, ReasoningContext, RespondResult}; /// Signal from the delegate indicating how the loop should proceed. pub enum LoopSignal { @@ -134,6 +134,9 @@ pub async fn run_agentic_loop( config: &AgenticLoopConfig, ) -> Result { let mut consecutive_tool_intent_nudges: u32 = 0; + // Accumulates across all iterations (not reset by text responses) so + // non-consecutive truncations still escalate to force_text. + let mut truncation_count: u32 = 0; for iteration in 1..=config.max_iterations { // Check for external signals (stop, cancellation, user messages) @@ -215,7 +218,35 @@ pub async fn run_agentic_loop( tool_calls, content, } => { + // If the response was truncated, tool call parameters are likely + // incomplete. Discard them and tell the LLM to try a different + // approach rather than executing malformed tool calls. + if output.finish_reason == FinishReason::Length { + truncation_count += 1; + let names: Vec<&str> = tool_calls.iter().map(|tc| tc.name.as_str()).collect(); + tracing::warn!( + iteration, + tools = ?names, + truncation_count, + "Discarding truncated tool calls (finish_reason=Length)" + ); + if let Some(ref text) = content { + reason_ctx.messages.push(ChatMessage::assistant(text)); + } + reason_ctx + .messages + .push(ChatMessage::user(crate::llm::TRUNCATED_TOOL_CALL_NOTICE)); + // After repeated truncations, force text-only mode so the LLM + // stops attempting tool calls it can't fit in the output budget. + if truncation_count >= 3 { + reason_ctx.force_text = true; + } + delegate.after_iteration(iteration).await; + continue; + } + consecutive_tool_intent_nudges = 0; + truncation_count = 0; if let Some(outcome) = delegate .execute_tool_calls(tool_calls, content, reason_ctx) @@ -271,6 +302,7 @@ mod tests { RespondOutput { result: RespondResult::Text(text.to_string()), usage: zero_usage(), + finish_reason: FinishReason::Stop, } } @@ -281,6 +313,7 @@ mod tests { content: None, }, usage: zero_usage(), + finish_reason: FinishReason::ToolUse, } } @@ -622,4 +655,95 @@ mod tests { let result = truncate_for_preview("café", 4); assert_eq!(result, "caf..."); } + + #[tokio::test] + async fn test_truncated_tool_calls_discarded_on_length() { + let truncated_tool_call = ToolCall { + id: "call_1".to_string(), + name: "memory_write".to_string(), + arguments: serde_json::json!({}), // empty — truncated + reasoning: None, + }; + let truncated_output = RespondOutput { + result: RespondResult::ToolCalls { + tool_calls: vec![truncated_tool_call], + content: Some("I'll write the report.".to_string()), + }, + usage: zero_usage(), + finish_reason: FinishReason::Length, // response was truncated + }; + let delegate = MockDelegate::new(vec![truncated_output, text_output("Summarized it.")]); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig { + max_iterations: 5, + ..Default::default() + }; + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + // Tool calls should NOT have been executed + assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0); + // The loop should have continued and returned the text response + assert!(matches!(outcome, LoopOutcome::Response(ref t) if t == "Summarized it.")); + // A truncation notice should have been injected into context + assert!( + ctx.messages + .iter() + .any(|m| m.role == crate::llm::Role::User && m.content.contains("truncated")), + "Should inject truncation notice into context" + ); + // The partial assistant content should have been preserved + assert!( + ctx.messages + .iter() + .any(|m| m.role == crate::llm::Role::Assistant + && m.content.contains("write the report")), + "Should preserve partial assistant content" + ); + } + + #[tokio::test] + async fn test_repeated_truncations_force_text_mode() { + let make_truncated = || RespondOutput { + result: RespondResult::ToolCalls { + tool_calls: vec![ToolCall { + id: "call_1".to_string(), + name: "memory_write".to_string(), + arguments: serde_json::json!({}), + reasoning: None, + }], + content: None, + }, + usage: zero_usage(), + finish_reason: FinishReason::Length, + }; + // Three truncated responses, then a text response + let delegate = MockDelegate::new(vec![ + make_truncated(), + make_truncated(), + make_truncated(), + text_output("Gave up on tool calls."), + ]); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig { + max_iterations: 5, + ..Default::default() + }; + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + assert!(matches!(outcome, LoopOutcome::Response(_))); + assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0); + // After 3 truncations, force_text should be set + assert!( + ctx.force_text, + "Should escalate to force_text after repeated truncations" + ); + } } diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 96bca197..a5f9cd6f 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -306,6 +306,8 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { // Update context for this iteration reason_ctx.available_tools = tool_defs; + // Preserve force_text if already set (e.g. by truncation escalation). + let force_text = force_text || reason_ctx.force_text; reason_ctx.system_prompt = Some(if force_text { self.cached_prompt_no_tools.clone() } else { diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 308b3983..d681547d 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -63,7 +63,8 @@ pub use provider::{ }; pub use reasoning::{ ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, SILENT_REPLY_TOKEN, - TOOL_INTENT_NUDGE, TokenUsage, ToolSelection, is_silent_reply, llm_signals_tool_intent, + TOOL_INTENT_NUDGE, TRUNCATED_TOOL_CALL_NOTICE, TokenUsage, ToolSelection, is_silent_reply, + llm_signals_tool_intent, }; pub use recording::RecordingLlm; pub use registry::{ProviderDefinition, ProviderProtocol, ProviderRegistry}; diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index 473eb16d..6e078ac7 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -8,8 +8,8 @@ use serde::{Deserialize, Serialize}; use crate::llm::error::LlmError; use crate::llm::{ - ChatMessage, CompletionRequest, LlmProvider, Role, ToolCall, ToolCompletionRequest, - ToolDefinition, + ChatMessage, CompletionRequest, FinishReason, LlmProvider, Role, ToolCall, + ToolCompletionRequest, ToolDefinition, }; /// Token the agent returns when it has nothing to say (e.g. in group chats). @@ -23,6 +23,13 @@ You said you would perform an action, but you did not include any tool calls.\n\ Do NOT describe what you intend to do — actually call the tool now.\n\ Use the tool_calls mechanism to invoke the appropriate tool."; +/// Notice injected when the LLM's response was truncated mid-tool-call, +/// causing incomplete parameters. Tells the LLM to try a different approach. +pub const TRUNCATED_TOOL_CALL_NOTICE: &str = "\ +Your previous response was truncated while generating tool call parameters. \ +The tool calls were discarded. Please try a different approach — \ +summarize or transform the data instead of echoing it verbatim in a tool call."; + /// Seed value used as the second argument to `generate_tool_call_id` when /// recovering tool calls from malformed LLM text responses. This must differ /// from the `0` seed used in `rig_adapter::normalized_tool_call_id` to avoid @@ -194,6 +201,8 @@ pub struct ReasoningContext { pub metadata: std::collections::HashMap, /// When true, force a text-only response (ignore available tools). /// Used by the agentic loop to guarantee termination near the iteration limit. + /// Sticky: once set, never cleared within a loop invocation. Callers must + /// create a fresh `ReasoningContext` per `run_agentic_loop()` call. pub force_text: bool, /// Pre-built system prompt. When set, `respond_with_tools` uses this directly /// instead of calling `build_system_prompt_with_tools`. Allows callers to build @@ -349,6 +358,7 @@ pub enum RespondResult { pub struct RespondOutput { pub result: RespondResult, pub usage: TokenUsage, + pub finish_reason: FinishReason, } /// Reasoning engine for the agent. @@ -530,6 +540,17 @@ impl Reasoning { let response = self.llm.complete_with_tools(request).await?; + // If the response was truncated, tool call parameters are likely incomplete. + // Return empty so the caller can fall through to respond_with_tools() which + // has a larger output token budget. + if response.finish_reason == FinishReason::Length { + tracing::warn!( + "select_tools response truncated (finish_reason=Length), \ + discarding potentially incomplete tool selections" + ); + return Ok(vec![]); + } + let shared_reasoning = response .content .map(|c| { @@ -722,6 +743,7 @@ Respond in JSON format: content: narrative, }, usage, + finish_reason: response.finish_reason, }); } @@ -749,6 +771,7 @@ Respond in JSON format: }, }, usage, + finish_reason: response.finish_reason, }); } @@ -774,6 +797,7 @@ Respond in JSON format: Ok(RespondOutput { result: RespondResult::Text(final_text), usage, + finish_reason: response.finish_reason, }) } else { // No tools, use simple completion @@ -805,6 +829,7 @@ Respond in JSON format: cache_read_input_tokens: response.cache_read_input_tokens, cache_creation_input_tokens: response.cache_creation_input_tokens, }, + finish_reason: response.finish_reason, }) } } @@ -3315,4 +3340,85 @@ That's my plan."#; let cleaned = clean_response(&pre_truncated); assert!(cleaned.trim().is_empty()); } + + // ---- select_tools truncation guard ---- + + /// Mock provider that returns tool calls with a configurable finish_reason. + struct TruncatingLlm { + finish_reason: crate::llm::FinishReason, + } + + #[async_trait::async_trait] + impl crate::llm::LlmProvider for TruncatingLlm { + fn model_name(&self) -> &str { + "truncating-stub" + } + fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) { + (rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO) + } + async fn complete( + &self, + _request: crate::llm::CompletionRequest, + ) -> Result { + unimplemented!() + } + async fn complete_with_tools( + &self, + _request: crate::llm::ToolCompletionRequest, + ) -> Result { + Ok(crate::llm::ToolCompletionResponse { + content: Some("I'll write the report.".to_string()), + tool_calls: vec![ToolCall { + id: "call_1".to_string(), + name: "memory_write".to_string(), + arguments: serde_json::json!({}), + reasoning: None, + }], + input_tokens: 5000, + output_tokens: 1024, + finish_reason: self.finish_reason, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }) + } + } + + #[tokio::test] + async fn test_select_tools_returns_empty_on_truncation() { + let llm = Arc::new(TruncatingLlm { + finish_reason: FinishReason::Length, + }); + let reasoning = Reasoning::new(llm); + let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report")); + ctx.available_tools.push(ToolDefinition { + name: "memory_write".to_string(), + description: "Write to memory".to_string(), + parameters: serde_json::json!({"type": "object"}), + }); + + let selections = reasoning.select_tools(&ctx).await.unwrap(); + assert!( + selections.is_empty(), + "Truncated tool selections should be discarded (got {} selections)", + selections.len() + ); + } + + #[tokio::test] + async fn test_select_tools_returns_selections_when_not_truncated() { + let llm = Arc::new(TruncatingLlm { + finish_reason: FinishReason::ToolUse, + }); + let reasoning = Reasoning::new(llm); + let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report")); + ctx.available_tools.push(ToolDefinition { + name: "memory_write".to_string(), + description: "Write to memory".to_string(), + parameters: serde_json::json!({"type": "object"}), + }); + + let selections = reasoning.select_tools(&ctx).await.unwrap(); + assert_eq!(selections.len(), 1); + assert_eq!(selections[0].tool_name, "memory_write"); + } } diff --git a/src/worker/job.rs b/src/worker/job.rs index 671b8864..f74d4ec8 100644 --- a/src/worker/job.rs +++ b/src/worker/job.rs @@ -1158,6 +1158,7 @@ impl<'a> JobDelegate<'a> { Ok(crate::llm::RespondOutput { result: RespondResult::Text(String::new()), usage: crate::llm::TokenUsage::default(), + finish_reason: crate::llm::FinishReason::Stop, }) } } @@ -1283,6 +1284,7 @@ impl<'a> LoopDelegate for JobDelegate<'a> { content: reasoning_text, }, usage: crate::llm::TokenUsage::default(), + finish_reason: crate::llm::FinishReason::ToolUse, }); } Ok(_) => {} // empty selections, fall through