fix(agent): discard truncated tool calls when finish_reason == Length (#1631) (#1632)

This commit is contained in:
rajulbhatnagar
2026-03-26 10:02:41 +03:00
committed by GitHub
parent b3fbef5287
commit ed4d92932a
5 changed files with 239 additions and 4 deletions
+125 -1
View File
@@ -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<LoopOutcome, Error> {
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"
);
}
}
+2
View File
@@ -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 {
+2 -1
View File
@@ -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};
+108 -2
View File
@@ -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<String, String>,
/// 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<crate::llm::CompletionResponse, crate::llm::error::LlmError> {
unimplemented!()
}
async fn complete_with_tools(
&self,
_request: crate::llm::ToolCompletionRequest,
) -> Result<crate::llm::ToolCompletionResponse, crate::llm::error::LlmError> {
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");
}
}
+2
View File
@@ -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