diff --git a/crates/tool_parser/src/factory.rs b/crates/tool_parser/src/factory.rs index 1c52c685d3..79d4488be0 100644 --- a/crates/tool_parser/src/factory.rs +++ b/crates/tool_parser/src/factory.rs @@ -126,7 +126,7 @@ impl ParserRegistry { /// Map a model name/pattern to a parser pub fn map_model(&self, model: impl Into, parser: impl Into) { let mut mapping = self.model_mapping.write(); - mapping.insert(model.into(), parser.into()); + mapping.insert(model.into().to_ascii_lowercase(), parser.into()); } /// Get a pooled parser by exact name. @@ -236,23 +236,33 @@ impl ParserRegistry { /// Returns None if no mapping matches (caller decides fallback). pub fn resolve_model_to_parser(&self, model: &str) -> Option { let mapping = self.model_mapping.read(); - // Try exact match - if let Some(parser_name) = mapping.get(model) { - return Some(parser_name.clone()); - } - // Case-insensitive substring matching (longest pattern wins) so namespaced - // and differently-cased ids (e.g. "org/Inkling-Chat") still resolve. - let model_lower = model.to_lowercase(); - mapping - .iter() - .filter_map(|(pattern, parser_name)| { - let stem = pattern.strip_suffix('*')?; - model_lower - .contains(&stem.to_lowercase()) - .then_some((stem, parser_name)) - }) - .max_by_key(|(stem, _)| stem.len()) - .map(|(_, parser_name)| parser_name.clone()) + // Case-insensitive resolution (longest pattern wins) so namespaced and + // differently-cased ids (e.g. "org/Inkling-Chat", "THUDM/GLM-5.2") still + // resolve. Match against both the full normalized id and the basename + // (last path segment) using prefix semantics — prefix (not substring) + // matching keeps resolution fail-closed and avoids false positives from + // a stem appearing mid-segment. + let normalized = model.to_ascii_lowercase(); + let basename = normalized + .rsplit(['/', '\\']) + .find(|part| !part.is_empty()) + .unwrap_or(normalized.as_str()); + + let resolved = std::iter::once(normalized.as_str()) + .chain((basename != normalized.as_str()).then_some(basename)) + .find_map(|candidate| { + mapping.get(candidate).cloned().or_else(|| { + mapping + .iter() + .filter(|(pattern, _)| { + pattern.ends_with('*') + && candidate.starts_with(&pattern[..pattern.len() - 1]) + }) + .max_by_key(|(pattern, _)| pattern.len()) + .map(|(_, parser_name)| parser_name.clone()) + }) + }); + resolved } /// Check if a parser can be created for a specific model without actually creating it. @@ -329,6 +339,8 @@ impl ParserFactory { registry.register_parser("deepseek_v4", || Box::new(DeepSeekDsmlParser::v4())); registry.register_parser("glm45_moe", || Box::new(Glm4MoeParser::glm45())); registry.register_parser("glm47_moe", || Box::new(Glm4MoeParser::glm47())); + registry.register_parser("glm45", || Box::new(Glm4MoeParser::glm45())); + registry.register_parser("glm47", || Box::new(Glm4MoeParser::glm47())); registry.register_parser("step3", || Box::new(Step3Parser::new())); registry.register_parser("sarashina", || Box::new(SarashinaParser::new())); registry.register_parser_with_structural_tag( diff --git a/crates/tool_parser/src/parsers/glm4_moe.rs b/crates/tool_parser/src/parsers/glm4_moe.rs index de86b3854f..dd3fbc7dfa 100644 --- a/crates/tool_parser/src/parsers/glm4_moe.rs +++ b/crates/tool_parser/src/parsers/glm4_moe.rs @@ -97,21 +97,48 @@ impl Glm4MoeParser { &self, args_text: &str, param_types: &HashMap, - ) -> serde_json::Map { + ) -> ParserResult> { let mut arguments = serde_json::Map::new(); + let mut consumed = 0; for capture in self.arg_extractor.captures_iter(args_text) { + let full_match = capture.get(0).ok_or_else(|| { + ParserError::ParsingFailed("malformed GLM argument block".to_string()) + })?; + if !args_text[consumed..full_match.start()].trim().is_empty() { + return Err(ParserError::ParsingFailed( + "malformed GLM argument tags".to_string(), + )); + } + let key = capture.get(1).map_or("", |m| m.as_str()).trim(); - let value_str = capture.get(2).map_or("", |m| m.as_str()).trim(); + if key.is_empty() { + return Err(ParserError::ParsingFailed( + "GLM argument key must not be empty".to_string(), + )); + } + let raw_value = capture.get(2).map_or("", |m| m.as_str()); + let declared_type = param_types.get(key).map(String::as_str); + let value_str = if declared_type == Some("string") { + raw_value + } else { + raw_value.trim() + }; - let value = - helpers::coerce_by_schema_type(value_str, param_types.get(key).map(String::as_str)) - .unwrap_or_else(|| infer_value(value_str)); + let value = helpers::coerce_by_schema_type(value_str, declared_type) + .unwrap_or_else(|| infer_value(value_str)); arguments.insert(key.to_string(), value); + consumed = full_match.end(); } - arguments + if !args_text[consumed..].trim().is_empty() { + return Err(ParserError::ParsingFailed( + "malformed GLM argument tags".to_string(), + )); + } + + Ok(arguments) } /// Parse a single tool call block @@ -124,7 +151,7 @@ impl Glm4MoeParser { let args_text = captures.get(2).map_or("", |m| m.as_str()); let param_types = helpers::param_types_for_function(tools, func_name); - let arguments = self.parse_arguments(args_text, ¶m_types); + let arguments = self.parse_arguments(args_text, ¶m_types)?; let arguments_str = serde_json::to_string(&arguments) .map_err(|e| ParserError::ParsingFailed(e.to_string()))?; @@ -141,34 +168,240 @@ impl Glm4MoeParser { } /// Parse all tool calls from text (shared logic for complete and incremental parsing) - fn parse_tool_calls_from_text(&self, text: &str, tools: &[Tool]) -> Vec { + fn parse_tool_calls_from_text( + &self, + text: &str, + tools: &[Tool], + ) -> ParserResult> { let mut parsed = Vec::new(); + let tool_indices = helpers::get_tool_indices(tools); for mat in self.tool_call_extractor.find_iter(text) { - match self.parse_tool_call(mat.as_str(), tools) { - Ok(Some(tool)) => parsed.push(tool), - Ok(None) => continue, - Err(e) => { - tracing::debug!("Failed to parse tool call: {}", e); - continue; - } + let Some(tool) = self.parse_tool_call(mat.as_str(), tools)? else { + return Err(ParserError::ParsingFailed( + "malformed GLM tool call".to_string(), + )); + }; + if !tools.is_empty() && !tool_indices.contains_key(&tool.function.name) { + return Err(ParserError::InvalidToolName(tool.function.name)); } + parsed.push(tool); } - parsed + Ok(parsed) } } impl Glm4MoeParser { + const STRUCTURAL_MARKERS: [&str; 6] = [ + "", + "", + "", + "", + "", + "", + ]; + + fn partial_marker_suffix_len(text: &str, marker: &str) -> usize { + (1..marker.len()) + .rev() + .find(|&len| text.ends_with(&marker[..len])) + .unwrap_or(0) + } + + fn partial_structural_marker_suffix_len(text: &str) -> usize { + Self::STRUCTURAL_MARKERS + .iter() + .map(|marker| Self::partial_marker_suffix_len(text, marker)) + .max() + .unwrap_or(0) + } + + fn contains_structural_fragment(text: &str) -> bool { + [" bool { + let mut saw_call = false; + let mut previous_end = 0; + + for block in self.tool_call_extractor.find_iter(text) { + if saw_call && Self::contains_structural_fragment(&text[previous_end..block.start()]) { + return true; + } + saw_call = true; + previous_end = block.end(); + } + + saw_call + && (Self::contains_structural_fragment(&text[previous_end..]) + || Self::partial_structural_marker_suffix_len(&text[previous_end..]) > 0) + } + + fn record_tool_call(&mut self, tool_call: ToolCall) -> ToolCallItem { + if self.current_tool_id == -1 { + self.current_tool_id = 0; + self.prev_tool_call_arr.clear(); + self.streamed_args_for_tool.clear(); + } + + let tool_id = self.current_tool_id as usize; + helpers::ensure_capacity( + self.current_tool_id, + &mut self.prev_tool_call_arr, + &mut self.streamed_args_for_tool, + ); + + if let Ok(args) = serde_json::from_str::(&tool_call.function.arguments) { + self.prev_tool_call_arr[tool_id] = serde_json::json!({ + "name": tool_call.function.name, + "arguments": args, + }); + } + self.streamed_args_for_tool[tool_id].clone_from(&tool_call.function.arguments); + self.current_tool_id += 1; + + ToolCallItem { + tool_index: tool_id, + name: Some(tool_call.function.name), + parameters: tool_call.function.arguments, + } + } + + fn drain_incremental( + &mut self, + tools: &[Tool], + end_of_input: bool, + ) -> ParserResult { + let mut normal_text = String::new(); + let mut calls = Vec::new(); + let tool_indices = helpers::get_tool_indices(tools); + + loop { + if let Some(start) = self.buffer.find(self.bot_token) { + if start > 0 { + if self.current_tool_id == -1 { + normal_text.push_str(&self.buffer[..start]); + } + self.buffer.drain(..start); + continue; + } + + let Some(end_pos) = self.buffer.find(self.eot_token) else { + break; + }; + let block_end = end_pos + self.eot_token.len(); + let block = self.buffer[..block_end].to_string(); + + let Some(tool_call) = self.parse_tool_call(&block, tools)? else { + return Err(ParserError::ParsingFailed( + "malformed GLM tool call".to_string(), + )); + }; + // Only validate tool names against a schema when tools are + // provided; no-schema callers (inference-only, native markup + // decoding) pass an empty slice and should accept any name. + if !tools.is_empty() && !tool_indices.contains_key(&tool_call.function.name) { + return Err(ParserError::InvalidToolName(tool_call.function.name)); + } + self.buffer.drain(..block_end); + calls.push(self.record_tool_call(tool_call)); + continue; + } + + if self.current_tool_id != -1 && Self::contains_structural_fragment(&self.buffer) { + return Err(ParserError::ParsingFailed( + "unexpected GLM structure after tool call".to_string(), + )); + } + + // If bot_token is already in the buffer (but eot_token hasn't arrived), + // we must hold everything from that index so the tool-call marker isn't + // drained and emitted as normal text. + let held_len = if let Some(start) = self.buffer.find(self.bot_token) { + self.buffer.len() - start + } else if self.current_tool_id == -1 { + Self::partial_marker_suffix_len(&self.buffer, self.bot_token) + } else { + Self::partial_structural_marker_suffix_len(&self.buffer) + }; + let emit_len = self.buffer.len() - held_len; + if self.current_tool_id == -1 { + normal_text.push_str(&self.buffer[..emit_len]); + } + self.buffer.drain(..emit_len); + break; + } + + if end_of_input && !self.buffer.is_empty() { + // At EOF, release short ambiguous prefixes (≤ 2 bytes) as normal + // text rather than failing with Incomplete. A lone "<" or " ParserResult<(String, Vec)> { - if !self.has_tool_markers(text) { + let has_complete_marker = self.has_tool_markers(text); + // Short ambiguous prefixes ("<", " 2; + if has_partial_marker { + return Err(ParserError::Incomplete); + } + // A standalone unmatched closing marker (`eot_token`) with no matching + // opening `bot_token` is malformed structured output. Without this + // guard `has_complete_marker` would be false and the raw GLM markup + // would leak to clients as ordinary assistant text, bypassing the + // parse-error/502 path. The streaming path already rejects this via + // its post-call structural-residue checks. + if !has_complete_marker && text.contains(self.eot_token) { + return Err(ParserError::ParsingFailed( + "unmatched GLM closing marker without opening marker".to_string(), + )); + } + if !has_complete_marker { return Ok((text.to_string(), vec![])); } + // Strip extracted tool-call blocks so orphan marker checks are not + // tripped by marker text that happens to appear inside argument values + // (e.g. a search query containing the literal ""). + let mut stripped = String::with_capacity(text.len()); + let mut last_end = 0; + for m in self.tool_call_extractor.find_iter(text) { + stripped.push_str(&text[last_end..m.start()]); + last_end = m.end(); + } + stripped.push_str(&text[last_end..]); + + if stripped.contains(self.bot_token) { + return Err(ParserError::Incomplete); + } + if stripped.contains(self.eot_token) || self.has_post_call_structural_residue(text) + { + return Err(ParserError::ParsingFailed( + "unexpected GLM structure outside tool call".to_string(), + )); + } // Find where tool calls begin // Safe: has_tool_markers() already confirmed the marker exists @@ -177,11 +410,13 @@ impl Glm4MoeParser { .ok_or_else(|| ParserError::ParsingFailed("tool call marker not found".to_string()))?; let normal_text = text[..idx].to_string(); - let parsed = self.parse_tool_calls_from_text(text, tools); + let parsed = self.parse_tool_calls_from_text(text, tools)?; - // If no tools were successfully parsed despite having markers, return entire text as fallback + // Structured output must never silently fall back to raw marker text. if parsed.is_empty() { - return Ok((text.to_string(), vec![])); + return Err(ParserError::ParsingFailed( + "malformed or incomplete GLM tool call".to_string(), + )); } Ok((normal_text, parsed)) @@ -236,126 +471,12 @@ impl ToolParser for Glm4MoeParser { chunk: &str, tools: &[Tool], ) -> ParserResult { - // Python logic: Wait for complete tool call, then parse it all at once self.buffer.push_str(chunk); - let current_text = &self.buffer.clone(); - - // Check if we have bot_token - let start = current_text.find(self.bot_token); - if start.is_none() { - self.buffer.clear(); - // If we're in the middle of streaming (current_tool_id > 0), don't return text - let normal_text = if self.current_tool_id > 0 { - String::new() - } else { - current_text.clone() - }; - return Ok(StreamingParseResult { - normal_text, - calls: vec![], - }); - } - - // Check if we have eot_token (end of tool call) - let end = current_text.find(self.eot_token); - if let Some(end_pos) = end { - // We have a complete tool call! - - // Initialize state if this is the first tool call - if self.current_tool_id == -1 { - self.current_tool_id = 0; - self.prev_tool_call_arr = Vec::new(); - self.streamed_args_for_tool = vec![String::new()]; - } - - // Ensure we have enough entries in our tracking arrays - helpers::ensure_capacity( - self.current_tool_id, - &mut self.prev_tool_call_arr, - &mut self.streamed_args_for_tool, - ); - - // Parse the complete block using shared helper - let block_end = end_pos + self.eot_token.len(); - let parsed_tools = self.parse_tool_calls_from_text(¤t_text[..block_end], tools); - - // Extract normal text before tool calls - let idx = current_text.find(self.bot_token); - let normal_text = if let Some(pos) = idx { - current_text[..pos].trim().to_string() - } else { - String::new() - }; - - // Build tool indices for validation - let tool_indices = helpers::get_tool_indices(tools); - - let mut calls = Vec::new(); - - if !parsed_tools.is_empty() { - // Take the first tool and convert to ToolCallItem - let tool_call = &parsed_tools[0]; - let tool_id = self.current_tool_id as usize; - - // Validate tool name - if !tool_indices.contains_key(&tool_call.function.name) { - // Invalid tool name - skip this tool, preserve indexing for next tool - tracing::debug!("Invalid tool name '{}' - skipping", tool_call.function.name); - helpers::reset_current_tool_state( - &mut self.buffer, - &mut false, // glm45_moe/glm47_moe doesn't track name_sent per tool - &mut self.streamed_args_for_tool, - &self.prev_tool_call_arr, - ); - return Ok(StreamingParseResult::default()); - } - - calls.push(ToolCallItem { - tool_index: tool_id, - name: Some(tool_call.function.name.clone()), - parameters: tool_call.function.arguments.clone(), - }); - - // Store in tracking arrays - if self.prev_tool_call_arr.len() <= tool_id { - self.prev_tool_call_arr - .resize_with(tool_id + 1, || Value::Null); - } - - // Parse parameters as JSON and store - if let Ok(args) = serde_json::from_str::(&tool_call.function.arguments) { - self.prev_tool_call_arr[tool_id] = serde_json::json!({ - "name": tool_call.function.name, - "arguments": args, - }); - } - - if self.streamed_args_for_tool.len() <= tool_id { - self.streamed_args_for_tool - .resize_with(tool_id + 1, String::new); - } - self.streamed_args_for_tool[tool_id].clone_from(&tool_call.function.arguments); - - self.current_tool_id += 1; - } - - // Remove processed portion from buffer - self.buffer = current_text[block_end..].to_string(); - return Ok(StreamingParseResult { normal_text, calls }); - } - - // No complete tool call yet - return normal text before start token - // Safe: start.is_none() case was handled above (early return) - let Some(start_pos) = start else { - return Ok(StreamingParseResult::default()); - }; - let normal_text = current_text[..start_pos].to_string(); - self.buffer = current_text[start_pos..].to_string(); + self.drain_incremental(tools, false) + } - Ok(StreamingParseResult { - normal_text, - calls: vec![], - }) + async fn finalize(&mut self, tools: &[Tool]) -> ParserResult { + self.drain_incremental(tools, true) } fn has_tool_markers(&self, text: &str) -> bool { @@ -438,4 +559,189 @@ mod tests { assert_eq!(args["limit"], Value::String("4".to_string())); assert_eq!(args["count"], Value::Number(5.into())); } + + #[tokio::test] + async fn test_non_streaming_rejects_incomplete_or_malformed_markers() { + let parser = Glm4MoeParser::glm47(); + + for text in [ + "plain text fquery", + "", + ] { + assert!( + parser.parse_complete(text).await.is_err(), + "structured marker must not fall back to raw text: {text}" + ); + } + } + + #[tokio::test] + async fn test_non_streaming_rejects_standalone_unmatched_closing_marker() { + let parser = Glm4MoeParser::glm47(); + // Build the closing marker via `format!` so the literal tag cannot be + // accidentally stripped from the source. A standalone `` + // with no matching `` is malformed structured output and + // must surface as a parse error rather than leaking as ordinary text. + let eot = format!("<{}tool_call>", "/"); + + for text in [ + eot.clone(), + format!("some preamble{eot}"), + format!("lorem ipsum {eot} trailing"), + ] { + assert!( + parser.parse_complete(&text).await.is_err(), + "a standalone unmatched GLM closing marker must not leak as ordinary text: {text}" + ); + } + } + + // Short (<=2 byte) ambiguous prefixes like "f"; + + for trailing in [ + "fqueryunfinished", + "", + ] { + let text = format!("{valid}{trailing}"); + assert!( + parser.parse_complete(&text).await.is_err(), + "a valid first call must not hide an incomplete second call: {text}" + ); + } + } + + #[tokio::test] + async fn test_non_streaming_rejects_unknown_tool_name() { + let tools = tool_with_props(serde_json::json!({})); + let text = "missing"; + + let error = Glm4MoeParser::glm47() + .parse_complete_with_tools(text, &tools) + .await + .unwrap_err(); + + assert!(matches!(error, ParserError::InvalidToolName(name) if name == "missing")); + } + + #[tokio::test] + async fn test_streaming_rejects_unknown_tool_name() { + let tools = tool_with_props(serde_json::json!({})); + let text = "missing"; + let mut parser = Glm4MoeParser::glm47(); + + let error = parser.parse_incremental(text, &tools).await.unwrap_err(); + + assert!(matches!(error, ParserError::InvalidToolName(name) if name == "missing")); + assert!(matches!( + parser.parse_incremental("", &tools).await.unwrap_err(), + ParserError::InvalidToolName(name) if name == "missing" + )); + } + + #[tokio::test] + async fn test_rejects_malformed_argument_tags() { + let tools = tool_with_props(serde_json::json!({"query": {"type": "string"}})); + + for text in [ + "fquerypeople", + "fquerypeople", + "fquerypeople", + "fpeople", + ] { + assert!( + Glm4MoeParser::glm47() + .parse_complete_with_tools(text, &tools) + .await + .is_err(), + "malformed argument tags must fail in complete parsing: {text}" + ); + + let mut parser = Glm4MoeParser::glm47(); + assert!( + parser.parse_incremental(text, &tools).await.is_err(), + "malformed argument tags must fail in incremental parsing: {text}" + ); + } + } + + #[tokio::test] + async fn test_rejects_structural_residue_after_valid_call() { + let tools = tool_with_props(serde_json::json!({"query": {"type": "string"}})); + let valid = "f"; + + for residue in [ + "", + "query", + "people", + ] { + let text = format!("{valid}{residue}"); + assert!( + Glm4MoeParser::glm47() + .parse_complete_with_tools(&text, &tools) + .await + .is_err(), + "complete parsing must reject post-call structural residue: {text}" + ); + } + } + + #[tokio::test] + async fn test_streaming_rejects_unmatched_end_marker_at_every_split() { + let tools = tool_with_props(serde_json::json!({})); + let valid = "f"; + let residue = ""; + + for split in 0..=residue.len() { + let mut parser = Glm4MoeParser::glm47(); + let first = format!("{valid}{}", &residue[..split]); + let mut failed = parser.parse_incremental(&first, &tools).await.is_err(); + if !failed { + failed = parser + .parse_incremental(&residue[split..], &tools) + .await + .is_err(); + } + if !failed { + failed = parser.finalize(&tools).await.is_err(); + } + + assert!( + failed, + "unmatched end marker must fail at split {split}: {first:?} + {:?}", + &residue[split..] + ); + } + } } diff --git a/crates/tool_parser/src/parsers/helpers.rs b/crates/tool_parser/src/parsers/helpers.rs index 9f232c6667..0291a95540 100644 --- a/crates/tool_parser/src/parsers/helpers.rs +++ b/crates/tool_parser/src/parsers/helpers.rs @@ -12,8 +12,9 @@ use crate::{ /// `param_name -> declared JSON-schema type` for the named function (empty if the /// function or its `properties` are absent). Lets XML-style parsers coerce by the /// declared type instead of guessing from text (e.g. keep a numeric-looking -/// `string` as a string). Only scalar `type` is read; unions/nullable schemas -/// have no single type, so they fall back to the caller's inference. +/// `string` as a string). For simple unions, a string branch takes precedence; +/// otherwise the first non-null scalar type is used. `$ref` resolution remains +/// the responsibility of the caller. pub fn param_types_for_function(tools: &[Tool], func_name: &str) -> HashMap { let mut types = HashMap::new(); let Some(tool) = tools.iter().find(|t| t.function.name == func_name) else { @@ -26,14 +27,41 @@ pub fn param_types_for_function(tools: &[Tool], func_name: &str) -> HashMap Option { + let mut candidates: Vec = Vec::new(); + + match schema.get("type") { + Some(Value::String(ty)) => candidates.push(ty.clone()), + Some(Value::Array(types)) => candidates.extend( + types + .iter() + .filter_map(Value::as_str) + .map(ToString::to_string), + ), + _ => {} + } + + for keyword in ["anyOf", "oneOf"] { + if let Some(branches) = schema.get(keyword).and_then(Value::as_array) { + candidates.extend(branches.iter().filter_map(declared_schema_type)); + } + } + + candidates + .iter() + .find(|ty| ty.as_str() == "string") + .or_else(|| candidates.iter().find(|ty| ty.as_str() != "null")) + .cloned() +} + /// Coerce a raw value by its declared JSON-schema type. For `string`, a JSON /// string literal (`"4"`) is unwrapped while bare text (`4`, `true`) is kept /// verbatim; other types are parsed. `None` (unknown type or parse failure) diff --git a/crates/tool_parser/src/traits.rs b/crates/tool_parser/src/traits.rs index 30ed8e8e0f..d704f3d9d9 100644 --- a/crates/tool_parser/src/traits.rs +++ b/crates/tool_parser/src/traits.rs @@ -36,6 +36,15 @@ pub trait ToolParser: Send + Sync { tools: &[Tool], ) -> ParserResult; + /// Finish a streaming parse at end-of-input. + /// + /// Stateful parsers can override this to drain complete buffered calls or + /// reject an incomplete structured marker. The default is intentionally a + /// no-op so existing parser implementations retain their current behavior. + async fn finalize(&mut self, _tools: &[Tool]) -> ParserResult { + Ok(StreamingParseResult::default()) + } + /// Check if text contains tool calls in this parser's format fn has_tool_markers(&self, text: &str) -> bool; diff --git a/crates/tool_parser/tests/tool_parser_glm47_moe.rs b/crates/tool_parser/tests/tool_parser_glm47_moe.rs index 3ace073179..cf096965b3 100644 --- a/crates/tool_parser/tests/tool_parser_glm47_moe.rs +++ b/crates/tool_parser/tests/tool_parser_glm47_moe.rs @@ -2,8 +2,22 @@ mod common; use common::create_test_tools; +use openai_protocol::common::{Function, Tool}; +use serde_json::{json, Value}; use tool_parser::{Glm4MoeParser, ParserFactory, ToolParser}; +fn schema_tool(properties: Value) -> Vec { + vec![Tool { + tool_type: "function".to_string(), + function: Function { + name: "lookup".to_string(), + description: None, + parameters: json!({"type": "object", "properties": properties}), + strict: None, + }, + }] +} + #[tokio::test] async fn test_glm47_complete_parsing() { let parser = Glm4MoeParser::glm47(); @@ -17,7 +31,7 @@ The weather will be..."; assert_eq!(normal_text, "Let me search for that.\n"); assert_eq!(tools[0].function.name, "get_weather"); - let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); + let args: Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); assert_eq!(args["city"], "Beijing"); assert_eq!(args["date"], "2024-12-25"); } @@ -45,11 +59,11 @@ async fn test_glm47_type_conversion() { assert_eq!(tools.len(), 1); assert_eq!(normal_text, ""); - let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); + let args: Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); assert_eq!(args["count"], 42); assert_eq!(args["rate"], 1.5); assert_eq!(args["enabled"], true); - assert_eq!(args["data"], serde_json::Value::Null); + assert_eq!(args["data"], Value::Null); assert_eq!(args["text"], "string value"); } @@ -86,6 +100,108 @@ async fn test_glm47_streaming() { assert!(found_name, "Should have found tool name during streaming"); } +#[tokio::test] +async fn test_glm47_streaming_holds_partial_tool_marker_at_every_split() { + let tools = create_test_tools(); + let input = "回答：get_weathercity杭州"; + let expected_prefix = "回答："; + + for split in input.char_indices().map(|(index, _)| index).skip(1) { + let mut parser = Glm4MoeParser::glm47(); + let mut normal_text = String::new(); + let mut calls = Vec::new(); + + for chunk in [&input[..split], &input[split..]] { + let result = parser.parse_incremental(chunk, &tools).await.unwrap(); + normal_text.push_str(&result.normal_text); + calls.extend(result.calls); + } + + assert_eq!(normal_text, expected_prefix, "split at byte {split}"); + assert_eq!(calls.len(), 1, "split at byte {split}"); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + let args: Value = serde_json::from_str(&calls[0].parameters).unwrap(); + assert_eq!(args["city"], "杭州", "split at byte {split}"); + } +} + +#[tokio::test] +async fn test_glm47_streaming_parses_two_calls_from_one_chunk() { + let mut parser = Glm4MoeParser::glm47(); + let tools = create_test_tools(); + let input = "get_weathercity杭州\ + translatetext天气target_langen"; + + let result = parser.parse_incremental(input, &tools).await.unwrap(); + + assert_eq!(result.calls.len(), 2); + assert_eq!(result.calls[0].tool_index, 0); + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls[1].tool_index, 1); + assert_eq!(result.calls[1].name.as_deref(), Some("translate")); +} + +#[tokio::test] +async fn test_glm47_finalize_accepts_complete_cjk_call_and_rejects_partial_eof() { + let tools = schema_tool(json!({"query": {"type": "string"}})); + let mut complete = Glm4MoeParser::glm47(); + let parsed = complete + .parse_incremental( + "lookupquery查询人事", + &tools, + ) + .await + .unwrap(); + assert_eq!(parsed.calls.len(), 1); + let final_result = complete.finalize(&tools).await.unwrap(); + assert!(final_result.normal_text.is_empty()); + assert!(final_result.calls.is_empty()); + + let mut unclosed = Glm4MoeParser::glm47(); + unclosed + .parse_incremental( + "lookupquery查询人事", + &tools, + ) + .await + .unwrap(); + assert!(unclosed.finalize(&tools).await.is_err()); + + let mut partial_marker = Glm4MoeParser::glm47(); + partial_marker + .parse_incremental("正文get_weathercity杭州"; + + for alias in ["glm45", "glm47"] { + assert!(factory.registry().create_parser(alias).is_some(), "{alias}"); + } + + for model in [ + "GLM-5.2", + "THUDM/GLM-5.2", + "zai-org/glm-5.2-fp8", + "/models/THUDM/GLM-5.2", + ] { + let parser = factory + .registry() + .create_for_model(model) + .unwrap_or_else(|| panic!("no parser for {model}")); + let (_, calls) = parser.parse_complete(input).await.unwrap(); + assert_eq!(calls.len(), 1, "{model}"); + assert_eq!(calls[0].function.name, "get_weather", "{model}"); + } +} + #[tokio::test] async fn test_python_literals() { let parser = Glm4MoeParser::glm47(); @@ -127,10 +269,10 @@ async fn test_python_literals() { assert_eq!(tools.len(), 1); assert_eq!(tools[0].function.name, "test_func"); - let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); + let args: Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); assert_eq!(args["bool_true"], true); assert_eq!(args["bool_false"], false); - assert_eq!(args["none_val"], serde_json::Value::Null); + assert_eq!(args["none_val"], Value::Null); } #[tokio::test] @@ -142,7 +284,7 @@ async fn test_glm47_nested_json_in_arg_values() { let (_normal_text, tools) = parser.parse_complete(input).await.unwrap(); assert_eq!(tools.len(), 1); - let args: serde_json::Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); + let args: Value = serde_json::from_str(&tools[0].function.arguments).unwrap(); assert!(args["data"].is_object()); assert!(args["list"].is_array()); } diff --git a/grpc_servicer/tests/test_sglang_generate_request_contract.py b/grpc_servicer/tests/test_sglang_generate_request_contract.py new file mode 100644 index 0000000000..1f3253a97d --- /dev/null +++ b/grpc_servicer/tests/test_sglang_generate_request_contract.py @@ -0,0 +1,89 @@ +"""Source-level contract tests for current SGLang request construction. + +These tests intentionally avoid importing SGLang so the compatibility contract +can be checked in the lightweight gRPC-servicer test environment. +""" + +import ast +import unittest +from pathlib import Path + +SERVICER_PATH = Path(__file__).resolve().parents[1] / "smg_grpc_servicer" / "sglang" / "servicer.py" +CURRENT_REQUIRED_FIELDS = { + "rid", + "input_text", + "input_ids", + "input_embeds", + "mm_inputs", + "token_type_ids", + "sampling_params", + "return_logprob", + "logprob_start_len", + "top_logprobs_num", + "token_ids_logprob", + "stream", +} + + +def _constructor_keywords(function_name: str) -> list[set[str]]: + return [ + {keyword.arg for keyword in call.keywords if keyword.arg is not None} + for call in _tokenized_generate_calls(function_name) + ] + + +def _tokenized_generate_calls(function_name: str) -> list[ast.Call]: + """Return every ``TokenizedGenerateReqInput(...)`` call in a function.""" + module = ast.parse(SERVICER_PATH.read_text()) + function = next( + node + for node in ast.walk(module) + if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef) and node.name == function_name + ) + return [ + call + for call in ast.walk(function) + if isinstance(call, ast.Call) + and isinstance(call.func, ast.Name) + and call.func.id == "TokenizedGenerateReqInput" + ] + + +class TestSGLangGenerateRequestContract(unittest.TestCase): + def test_generate_converter_supplies_current_required_fields(self): + calls = _constructor_keywords("_convert_generate_request") + self.assertEqual(len(calls), 1) + self.assertEqual(CURRENT_REQUIRED_FIELDS - calls[0], set()) + + def test_health_probe_supplies_current_required_fields(self): + calls = _constructor_keywords("HealthCheck") + self.assertEqual(len(calls), 1) + self.assertEqual(CURRENT_REQUIRED_FIELDS - calls[0], set()) + + def test_generate_converter_sets_optional_fields_to_none(self): + self._assert_optional_fields_none("_convert_generate_request") + + def test_health_probe_sets_optional_fields_to_none(self): + self._assert_optional_fields_none("HealthCheck") + + def _assert_optional_fields_none(self, function_name: str) -> None: + """``input_embeds`` and ``token_type_ids`` must be literal ``None``.""" + calls = _tokenized_generate_calls(function_name) + self.assertEqual(len(calls), 1, f"expected one constructor call in {function_name}") + keywords = {kw.arg: kw.value for kw in calls[0].keywords if kw.arg is not None} + for field in ("input_embeds", "token_type_ids"): + self.assertIn(field, keywords, f"{field} missing in {function_name}") + value = keywords[field] + self.assertIsInstance( + value, + ast.Constant, + f"{field} must be a literal (ast.Constant) in {function_name}", + ) + self.assertIsNone( + value.value, + f"{field} must be literal None in {function_name}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/model_gateway/src/routers/grpc/regular/mod.rs b/model_gateway/src/routers/grpc/regular/mod.rs index 38311b10a1..f12598960b 100644 --- a/model_gateway/src/routers/grpc/regular/mod.rs +++ b/model_gateway/src/routers/grpc/regular/mod.rs @@ -7,3 +7,6 @@ pub(crate) mod processor; pub(crate) mod responses; pub(crate) mod stages; pub(crate) mod streaming; + +#[cfg(test)] +pub(crate) mod test_fakes; diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index 2513d1e6cd..bf953913d6 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -128,17 +128,13 @@ impl ResponseProcessor { parser.mark_reasoning_started(); } - match parser.detect_and_parse_reasoning(&processed_text) { - Ok(result) => { - if !result.reasoning_text.is_empty() { - reasoning_text = Some(result.reasoning_text); - } - processed_text = result.normal_text; - } - Err(e) => { - warn!("Reasoning parsing error, skipping parsing: {e}"); - } + let result = parser + .detect_and_parse_reasoning(&processed_text) + .map_err(|e| format!("Reasoning parsing failed: {e}"))?; + if !result.reasoning_text.is_empty() { + reasoning_text = Some(result.reasoning_text); } + processed_text = result.normal_text; } } @@ -182,7 +178,8 @@ impl ResponseProcessor { original_request.tools.as_deref().unwrap_or(&[]), history_tool_calls_count, ) - .await; + .await + .map_err(|e| format!("Tool call parsing failed: {e}"))?; } } @@ -309,10 +306,7 @@ impl ResponseProcessor { { Ok(choice) => choices.push(choice), Err(e) => { - return Err(error::internal_error( - "process_choice_failed", - format!("Failed to process choice {index}: {e}"), - )); + return Err(process_choice_error_response(index, e)); } } } @@ -340,7 +334,7 @@ impl ResponseProcessor { tool_parser_name: Option<&str>, tools: &[Tool], history_tool_calls_count: usize, - ) -> (Option>, String) { + ) -> tool_parser::errors::ParserResult<(Option>, String)> { // Get pooled parser for this model let pooled_parser = utils::get_tool_parser(&self.tool_parser_factory, tool_parser_name, model); @@ -356,40 +350,33 @@ impl ResponseProcessor { // Lock is dropped here }; - match result { - Ok((normal_text, parsed_tool_calls)) => { - if parsed_tool_calls.is_empty() { - return (None, normal_text); - } - - let spec_tool_calls = parsed_tool_calls - .into_iter() - .enumerate() - .map(|(index, tc)| { - // Generate ID for this tool call - let id = utils::generate_tool_call_id( - model, - &tc.function.name, - index, - history_tool_calls_count, - ); - ToolCall { - id, - tool_type: "function".to_string(), - function: FunctionCallResponse { - name: tc.function.name, - arguments: Some(tc.function.arguments), - }, - } - }) - .collect(); - (Some(spec_tool_calls), normal_text) - } - Err(e) => { - error!("Tool call parsing error: {}", e); - (None, processed_text.to_string()) - } + let (normal_text, parsed_tool_calls) = result?; + if parsed_tool_calls.is_empty() { + return Ok((None, normal_text)); } + + let spec_tool_calls = parsed_tool_calls + .into_iter() + .enumerate() + .map(|(index, tc)| { + // Generate ID for this tool call + let id = utils::generate_tool_call_id( + model, + &tc.function.name, + index, + history_tool_calls_count, + ); + ToolCall { + id, + tool_type: "function".to_string(), + function: FunctionCallResponse { + name: tc.function.name, + arguments: Some(tc.function.arguments), + }, + } + }) + .collect(); + Ok((Some(spec_tool_calls), normal_text)) } /// Process non-streaming generate response (collects all responses and builds final response array) @@ -412,15 +399,7 @@ impl ResponseProcessor { stop_decoder.reset(); // Process tokens through stop decoder - let outputs = match stop_decoder.process_tokens(complete.output_ids()) { - Ok(outputs) => outputs, - Err(e) => { - return Err(error::internal_error( - "process_tokens_failed", - format!("Failed to process tokens: {e}"), - )) - } - }; + let outputs = process_non_streaming_output_tokens(stop_decoder, complete.output_ids())?; // Accumulate text with early breaks let mut decoded_text = String::new(); @@ -701,7 +680,7 @@ impl ResponseProcessor { .as_deref() .map(utils::message_utils::extract_chat_tools) .unwrap_or_default(); - (tool_calls, processed_text) = self + match self .parse_tool_calls( &processed_text, &messages_request.model, @@ -711,7 +690,16 @@ impl ResponseProcessor { &messages_request, ), ) - .await; + .await + { + Ok(parsed) => (tool_calls, processed_text) = parsed, + Err(e) => { + return Err(error::bad_gateway( + "upstream_output_parse_failed", + format!("Tool call parsing failed: {e}"), + )); + } + } } } @@ -864,15 +852,8 @@ impl ResponseProcessor { for (i, complete) in all_responses.into_iter().enumerate() { stop_decoder.reset(); - let outputs = match stop_decoder.process_tokens(complete.output_ids()) { - Ok(outputs) => outputs, - Err(e) => { - return Err(error::internal_error( - "process_tokens_failed", - format!("Failed to process tokens: {e}"), - )) - } - }; + let outputs = + process_non_streaming_output_tokens(stop_decoder, complete.output_ids())?; let mut decoded_text = String::new(); let mut stopped = false; @@ -968,6 +949,32 @@ impl ResponseProcessor { } } +#[expect( + clippy::result_large_err, + reason = "router errors are returned as complete HTTP responses" +)] +fn process_non_streaming_output_tokens( + stop_decoder: &mut StopSequenceDecoder, + output_ids: &[u32], +) -> Result, axum::response::Response> { + stop_decoder.process_tokens(output_ids).map_err(|e| { + error::bad_gateway( + "upstream_output_parse_failed", + format!("Failed to process tokens: {e}"), + ) + }) +} + +fn process_choice_error_response( + index: usize, + error_message: impl std::fmt::Display, +) -> axum::response::Response { + error::bad_gateway( + "upstream_output_parse_failed", + format!("Failed to process choice {index}: {error_message}"), + ) +} + /// Residual assistant text → OpenAI `content`. Whitespace-only (the `"\n\n"` left /// after reasoning + tool-call extraction) becomes `None`, not `Some("\n\n")`, which /// would otherwise diverge multi-turn conversations. Real content is kept verbatim. @@ -980,8 +987,34 @@ fn normalize_assistant_content(text: String) -> Option { } #[cfg(test)] -mod content_normalization_tests { - use super::normalize_assistant_content; +mod tests { + use axum::http::StatusCode; + use llm_tokenizer::{mock::MockTokenizer, stop::StopSequenceConfig}; + use openai_protocol::common::Function; + use smg_grpc_client::sglang_proto; + + use super::*; + use crate::routers::grpc::regular::test_fakes::{FailingReasoningParser, FailingToolParser, FailingTokenizer}; + + fn complete_with_text_token() -> ProtoGenerateComplete { + ProtoGenerateComplete::Sglang(sglang_proto::GenerateComplete { + output_ids: vec![1], + finish_reason: "stop".to_string(), + ..Default::default() + }) + } + + fn lookup_tool() -> Tool { + Tool { + tool_type: "function".to_string(), + function: Function { + name: "lookup".to_string(), + description: None, + parameters: serde_json::json!({"type": "object"}), + strict: None, + }, + } + } #[test] fn whitespace_only_is_none_real_text_kept_verbatim() { @@ -992,4 +1025,141 @@ mod content_normalization_tests { Some("\n\nDone.".to_string()) ); } + + #[tokio::test] + async fn non_streaming_reasoning_parse_error_is_returned() { + let reasoning_factory = ReasoningParserFactory::new(); + reasoning_factory + .registry() + .register_parser("failing", || Box::new(FailingReasoningParser)); + let processor = ResponseProcessor::new( + ToolParserFactory::new(), + reasoning_factory, + utils::ParserResolver::configured_only(None, Some("failing".to_string())), + ); + let tokenizer: Arc = Arc::new(MockTokenizer::new()); + let mut decoder = + StopSequenceDecoder::new(Arc::clone(&tokenizer), StopSequenceConfig::default(), false); + let request = ChatCompletionRequest { + model: "test-model".to_string(), + separate_reasoning: true, + ..Default::default() + }; + + let error = processor + .process_single_choice( + &complete_with_text_token(), + 0, + &request, + &tokenizer, + &mut decoder, + 0, + true, + false, + Some("failing"), + None, + ) + .await + .unwrap_err(); + + assert!(error.contains("Reasoning parsing failed")); + assert!(error.contains("fake failure")); + } + + #[tokio::test] + async fn non_streaming_tool_parse_error_is_returned() { + let tool_factory = ToolParserFactory::new(); + tool_factory + .registry() + .register_parser("failing", || Box::new(FailingToolParser)); + let processor = ResponseProcessor::new( + tool_factory, + ReasoningParserFactory::new(), + utils::ParserResolver::configured_only(Some("failing".to_string()), None), + ); + let tokenizer: Arc = Arc::new(MockTokenizer::new()); + let mut decoder = + StopSequenceDecoder::new(Arc::clone(&tokenizer), StopSequenceConfig::default(), false); + let request = ChatCompletionRequest { + model: "test-model".to_string(), + tools: Some(vec![lookup_tool()]), + separate_reasoning: false, + ..Default::default() + }; + + let error = processor + .process_single_choice( + &complete_with_text_token(), + 0, + &request, + &tokenizer, + &mut decoder, + 0, + false, + true, + None, + Some("failing"), + ) + .await + .unwrap_err(); + + assert!(error.contains("Tool call parsing failed")); + assert!(error.contains("fake failure")); + } + + #[tokio::test] + async fn non_streaming_parse_error_maps_to_bad_gateway() { + let response = process_choice_error_response(2, "fake parse failure"); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + assert_eq!( + response + .headers() + .get(error::HEADER_X_SMG_ERROR_CODE) + .and_then(|value| value.to_str().ok()), + Some("upstream_output_parse_failed") + ); + + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["error"]["code"], "upstream_output_parse_failed"); + assert!(json["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("fake parse failure"))); + } + + #[tokio::test] + async fn non_streaming_generate_and_completion_decode_errors_map_to_bad_gateway() { + for endpoint in ["generate", "completion"] { + let tokenizer: Arc = Arc::new(FailingTokenizer::default()); + let mut decoder = + StopSequenceDecoder::new(tokenizer, StopSequenceConfig::default(), false); + + let response = process_non_streaming_output_tokens(&mut decoder, &[7]).unwrap_err(); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY, "{endpoint}"); + assert_eq!( + response + .headers() + .get(error::HEADER_X_SMG_ERROR_CODE) + .and_then(|value| value.to_str().ok()), + Some("upstream_output_parse_failed"), + "{endpoint}" + ); + + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!( + json["error"]["code"], "upstream_output_parse_failed", + "{endpoint}" + ); + assert!(json["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("fake decode failure"))); + } + } } diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index 43bd04eae8..a07a141c0f 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, @@ -95,6 +96,32 @@ struct GenerateStreamContext { expected_choices: u32, } +type PooledReasoningParser = Arc>>; +type PooledToolParser = Arc>>; + +struct ChatTextProcessingContext<'a> { + separate_reasoning: bool, + reasoning_parser_available: bool, + thinking_override: bool, + think_in_prefill: bool, + // Resolved once per request by the caller: re-resolving inside the helper + // could disagree with the upfront availability check if the worker + // registry changed mid-stream. + reasoning_parser_name: Option<&'a str>, + tool_choice: Option<&'a ToolChoice>, + tools: Option<&'a [Tool]>, + tool_choice_enabled: bool, + tool_parser_available: bool, + tool_parser_name: Option<&'a str>, + used_json_schema: bool, + is_specific_function: bool, + request_id: &'a str, + model: &'a str, + created: u64, + system_fingerprint: Option<&'a str>, + history_tool_calls_count: usize, +} + impl StreamingProcessor { pub fn new( tool_parser_factory: ToolParserFactory, @@ -110,6 +137,31 @@ impl StreamingProcessor { } } + async fn finish_openai_streaming_task( + tx: &SseSender, + result: Result<(), String>, + ) { + match result { + Ok(()) => { + let _ = tx + .send(Ok(Bytes::from("data: [DONE]\n\n"))) + .await; + } + Err(error) => { + utils::send_error_sse(tx, error, "internal_error").await; + } + } + } + + fn decode_generate_chunk( + tokenizer: &dyn Tokenizer, + token_ids: &[u32], + ) -> Result { + tokenizer + .decode(token_ids, true) + .map_err(|e| format!("Failed to decode generate chunk: {e}")) + } + /// Process streaming chat response and return SSE response /// /// This is the high-level entry point for streaming responses, handling: @@ -164,11 +216,8 @@ impl StreamingProcessor { ) .await; - if let Err(e) = result { - utils::send_error_sse(&tx, &e, "internal_error").await; - } + Self::finish_openai_streaming_task(&tx, result).await; - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))).await; }); } context::ExecutionResult::PrefillDecode { @@ -197,11 +246,8 @@ impl StreamingProcessor { ) .await; - if let Err(e) = result { - utils::send_error_sse(&tx, &e, "internal_error").await; - } + Self::finish_openai_streaming_task(&tx, result).await; - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))).await; }); } context::ExecutionResult::Embedding { .. } => { @@ -282,7 +328,6 @@ impl StreamingProcessor { // Phase 1: Initialize state tracking (per-index for n>1 support) let mut is_firsts: HashMap = HashMap::new(); - let mut stream_buffers: HashMap = HashMap::new(); let mut finish_reasons: HashMap = HashMap::new(); let mut matched_stops: HashMap> = HashMap::new(); // Indices whose local stop decoder fired: their finish reason is pinned @@ -294,10 +339,8 @@ impl StreamingProcessor { let mut reasoning_tokens: HashMap = HashMap::new(); // Parser state (lazy initialization per index) - type PooledReasoningParser = Arc>>; let mut reasoning_parsers: HashMap = HashMap::new(); - type PooledToolParser = Arc>>; let mut tool_parsers: HashMap = HashMap::new(); let mut has_tool_calls: HashMap = HashMap::new(); @@ -369,6 +412,8 @@ impl StreamingProcessor { tool_parser_name.as_deref(), model, ); + let tool_choice_enabled = + !matches!(tool_choice, Some(ToolChoice::Value(ToolChoiceValue::None))); if separate_reasoning && !reasoning_parser_available { debug!( @@ -384,6 +429,26 @@ impl StreamingProcessor { ); } + let text_processing_context = ChatTextProcessingContext { + separate_reasoning, + reasoning_parser_available, + thinking_override, + think_in_prefill, + reasoning_parser_name: reasoning_parser_name.as_deref(), + tool_choice: tool_choice.as_ref(), + tools: tools.as_deref(), + tool_choice_enabled, + tool_parser_available, + tool_parser_name: tool_parser_name.as_deref(), + used_json_schema, + is_specific_function, + request_id, + model, + created, + system_fingerprint, + history_tool_calls_count, + }; + // Phase 2: Main streaming loop while let Some(response) = grpc_stream.next().await { let gen_response = response.map_err(|e| format!("Stream error: {}", e.message()))?; @@ -434,7 +499,8 @@ impl StreamingProcessor { // Process tokens through stop decoder let (chunk_text, should_stop) = - Self::process_chunk_tokens(stop_decoder, chunk.token_ids()); + Self::process_chunk_tokens(stop_decoder, chunk.token_ids())?; + if should_stop { // Stop-decoder match takes precedence: pin "stop" even if @@ -462,9 +528,6 @@ 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) @@ -479,109 +542,23 @@ impl StreamingProcessor { 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); + let chunks = self + .process_chat_text_delta( + &chunk_text, + index, + choice_logprobs, + &mut reasoning_parsers, + &mut tool_parsers, + &mut has_tool_calls, + &text_processing_context, + ) + .await?; + for chunk in chunks { + Self::format_sse_chunk_into(&mut sse_buffer, &chunk); tx.send(Ok(Bytes::from(sse_buffer.clone()))) .await - .map_err(|_| "Failed to send content chunk".to_string())?; + .map_err(|_| "Failed to send parsed chunk".to_string())?; + } } ProtoResponseVariant::Complete(complete) => { @@ -591,23 +568,27 @@ impl StreamingProcessor { 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())?; + let chunks = self + .process_chat_text_delta( + &text, + index, + None, + &mut reasoning_parsers, + &mut tool_parsers, + &mut has_tool_calls, + &text_processing_context, + ) + .await?; + for chunk in chunks { + let sse_chunk = + sse_encoder.encode_data(&chunk).map_err(|e| { + format!("Failed to serialize flushed chunk: {e}") + })?; + tx.send(Ok(sse_chunk)) + .await + .map_err(|_| "Failed to send flushed chunk".to_string())?; + } + } } } @@ -635,7 +616,29 @@ impl StreamingProcessor { // Phase 3: Check unstreamed tool args for (index, parser) in &tool_parsers { - let parser_guard = parser.lock().await; + let mut parser_guard = parser.lock().await; + let final_result = parser_guard + .finalize(tools.as_deref().unwrap_or_default()) + .await + .map_err(|e| format!("Tool parser finalization failed: {e}"))?; + for chunk in Self::tool_result_to_chat_chunks( + final_result, + *index, + &mut has_tool_calls, + request_id, + model, + created, + system_fingerprint, + history_tool_calls_count, + None, + ) { + let sse_chunk = sse_encoder + .encode_data(&chunk) + .map_err(|e| format!("Failed to serialize finalized tool chunk: {e}"))?; + tx.send(Ok(sse_chunk)) + .await + .map_err(|_| "Failed to send finalized tool output".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 { @@ -866,11 +869,8 @@ impl StreamingProcessor { Self::process_generate_streaming(tokenizer, stream, ctx, &tx, reservation) .await; - if let Err(e) = result { - utils::send_error_sse(&tx, &e, "internal_error").await; - } + Self::finish_openai_streaming_task(&tx, result).await; - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))).await; }); } context::ExecutionResult::PrefillDecode { @@ -896,11 +896,8 @@ impl StreamingProcessor { ) .await; - if let Err(e) = result { - utils::send_error_sse(&tx, &e, "internal_error").await; - } + Self::finish_openai_streaming_task(&tx, result).await; - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))).await; }); } context::ExecutionResult::Embedding { .. } => { @@ -970,9 +967,8 @@ impl StreamingProcessor { let current_completion_tokens = *completion_tokens; // Decode tokens to text (skip_special_tokens=true to handle newlines correctly) - let chunk_text = tokenizer - .decode(chunk.token_ids(), true) - .unwrap_or_default(); + let chunk_text = + Self::decode_generate_chunk(tokenizer.as_ref(), chunk.token_ids())?; // Accumulate text for this index let accumulated_text = accumulated_texts.entry(index).or_default(); @@ -1175,9 +1171,8 @@ impl StreamingProcessor { let current_completion_tokens = *completion_tokens; // Decode tokens to text - let chunk_text = tokenizer - .decode(chunk.token_ids(), true) - .unwrap_or_default(); + let chunk_text = + Self::decode_generate_chunk(tokenizer.as_ref(), chunk.token_ids())?; // Accumulate text for this index let accumulated_text = accumulated_texts.entry(index).or_default(); @@ -1338,31 +1333,28 @@ impl StreamingProcessor { fn process_chunk_tokens( stop_decoder: &mut StopSequenceDecoder, token_ids: &[u32], - ) -> (String, bool) { + ) -> Result<(String, bool), String> { let mut chunk_text = String::new(); for &token_id in token_ids { - match stop_decoder.process_token(token_id).unwrap_or_else(|e| { - debug!( - "Error processing token {}: {}. Treating as Held.", - token_id, e - ); - SequenceDecoderOutput::Held - }) { + let output = stop_decoder + .process_token(token_id) + .map_err(|e| format!("Failed to decode token {token_id}: {e}"))?; + match output { SequenceDecoderOutput::Text(text) => { chunk_text.push_str(&text); } SequenceDecoderOutput::StoppedWithText(text) => { chunk_text.push_str(&text); - return (chunk_text, true); + return Ok((chunk_text, true)); } SequenceDecoderOutput::Stopped => { - return (chunk_text, true); + return Ok((chunk_text, true)); } SequenceDecoderOutput::Held => {} } } - (chunk_text, false) + Ok((chunk_text, false)) } /// Helper: Process reasoning content in streaming mode @@ -1382,7 +1374,7 @@ impl StreamingProcessor { model: &str, created: u64, system_fingerprint: Option<&str>, - ) -> (String, Option, bool) { + ) -> Result<(String, Option, bool), String> { // Create fresh parser for this index (not pooled, to avoid state pollution) #[expect( clippy::expect_used, @@ -1428,15 +1420,111 @@ impl StreamingProcessor { .build(), ) }; - return (normal_text, chunk, in_reasoning); + return Ok((normal_text, chunk, in_reasoning)); } Err(e) => { - warn!("Reasoning parsing error: {}", e); + return Err(format!("Reasoning parsing failed: {e}")); } } } - (delta.to_string(), None, false) + Ok((delta.to_string(), None, false)) + } + + #[expect(clippy::too_many_arguments)] + async fn process_chat_text_delta( + &self, + delta: &str, + index: u32, + choice_logprobs: Option, + reasoning_parsers: &mut HashMap, + tool_parsers: &mut HashMap, + has_tool_calls: &mut HashMap, + context: &ChatTextProcessingContext<'_>, + ) -> Result, String> { + let mut chunks = Vec::new(); + let mut normal_text = delta.to_string(); + + let in_reasoning = if context.separate_reasoning && context.reasoning_parser_available { + let (parsed_text, reasoning_chunk, in_reasoning) = self + .process_reasoning_stream( + &normal_text, + index, + reasoning_parsers, + context.thinking_override, + context.think_in_prefill, + context.reasoning_parser_name, + context.request_id, + context.model, + context.created, + context.system_fingerprint, + ) + .await?; + if let Some(chunk) = reasoning_chunk { + chunks.push(chunk); + } + normal_text = parsed_text; + in_reasoning + } else { + false + }; + + if let Some(tools) = context.tools { + if !in_reasoning + && context.tool_choice_enabled + && (context.tool_parser_available || context.used_json_schema) + { + let tool_chunks = if context.is_specific_function { + Self::process_specific_function_stream( + &normal_text, + index, + has_tool_calls, + context.tool_choice, + context.request_id, + context.model, + context.created, + context.system_fingerprint, + context.history_tool_calls_count, + ) + } else { + self.process_tool_calls_stream( + &normal_text, + index, + tool_parsers, + has_tool_calls, + tools, + context.tool_parser_name, + context.request_id, + context.model, + context.created, + context.system_fingerprint, + context.history_tool_calls_count, + context.used_json_schema, + choice_logprobs, + ) + .await? + }; + chunks.extend(tool_chunks); + return Ok(chunks); + } + } + + if !normal_text.is_empty() { + chunks.push( + ChatCompletionStreamResponse::builder(context.request_id, context.model) + .created(context.created) + .add_choice_content_with_logprobs( + index, + "assistant", + normal_text, + choice_logprobs, + ) + .maybe_system_fingerprint(context.system_fingerprint) + .build(), + ); + } + + Ok(chunks) } /// Helper: Process specific function case - emit tool call deltas with arguments @@ -1509,9 +1597,8 @@ impl StreamingProcessor { system_fingerprint: Option<&str>, history_tool_calls_count: usize, use_json_parser: bool, - ) -> Vec { - let mut chunks = Vec::new(); - + choice_logprobs: Option, + ) -> Result, String> { // Create fresh parser for this index (not pooled, to avoid state pollution) #[expect( clippy::expect_used, @@ -1532,68 +1619,84 @@ impl StreamingProcessor { let mut parser = pooled_parser.lock().await; match parser.parse_incremental(delta, tools).await { - Ok(StreamingParseResult { normal_text, calls }) => { - // Emit normal text if present - if !normal_text.is_empty() { - chunks.push( - ChatCompletionStreamResponse::builder(request_id, model) - .created(created) - .add_choice_content(index, "assistant", normal_text) - .maybe_system_fingerprint(system_fingerprint) - .build(), - ); - } - - // Emit tool call chunks - for tool_call_item in calls { - has_tool_calls.insert(index, true); - - let tool_call_id = if let Some(ref name) = tool_call_item.name { - Some(utils::generate_tool_call_id( - model, - name, - tool_call_item.tool_index, - history_tool_calls_count, - )) - } else { - None - }; - - let tool_call_delta = ToolCallDelta { - index: tool_call_item.tool_index as u32, - id: tool_call_id, - tool_type: if tool_call_item.name.is_some() { - Some("function".to_string()) - } else { - None - }, - function: Some(FunctionCallDelta { - name: tool_call_item.name, - arguments: if tool_call_item.parameters.is_empty() { - None - } else { - Some(tool_call_item.parameters) - }, - }), - }; - - chunks.push( - ChatCompletionStreamResponse::builder(request_id, model) - .created(created) - .add_choice_tool_call_delta(index, tool_call_delta) - .maybe_system_fingerprint(system_fingerprint) - .build(), - ); - } - - return chunks; + Ok(result) => { + return Ok(Self::tool_result_to_chat_chunks( + result, + index, + has_tool_calls, + request_id, + model, + created, + system_fingerprint, + history_tool_calls_count, + choice_logprobs, + )); } Err(e) => { - error!("Tool call parsing error: {}", e); + return Err(format!("Tool call parsing failed: {e}")); } } } + Ok(Vec::new()) + } + + #[expect(clippy::too_many_arguments)] + fn tool_result_to_chat_chunks( + result: StreamingParseResult, + index: u32, + has_tool_calls: &mut HashMap, + request_id: &str, + model: &str, + created: u64, + system_fingerprint: Option<&str>, + history_tool_calls_count: usize, + choice_logprobs: Option, + ) -> Vec { + let mut chunks = Vec::new(); + if !result.normal_text.is_empty() { + chunks.push( + ChatCompletionStreamResponse::builder(request_id, model) + .created(created) + .add_choice_content_with_logprobs( + index, + "assistant", + result.normal_text, + choice_logprobs, + ) + .maybe_system_fingerprint(system_fingerprint) + .build(), + ); + } + + for tool_call_item in result.calls { + has_tool_calls.insert(index, true); + let tool_call_id = tool_call_item.name.as_ref().map(|name| { + utils::generate_tool_call_id( + model, + name, + tool_call_item.tool_index, + history_tool_calls_count, + ) + }); + let tool_call_delta = ToolCallDelta { + index: tool_call_item.tool_index as u32, + id: tool_call_id, + tool_type: tool_call_item.name.as_ref().map(|_| "function".to_string()), + function: Some(FunctionCallDelta { + name: tool_call_item.name, + arguments: (!tool_call_item.parameters.is_empty()) + .then_some(tool_call_item.parameters), + }), + }; + chunks.push( + ChatCompletionStreamResponse::builder(request_id, model) + .created(created) + .add_choice_tool_call_delta(index, tool_call_delta) + .maybe_system_fingerprint(system_fingerprint) + .build(), + ); + } chunks } @@ -2046,7 +2149,8 @@ impl StreamingProcessor { completion_tokens.record_chunk(&chunk); let (chunk_text, should_stop) = - Self::process_chunk_tokens(&mut stop_decoder, chunk.token_ids()); + Self::process_chunk_tokens(&mut stop_decoder, chunk.token_ids())?; + if should_stop { // Stop-decoder match takes precedence over the engine's @@ -2288,7 +2392,19 @@ impl StreamingProcessor { } } Err(e) => { - error!("Tool call parsing error in messages streaming: {}", e); + let error_event = MessageStreamEvent::Error { + error: messages::ErrorResponse { + error_type: "api_error".to_string(), + message: format!("Tool call parsing failed: {e}"), + }, + }; + let _ = Self::send_messages_event( + tx, + &mut sse_buffer, + &error_event, + ) + .await; + return Ok(()); } } } @@ -2897,7 +3013,7 @@ impl StreamingProcessor { }); let (decoded_text, stopped) = - Self::process_chunk_tokens(stop_decoder, chunk.token_ids()); + Self::process_chunk_tokens(stop_decoder, chunk.token_ids())?; chunk_text.clear(); chunk_text.push_str(&decoded_text); @@ -3211,7 +3327,146 @@ impl StreamingProcessor { #[cfg(test)] mod tests { + use openai_protocol::common::Function; + use super::*; + use crate::routers::grpc::regular::test_fakes::{FailingReasoningParser, FailingToolParser, FailingTokenizer}; + + fn test_processor( + configured_tool_parser: Option<&str>, + configured_reasoning_parser: Option<&str>, + ) -> StreamingProcessor { + StreamingProcessor::new( + ToolParserFactory::new(), + ReasoningParserFactory::new(), + utils::ParserResolver::configured_only( + configured_tool_parser.map(ToString::to_string), + configured_reasoning_parser.map(ToString::to_string), + ), + "sglang", + ) + } + + fn lookup_tools() -> Vec { + vec![Tool { + tool_type: "function".to_string(), + function: Function { + name: "lookup".to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": {"query": {"type": "string"}} + }), + strict: None, + }, + }] + } + + fn text_context<'a>( + tools: Option<&'a [Tool]>, + separate_reasoning: bool, + tool_parser_name: Option<&'a str>, + reasoning_parser_name: Option<&'a str>, + ) -> ChatTextProcessingContext<'a> { + ChatTextProcessingContext { + separate_reasoning, + reasoning_parser_available: separate_reasoning, + thinking_override: false, + think_in_prefill: false, + reasoning_parser_name, + tool_choice: None, + tools, + tool_choice_enabled: true, + tool_parser_available: tools.is_some(), + tool_parser_name, + used_json_schema: false, + is_specific_function: false, + request_id: "chatcmpl-test", + model: if separate_reasoning { + "glm47" + } else { + "test-model" + }, + created: 1, + system_fingerprint: None, + history_tool_calls_count: 0, + } + } + + #[tokio::test] + async fn failed_chat_stream_sends_error_without_done() { + let (tx, mut rx) = sse_channel(); + + StreamingProcessor::finish_openai_streaming_task( + &tx, + Err("fake parse failure".to_string()), + ) + .await; + drop(tx); + + let mut events: Vec = Vec::new(); + while let Some(event) = rx.recv().await { + events.push( + String::from_utf8(event.expect("SSE event").to_vec()).expect("valid UTF-8 SSE"), + ); + } + + assert_eq!(events.len(), 1); + assert!(events[0].contains("fake parse failure")); + assert!(!events.iter().any(|event| event.contains("[DONE]"))); + } + + #[tokio::test] + async fn failed_generate_and_completion_streams_send_error_without_done() { + for endpoint in ["generate", "completion"] { + let (tx, mut rx) = sse_channel(); + + StreamingProcessor::finish_openai_streaming_task( + &tx, + Err(format!("fake {endpoint} failure")), + ) + .await; + drop(tx); + + let mut events: Vec = Vec::new(); + while let Some(event) = rx.recv().await { + events.push( + String::from_utf8(event.expect("SSE event").to_vec()).expect("valid UTF-8 SSE"), + ); + } + + assert_eq!(events.len(), 1, "{endpoint}"); + assert!(events[0].contains(&format!("fake {endpoint} failure"))); + assert!( + !events.iter().any(|event| event.contains("[DONE]")), + "{endpoint}" + ); + } + } + + #[test] + fn token_decode_error_is_returned_instead_of_held() { + let tokenizer: Arc = Arc::new(FailingTokenizer::default()); + let mut decoder = StopSequenceDecoder::new(tokenizer, Default::default(), false); + + let error = StreamingProcessor::process_chunk_tokens(&mut decoder, &[7]).unwrap_err(); + + assert_eq!(error, "Failed to decode token 7: fake decode failure"); + } + + #[test] + fn generate_decode_error_is_returned_for_regular_and_pd_paths() { + let tokenizer = FailingTokenizer::default(); + + for path in ["regular", "prefill-decode"] { + let error = StreamingProcessor::decode_generate_chunk(&tokenizer, &[7]).unwrap_err(); + + assert_eq!( + error, "Failed to decode generate chunk: fake decode failure", + "{path}" + ); + } + } #[test] fn completion_streaming_usage_includes_reasoning_tokens() { @@ -3235,4 +3490,144 @@ mod tests { Some(3) ); } + + #[tokio::test] + async fn tool_parse_error_is_returned_without_success_chunks() { + let processor = test_processor(None, None); + let tools = lookup_tools(); + let mut reasoning_parsers = HashMap::new(); + let mut tool_parsers: HashMap = HashMap::new(); + tool_parsers.insert( + 0, + Arc::new(tokio::sync::Mutex::new(Box::new(FailingToolParser))), + ); + let mut has_tool_calls = HashMap::new(); + + let result = processor + .process_chat_text_delta( + "ignored", + 0, + None, + &mut reasoning_parsers, + &mut tool_parsers, + &mut has_tool_calls, + &text_context(Some(&tools), false, None, None), + ) + .await; + + assert!(result.is_err()); + assert!(has_tool_calls.is_empty()); + } + + #[tokio::test] + async fn reasoning_parse_error_is_returned_without_raw_fallback() { + let processor = test_processor(None, None); + let mut reasoning_parsers: HashMap = HashMap::new(); + reasoning_parsers.insert( + 0, + Arc::new(tokio::sync::Mutex::new(Box::new(FailingReasoningParser))), + ); + let mut tool_parsers = HashMap::new(); + let mut has_tool_calls = HashMap::new(); + + let result = processor + .process_chat_text_delta( + "must not leak", + 0, + None, + &mut reasoning_parsers, + &mut tool_parsers, + &mut has_tool_calls, + &text_context(None, true, None, None), + ) + .await; + + assert_eq!( + result.unwrap_err(), + "Reasoning parsing failed: Parser configuration error: fake failure" + ); + } + + #[tokio::test] + async fn flush_only_cjk_tool_call_emits_tool_delta() { + let processor = test_processor(Some("glm47"), None); + let tools = lookup_tools(); + let mut reasoning_parsers = HashMap::new(); + let mut tool_parsers = HashMap::new(); + let mut has_tool_calls = HashMap::new(); + + let chunks = processor + .process_chat_text_delta( + "lookupquery查询人事", + 0, + None, + &mut reasoning_parsers, + &mut tool_parsers, + &mut has_tool_calls, + &text_context(Some(&tools), false, Some("glm47"), None), + ) + .await + .unwrap(); + + assert_eq!(chunks.len(), 1); + let tool_calls = chunks[0].choices[0] + .delta + .tool_calls + .as_ref() + .expect("tool delta"); + let arguments = tool_calls[0] + .function + .as_ref() + .and_then(|function| function.arguments.as_deref()) + .expect("arguments"); + assert_eq!( + serde_json::from_str::(arguments).unwrap()["query"], + "查询人事" + ); + } + + #[tokio::test] + async fn flush_only_reasoning_close_tag_does_not_leak_as_content() { + let processor = test_processor(None, Some("glm45")); + let mut reasoning_parsers = HashMap::new(); + let mut tool_parsers = HashMap::new(); + let mut has_tool_calls = HashMap::new(); + let context = text_context(None, true, None, Some("glm45")); + + processor + .process_chat_text_delta( + "内部推理", + 0, + None, + &mut reasoning_parsers, + &mut tool_parsers, + &mut has_tool_calls, + &context, + ) + .await + .unwrap(); + let flushed = processor + .process_chat_text_delta( + "", + 0, + None, + &mut reasoning_parsers, + &mut tool_parsers, + &mut has_tool_calls, + &context, + ) + .await + .unwrap(); + + assert!(flushed.iter().all(|chunk| { + chunk.choices.iter().all(|choice| { + choice + .delta + .content + .as_deref() + .unwrap_or_default() + .is_empty() + }) + })); + } } diff --git a/model_gateway/src/routers/grpc/regular/test_fakes.rs b/model_gateway/src/routers/grpc/regular/test_fakes.rs new file mode 100644 index 0000000000..d8839d7a76 --- /dev/null +++ b/model_gateway/src/routers/grpc/regular/test_fakes.rs @@ -0,0 +1,121 @@ +//! Shared test fakes for the regular (non-harmony) gRPC router tests. +//! +//! These fakes exist solely to exercise error-propagation paths in +//! [`processor`] and [`streaming`]. They are compiled only under +//! `#[cfg(test)]` and intentionally fail every operation they implement. +//! +//! [`processor`]: super::processor +//! [`streaming`]: super::streaming + +use async_trait::async_trait; +use llm_tokenizer::traits::{Decoder, Encoder, Encoding, SpecialTokens, Tokenizer}; +use openai_protocol::common::Tool; +use reasoning_parser::{ParseError, ParserResult, ReasoningParser}; +use tool_parser::{ + errors::{ParserError as ToolParserError, ParserResult as ToolParserResult}, + types::{StreamingParseResult, ToolCall as ParsedToolCall}, + ToolParser, +}; + +pub(crate) struct FailingReasoningParser; + +pub(crate) struct FailingToolParser; + +#[derive(Default)] +pub(crate) struct FailingTokenizer { + special_tokens: SpecialTokens, +} + +impl Encoder for FailingTokenizer { + fn encode(&self, _input: &str, _add_special_tokens: bool) -> anyhow::Result { + Ok(Encoding::Plain(Vec::new())) + } + + fn encode_batch( + &self, + inputs: &[&str], + _add_special_tokens: bool, + ) -> anyhow::Result> { + Ok(inputs.iter().map(|_| Encoding::Plain(Vec::new())).collect()) + } +} + +impl Decoder for FailingTokenizer { + fn decode(&self, _token_ids: &[u32], _skip_special_tokens: bool) -> anyhow::Result { + Err(anyhow::anyhow!("fake decode failure")) + } +} + +impl Tokenizer for FailingTokenizer { + fn vocab_size(&self) -> usize { + 0 + } + + fn get_special_tokens(&self) -> &SpecialTokens { + &self.special_tokens + } + + fn token_to_id(&self, _token: &str) -> Option { + None + } + + fn id_to_token(&self, _id: u32) -> Option { + None + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +impl ReasoningParser for FailingReasoningParser { + fn detect_and_parse_reasoning( + &mut self, + _text: &str, + ) -> Result { + Err(ParseError::ConfigError("fake failure".to_string())) + } + + fn parse_reasoning_streaming_incremental( + &mut self, + _text: &str, + ) -> Result { + Err(ParseError::ConfigError("fake failure".to_string())) + } + + fn reset(&mut self) {} + + fn model_type(&self) -> &str { + "failing" + } + + fn is_in_reasoning(&self) -> bool { + false + } + + fn mark_reasoning_started(&mut self) {} + + fn mark_think_start_stripped(&mut self) {} +} + +#[async_trait] +impl ToolParser for FailingToolParser { + async fn parse_complete( + &self, + _output: &str, + ) -> ToolParserResult<(String, Vec)> { + Err(ToolParserError::ParsingFailed("fake failure".to_string())) + } + + async fn parse_incremental( + &mut self, + _chunk: &str, + _tools: &[Tool], + ) -> ToolParserResult { + Err(ToolParserError::ParsingFailed("fake failure".to_string())) + } + + fn has_tool_markers(&self, _text: &str) -> bool { + false + } +} diff --git a/model_gateway/src/routers/grpc/utils/parsers.rs b/model_gateway/src/routers/grpc/utils/parsers.rs index 51ea413ed8..59e6a9afc5 100644 --- a/model_gateway/src/routers/grpc/utils/parsers.rs +++ b/model_gateway/src/routers/grpc/utils/parsers.rs @@ -57,6 +57,21 @@ impl ParserResolver { } } + /// Test-only resolver that carries configured parser names but skips + /// model-card lookups (no worker registry). Mirrors `disabled()` for + /// endpoints that rely solely on the gateway-configured parser names. + #[cfg(test)] + pub(crate) fn configured_only( + configured_tool_parser: Option, + configured_reasoning_parser: Option, + ) -> Self { + Self { + worker_registry: None, + configured_tool_parser, + configured_reasoning_parser, + } + } + /// Effective tool-parser name for `model`, if any. pub(crate) fn tool_parser(&self, model: &str) -> Option { self.card_parser(model, |card| card.tool_parser.as_ref())