diff --git a/crates/buzz-agent/src/agent.rs b/crates/buzz-agent/src/agent.rs index b121564d40..f3fbabdcda 100644 --- a/crates/buzz-agent/src/agent.rs +++ b/crates/buzz-agent/src/agent.rs @@ -24,6 +24,17 @@ const ERROR_REFLECTION_SUFFIX: &str = const UNSUPPORTED_IMAGE_TOOL_MESSAGE: &str = "The current model does not support image input. The image was removed from conversation history so this turn can continue. Use a text-based inspection tool or ask the user for a textual description instead."; +/// Model-visible feedback after the provider truncates an assistant response at +/// its output-token limit. This is a user message rather than a synthetic tool +/// result because truncation can happen without a tool call (and an unpaired +/// tool result is invalid on every provider wire format). +const MAX_TOKENS_RECOVERY_MESSAGE: &str = "Your previous response exceeded the model's output token limit and was truncated. Any incomplete tool call was not run. Continue the task, breaking the work or tool call into smaller steps and keeping the response concise."; + +/// A provider can repeatedly spend its entire output allowance without making +/// progress, while `max_rounds` is unbounded by default. Keep the in-turn rescue +/// finite so a persistently truncating model eventually surfaces `max_tokens`. +const MAX_TOKENS_RECOVERIES_PER_RUN: u32 = 2; + /// Remove image blocks that the provider has explicitly rejected while keeping /// their surrounding tool result (and therefore the tool-call/result pairing) /// intact. Returns the number of images removed; zero means the provider error @@ -262,6 +273,10 @@ impl RunCtx<'_> { // per-session: a fresh prompt deserves a fresh chance to recover, and // `max_rounds` defaults to 0 (unbounded) so it cannot bound this. let mut context_recoveries = 0u32; + // Per-run output-truncation recovery budget. Unlike context recovery, + // these successful provider requests consume a real round and are not + // refunded; this counter only bounds the default-unlimited case. + let mut max_tokens_recoveries = 0u32; loop { if self.cfg.max_rounds > 0 && round >= self.cfg.max_rounds { return Ok(StopReason::MaxTurnRequests); @@ -515,6 +530,37 @@ impl RunCtx<'_> { .await; } + // `max_tokens` describes a truncated assistant response, not turn + // completion. Never execute tool calls from it: although one may + // parse as valid, a later call (or surrounding instructions) may + // have been cut off. Replay only the text, with no tool calls, so + // the history remains valid without fabricated tool results; then + // add actionable user-role feedback and ask the model to continue. + if response.stop == ProviderStop::MaxTokens { + self.history.push(HistoryItem::Assistant { + text: response.text, + tool_calls: Vec::new(), + reasoning_details: response.reasoning_details, + }); + if max_tokens_recoveries >= MAX_TOKENS_RECOVERIES_PER_RUN { + tracing::warn!( + recoveries = max_tokens_recoveries, + "provider repeatedly hit output token limit; recovery budget exhausted" + ); + return Ok(StopReason::MaxTokens); + } + max_tokens_recoveries = max_tokens_recoveries.saturating_add(1); + tracing::warn!( + recovery = max_tokens_recoveries, + max_recoveries = MAX_TOKENS_RECOVERIES_PER_RUN, + discarded_tool_calls = response.tool_calls.len(), + "provider hit output token limit; asking model to continue in smaller steps" + ); + self.history + .push(HistoryItem::User(MAX_TOKENS_RECOVERY_MESSAGE.to_string())); + continue; + } + if response.tool_calls.is_empty() { if response.stop == ProviderStop::ToolUse { return Err(AgentError::Llm( diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index c452775c9e..a091790425 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -1229,6 +1229,11 @@ fn databricks_v2_path(route: DatabricksV2Route) -> &'static str { } fn parse_responses(v: Value) -> Result { + let max_tokens = v.get("status").and_then(Value::as_str) == Some("incomplete") + && v.get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(Value::as_str) + == Some("max_output_tokens"); let mut text = String::new(); let mut reasoning = String::new(); let mut tool_calls = Vec::new(); @@ -1259,7 +1264,7 @@ fn parse_responses(v: Value) -> Result { } } } - Some("function_call") => { + Some("function_call") if !max_tokens => { saw_function_call = true; let raw = item .get("arguments") @@ -1279,6 +1284,9 @@ fn parse_responses(v: Value) -> Result { Default::default(), )?); } + // Incomplete Responses output can carry a partial function call. + // It is intentionally discarded by the in-turn recovery path. + Some("function_call") => {} Some("reasoning") => { // Reasoning summary items from the Responses API. Each item has a // `summary` array of `{"type": "summary_text", "text": "..."}` objects. @@ -1551,13 +1559,19 @@ fn parse_anthropic(v: Value) -> Result { reasoning.push_str(t); } } - // Anthropic's replay shape is fully modelled, so nothing to keep. - Some("tool_use") => tool_calls.push(make_tool_call( - str_field(b, "id"), - str_field(b, "name"), - b.get("input").cloned().unwrap_or(Value::Null), - Default::default(), - )?), + // A max-token response may end in the middle of a tool input. + // The agent discards all calls from truncated responses, so do + // not reject the whole response trying to parse an input that + // can never be executed. + Some("tool_use") if stop != ProviderStop::MaxTokens => { + tool_calls.push(make_tool_call( + str_field(b, "id"), + str_field(b, "name"), + b.get("input").cloned().unwrap_or(Value::Null), + Default::default(), + )?) + } + Some("tool_use") => {} _ => {} } } @@ -1645,30 +1659,33 @@ fn parse_openai(v: Value) -> Result { } }; let mut tool_calls = Vec::new(); - if let Some(arr) = msg.get("tool_calls").and_then(Value::as_array) { - for tc in arr { - let f = tc - .get("function") - .ok_or_else(|| AgentError::Llm("tool_call missing function".into()))?; - let raw = f.get("arguments").and_then(Value::as_str).unwrap_or("{}"); - let args: Value = serde_json::from_str(raw) - .map_err(|e| AgentError::Llm(format!("tool_call.arguments not valid JSON: {e}")))?; - // Everything on the wire object we do not model, kept for replay. - let extra = tc - .as_object() - .map(|o| { - o.iter() - .filter(|(k, _)| !matches!(k.as_str(), "id" | "type" | "function")) - .map(|(k, v)| (k.clone(), v.clone())) - .collect() - }) - .unwrap_or_default(); - tool_calls.push(make_tool_call( - str_field(tc, "id"), - str_field(f, "name"), - args, - extra, - )?); + if stop != ProviderStop::MaxTokens { + if let Some(arr) = msg.get("tool_calls").and_then(Value::as_array) { + for tc in arr { + let f = tc + .get("function") + .ok_or_else(|| AgentError::Llm("tool_call missing function".into()))?; + let raw = f.get("arguments").and_then(Value::as_str).unwrap_or("{}"); + let args: Value = serde_json::from_str(raw).map_err(|e| { + AgentError::Llm(format!("tool_call.arguments not valid JSON: {e}")) + })?; + // Everything on the wire object we do not model, kept for replay. + let extra = tc + .as_object() + .map(|o| { + o.iter() + .filter(|(k, _)| !matches!(k.as_str(), "id" | "type" | "function")) + .map(|(k, v)| (k.clone(), v.clone())) + .collect() + }) + .unwrap_or_default(); + tool_calls.push(make_tool_call( + str_field(tc, "id"), + str_field(f, "name"), + args, + extra, + )?); + } } } dedupe_provider_ids(&mut tool_calls); @@ -3527,10 +3544,51 @@ mod tests { let v = serde_json::json!({ "status": "incomplete", "incomplete_details": {"reason": "max_output_tokens"}, - "output": [], + "output": [{ + "type": "function_call", + "call_id": "partial", + "name": "dev__shell", + "arguments": "{\"command\":\"unterminated", + }], }); let r = parse_responses(v).unwrap(); assert_eq!(r.stop, ProviderStop::MaxTokens); + assert!(r.tool_calls.is_empty()); + } + + #[test] + fn truncated_openai_tool_arguments_are_discarded_not_rejected() { + let v = serde_json::json!({ + "choices": [{ + "finish_reason": "length", + "message": { + "content": "partial text", + "tool_calls": [{ + "id": "partial", + "type": "function", + "function": { + "name": "dev__shell", + "arguments": "{\"command\":\"unterminated", + }, + }], + }, + }], + }); + let r = parse_openai(v).unwrap(); + assert_eq!(r.stop, ProviderStop::MaxTokens); + assert_eq!(r.text, "partial text"); + assert!(r.tool_calls.is_empty()); + } + + #[test] + fn truncated_anthropic_tool_use_is_discarded_not_rejected() { + let v = serde_json::json!({ + "stop_reason": "max_tokens", + "content": [{"type": "tool_use", "id": "", "name": "", "input": null}], + }); + let r = parse_anthropic(v).unwrap(); + assert_eq!(r.stop, ProviderStop::MaxTokens); + assert!(r.tool_calls.is_empty()); } #[test] diff --git a/crates/buzz-agent/tests/regressions.rs b/crates/buzz-agent/tests/regressions.rs index c82be76dc0..a972f9b242 100644 --- a/crates/buzz-agent/tests/regressions.rs +++ b/crates/buzz-agent/tests/regressions.rs @@ -233,6 +233,26 @@ fn openai_text_with_usage(content: &str, prompt_tokens: u64) -> Value { v } +fn openai_max_tokens(content: &str, tool_calls: Value) -> Value { + json!({ + "id": "cc-max", "object": "chat.completion", "model": "fake-model", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": content, + "tool_calls": tool_calls, + }, + "finish_reason": "length", + }], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 100, + "total_tokens": 110, + }, + }) +} + fn openai_tool_call(id: &str, name: &str, args: Value) -> Value { json!({ "id": "cc-2", "object": "chat.completion", "model": "fake-model", @@ -2437,6 +2457,110 @@ async fn context_window_400_recovers_instead_of_sticking() { h.shutdown().await; } +/// A provider output-token stop is an interrupted round, not completion. The +/// agent must preserve any text, discard a possibly partial tool call, provide +/// actionable feedback, and let the same prompt finish normally. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn max_tokens_recovers_in_turn_without_running_partial_tool_call() { + const PARTIAL: &str = "partial-before-limit"; + let partial_call = json!([{ + "id": "partial-call", "type": "function", + "function": { "name": "dev__shell", "arguments": "{\"command\":\"echo" }, + }]); + let llm = spawn_capturing_llm(vec![ + openai_max_tokens(PARTIAL, partial_call), + openai_text("done after truncation"), + ]) + .await; + let mut h = Harness::spawn(&llm.url).await; + let sid = init_session(&mut h, json!([])).await; + + let prompt_id = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"do the task"}]}), + ) + .await; + let reply = h.recv_until(|v| v["id"] == json!(prompt_id)).await; + assert_eq!(reply["result"]["stopReason"], "end_turn", "{reply}"); + + let requests = llm.captured.lock().await; + assert_eq!(requests.len(), 2, "truncation should trigger one retry"); + let retry = &requests[1]["messages"]; + let serialized = retry.to_string(); + assert!( + serialized.contains(PARTIAL), + "partial text was lost: {retry}" + ); + assert!( + serialized.contains("output token limit") + && serialized.contains("smaller steps") + && serialized.contains("tool call"), + "retry lacks actionable truncation feedback: {retry}" + ); + assert!( + !serialized.contains("partial-call") && !serialized.contains("tool_call_id"), + "partial tool call must not be replayed or executed: {retry}" + ); + drop(requests); + h.shutdown().await; +} + +/// `max_rounds` counts max-token responses because they are successful, billed +/// provider requests. Recovery must not grant them the refund reserved for a +/// rejected context-overflow request. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn max_tokens_recovery_respects_finite_round_cap() { + let llm = spawn_capturing_llm(vec![ + openai_max_tokens("cut off", json!([])), + openai_text("must not be requested"), + ]) + .await; + let mut h = Harness::spawn_with_env(&llm.url, &[("BUZZ_AGENT_MAX_ROUNDS", "1")]).await; + let sid = init_session(&mut h, json!([])).await; + let prompt_id = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), + ) + .await; + let reply = h.recv_until(|v| v["id"] == json!(prompt_id)).await; + assert_eq!( + reply["result"]["stopReason"], "max_turn_requests", + "{reply}" + ); + assert_eq!(llm.captured.lock().await.len(), 1); + h.shutdown().await; +} + +/// With the production-unbounded round setting, a model that always fills its +/// output allowance still has to return. Two recovery prompts are allowed; the +/// third truncation surfaces the original stop reason. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn repeated_max_tokens_is_bounded() { + let responses = (0..4) + .map(|_| openai_max_tokens("still truncated", json!([]))) + .collect(); + let llm = spawn_capturing_llm(responses).await; + let mut h = Harness::spawn_with_env(&llm.url, &[("BUZZ_AGENT_MAX_ROUNDS", "0")]).await; + let sid = init_session(&mut h, json!([])).await; + let prompt_id = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), + ) + .await; + let reply = tokio::time::timeout( + Duration::from_secs(10), + h.recv_until(|v| v["id"] == json!(prompt_id)), + ) + .await + .expect("max-token recovery must be bounded"); + assert_eq!(reply["result"]["stopReason"], "max_tokens", "{reply}"); + assert_eq!(llm.captured.lock().await.len(), 3); + h.shutdown().await; +} + /// A successful recovery must actually send the recovered completion, even when /// `max_rounds` is finite. `round` is incremented BEFORE the completion that /// gets rejected, so a naive `continue` after recovery re-enters the loop with