diff --git a/crates/tokenizer/src/stop.rs b/crates/tokenizer/src/stop.rs index 3b0ed3e6fa..e285d93f7b 100644 --- a/crates/tokenizer/src/stop.rs +++ b/crates/tokenizer/src/stop.rs @@ -72,9 +72,6 @@ pub struct StopSequenceDecoder { visible_boundary_idx: usize, /// Buffer for partial matches (the "jail") jail_buffer: String, - /// Maximum bytes to retain in jail_buffer — equal to the longest stop sequence. - /// Text beyond this window cannot participate in a future match and is safe to drain. - jail_max_bytes: usize, /// Whether we've stopped stopped: bool, /// The string stop sequence that triggered the stop, if any. Set only for @@ -109,17 +106,6 @@ impl StopSequenceDecoder { .map(|s| s.as_str()), ); - // Precompute the maximum stop sequence length in bytes. - // The jail buffer only needs to retain this many bytes — any text older than - // this window cannot be part of a future match and is safe to emit. - let jail_max_bytes = config - .stop_sequences - .iter() - .chain(&config.visible_stop_sequences) - .map(|s| s.len()) - .max() - .unwrap_or(0); - let aho_corasick = if patterns.is_empty() { None } else { @@ -140,13 +126,52 @@ impl StopSequenceDecoder { aho_corasick, visible_boundary_idx, jail_buffer: String::new(), - jail_max_bytes, stopped: false, matched_stop: None, token_only, } } + /// Byte length of the longest suffix of the jail buffer that is a *proper* + /// prefix of some stop sequence. + /// + /// That suffix is the only part of the buffer that can still grow into a stop + /// sequence once more tokens arrive; everything before it can never take part + /// in a match. A full-length match is excluded because the Aho-Corasick scan + /// has already ruled one out by the time this runs. + fn pending_match_len(&self) -> usize { + let buf = self.jail_buffer.as_bytes(); + let Some(&last) = buf.last() else { + return 0; + }; + + let mut longest = 0; + for pattern in self + .config + .stop_sequences + .iter() + .chain(&self.config.visible_stop_sequences) + { + let pat = pattern.as_bytes(); + // Only proper prefixes count, and only ones longer than the best so far. + let max_n = pat.len().saturating_sub(1).min(buf.len()); + for n in (longest + 1..=max_n).rev() { + // Cheap necessary condition first: the prefix has to end on the + // byte the buffer ends on, which skips most of the comparisons. + if pat[n - 1] != last { + continue; + } + // A prefix ending mid-character would split a multi-byte codepoint + // in the buffer, so only consider character-aligned prefixes. + if pattern.is_char_boundary(n) && buf.ends_with(&pat[..n]) { + longest = n; + break; + } + } + } + longest + } + /// Process a single token pub fn process_token(&mut self, token_id: TokenIdType) -> Result { if self.stopped { @@ -190,27 +215,13 @@ impl StopSequenceDecoder { return Ok(SequenceDecoderOutput::Text(new_text)); } - let old_len = self.jail_buffer.len(); self.jail_buffer.push_str(&new_text); - // Check for stop sequences using Aho-Corasick, scoped to avoid rescanning - // old text: a match can start no earlier than `old_len - jail_max_bytes + 1` - // because any earlier match would have been found on a previous call. + // Check for stop sequences using Aho-Corasick. The jail only ever retains a + // proper prefix of some pattern, so the buffer is bounded by + // `longest_pattern + one token of text` and scanning it whole is cheap. if let Some(ac) = &self.aho_corasick { - let search_start = if old_len >= self.jail_max_bytes { - // Walk forward to a char boundary (we must not start mid-codepoint) - let raw = old_len + 1 - self.jail_max_bytes; - let mut start = raw; - while start < self.jail_buffer.len() && !self.jail_buffer.is_char_boundary(start) { - start += 1; - } - start - } else { - 0 - }; - - let input = Input::new(&self.jail_buffer).span(search_start..self.jail_buffer.len()); - if let Some(mat) = ac.find(input) { + if let Some(mat) = ac.find(Input::new(&self.jail_buffer)) { self.stopped = true; self.matched_stop = Some(self.jail_buffer[mat.start()..mat.end()].to_string()); let is_visible = mat.pattern().as_usize() >= self.visible_boundary_idx; @@ -233,25 +244,17 @@ impl StopSequenceDecoder { } } - // Drain the jail buffer down to at most jail_max_bytes, emitting safe text. - // Any text older than the window cannot be part of a future stop sequence match. - if self.jail_buffer.len() > self.jail_max_bytes { - // Find a char-safe drain point: we want to keep the last jail_max_bytes, - // but must not split a multi-byte UTF-8 character. - let mut drain_to = self.jail_buffer.len() - self.jail_max_bytes; - while drain_to > 0 && !self.jail_buffer.is_char_boundary(drain_to) { - // Move backward to retain at least jail_max_bytes (safe: retains more, not less) - drain_to -= 1; - } - - if drain_to > 0 { - let suffix = self.jail_buffer.split_off(drain_to); - let to_output = std::mem::replace(&mut self.jail_buffer, suffix); - return Ok(SequenceDecoderOutput::Text(to_output)); - } + // Withhold only the longest suffix that is still a *partial* stop sequence. + // Everything before it can never take part in a match, so emit it now + // rather than trailing the stream by the length of the longest stop word. + let drain_to = self.jail_buffer.len() - self.pending_match_len(); + if drain_to > 0 { + let suffix = self.jail_buffer.split_off(drain_to); + let to_output = std::mem::replace(&mut self.jail_buffer, suffix); + return Ok(SequenceDecoderOutput::Text(to_output)); } - // Buffer is within the window — hold everything for potential partial match + // The whole buffer is a partial stop sequence — hold it. Ok(SequenceDecoderOutput::Held) } @@ -502,19 +505,85 @@ mod tests { } #[test] - fn test_flush_after_partial() { + fn test_flush_returns_pending_partial_match() { + let tokenizer = Arc::new(MockTokenizer::new()); + // "Hello" is a proper prefix of the stop sequence, so it must be held + // back: the next token could complete the match. + let config = StopSequenceConfig::default().with_stop_sequence("Hello world"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + assert_eq!( + decoder.process_token(1).unwrap(), // "Hello" + SequenceDecoderOutput::Held, + "a proper prefix of the stop sequence must be withheld" + ); + + // The stream ended without completing the match, so flush releases it. + assert_eq!( + decoder.flush(), + SequenceDecoderOutput::Text("Hello".to_string()) + ); + assert_eq!(decoder.flush(), SequenceDecoderOutput::Held); + } + + #[test] + fn test_text_that_cannot_match_is_emitted_immediately() { let tokenizer = Arc::new(MockTokenizer::new()); let config = StopSequenceConfig::default().with_stop_sequence("NEVER_MATCH"); let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); - // Process a token - decoder.process_token(1).unwrap(); // "Hello" + // No suffix of "Hello" can grow into "NEVER_MATCH", so nothing is jailed + // and the text goes out with the token that produced it. + assert_eq!( + decoder.process_token(1).unwrap(), + SequenceDecoderOutput::Text("Hello".to_string()) + ); - // Flush should return any remaining text in jail - let result = decoder.flush(); + // Nothing is left behind for the caller to flush unparsed at end of stream. + assert_eq!(decoder.flush(), SequenceDecoderOutput::Held); + } - // After processing, flush should work - assert!(matches!(result, SequenceDecoderOutput::Text(_))); + #[test] + fn test_control_token_is_not_sliced_by_the_jail() { + let tokenizer = Arc::new(MockTokenizer::new()); + // A stop sequence long enough to span a control token, and sharing its + // "<|" opening. This is the shape that leaked: the jail used to retain + // the last `len(stop)` bytes unconditionally, so a control token landing + // on that boundary was cut in half and its tail escaped through `flush()`. + let config = StopSequenceConfig::default().with_stop_sequence("<|im_end|>"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // "<|im_start|>" cannot become "<|im_end|>" — no suffix of it is a prefix + // of the stop sequence — so it must be emitted whole and at once. + assert_eq!( + decoder.process_token(1001).unwrap(), + SequenceDecoderOutput::Text("<|im_start|>".to_string()) + ); + assert_eq!( + decoder.flush(), + SequenceDecoderOutput::Held, + "no fragment of the control token may be left for an unparsed flush" + ); + assert!(!decoder.is_stopped()); + } + + #[test] + fn test_partial_match_is_released_when_it_diverges() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_sequence("Hello world"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // "Hello" looks like the start of the stop sequence — hold it. + assert_eq!( + decoder.process_token(1).unwrap(), + SequenceDecoderOutput::Held + ); + + // "test" makes the match impossible; the held text must come back out + // in full rather than being dropped or trimmed. + let out = decoder.process_token(3).unwrap(); + assert_eq!(out, SequenceDecoderOutput::Text("Hello test".to_string())); + assert_eq!(decoder.flush(), SequenceDecoderOutput::Held); } #[test] @@ -663,12 +732,14 @@ mod tests { /// Tokens: 3 ("test"), 1 ("Hello"), 2 ("world") /// Stop sequence: "Hello world" (11 bytes) /// - /// With the bounded jail window, all text is held until the jail exceeds - /// jail_max_bytes (11). The jail accumulates: - /// - Token 3: jail = "test" (4 bytes ≤ 11) → Held - /// - Token 1: jail = "test Hello" (10 bytes ≤ 11) → Held - /// - Token 2: jail = "test Hello world" → Aho-Corasick matches "Hello world" - /// → StoppedWithText("test ") (text before the hidden stop sequence) + /// Only a suffix that is a proper prefix of the stop sequence is withheld, + /// so text flows out as soon as it can no longer be part of a match: + /// - Token 3: "test" cannot start "Hello world" → Text("test") + /// - Token 1: jail = " Hello", holds "Hello" → Text(" ") + /// - Token 2: jail = "Hello world" — Aho-Corasick matches → Stopped + /// + /// The caller sees "test " either way; the difference is that it arrives + /// with the tokens that produced it instead of trailing the stream. #[test] fn test_stop_sequence_spanning_tokens_with_preceding_text() { let tokenizer = Arc::new(MockTokenizer::new()); @@ -676,39 +747,38 @@ mod tests { let config = StopSequenceConfig::default().with_stop_sequence("Hello world"); let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); - // Token 3 ("test"): jail = "test" (4 bytes), within the 11-byte window → Held - let result1 = decoder.process_token(3).unwrap(); - assert!( - matches!(result1, SequenceDecoderOutput::Held), - "Expected Held for token within jail window, got {result1:?}" - ); + let mut emitted = String::new(); + let mut stopped = false; + for token in [3u32, 1, 2] { + match decoder.process_token(token).unwrap() { + SequenceDecoderOutput::Text(t) => emitted.push_str(&t), + SequenceDecoderOutput::StoppedWithText(t) => { + emitted.push_str(&t); + stopped = true; + } + SequenceDecoderOutput::Stopped => stopped = true, + SequenceDecoderOutput::Held => {} + } + } - // Token 1 ("Hello"): jail = "test Hello" (10 bytes), still within window → Held - let result2 = decoder.process_token(1).unwrap(); assert!( - matches!(result2, SequenceDecoderOutput::Held), - "Expected Held for token within jail window, got {result2:?}" + stopped, + "the stop sequence spanning tokens should have fired" ); - - // Token 2 ("world"): jail = "test Hello world" — Aho-Corasick matches - // "Hello world", so we stop. Text before the match ("test ") is emitted. - let result3 = decoder.process_token(2).unwrap(); - assert!( - matches!( - result3, - SequenceDecoderOutput::Stopped | SequenceDecoderOutput::StoppedWithText(_) - ), - "Expected Stopped or StoppedWithText when stop sequence completes, got {result3:?}" + assert_eq!( + emitted, "test ", + "everything before the hidden stop sequence is emitted, and nothing more" + ); + assert_eq!( + decoder.flush(), + SequenceDecoderOutput::Held, + "a completed stop must leave nothing jailed for an unparsed flush" ); assert!(decoder.is_stopped()); - - // Verify that any text before the stop sequence is preserved - if let SequenceDecoderOutput::StoppedWithText(text) = &result3 { - assert!( - !text.contains("Hello world"), - "Hidden stop sequence should not appear in output, got: {text:?}" - ); - } + assert!( + !emitted.contains("Hello world"), + "the hidden stop sequence must not appear in the output, got: {emitted:?}" + ); } #[test] diff --git a/e2e_test/chat_completions/test_reasoning_content.py b/e2e_test/chat_completions/test_reasoning_content.py index b57f27e174..370116785a 100644 --- a/e2e_test/chat_completions/test_reasoning_content.py +++ b/e2e_test/chat_completions/test_reasoning_content.py @@ -121,6 +121,53 @@ def test_streaming_separate_reasoning_true_stream_reasoning_false(self, model, a assert len(reasoning_content) > 0 assert len(content) > 0 + @pytest.mark.parametrize("stop", [["wtf"], ["wtfx"], ["0123456789"]]) + def test_streaming_unmatched_stop_word_does_not_change_output(self, model, api_client, stop): + """A `stop` word that never fires must not change what the client sees. + + The stop decoder withholds text that could still complete a stop + sequence. When generation is cut short — here by `max_tokens`, before + the model ever emits the stop word — whatever is still held is released + at end of stream. That release used to skip the reasoning parser, so the + held bytes surfaced as assistant `content`: a tail of the reasoning + text, or a fragment of the model's own structural tokens. The leak was + exactly as long as the stop word, which is what the parametrize covers. + """ + + def run(stop_words): + kwargs = {"stop": stop_words} if stop_words else {} + response = api_client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "What is 1+3?"}], + # Cut generation off early so the stream ends mid-reasoning, + # with text still held back by the stop decoder. + max_tokens=24, + temperature=0, + stream=True, + extra_body={"separate_reasoning": True}, + **kwargs, + ) + reasoning, content = "", "" + for chunk in response: + delta = chunk.choices[0].delta + if delta.content: + content += delta.content + if delta.reasoning_content: + reasoning += delta.reasoning_content + return reasoning, content + + baseline_reasoning, baseline_content = run(None) + stopped_reasoning, stopped_content = run(stop) + + assert stopped_content == baseline_content, ( + f"stop={stop!r} leaked {stopped_content[len(baseline_content) :]!r} " + "into content; text held by the stop decoder must still be parsed" + ) + assert stopped_reasoning == baseline_reasoning, ( + f"stop={stop!r} changed reasoning_content; a stop word that never " + "matches must not affect the reasoning/content split" + ) + def test_nonstreaming_separate_reasoning_false(self, model, api_client): """Test non-streaming with separate_reasoning=False, reasoning_content should be empty.""" diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index 43bd04eae8..7d210a47be 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -18,7 +18,8 @@ use llm_tokenizer::{ use openai_protocol::{ chat::{ChatCompletionRequest, ChatCompletionStreamResponse}, common::{ - FunctionCallDelta, StringOrArray, Tool, ToolCallDelta, ToolChoice, ToolChoiceValue, Usage, + ChatLogProbs, FunctionCallDelta, StringOrArray, Tool, ToolCallDelta, ToolChoice, + ToolChoiceValue, Usage, }, completion::{CompletionRequest, CompletionStreamChoice, CompletionStreamResponse}, generate::GenerateRequest, @@ -388,7 +389,12 @@ impl StreamingProcessor { while let Some(response) = grpc_stream.next().await { let gen_response = response.map_err(|e| format!("Stream error: {}", e.message()))?; - match gen_response.into_response() { + // Text the stop decoder produced for this response, if any. Per-chunk + // text and the end-of-stream flush both funnel into the shared emission + // below, so neither can reach the client without being parsed. + let pending: Option<(u32, String, Option)> = match gen_response + .into_response() + { ProtoResponseVariant::Chunk(chunk) => { // Track TTFT immediately on first chunk received from backend if first_token_time.is_none() { @@ -462,155 +468,21 @@ impl StreamingProcessor { utils::convert_proto_to_openai_logprobs(proto_logprobs, &tokenizer) }); - // Initialize stream buffer if first time - let stream_buffer = stream_buffers.entry(index).or_default(); - - // Send first chunk with role - if is_firsts.get(&index).copied().unwrap_or(true) { - let first_chunk = ChatCompletionStreamResponse::builder(request_id, model) - .created(created) - .add_choice_role(index, "assistant") - .maybe_system_fingerprint(system_fingerprint) - .build(); - Self::format_sse_chunk_into(&mut sse_buffer, &first_chunk); - tx.send(Ok(Bytes::from(sse_buffer.clone()))) - .await - .map_err(|_| "Failed to send first chunk".to_string())?; - is_firsts.insert(index, false); - } - - // Calculate delta - let mut delta = chunk_text; - stream_buffer.push_str(&delta); - - // Reasoning content handling - let in_reasoning = if separate_reasoning && reasoning_parser_available { - let (normal_text, reasoning_chunk, in_reasoning) = self - .process_reasoning_stream( - &delta, - index, - &mut reasoning_parsers, - thinking_override, - think_in_prefill, - reasoning_parser_name.as_deref(), - request_id, - model, - created, - system_fingerprint, - ) - .await; - if let Some(chunk) = reasoning_chunk { - Self::format_sse_chunk_into(&mut sse_buffer, &chunk); - tx.send(Ok(Bytes::from(sse_buffer.clone()))) - .await - .map_err(|_| "Failed to send reasoning chunk".to_string())?; - } - delta = normal_text; - in_reasoning - } else { - false - }; - - // Tool call handling - let tool_choice_enabled = - !matches!(tool_choice, Some(ToolChoice::Value(ToolChoiceValue::None))); - - if let Some(tools_ref) = tools.as_ref() { - if !in_reasoning - && tool_choice_enabled - && (tool_parser_available || used_json_schema) - { - let tool_chunks = if is_specific_function { - // Handle specific function case - emit tool call deltas with arguments - Self::process_specific_function_stream( - &delta, - index, - &mut has_tool_calls, - tool_choice.as_ref(), - request_id, - model, - created, - system_fingerprint, - history_tool_calls_count, - ) - } else { - // Use incremental parser for regular/required modes - self.process_tool_calls_stream( - &delta, - index, - &mut tool_parsers, - &mut has_tool_calls, - tools_ref, - tool_parser_name.as_deref(), - request_id, - model, - created, - system_fingerprint, - history_tool_calls_count, - used_json_schema, - ) - .await - }; - - for chunk in tool_chunks { - Self::format_sse_chunk_into(&mut sse_buffer, &chunk); - tx.send(Ok(Bytes::from(sse_buffer.clone()))) - .await - .map_err(|_| "Failed to send tool call chunk".to_string())?; - } - - // Always skip regular content when tool parsing is active - // Parser either emitted chunks or buffered content - continue; - } - } - - // Regular content emission - if !delta.is_empty() { - let content_chunk = - ChatCompletionStreamResponse::builder(request_id, model) - .created(created) - .add_choice_content_with_logprobs( - index, - "assistant", - delta, - choice_logprobs, - ) - .maybe_system_fingerprint(system_fingerprint) - .build(); - Self::format_sse_chunk_into(&mut sse_buffer, &content_chunk); - tx.send(Ok(Bytes::from(sse_buffer.clone()))) - .await - .map_err(|_| "Failed to send content chunk".to_string())?; - } + Some((index, chunk_text, choice_logprobs)) } ProtoResponseVariant::Complete(complete) => { let index = complete.index(); - // Flush any remaining text for this index's stop_decoder - if let Some(decoder) = stop_decoders.get_mut(&index) { - if let SequenceDecoderOutput::Text(text) = decoder.flush() { - if !text.is_empty() { - let stream_buffer = stream_buffers.entry(index).or_default(); - stream_buffer.push_str(&text); - - let content_chunk = - ChatCompletionStreamResponse::builder(request_id, model) - .created(created) - .add_choice_content(index, "assistant", 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".to_string())?; - } - } - } + // Release whatever the stop decoder still holds. It only ever + // retains a partial stop-sequence match, and it is routed through + // the same parsers as every other chunk rather than straight out. + let flushed = + stop_decoders + .get_mut(&index) + .and_then(|decoder| match decoder.flush() { + SequenceDecoderOutput::Text(text) if !text.is_empty() => Some(text), + _ => None, + }); // Store metadata prompt_tokens.insert(index, complete.prompt_tokens()); @@ -628,8 +500,129 @@ impl StreamingProcessor { } // Don't break - continue reading all Complete messages for n>1 + flushed.map(|text| (index, text, None)) } ProtoResponseVariant::None => continue, + }; + + let Some((index, text, choice_logprobs)) = pending else { + continue; + }; + + // Initialize stream buffer if first time + let stream_buffer = stream_buffers.entry(index).or_default(); + + // Send first chunk with role + if is_firsts.get(&index).copied().unwrap_or(true) { + let first_chunk = ChatCompletionStreamResponse::builder(request_id, model) + .created(created) + .add_choice_role(index, "assistant") + .maybe_system_fingerprint(system_fingerprint) + .build(); + Self::format_sse_chunk_into(&mut sse_buffer, &first_chunk); + tx.send(Ok(Bytes::from(sse_buffer.clone()))) + .await + .map_err(|_| "Failed to send first chunk".to_string())?; + is_firsts.insert(index, false); + } + + // Calculate delta + let mut delta = text; + stream_buffer.push_str(&delta); + + // Reasoning content handling + let in_reasoning = if separate_reasoning && reasoning_parser_available { + let (normal_text, reasoning_chunk, in_reasoning) = self + .process_reasoning_stream( + &delta, + index, + &mut reasoning_parsers, + thinking_override, + think_in_prefill, + reasoning_parser_name.as_deref(), + request_id, + model, + created, + system_fingerprint, + ) + .await; + if let Some(chunk) = reasoning_chunk { + Self::format_sse_chunk_into(&mut sse_buffer, &chunk); + tx.send(Ok(Bytes::from(sse_buffer.clone()))) + .await + .map_err(|_| "Failed to send reasoning chunk".to_string())?; + } + delta = normal_text; + in_reasoning + } else { + false + }; + + // Tool call handling + let tool_choice_enabled = + !matches!(tool_choice, Some(ToolChoice::Value(ToolChoiceValue::None))); + + if let Some(tools_ref) = tools.as_ref() { + if !in_reasoning + && tool_choice_enabled + && (tool_parser_available || used_json_schema) + { + let tool_chunks = if is_specific_function { + // Handle specific function case - emit tool call deltas with arguments + Self::process_specific_function_stream( + &delta, + index, + &mut has_tool_calls, + tool_choice.as_ref(), + request_id, + model, + created, + system_fingerprint, + history_tool_calls_count, + ) + } else { + // Use incremental parser for regular/required modes + self.process_tool_calls_stream( + &delta, + index, + &mut tool_parsers, + &mut has_tool_calls, + tools_ref, + tool_parser_name.as_deref(), + request_id, + model, + created, + system_fingerprint, + history_tool_calls_count, + used_json_schema, + ) + .await + }; + + for chunk in tool_chunks { + Self::format_sse_chunk_into(&mut sse_buffer, &chunk); + tx.send(Ok(Bytes::from(sse_buffer.clone()))) + .await + .map_err(|_| "Failed to send tool call chunk".to_string())?; + } + + // Always skip regular content when tool parsing is active + // Parser either emitted chunks or buffered content + continue; + } + } + + // Regular content emission + if !delta.is_empty() { + let content_chunk = ChatCompletionStreamResponse::builder(request_id, model) + .created(created) + .add_choice_content_with_logprobs(index, "assistant", delta, choice_logprobs) + .maybe_system_fingerprint(system_fingerprint) + .build(); + Self::format_sse_chunk_into(&mut sse_buffer, &content_chunk); + tx.send(Ok(Bytes::from(sse_buffer.clone()))) + .await + .map_err(|_| "Failed to send content chunk".to_string())?; } } @@ -2031,7 +2024,10 @@ impl StreamingProcessor { while let Some(response) = grpc_stream.next().await { let gen_response = response.map_err(|e| format!("Stream error: {}", e.message()))?; - match gen_response.into_response() { + // Text the stop decoder produced for this response, if any. Per-chunk + // text and the end-of-stream flush both funnel into the shared emission + // below, so neither can reach the client without being parsed. + let pending: Option = match gen_response.into_response() { ProtoResponseVariant::Chunk(chunk) => { if first_token_time.is_none() { first_token_time = Some(Instant::now()); @@ -2064,307 +2060,292 @@ impl StreamingProcessor { continue; } - // Apply reasoning parser - let (normal_text, reasoning_chunk_text, in_reasoning) = - if reasoning_parser_available { - self.process_messages_reasoning( - &chunk_text, - &mut reasoning_parser, - thinking_override, - think_in_prefill, - reasoning_parser_name.as_deref(), - model, - ) - .await - } else { - (chunk_text, String::new(), false) - }; + Some(chunk_text) + } + ProtoResponseVariant::Complete(complete) => { + // Release whatever the stop decoder still holds. It only ever + // retains a partial stop-sequence match, and it is routed through + // the same parsers as every other chunk rather than straight out. + let flushed = match stop_decoder.flush() { + SequenceDecoderOutput::Text(text) if !text.is_empty() => Some(text), + _ => None, + }; + + prompt_tokens = complete.prompt_tokens(); + saw_complete = true; + completion_tokens.record_complete(&complete); + // A local stop-decoder match already pinned "stop"; don't let + // the engine's finish reason overwrite it. + if !stopped { + finish_reason_str = complete.finish_reason().to_string(); + matched_stop = complete.matched_stop_json(); + } + flushed + } + ProtoResponseVariant::None => continue, + }; + + let Some(chunk_text) = pending else { + continue; + }; + + // Apply reasoning parser + let (normal_text, reasoning_chunk_text, in_reasoning) = if reasoning_parser_available { + self.process_messages_reasoning( + &chunk_text, + &mut reasoning_parser, + thinking_override, + think_in_prefill, + reasoning_parser_name.as_deref(), + model, + ) + .await + } else { + (chunk_text, String::new(), false) + }; + + // Emit thinking content block deltas + if !reasoning_chunk_text.is_empty() { + if !thinking_block_open { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStart { + index: current_block_index, + content_block: ContentBlock::Thinking { + thinking: String::new(), + signature: String::new(), + }, + }, + ) + .await?; + thinking_block_open = true; + } + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockDelta { + index: current_block_index, + delta: ContentBlockDelta::ThinkingDelta { + thinking: reasoning_chunk_text, + }, + }, + ) + .await?; + } - // Emit thinking content block deltas - if !reasoning_chunk_text.is_empty() { - if !thinking_block_open { + // Transition: reasoning ended, close thinking block + if thinking_block_open && !in_reasoning && !normal_text.is_empty() { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStop { + index: current_block_index, + }, + ) + .await?; + thinking_block_open = false; + current_block_index += 1; + } + + // Tool call handling: incremental streaming parser + if !in_reasoning && streaming_tool_parser.is_some() { + if is_specific_function { + // Specific function: entire output is arguments for one tool + if !has_tool_calls { + has_tool_calls = true; + // Close text block if open before starting tool block + if text_block_open { Self::send_messages_event( tx, &mut sse_buffer, - &MessageStreamEvent::ContentBlockStart { + &MessageStreamEvent::ContentBlockStop { index: current_block_index, - content_block: ContentBlock::Thinking { - thinking: String::new(), - signature: String::new(), - }, }, ) .await?; - thinking_block_open = true; + text_block_open = false; + current_block_index += 1; } + // Emit content_block_start for the tool_use + let tool_name = match &original_request.tool_choice { + Some(messages::ToolChoice::Tool { name, .. }) => name.clone(), + _ => String::new(), + }; + let tool_call_id = utils::generate_tool_call_id( + model, + &tool_name, + 0, + history_tool_calls_count, + ); Self::send_messages_event( tx, &mut sse_buffer, - &MessageStreamEvent::ContentBlockDelta { + &MessageStreamEvent::ContentBlockStart { index: current_block_index, - delta: ContentBlockDelta::ThinkingDelta { - thinking: reasoning_chunk_text, + content_block: ContentBlock::ToolUse { + id: message_utils::anthropic_tool_use_id(&tool_call_id), + name: tool_name, + input: Value::Object(serde_json::Map::new()), }, }, ) .await?; + tool_block_open = true; } - - // Transition: reasoning ended, close thinking block - if thinking_block_open && !in_reasoning && !normal_text.is_empty() { + // Emit arguments delta + if !normal_text.is_empty() { Self::send_messages_event( tx, &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { + &MessageStreamEvent::ContentBlockDelta { index: current_block_index, + delta: ContentBlockDelta::InputJsonDelta { + partial_json: normal_text, + }, }, ) .await?; - thinking_block_open = false; - current_block_index += 1; } - - // Tool call handling: incremental streaming parser - if !in_reasoning && streaming_tool_parser.is_some() { - if is_specific_function { - // Specific function: entire output is arguments for one tool - if !has_tool_calls { - has_tool_calls = true; - // Close text block if open before starting tool block - if text_block_open { + } else if let Some(ref mut parser) = streaming_tool_parser { + // Regular/required tool choice: use incremental parser + match parser.parse_incremental(&normal_text, &chat_tools).await { + Ok(StreamingParseResult { + normal_text: text, + calls, + }) => { + // Emit normal text from parser as text content blocks + if !text.is_empty() { + if !text_block_open { Self::send_messages_event( tx, &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { + &MessageStreamEvent::ContentBlockStart { index: current_block_index, + content_block: ContentBlock::Text { + text: String::new(), + citations: None, + }, }, ) .await?; - text_block_open = false; - current_block_index += 1; + text_block_open = true; } - // Emit content_block_start for the tool_use - let tool_name = match &original_request.tool_choice { - Some(messages::ToolChoice::Tool { name, .. }) => name.clone(), - _ => String::new(), - }; - let tool_call_id = utils::generate_tool_call_id( - model, - &tool_name, - 0, - history_tool_calls_count, - ); - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStart { - index: current_block_index, - content_block: ContentBlock::ToolUse { - id: message_utils::anthropic_tool_use_id(&tool_call_id), - name: tool_name, - input: Value::Object(serde_json::Map::new()), - }, - }, - ) - .await?; - tool_block_open = true; - } - // Emit arguments delta - if !normal_text.is_empty() { Self::send_messages_event( tx, &mut sse_buffer, &MessageStreamEvent::ContentBlockDelta { index: current_block_index, - delta: ContentBlockDelta::InputJsonDelta { - partial_json: normal_text, - }, + delta: ContentBlockDelta::TextDelta { text }, }, ) .await?; } - } else if let Some(ref mut parser) = streaming_tool_parser { - // Regular/required tool choice: use incremental parser - match parser.parse_incremental(&normal_text, &chat_tools).await { - Ok(StreamingParseResult { - normal_text: text, - calls, - }) => { - // Emit normal text from parser as text content blocks - if !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; - } + + // Emit tool call events + for tool_call_item in calls { + has_tool_calls = true; + + if let Some(ref name) = tool_call_item.name { + // New tool call: close previous blocks, emit start + if text_block_open { Self::send_messages_event( tx, &mut sse_buffer, - &MessageStreamEvent::ContentBlockDelta { + &MessageStreamEvent::ContentBlockStop { index: current_block_index, - delta: ContentBlockDelta::TextDelta { text }, }, ) .await?; + text_block_open = false; + current_block_index += 1; } - - // Emit tool call events - for tool_call_item in calls { - has_tool_calls = true; - - if let Some(ref name) = tool_call_item.name { - // New tool call: close previous blocks, emit start - if text_block_open { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { - index: current_block_index, - }, - ) - .await?; - text_block_open = false; - current_block_index += 1; - } - if tool_block_open { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStop { - index: current_block_index, - }, - ) - .await?; - current_block_index += 1; - } - - let tool_call_id = utils::generate_tool_call_id( - model, - name, - tool_call_item.tool_index, - history_tool_calls_count, - ); - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockStart { - index: current_block_index, - content_block: ContentBlock::ToolUse { - id: message_utils::anthropic_tool_use_id( - &tool_call_id, - ), - name: name.clone(), - input: Value::Object(serde_json::Map::new()), - }, - }, - ).await?; - tool_block_open = true; - } - - // Emit incremental arguments - if !tool_call_item.parameters.is_empty() { - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockDelta { - index: current_block_index, - delta: ContentBlockDelta::InputJsonDelta { - partial_json: tool_call_item.parameters, - }, - }, - ) - .await?; - } + if tool_block_open { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStop { + index: current_block_index, + }, + ) + .await?; + current_block_index += 1; } + + let tool_call_id = utils::generate_tool_call_id( + model, + name, + tool_call_item.tool_index, + history_tool_calls_count, + ); + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStart { + index: current_block_index, + content_block: ContentBlock::ToolUse { + id: message_utils::anthropic_tool_use_id( + &tool_call_id, + ), + name: name.clone(), + input: Value::Object(serde_json::Map::new()), + }, + }, + ) + .await?; + tool_block_open = true; } - Err(e) => { - error!("Tool call parsing error in messages streaming: {}", e); + + // Emit incremental arguments + if !tool_call_item.parameters.is_empty() { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockDelta { + index: current_block_index, + delta: ContentBlockDelta::InputJsonDelta { + partial_json: tool_call_item.parameters, + }, + }, + ) + .await?; } } } - continue; - } - - // Regular text emission (no tools active) - if !normal_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; + Err(e) => { + error!("Tool call parsing error in messages streaming: {}", e); } - Self::send_messages_event( - tx, - &mut sse_buffer, - &MessageStreamEvent::ContentBlockDelta { - index: current_block_index, - delta: ContentBlockDelta::TextDelta { text: normal_text }, - }, - ) - .await?; } } - ProtoResponseVariant::Complete(complete) => { - // Flush stop decoder - if let SequenceDecoderOutput::Text(text) = stop_decoder.flush() { - if !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 }, - }, - ) - .await?; - } - } + continue; + } - prompt_tokens = complete.prompt_tokens(); - saw_complete = true; - completion_tokens.record_complete(&complete); - // A local stop-decoder match already pinned "stop"; don't let - // the engine's finish reason overwrite it. - if !stopped { - finish_reason_str = complete.finish_reason().to_string(); - matched_stop = complete.matched_stop_json(); - } + // Regular text emission (no tools active) + if !normal_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; } - ProtoResponseVariant::None => continue, + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockDelta { + index: current_block_index, + delta: ContentBlockDelta::TextDelta { text: normal_text }, + }, + ) + .await?; } }