From 751eb2b0256985eadf515f132d26069dbae8fc5b Mon Sep 17 00:00:00 2001 From: yifeng liu <31553858+pallasathena92@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:40:13 -0700 Subject: [PATCH 1/2] fix(tool_parser): flush unconsumed streaming buffer as content instead of swallowing it Streaming tool parsing could swallow an entire completion: any text the incremental parser buffered as a prospective tool call that never became a valid declared tool call was silently dropped, producing a stream with no tool_call deltas and no content deltas, while the non-streaming path correctly returned the same text as content. Observed deterministically on PR #2261's e2e runs (all four engines): Llama-3.2-1B with --tool-call-parser llama, tools present, tool_choice=auto, streaming emitted a {-prefixed JSON answer that LlamaParser::has_tool_markers() routed into handle_json_tool_streaming(), where it buffered forever. Two-level fix: - handle_json_tool_streaming() now bails out to normal_text when the buffered JSON is definitively not a tool call: a complete JSON value with no name field, or a (by construction complete) name that is not among the declared tools. Previously the former buffered forever and the latter cleared the buffer silently, losing the text. - New ToolParser::take_unstreamed_normal_text() (default: empty) drains text still buffered at end of stream - truncated tool JSON, partial start markers - so routers can emit it as content. Implemented via a shared helper for the five parsers that use the shared JSON streaming helper (llama, json, mistral, qwen, cohere); parsers that announced a tool call from the buffer return nothing (the remaining arguments are recovered via get_unstreamed_tool_args, matching the non-streaming path which drops trailing text once tool calls were extracted). Wired into both stream-finish sites: Chat Completions (shared by single and PD mode, gRPC and ZMQ transports) and Messages API. Also: QwenParser::reset() now clears normal_text_buffer, and the end-of-stream flush drains it (a partial suffix held there is real text when no tool call was ever announced). Legit buffering is preserved: partial bot_token suffixes and incomplete-but-potentially-valid tool JSON keep buffering mid-stream; new regression tests cover chunk boundaries inside JSON, undeclared names, truncated JSON at end of stream, and the announced-tool guard. Signed-off-by: yifeng liu <31553858+pallasathena92@users.noreply.github.com> --- crates/tool_parser/src/parsers/cohere.rs | 6 + crates/tool_parser/src/parsers/helpers.rs | 52 ++- crates/tool_parser/src/parsers/json.rs | 4 + crates/tool_parser/src/parsers/llama.rs | 4 + crates/tool_parser/src/parsers/mistral.rs | 4 + crates/tool_parser/src/parsers/qwen.rs | 14 + crates/tool_parser/src/traits.rs | 17 + .../tests/tool_parser_streaming_flush.rs | 384 ++++++++++++++++++ .../src/routers/grpc/regular/streaming.rs | 61 ++- 9 files changed, 540 insertions(+), 6 deletions(-) create mode 100644 crates/tool_parser/tests/tool_parser_streaming_flush.rs diff --git a/crates/tool_parser/src/parsers/cohere.rs b/crates/tool_parser/src/parsers/cohere.rs index e56095085..40afc343d 100644 --- a/crates/tool_parser/src/parsers/cohere.rs +++ b/crates/tool_parser/src/parsers/cohere.rs @@ -327,6 +327,12 @@ impl ToolParser for CohereParser { helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool) } + fn take_unstreamed_normal_text(&mut self) -> String { + // Covers both a partial START_ACTION held in Text state and an action + // block whose END_ACTION never arrived (truncated stream). + helpers::take_unstreamed_normal_text(&mut self.buffer, self.current_tool_id) + } + fn reset(&mut self) { self.state = ParseState::Text; helpers::reset_parser_state( diff --git a/crates/tool_parser/src/parsers/helpers.rs b/crates/tool_parser/src/parsers/helpers.rs index 9f232c666..519600e77 100644 --- a/crates/tool_parser/src/parsers/helpers.rs +++ b/crates/tool_parser/src/parsers/helpers.rs @@ -124,6 +124,27 @@ pub fn get_unstreamed_args( }]) } +/// End-of-stream flush: take any text still buffered as a *prospective* tool +/// call so the caller can emit it as normal content instead of dropping it. +/// +/// Returns the buffer verbatim when no tool call was ever announced this +/// request (`current_tool_id == -1`): the buffered text was only ever a +/// tool-call *candidate* (e.g. a bare `{` prefix or a partial start marker) +/// that never materialized, so it is content. Once a tool call has been +/// announced (`current_tool_id >= 0`), the buffered tail is tool-call syntax +/// (separators, closing brackets, or a partially-streamed call whose +/// remaining arguments are recovered via [`get_unstreamed_args`]) and is +/// discarded — matching the non-streaming path, which drops trailing +/// non-JSON text once tool calls were extracted. +pub fn take_unstreamed_normal_text(buffer: &mut String, current_tool_id: i32) -> String { + let text = std::mem::take(buffer); + if current_tool_id == -1 { + text + } else { + String::new() + } +} + /// Check if a buffer ends with a partial occurrence of a token /// Returns Some(length) if there's a partial match, None otherwise pub fn ends_with_partial_token(buffer: &str, token: &str) -> Option { @@ -284,16 +305,41 @@ pub(crate) fn handle_json_tool_streaming( // Validate tool name if present if let Some(name) = current_tool_call.get("name").and_then(|v| v.as_str()) { if !tool_indices.contains_key(name) { - // Invalid tool name - skip this tool, preserve indexing for next tool - tracing::debug!("Invalid tool name '{}' - skipping", name); + // The name string is complete (partial strings are disallowed + // until the name has been sent), so this JSON can never become a + // declared tool call. Surface the buffered text as normal content + // instead of dropping it: silently clearing the buffer here is + // how streams ended up empty while the non-streaming path + // returned the same text as content. Any remainder of the JSON + // still arriving flows through as normal text on later chunks. + tracing::debug!( + "Undeclared tool name '{}' - emitting buffered text as content", + name + ); + let normal_text = current_text.to_string(); reset_current_tool_state( buffer, current_tool_name_sent, streamed_args_for_tool, prev_tool_call_arr, ); - return Ok(StreamingParseResult::default()); + return Ok(StreamingParseResult { + normal_text, + calls: vec![], + }); } + } else if is_complete { + // A complete JSON value with no tool name is definitively not a tool + // call. Emit the consumed text as normal content instead of buffering + // it forever (which swallowed the whole stream), keeping any tail for + // further parsing. + let consumed = start_idx + safe_end_idx; + let normal_text = current_text[..consumed].to_string(); + *buffer = current_text[consumed..].to_string(); + return Ok(StreamingParseResult { + normal_text, + calls: vec![], + }); } let mut result = StreamingParseResult::default(); diff --git a/crates/tool_parser/src/parsers/json.rs b/crates/tool_parser/src/parsers/json.rs index 51f53b1df..b1d05e477 100644 --- a/crates/tool_parser/src/parsers/json.rs +++ b/crates/tool_parser/src/parsers/json.rs @@ -294,6 +294,10 @@ impl ToolParser for JsonParser { helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool) } + fn take_unstreamed_normal_text(&mut self) -> String { + helpers::take_unstreamed_normal_text(&mut self.buffer, self.current_tool_id) + } + fn reset(&mut self) { helpers::reset_parser_state( &mut self.buffer, diff --git a/crates/tool_parser/src/parsers/llama.rs b/crates/tool_parser/src/parsers/llama.rs index 7bdadb897..717a4bfaa 100644 --- a/crates/tool_parser/src/parsers/llama.rs +++ b/crates/tool_parser/src/parsers/llama.rs @@ -231,6 +231,10 @@ impl ToolParser for LlamaParser { helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool) } + fn take_unstreamed_normal_text(&mut self) -> String { + helpers::take_unstreamed_normal_text(&mut self.buffer, self.current_tool_id) + } + fn reset(&mut self) { helpers::reset_parser_state( &mut self.buffer, diff --git a/crates/tool_parser/src/parsers/mistral.rs b/crates/tool_parser/src/parsers/mistral.rs index 1d74b9220..1c2e6a257 100644 --- a/crates/tool_parser/src/parsers/mistral.rs +++ b/crates/tool_parser/src/parsers/mistral.rs @@ -315,6 +315,10 @@ impl ToolParser for MistralParser { helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool) } + fn take_unstreamed_normal_text(&mut self) -> String { + helpers::take_unstreamed_normal_text(&mut self.buffer, self.current_tool_id) + } + fn reset(&mut self) { helpers::reset_parser_state( &mut self.buffer, diff --git a/crates/tool_parser/src/parsers/qwen.rs b/crates/tool_parser/src/parsers/qwen.rs index 2c8381afb..e34d30b60 100644 --- a/crates/tool_parser/src/parsers/qwen.rs +++ b/crates/tool_parser/src/parsers/qwen.rs @@ -251,7 +251,21 @@ impl ToolParser for QwenParser { helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool) } + fn take_unstreamed_normal_text(&mut self) -> String { + // `normal_text_buffer` may be holding back a suffix that looked like a + // partial "" tag; at end of stream it is real text when no + // tool call was ever announced, and marker debris otherwise. + let held = std::mem::take(&mut self.normal_text_buffer); + let tail = helpers::take_unstreamed_normal_text(&mut self.buffer, self.current_tool_id); + if self.current_tool_id == -1 { + held + &tail + } else { + String::new() + } + } + fn reset(&mut self) { + self.normal_text_buffer.clear(); helpers::reset_parser_state( &mut self.buffer, &mut self.prev_tool_call_arr, diff --git a/crates/tool_parser/src/traits.rs b/crates/tool_parser/src/traits.rs index 30ed8e8e0..a0c751951 100644 --- a/crates/tool_parser/src/traits.rs +++ b/crates/tool_parser/src/traits.rs @@ -45,6 +45,23 @@ pub trait ToolParser: Send + Sync { None } + /// Take any text still buffered by the streaming parser that never became + /// a tool call, transferring ownership to the caller. + /// + /// Streaming consumers call this once at end of stream, alongside + /// [`Self::get_unstreamed_tool_args`]: text held back as a *prospective* + /// tool call (a bare `{` prefix, a partial start marker, or tool JSON + /// that never completed) must be surfaced as normal content instead of + /// being silently dropped — mirroring the non-streaming fallback that + /// returns unparseable tool text verbatim. Parsers that announced a tool + /// call from the buffered text return an empty string (the remaining + /// arguments are recovered via `get_unstreamed_tool_args`). + /// + /// The default returns an empty string (parser holds no buffered text). + fn take_unstreamed_normal_text(&mut self) -> String { + String::new() + } + /// Reset the parser state for reuse across requests. /// This should clear all buffers and reset state to initial values. fn reset(&mut self) { diff --git a/crates/tool_parser/tests/tool_parser_streaming_flush.rs b/crates/tool_parser/tests/tool_parser_streaming_flush.rs new file mode 100644 index 000000000..10fb09d92 --- /dev/null +++ b/crates/tool_parser/tests/tool_parser_streaming_flush.rs @@ -0,0 +1,384 @@ +//! Streaming buffered-content flush tests +//! +//! Regression tests for the streaming content-loss bug observed on PR #2261's +//! e2e runs (all four engines, deterministic): Llama-3.2-1B with +//! `--tool-call-parser llama`, tools present, `tool_choice=auto`, streaming. +//! When the model emits `{`-prefixed text that never becomes a valid declared +//! tool call, the incremental parser buffered everything while waiting for a +//! tool call that never materialized and the client received a completely +//! empty stream — no tool_call deltas AND no content deltas — while the +//! non-streaming path correctly fell back to returning the text as content. +//! +//! Two properties are covered here: +//! 1. Text that is definitively not a declared tool call (complete JSON with +//! a missing or undeclared name) is surfaced as normal text mid-stream +//! instead of being buffered forever or silently dropped. +//! 2. Text still buffered at end of stream (truncated JSON, partial start +//! markers) is recoverable via the end-of-stream flush instead of being +//! swallowed. +mod common; + +use common::{create_test_tools, streaming_helpers}; +use tool_parser::{ + types::ToolCallItem, CohereParser, JsonParser, LlamaParser, MistralParser, QwenParser, + ToolParser, +}; + +/// Representative Llama-3.2-1B output for "What's the weather in Tokyo?" with +/// declared tools: `{`-prefixed JSON that never becomes a valid declared tool +/// call because the function name is under a non-standard key (the parser +/// only recognizes `name`). The CI worker-log artifacts do not record raw +/// generations, so this reproduces the failure class rather than the literal +/// string. +const NON_TOOL_JSON: &str = + r#"{"type": "function", "function": "get_weather", "parameters": {"city": "Tokyo"}}"#; + +/// Valid, complete JSON in llama tool shape, but the name is not among the +/// declared tools. +const UNDECLARED_TOOL_JSON: &str = + r#"{"name": "get_stock_price", "parameters": {"ticker": "AAPL"}}"#; + +/// Drive a parser through chunked streaming and collect everything it emits. +#[expect( + clippy::unwrap_used, + reason = "test helper; allow-unwrap-in-tests only covers #[test] fns" +)] +async fn stream_chunks( + parser: &mut dyn ToolParser, + chunks: &[&str], +) -> (String, Vec) { + let tools = create_test_tools(); + let mut normal_text = String::new(); + let mut calls = Vec::new(); + for chunk in chunks { + let result = parser.parse_incremental(chunk, &tools).await.unwrap(); + normal_text.push_str(&result.normal_text); + calls.extend(result.calls); + } + (normal_text, calls) +} + +// ============================================================================ +// Property 1: definitively-not-a-tool JSON must surface as normal text +// ============================================================================ + +#[tokio::test] +async fn test_llama_streaming_non_tool_json_not_swallowed() { + let mut parser = LlamaParser::new(); + // Realistic 2-3 char chunks: many chunk boundaries inside the JSON. + let chunks = streaming_helpers::create_realistic_chunks(NON_TOOL_JSON); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + + let (normal_text, calls) = stream_chunks(&mut parser, &chunk_refs).await; + + assert!( + calls.is_empty(), + "no declared tool call was made: {calls:?}" + ); + assert_eq!( + normal_text, NON_TOOL_JSON, + "non-tool JSON text must be streamed as content, not swallowed" + ); +} + +#[tokio::test] +async fn test_llama_streaming_undeclared_tool_name_surfaces_as_text() { + let mut parser = LlamaParser::new(); + // Chunk boundaries inside the JSON, including inside the (undeclared) name. + let chunks = [ + r#"{"name": "get_st"#, + r#"ock_price", "para"#, + r#"meters": {"ticker": "AAPL"}}"#, + ]; + + let (normal_text, calls) = stream_chunks(&mut parser, &chunks).await; + + assert!( + calls.is_empty(), + "undeclared tool must not be emitted: {calls:?}" + ); + assert_eq!( + normal_text, UNDECLARED_TOOL_JSON, + "undeclared-tool JSON must be surfaced as content, not dropped" + ); +} + +#[tokio::test] +async fn test_json_streaming_non_tool_json_not_swallowed() { + let mut parser = JsonParser::new(); + let input = r#"{"result": 42, "status": "ok"}"#; + let chunks = streaming_helpers::create_realistic_chunks(input); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + + let (normal_text, calls) = stream_chunks(&mut parser, &chunk_refs).await; + + assert!(calls.is_empty()); + assert_eq!( + normal_text, input, + "non-tool JSON text must be streamed as content, not swallowed" + ); +} + +#[tokio::test] +async fn test_mistral_streaming_undeclared_tool_name_not_dropped() { + let mut parser = MistralParser::new(); + let input = r#"[TOOL_CALLS] [{"name": "frobnicate", "arguments": {"x": 1}}]"#; + let chunks = [ + "[TOOL_CALLS] ", + r#"[{"name": "frobni"#, + r#"cate", "arguments"#, + r#"": {"x": 1}}]"#, + ]; + + let (normal_text, calls) = stream_chunks(&mut parser, &chunks).await; + + assert!( + calls.is_empty(), + "undeclared tool must not be emitted: {calls:?}" + ); + // Content-preserving: everything the model emitted must reach the client + // as text (the non-streaming path never drops it silently either). + let compact: String = normal_text.chars().filter(|c| !c.is_whitespace()).collect(); + let expected: String = input.chars().filter(|c| !c.is_whitespace()).collect(); + assert_eq!(compact, expected, "undeclared-tool text was dropped"); +} + +#[tokio::test] +async fn test_qwen_streaming_non_tool_json_in_markers_not_swallowed() { + let mut parser = QwenParser::new(); + // Complete JSON inside qwen markers but with no "name" field at all. + let chunks = ["\n", r#"{"foo": "#, "1}", "\n"]; + + let (normal_text, calls) = stream_chunks(&mut parser, &chunks).await; + + assert!(calls.is_empty()); + assert!( + normal_text.contains(r#"{"foo": 1}"#), + "non-tool JSON inside markers must be surfaced as content, got {normal_text:?}" + ); +} + +// ============================================================================ +// Property 1b: legitimate tool calls must be unaffected +// ============================================================================ + +#[tokio::test] +async fn test_llama_streaming_valid_tool_call_still_works() { + let mut parser = LlamaParser::new(); + let chunks = [ + r#"{"name": "get_we"#, + r#"ather", "parameters"#, + r#"": {"city": "Tokyo"}}"#, + ]; + + let (normal_text, calls) = stream_chunks(&mut parser, &chunks).await; + + assert_eq!(normal_text, "", "no content expected for a valid tool call"); + assert!( + calls + .iter() + .any(|c| c.name.as_deref() == Some("get_weather")), + "declared tool call must still be announced: {calls:?}" + ); + let args: String = calls.iter().map(|c| c.parameters.as_str()).collect(); + assert!( + args.contains("Tokyo"), + "arguments must be streamed: {args:?}" + ); +} + +#[tokio::test] +async fn test_llama_streaming_partial_declared_name_keeps_buffering() { + let mut parser = LlamaParser::new(); + let tools = create_test_tools(); + + // "get_we" is a prefix of the declared "get_weather": the parser must NOT + // bail out to normal text while the name string is still open. + let result = parser + .parse_incremental(r#"{"name": "get_we"#, &tools) + .await + .unwrap(); + assert_eq!(result.normal_text, ""); + assert!(result.calls.is_empty()); + + let result = parser + .parse_incremental(r#"ather", "parameters": {"city": "Paris"}}"#, &tools) + .await + .unwrap(); + let mut calls = result.calls; + calls.extend(parser.parse_incremental("", &tools).await.unwrap().calls); + assert!( + calls + .iter() + .any(|c| c.name.as_deref() == Some("get_weather")), + "declared tool call must be announced after the name completes: {calls:?}" + ); +} + +// ============================================================================ +// Property 2: end-of-stream flush recovers buffered text +// ============================================================================ + +#[tokio::test] +async fn test_llama_truncated_json_flushed_at_end_of_stream() { + let mut parser = LlamaParser::new(); + let tools = create_test_tools(); + + // Stream ends mid-name: the parser must hold it while streaming (it could + // still become "get_weather"), but surface it at end of stream. + let truncated = r#"{"name": "get_wea"#; + let result = parser.parse_incremental(truncated, &tools).await.unwrap(); + assert_eq!(result.normal_text, ""); + assert!(result.calls.is_empty()); + + assert_eq!( + parser.take_unstreamed_normal_text(), + truncated, + "text buffered at end of stream must be flushed as content" + ); + // Drained: a second flush returns nothing. + assert_eq!(parser.take_unstreamed_normal_text(), ""); +} + +#[tokio::test] +async fn test_llama_partial_bot_token_flushed_at_end_of_stream() { + let mut parser = LlamaParser::new(); + let tools = create_test_tools(); + + let text = "Sure, let me check <|py"; + let result = parser.parse_incremental(text, &tools).await.unwrap(); + assert_eq!(result.normal_text, ""); + assert!(result.calls.is_empty()); + + assert_eq!(parser.take_unstreamed_normal_text(), text); +} + +#[tokio::test] +async fn test_llama_flush_empty_after_completed_tool_call() { + let mut parser = LlamaParser::new(); + let tools = create_test_tools(); + + let result = parser + .parse_incremental( + r#"{"name": "get_weather", "parameters": {"city": "Tokyo"}}"#, + &tools, + ) + .await + .unwrap(); + let mut calls = result.calls; + calls.extend(parser.parse_incremental("", &tools).await.unwrap().calls); + assert!(calls + .iter() + .any(|c| c.name.as_deref() == Some("get_weather"))); + + assert_eq!( + parser.take_unstreamed_normal_text(), + "", + "no content must be invented after a real tool call" + ); +} + +#[tokio::test] +async fn test_llama_flush_empty_after_announced_tool_truncated_args() { + let mut parser = LlamaParser::new(); + let tools = create_test_tools(); + + // Name announced, arguments truncated at end of stream: the buffered tail + // is tool syntax, not content — flushing it as text would duplicate the + // tool call. Remaining args are recovered via get_unstreamed_tool_args. + let result = parser + .parse_incremental( + r#"{"name": "get_weather", "parameters": {"city": "Par"#, + &tools, + ) + .await + .unwrap(); + assert!(result + .calls + .iter() + .any(|c| c.name.as_deref() == Some("get_weather"))); + + assert_eq!( + parser.take_unstreamed_normal_text(), + "", + "announced tool call tail must not be re-emitted as content" + ); +} + +#[tokio::test] +async fn test_json_truncated_json_flushed_at_end_of_stream() { + let mut parser = JsonParser::new(); + let tools = create_test_tools(); + + let truncated = r#"{"name": "sea"#; + let result = parser.parse_incremental(truncated, &tools).await.unwrap(); + assert_eq!(result.normal_text, ""); + assert!(result.calls.is_empty()); + + assert_eq!(parser.take_unstreamed_normal_text(), truncated); +} + +#[tokio::test] +async fn test_mistral_truncated_tool_call_flushed_at_end_of_stream() { + let mut parser = MistralParser::new(); + let tools = create_test_tools(); + + let truncated = r#"[TOOL_CALLS] [{"unknown"#; + let result = parser.parse_incremental(truncated, &tools).await.unwrap(); + assert_eq!(result.normal_text, ""); + assert!(result.calls.is_empty()); + + assert_eq!(parser.take_unstreamed_normal_text(), truncated); +} + +#[tokio::test] +async fn test_qwen_truncated_tool_call_flushed_at_end_of_stream() { + let mut parser = QwenParser::new(); + let tools = create_test_tools(); + + let truncated = "\n{\"a\": 1"; + let result = parser.parse_incremental(truncated, &tools).await.unwrap(); + assert_eq!(result.normal_text, ""); + assert!(result.calls.is_empty()); + + assert_eq!(parser.take_unstreamed_normal_text(), truncated); +} + +#[tokio::test] +async fn test_cohere_unterminated_action_flushed_at_end_of_stream() { + let mut parser = CohereParser::new(); + let tools = create_test_tools(); + + // START_ACTION seen, END_ACTION never arrives (truncated stream). + let json_part = r#"{"tool_name": "search", "parameters": {"query": "x"#; + let result = parser + .parse_incremental(&format!("<|START_ACTION|>{json_part}"), &tools) + .await + .unwrap(); + assert_eq!(result.normal_text, ""); + assert!(result.calls.is_empty()); + + assert_eq!( + parser.take_unstreamed_normal_text(), + json_part, + "unterminated action block must be flushed as content" + ); +} + +#[tokio::test] +async fn test_flush_empty_after_reset() { + let mut parser = LlamaParser::new(); + let tools = create_test_tools(); + + parser + .parse_incremental(r#"{"name": "get_wea"#, &tools) + .await + .unwrap(); + parser.reset(); + + assert_eq!( + parser.take_unstreamed_normal_text(), + "", + "reset must clear the streaming buffer" + ); +} diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index c08d78a17..55a79eefc 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -626,9 +626,29 @@ impl StreamingProcessor { } } - // Phase 3: Check unstreamed tool args + // Phase 3: End-of-stream parser flush: first any text still buffered + // as a prospective tool call that never materialized (dropping it + // produced fully-empty streams), then any parsed-but-unstreamed tool + // arguments. for (index, parser) in &tool_parsers { - let parser_guard = parser.lock().await; + let mut parser_guard = parser.lock().await; + + let leftover_text = parser_guard.take_unstreamed_normal_text(); + if !leftover_text.is_empty() { + let content_chunk = ChatCompletionStreamResponse::builder(request_id, model) + .created(created) + .add_choice_content(*index, "assistant", leftover_text) + .maybe_system_fingerprint(system_fingerprint) + .build(); + + let sse_chunk = sse_encoder + .encode_data(&content_chunk) + .map_err(|e| format!("Failed to serialize content chunk: {e}"))?; + tx.send(Ok(sse_chunk)) + .await + .map_err(|_| "Failed to send flushed content chunk".to_string())?; + } + if let Some(unstreamed_items) = parser_guard.get_unstreamed_tool_args() { for tool_call_item in unstreamed_items { let tool_call_delta = ToolCallDelta { @@ -2348,7 +2368,42 @@ impl StreamingProcessor { } } - // Phase 3: Flush unstreamed tool args from the incremental parser + // Phase 3: End-of-stream parser flush: first any text still buffered + // as a prospective tool call that never materialized (dropping it + // produced fully-empty streams), then any parsed-but-unstreamed tool + // arguments. + if let Some(ref mut parser) = streaming_tool_parser { + let leftover_text = parser.take_unstreamed_normal_text(); + if !leftover_text.is_empty() { + if !text_block_open { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStart { + index: current_block_index, + content_block: ContentBlock::Text { + text: String::new(), + citations: None, + }, + }, + ) + .await?; + text_block_open = true; + } + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockDelta { + index: current_block_index, + delta: ContentBlockDelta::TextDelta { + text: leftover_text, + }, + }, + ) + .await?; + } + } + if let Some(ref parser) = streaming_tool_parser { if let Some(unstreamed_items) = parser.get_unstreamed_tool_args() { for tool_call_item in unstreamed_items { From 047c9a4383b03ba2dc248e503e5ed4a547c32420 Mon Sep 17 00:00:00 2001 From: yifeng liu <31553858+pallasathena92@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:55:59 -0700 Subject: [PATCH 2/2] fix(tool_parser): drain adjacent JSON values; bundle streaming state A declared tool call trailing a non-tool JSON value in the same (possibly final) chunk is now re-parsed into tool-call deltas instead of stranding for the end-of-stream text flush. The threaded parser state moves into JsonToolStreamState, dropping both too_many_arguments suppressions and making call sites transposition-safe. Drain runs only on the cold bail-out paths with one bounded buffer copy; recursion is bounded by the number of adjacent values in a chunk. Also fixes the codespell typo and marks intentional mid-word chunk-boundary test data. Signed-off-by: yifeng liu <31553858+pallasathena92@users.noreply.github.com> --- crates/tool_parser/src/parsers/cohere.rs | 16 +- crates/tool_parser/src/parsers/helpers.rs | 157 ++++++++++++------ crates/tool_parser/src/parsers/json.rs | 16 +- crates/tool_parser/src/parsers/llama.rs | 16 +- crates/tool_parser/src/parsers/mistral.rs | 16 +- crates/tool_parser/src/parsers/qwen.rs | 16 +- crates/tool_parser/src/traits.rs | 2 +- .../tests/tool_parser_streaming_flush.rs | 59 ++++++- 8 files changed, 210 insertions(+), 88 deletions(-) diff --git a/crates/tool_parser/src/parsers/cohere.rs b/crates/tool_parser/src/parsers/cohere.rs index 40afc343d..96b3ed3d6 100644 --- a/crates/tool_parser/src/parsers/cohere.rs +++ b/crates/tool_parser/src/parsers/cohere.rs @@ -295,13 +295,15 @@ impl ToolParser for CohereParser { let result = helpers::handle_json_tool_streaming( &json_content, 0, - &mut self.partial_json, - &tool_indices, - &mut temp_buffer, - &mut self.current_tool_id, - &mut self.current_tool_name_sent, - &mut self.streamed_args_for_tool, - &mut self.prev_tool_call_arr, + &mut helpers::JsonToolStreamState { + partial_json: &mut self.partial_json, + tool_indices: &tool_indices, + buffer: &mut temp_buffer, + current_tool_id: &mut self.current_tool_id, + current_tool_name_sent: &mut self.current_tool_name_sent, + streamed_args_for_tool: &mut self.streamed_args_for_tool, + prev_tool_call_arr: &mut self.prev_tool_call_arr, + }, )?; // Move past END_ACTION and switch back to Text state diff --git a/crates/tool_parser/src/parsers/helpers.rs b/crates/tool_parser/src/parsers/helpers.rs index 519600e77..a346d39d6 100644 --- a/crates/tool_parser/src/parsers/helpers.rs +++ b/crates/tool_parser/src/parsers/helpers.rs @@ -145,6 +145,56 @@ pub fn take_unstreamed_normal_text(buffer: &mut String, current_tool_id: i32) -> } } +/// The mutable parser state `handle_json_tool_streaming` threads through the +/// JSON-tool streaming flow. Grouping it keeps call sites transposition-safe +/// and lets the drain path recurse without re-listing nine arguments. +pub(crate) struct JsonToolStreamState<'a> { + pub partial_json: &'a mut crate::partial_json::PartialJson, + pub tool_indices: &'a HashMap, + pub buffer: &'a mut String, + pub current_tool_id: &'a mut i32, + pub current_tool_name_sent: &'a mut bool, + pub streamed_args_for_tool: &'a mut Vec, + pub prev_tool_call_arr: &'a mut Vec, +} + +/// After a non-tool JSON value was emitted as content, re-parse any adjacent +/// JSON value left in the buffer instead of stranding it: a declared tool +/// call trailing the emitted value in the same (possibly final) chunk must +/// become tool-call deltas, not an end-of-stream text flush. Separator +/// characters between adjacent values join the emitted text; a +/// non-JSON-looking tail (e.g. a partial marker) stays buffered for later +/// chunks, as before. Runs only on the cold bail-out paths; recursion is +/// bounded by the number of adjacent complete values in one chunk. +fn drain_adjacent_values( + mut normal_text: String, + state: &mut JsonToolStreamState<'_>, +) -> ParserResult { + let json_start = state + .buffer + .char_indices() + .find(|(_, c)| !c.is_whitespace() && *c != ',' && *c != ';') + .map(|(i, _)| i) + .unwrap_or(state.buffer.len()); + if !state.buffer[json_start..].starts_with('{') && !state.buffer[json_start..].starts_with('[') + { + return Ok(StreamingParseResult { + normal_text, + calls: vec![], + }); + } + // Separators between adjacent values belong to the emitted text. + normal_text.push_str(&state.buffer[..json_start]); + let remainder = state.buffer.split_off(json_start); + *state.buffer = remainder; + // One bounded copy: the flow reads `current_text` while mutating the + // buffer, so they cannot alias. + let text = state.buffer.clone(); + let mut follow = handle_json_tool_streaming(&text, 0, state)?; + follow.normal_text = format!("{normal_text}{}", follow.normal_text); + Ok(follow) +} + /// Check if a buffer ends with a partial occurrence of a token /// Returns Some(length) if there's a partial match, None otherwise pub fn ends_with_partial_token(buffer: &str, token: &str) -> Option { @@ -253,17 +303,10 @@ pub fn normalize_tool_call_fields(obj: Value) -> Value { /// name then arguments, advance the buffer) for the JSON, Llama, Mistral, and Qwen /// parsers. `start_idx` is where JSON begins in `current_text`; `current_tool_id == /// -1` means no active tool. -#[expect(clippy::too_many_arguments)] pub(crate) fn handle_json_tool_streaming( current_text: &str, start_idx: usize, - partial_json: &mut crate::partial_json::PartialJson, - tool_indices: &HashMap, - buffer: &mut String, - current_tool_id: &mut i32, - current_tool_name_sent: &mut bool, - streamed_args_for_tool: &mut Vec, - prev_tool_call_arr: &mut Vec, + state: &mut JsonToolStreamState<'_>, ) -> ParserResult { // Check if we have content to parse if start_idx >= current_text.len() { @@ -273,12 +316,15 @@ pub(crate) fn handle_json_tool_streaming( // Extract JSON string from current position let json_str = ¤t_text[start_idx..]; - // When current_tool_name_sent is false, don't allow partial strings to avoid + // When state.current_tool_name_sent is false, don't allow partial strings to avoid // parsing incomplete tool names as empty strings - let allow_partial_strings = *current_tool_name_sent; + let allow_partial_strings = *state.current_tool_name_sent; // Parse partial JSON - let (obj, end_idx) = match partial_json.parse_value(json_str, allow_partial_strings) { + let (obj, end_idx) = match state + .partial_json + .parse_value(json_str, allow_partial_strings) + { Ok(result) => result, Err(_) => { return Ok(StreamingParseResult::default()); @@ -304,11 +350,11 @@ pub(crate) fn handle_json_tool_streaming( // Validate tool name if present if let Some(name) = current_tool_call.get("name").and_then(|v| v.as_str()) { - if !tool_indices.contains_key(name) { + if !state.tool_indices.contains_key(name) { // The name string is complete (partial strings are disallowed // until the name has been sent), so this JSON can never become a // declared tool call. Surface the buffered text as normal content - // instead of dropping it: silently clearing the buffer here is + // instead of dropping it: silently clearing the state.buffer here is // how streams ended up empty while the non-streaming path // returned the same text as content. Any remainder of the JSON // still arriving flows through as normal text on later chunks. @@ -316,17 +362,22 @@ pub(crate) fn handle_json_tool_streaming( "Undeclared tool name '{}' - emitting buffered text as content", name ); - let normal_text = current_text.to_string(); + // Emit the undeclared call (with any marker prefix) as content; + // the tail may hold a declared call and gets drained below. + let consumed = if is_complete { + start_idx + safe_end_idx + } else { + current_text.len() + }; + let normal_text = current_text[..consumed].to_string(); reset_current_tool_state( - buffer, - current_tool_name_sent, - streamed_args_for_tool, - prev_tool_call_arr, + state.buffer, + state.current_tool_name_sent, + state.streamed_args_for_tool, + state.prev_tool_call_arr, ); - return Ok(StreamingParseResult { - normal_text, - calls: vec![], - }); + *state.buffer = current_text[consumed..].to_string(); + return drain_adjacent_values(normal_text, state); } } else if is_complete { // A complete JSON value with no tool name is definitively not a tool @@ -335,32 +386,33 @@ pub(crate) fn handle_json_tool_streaming( // further parsing. let consumed = start_idx + safe_end_idx; let normal_text = current_text[..consumed].to_string(); - *buffer = current_text[consumed..].to_string(); - return Ok(StreamingParseResult { - normal_text, - calls: vec![], - }); + *state.buffer = current_text[consumed..].to_string(); + return drain_adjacent_values(normal_text, state); } let mut result = StreamingParseResult::default(); // Case 1: Handle tool name streaming - if !*current_tool_name_sent { + if !*state.current_tool_name_sent { if let Some(function_name) = current_tool_call.get("name").and_then(|v| v.as_str()) { - if tool_indices.contains_key(function_name) { + if state.tool_indices.contains_key(function_name) { // Initialize if first tool - if *current_tool_id == -1 { - *current_tool_id = 0; - streamed_args_for_tool.push(String::new()); - } else if *current_tool_id as usize >= streamed_args_for_tool.len() { + if *state.current_tool_id == -1 { + *state.current_tool_id = 0; + state.streamed_args_for_tool.push(String::new()); + } else if *state.current_tool_id as usize >= state.streamed_args_for_tool.len() { // Ensure capacity for subsequent tools - ensure_capacity(*current_tool_id, prev_tool_call_arr, streamed_args_for_tool); + ensure_capacity( + *state.current_tool_id, + state.prev_tool_call_arr, + state.streamed_args_for_tool, + ); } // Send tool name with empty parameters - *current_tool_name_sent = true; + *state.current_tool_name_sent = true; result.calls.push(ToolCallItem { - tool_index: *current_tool_id as usize, + tool_index: *state.current_tool_id as usize, name: Some(function_name.to_string()), parameters: String::new(), }); @@ -369,8 +421,9 @@ pub(crate) fn handle_json_tool_streaming( } // Case 2: Handle streaming arguments else if let Some(cur_arguments) = current_tool_call.get("arguments") { - let tool_id = *current_tool_id as usize; - let sent = streamed_args_for_tool + let tool_id = *state.current_tool_id as usize; + let sent = state + .streamed_args_for_tool .get(tool_id) .map(|s| s.len()) .unwrap_or(0); @@ -378,8 +431,8 @@ pub(crate) fn handle_json_tool_streaming( .map_err(|e| ParserError::ParsingFailed(e.to_string()))?; // Get prev_arguments (matches Python's structure) - let prev_arguments = if tool_id < prev_tool_call_arr.len() { - prev_tool_call_arr[tool_id].get("arguments") + let prev_arguments = if tool_id < state.prev_tool_call_arr.len() { + state.prev_tool_call_arr[tool_id].get("arguments") } else { None }; @@ -412,8 +465,8 @@ pub(crate) fn handle_json_tool_streaming( // Send diff if present if let Some(diff) = argument_diff { if !diff.is_empty() { - if tool_id < streamed_args_for_tool.len() { - streamed_args_for_tool[tool_id].push_str(&diff); + if tool_id < state.streamed_args_for_tool.len() { + state.streamed_args_for_tool[tool_id].push_str(&diff); } result.calls.push(ToolCallItem { tool_index: tool_id, @@ -423,20 +476,24 @@ pub(crate) fn handle_json_tool_streaming( } } - // Update prev_tool_call_arr with current state - if *current_tool_id >= 0 { - ensure_capacity(*current_tool_id, prev_tool_call_arr, streamed_args_for_tool); + // Update state.prev_tool_call_arr with current state + if *state.current_tool_id >= 0 { + ensure_capacity( + *state.current_tool_id, + state.prev_tool_call_arr, + state.streamed_args_for_tool, + ); - if tool_id < prev_tool_call_arr.len() { - prev_tool_call_arr[tool_id] = current_tool_call; + if tool_id < state.prev_tool_call_arr.len() { + state.prev_tool_call_arr[tool_id] = current_tool_call; } } // If complete, advance to next tool if is_complete { - *buffer = current_text[start_idx + end_idx..].to_string(); - *current_tool_name_sent = false; - *current_tool_id += 1; + *state.buffer = current_text[start_idx + end_idx..].to_string(); + *state.current_tool_name_sent = false; + *state.current_tool_id += 1; } } diff --git a/crates/tool_parser/src/parsers/json.rs b/crates/tool_parser/src/parsers/json.rs index b1d05e477..78da5a297 100644 --- a/crates/tool_parser/src/parsers/json.rs +++ b/crates/tool_parser/src/parsers/json.rs @@ -275,13 +275,15 @@ impl ToolParser for JsonParser { helpers::handle_json_tool_streaming( current_text, start_idx, - &mut self.partial_json, - &tool_indices, - &mut self.buffer, - &mut self.current_tool_id, - &mut self.current_tool_name_sent, - &mut self.streamed_args_for_tool, - &mut self.prev_tool_call_arr, + &mut helpers::JsonToolStreamState { + partial_json: &mut self.partial_json, + tool_indices: &tool_indices, + buffer: &mut self.buffer, + current_tool_id: &mut self.current_tool_id, + current_tool_name_sent: &mut self.current_tool_name_sent, + streamed_args_for_tool: &mut self.streamed_args_for_tool, + prev_tool_call_arr: &mut self.prev_tool_call_arr, + }, ) } diff --git a/crates/tool_parser/src/parsers/llama.rs b/crates/tool_parser/src/parsers/llama.rs index 717a4bfaa..24af77c1b 100644 --- a/crates/tool_parser/src/parsers/llama.rs +++ b/crates/tool_parser/src/parsers/llama.rs @@ -212,13 +212,15 @@ impl ToolParser for LlamaParser { helpers::handle_json_tool_streaming( current_text, start_idx, - &mut self.partial_json, - &tool_indices, - &mut self.buffer, - &mut self.current_tool_id, - &mut self.current_tool_name_sent, - &mut self.streamed_args_for_tool, - &mut self.prev_tool_call_arr, + &mut helpers::JsonToolStreamState { + partial_json: &mut self.partial_json, + tool_indices: &tool_indices, + buffer: &mut self.buffer, + current_tool_id: &mut self.current_tool_id, + current_tool_name_sent: &mut self.current_tool_name_sent, + streamed_args_for_tool: &mut self.streamed_args_for_tool, + prev_tool_call_arr: &mut self.prev_tool_call_arr, + }, ) } diff --git a/crates/tool_parser/src/parsers/mistral.rs b/crates/tool_parser/src/parsers/mistral.rs index 1c2e6a257..c35806196 100644 --- a/crates/tool_parser/src/parsers/mistral.rs +++ b/crates/tool_parser/src/parsers/mistral.rs @@ -297,13 +297,15 @@ impl ToolParser for MistralParser { helpers::handle_json_tool_streaming( current_text, start_idx, - &mut self.partial_json, - &tool_indices, - &mut self.buffer, - &mut self.current_tool_id, - &mut self.current_tool_name_sent, - &mut self.streamed_args_for_tool, - &mut self.prev_tool_call_arr, + &mut helpers::JsonToolStreamState { + partial_json: &mut self.partial_json, + tool_indices: &tool_indices, + buffer: &mut self.buffer, + current_tool_id: &mut self.current_tool_id, + current_tool_name_sent: &mut self.current_tool_name_sent, + streamed_args_for_tool: &mut self.streamed_args_for_tool, + prev_tool_call_arr: &mut self.prev_tool_call_arr, + }, ) } diff --git a/crates/tool_parser/src/parsers/qwen.rs b/crates/tool_parser/src/parsers/qwen.rs index e34d30b60..902bf13bd 100644 --- a/crates/tool_parser/src/parsers/qwen.rs +++ b/crates/tool_parser/src/parsers/qwen.rs @@ -199,13 +199,15 @@ impl ToolParser for QwenParser { let mut result = helpers::handle_json_tool_streaming( current_text, start_idx, - &mut self.partial_json, - &tool_indices, - &mut self.buffer, - &mut self.current_tool_id, - &mut self.current_tool_name_sent, - &mut self.streamed_args_for_tool, - &mut self.prev_tool_call_arr, + &mut helpers::JsonToolStreamState { + partial_json: &mut self.partial_json, + tool_indices: &tool_indices, + buffer: &mut self.buffer, + current_tool_id: &mut self.current_tool_id, + current_tool_name_sent: &mut self.current_tool_name_sent, + streamed_args_for_tool: &mut self.streamed_args_for_tool, + prev_tool_call_arr: &mut self.prev_tool_call_arr, + }, )?; // Qwen-specific: Handle partial end tokens in normal text diff --git a/crates/tool_parser/src/traits.rs b/crates/tool_parser/src/traits.rs index a0c751951..e11ef4def 100644 --- a/crates/tool_parser/src/traits.rs +++ b/crates/tool_parser/src/traits.rs @@ -53,7 +53,7 @@ pub trait ToolParser: Send + Sync { /// tool call (a bare `{` prefix, a partial start marker, or tool JSON /// that never completed) must be surfaced as normal content instead of /// being silently dropped — mirroring the non-streaming fallback that - /// returns unparseable tool text verbatim. Parsers that announced a tool + /// returns unparsable tool text verbatim. Parsers that announced a tool /// call from the buffered text return an empty string (the remaining /// arguments are recovered via `get_unstreamed_tool_args`). /// diff --git a/crates/tool_parser/tests/tool_parser_streaming_flush.rs b/crates/tool_parser/tests/tool_parser_streaming_flush.rs index 10fb09d92..e2dd06aab 100644 --- a/crates/tool_parser/tests/tool_parser_streaming_flush.rs +++ b/crates/tool_parser/tests/tool_parser_streaming_flush.rs @@ -167,7 +167,7 @@ async fn test_llama_streaming_valid_tool_call_still_works() { let mut parser = LlamaParser::new(); let chunks = [ r#"{"name": "get_we"#, - r#"ather", "parameters"#, + r#"ather", "parameters"#, // codespell:ignore ather r#"": {"city": "Tokyo"}}"#, ]; @@ -202,7 +202,7 @@ async fn test_llama_streaming_partial_declared_name_keeps_buffering() { assert!(result.calls.is_empty()); let result = parser - .parse_incremental(r#"ather", "parameters": {"city": "Paris"}}"#, &tools) + .parse_incremental(r#"ather", "parameters": {"city": "Paris"}}"#, &tools) // codespell:ignore ather .await .unwrap(); let mut calls = result.calls; @@ -382,3 +382,58 @@ async fn test_flush_empty_after_reset() { "reset must clear the streaming buffer" ); } + +// ============================================================================ +// Property 4: adjacent values in one chunk — a declared tool call trailing a +// non-tool JSON value must become tool-call deltas, never stranded text +// ============================================================================ + +#[tokio::test] +async fn test_json_adjacent_non_tool_then_declared_call_in_final_chunk() { + let mut parser = JsonParser::new(); + let (normal_text, calls) = stream_chunks( + &mut parser, + &[r#"{"note": "checking"} {"name": "get_weather", "arguments": {"city": "Tokyo"}}"#], + ) + .await; + + assert!( + normal_text.contains(r#"{"note": "checking"}"#), + "leading non-tool JSON must surface as content, got {normal_text:?}" + ); + assert!( + calls + .iter() + .any(|c| c.name.as_deref() == Some("get_weather")), + "trailing declared call must be announced as a tool call, got {calls:?}" + ); + assert_eq!( + parser.take_unstreamed_normal_text(), + "", + "nothing may be stranded for the end-of-stream flush" + ); +} + +#[tokio::test] +async fn test_json_undeclared_then_declared_call_in_final_chunk() { + let mut parser = JsonParser::new(); + let (normal_text, calls) = stream_chunks( + &mut parser, + &[concat!( + r#"{"name": "bogus_tool", "arguments": {}}"#, + r#"{"name": "get_weather", "arguments": {"city": "Paris"}}"# + )], + ) + .await; + + assert!( + normal_text.contains("bogus_tool"), + "undeclared call must surface as content, got {normal_text:?}" + ); + assert!( + calls + .iter() + .any(|c| c.name.as_deref() == Some("get_weather")), + "declared call after an undeclared one must still parse, got {calls:?}" + ); +}