From 017a570d7dec22a4dcfd989084f08f3b29293993 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sun, 13 Sep 2026 07:02:05 -0400 Subject: [PATCH] fix(sampler): a max_tokens stop after complete tool calls runs them The Messages stream failed every max_tokens stop with MaxTokensTruncation, so a response that had already emitted complete tool calls before the output budget ran out lost those calls and failed the turn. The stream now completes with StopReason::ToolCalls when every collected call's arguments parse (or are empty, the zero-argument convention); a cut call or a text-only truncation still fails so compact-on-error keeps working. The ChatCompletions and Responses streams already behaved this way. Ported from upstream 1.0.1x LengthPolicy::CompleteToolCalls, reshaped into the stream layer because Kigi has no caller that chooses a policy. Claude-Session: https://claude.ai/code/session_01RUqV8HVGXcAETyYbZ2tPQ4 --- .../kigi-sampler/src/stream/messages.rs | 26 ++++-- .../kigi-sampler/src/stream/messages_tests.rs | 87 +++++++++++++++++++ 2 files changed, 108 insertions(+), 5 deletions(-) diff --git a/crates/codegen/kigi-sampler/src/stream/messages.rs b/crates/codegen/kigi-sampler/src/stream/messages.rs index 4ba45e1..729e053 100644 --- a/crates/codegen/kigi-sampler/src/stream/messages.rs +++ b/crates/codegen/kigi-sampler/src/stream/messages.rs @@ -55,6 +55,15 @@ enum BlockType { Thinking, } +/// Executable after a `max_tokens` stop: every call's arguments parse or are empty (zero-arg convention). +fn tool_calls_complete(tool_calls: &[ToolCall]) -> bool { + !tool_calls.is_empty() + && tool_calls.iter().all(|tc| { + tc.arguments.trim().is_empty() + || serde_json::from_str::(&tc.arguments).is_ok() + }) +} + /// Transform a raw Anthropic Messages API stream into a stream of /// [`SamplingEvent`]s. /// @@ -454,11 +463,18 @@ pub fn stream_messages<'a>( } if final_stop_reason == Some(StopReason::Length) { - yield SamplingEvent::Failed { - request_id: request_id.clone(), - error: SamplingErrorInfo::from(&SamplingError::MaxTokensTruncation), - }; - return; + if !tool_calls_complete(&assistant_tool_calls) { + yield SamplingEvent::Failed { + request_id: request_id.clone(), + error: SamplingErrorInfo::from(&SamplingError::MaxTokensTruncation), + }; + return; + } + tracing::info!( + request_id = %request_id, + tool_calls = assistant_tool_calls.len(), + "max_tokens stop after complete tool calls; completing the turn so they run" + ); } // Build the final response diff --git a/crates/codegen/kigi-sampler/src/stream/messages_tests.rs b/crates/codegen/kigi-sampler/src/stream/messages_tests.rs index 6fbf52e..b6c8f3f 100644 --- a/crates/codegen/kigi-sampler/src/stream/messages_tests.rs +++ b/crates/codegen/kigi-sampler/src/stream/messages_tests.rs @@ -653,3 +653,90 @@ async fn pure_cache_hit_with_zero_uncached_still_emits_usage() { assert_eq!(usage.cached_prompt_tokens, 2500); assert_eq!(usage.total_tokens, 2501); } + +fn tool_use_start(index: u32, id: &str) -> MessageStreamEvent { + MessageStreamEvent::ContentBlockStart { + index, + content_block: ContentBlock::ToolUse { + id: id.into(), + name: "do_thing".into(), + input: serde_json::json!({}), + }, + } +} + +fn arg_delta(index: u32, partial_json: &str) -> MessageStreamEvent { + MessageStreamEvent::ContentBlockDelta { + index, + delta: StreamDelta::InputJsonDelta { + partial_json: partial_json.into(), + }, + } +} + +async fn max_tokens_after(events: Vec) -> Vec { + let mut all = vec![Ok(message_start())]; + all.extend(events.into_iter().map(Ok)); + all.push(Ok(message_delta_with_stop(messages::StopReason::MaxTokens))); + all.push(Ok(MessageStreamEvent::MessageStop)); + let raw = stream::iter(all).boxed(); + collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await +} + +/// Complete tool calls before a max_tokens stop complete the turn as ToolCalls. +#[tokio::test] +async fn max_tokens_after_complete_tool_calls_completes_as_tool_calls() { + let evs = max_tokens_after(vec![ + tool_use_start(0, "call_done"), + arg_delta(0, "{\"x\": 1}"), + block_stop(0), + ]) + .await; + match evs.last().unwrap() { + SamplingEvent::Completed { response, .. } => { + assert_eq!(response.stop_reason, Some(StopReason::ToolCalls)); + assert_eq!(response.tool_calls().len(), 1); + assert_eq!(response.tool_calls()[0].arguments.as_ref(), "{\"x\": 1}"); + } + other => panic!("expected Completed(ToolCalls), got {other:?}"), + } +} + +/// A tool_use block with no argument deltas is the zero-argument convention, not truncation. +#[tokio::test] +async fn max_tokens_after_tool_call_without_arguments_completes() { + let evs = max_tokens_after(vec![tool_use_start(0, "call_no_args"), block_stop(0)]).await; + match evs.last().unwrap() { + SamplingEvent::Completed { response, .. } => { + assert_eq!(response.stop_reason, Some(StopReason::ToolCalls)); + assert_eq!(response.tool_calls()[0].arguments.as_ref(), ""); + } + other => panic!("expected Completed(ToolCalls), got {other:?}"), + } +} + +/// A call cut mid-arguments is truncation: the turn fails so the caller can compact. +#[tokio::test] +async fn max_tokens_with_truncated_tool_arguments_fails() { + let evs = max_tokens_after(vec![ + tool_use_start(0, "call_ok"), + arg_delta(0, "{}"), + block_stop(0), + tool_use_start(1, "call_cut"), + arg_delta(1, "{\"x\": \"trunc"), + block_stop(1), + ]) + .await; + assert!( + !evs.iter() + .any(|e| matches!(e, SamplingEvent::Completed { .. })), + "a cut tool call must not complete: {evs:?}" + ); + match evs.last().unwrap() { + SamplingEvent::Failed { error, .. } => assert_eq!( + error.kind, + crate::events::SamplingErrorKind::MaxTokensTruncation + ), + other => panic!("expected Failed(MaxTokensTruncation), got {other:?}"), + } +}