[router][grpc] Fix streaming bugs: empty tool names, state pollution, and panics (#11373)

This commit is contained in:
Chang Su
2025-10-09 03:53:23 -07:00
committed by GitHub
parent a4b424c632
commit ab926dd697
33 changed files with 1145 additions and 534 deletions

View File

@@ -1,5 +1,5 @@
use crate::tool_parser::{
errors::{ToolParserError, ToolParserResult},
errors::{ParserError, ParserResult},
traits::PartialJsonParser,
};
use serde_json::{Map, Value};
@@ -22,8 +22,22 @@ impl PartialJson {
}
/// Parse potentially incomplete JSON, returning parsed value and consumed bytes
pub fn parse_value(&self, input: &str) -> ToolParserResult<(Value, usize)> {
let mut parser = Parser::new(input, self.max_depth, self.allow_incomplete);
///
/// # Arguments
/// * `input` - The JSON string to parse
/// * `allow_partial_strings` - When false, incomplete strings cause parsing to stop
/// (matches Python's Allow.ALL & ~Allow.STR behavior)
pub fn parse_value(
&self,
input: &str,
allow_partial_strings: bool,
) -> ParserResult<(Value, usize)> {
let mut parser = Parser::new(
input,
self.max_depth,
self.allow_incomplete,
allow_partial_strings,
);
let value = parser.parse_value(0)?;
Ok((value, parser.position))
}
@@ -36,8 +50,9 @@ impl Default for PartialJson {
}
impl PartialJsonParser for PartialJson {
fn parse(&self, input: &str) -> ToolParserResult<(Value, usize)> {
self.parse_value(input)
fn parse(&self, input: &str) -> ParserResult<(Value, usize)> {
// Default to allowing partial strings
self.parse_value(input, true)
}
fn is_complete(&self, input: &str) -> bool {
@@ -56,15 +71,22 @@ struct Parser<'a> {
position: usize,
max_depth: usize,
allow_incomplete: bool,
allow_partial_strings: bool,
}
impl<'a> Parser<'a> {
fn new(input: &'a str, max_depth: usize, allow_incomplete: bool) -> Self {
fn new(
input: &'a str,
max_depth: usize,
allow_incomplete: bool,
allow_partial_strings: bool,
) -> Self {
Self {
chars: input.chars().peekable(),
position: 0,
max_depth,
allow_incomplete,
allow_partial_strings,
}
}
@@ -88,9 +110,9 @@ impl<'a> Parser<'a> {
}
}
fn parse_value(&mut self, depth: usize) -> ToolParserResult<Value> {
fn parse_value(&mut self, depth: usize) -> ParserResult<Value> {
if depth > self.max_depth {
return Err(ToolParserError::DepthExceeded(self.max_depth));
return Err(ParserError::DepthExceeded(self.max_depth));
}
self.skip_whitespace();
@@ -106,17 +128,15 @@ impl<'a> Parser<'a> {
if self.allow_incomplete {
Ok(Value::Null)
} else {
Err(ToolParserError::ParsingFailed(
"Unexpected character".into(),
))
Err(ParserError::ParsingFailed("Unexpected character".into()))
}
}
}
}
fn parse_object(&mut self, depth: usize) -> ToolParserResult<Value> {
fn parse_object(&mut self, depth: usize) -> ParserResult<Value> {
if depth > self.max_depth {
return Err(ToolParserError::DepthExceeded(self.max_depth));
return Err(ParserError::DepthExceeded(self.max_depth));
}
let mut object = Map::new();
@@ -140,7 +160,7 @@ impl<'a> Parser<'a> {
return Ok(Value::Object(object));
}
Err(e) => return Err(e),
_ => return Err(ToolParserError::ParsingFailed("Expected string key".into())),
_ => return Err(ParserError::ParsingFailed("Expected string key".into())),
};
self.skip_whitespace();
@@ -152,7 +172,7 @@ impl<'a> Parser<'a> {
object.insert(key, Value::Null);
return Ok(Value::Object(object));
}
return Err(ToolParserError::ParsingFailed("Expected ':'".into()));
return Err(ParserError::ParsingFailed("Expected ':'".into()));
}
self.advance();
self.skip_whitespace();
@@ -161,8 +181,13 @@ impl<'a> Parser<'a> {
let value = match self.parse_value(depth) {
Ok(v) => v,
Err(_) if self.allow_incomplete => {
// Add null for incomplete value
object.insert(key, Value::Null);
// When allow_partial_strings is false, don't add the key with Null
// Just return the object without this incomplete key-value pair
// This matches Python's behavior: Allow.ALL & ~Allow.STR
if self.allow_partial_strings {
// Add null for incomplete value
object.insert(key, Value::Null);
}
return Ok(Value::Object(object));
}
Err(e) => return Err(e),
@@ -192,15 +217,15 @@ impl<'a> Parser<'a> {
if self.allow_incomplete {
return Ok(Value::Object(object));
}
return Err(ToolParserError::ParsingFailed("Expected ',' or '}'".into()));
return Err(ParserError::ParsingFailed("Expected ',' or '}'".into()));
}
}
}
}
fn parse_array(&mut self, depth: usize) -> ToolParserResult<Value> {
fn parse_array(&mut self, depth: usize) -> ParserResult<Value> {
if depth > self.max_depth {
return Err(ToolParserError::DepthExceeded(self.max_depth));
return Err(ParserError::DepthExceeded(self.max_depth));
}
let mut array = Vec::new();
@@ -249,15 +274,15 @@ impl<'a> Parser<'a> {
if self.allow_incomplete {
return Ok(Value::Array(array));
}
return Err(ToolParserError::ParsingFailed("Expected ',' or ']'".into()));
return Err(ParserError::ParsingFailed("Expected ',' or ']'".into()));
}
}
}
}
fn parse_string(&mut self) -> ToolParserResult<Value> {
fn parse_string(&mut self) -> ParserResult<Value> {
if self.peek() != Some('"') {
return Err(ToolParserError::ParsingFailed("Expected '\"'".into()));
return Err(ParserError::ParsingFailed("Expected '\"'".into()));
}
// Consume opening quote
@@ -301,14 +326,14 @@ impl<'a> Parser<'a> {
}
// Incomplete string
if self.allow_incomplete {
if self.allow_incomplete && self.allow_partial_strings {
Ok(Value::String(string))
} else {
Err(ToolParserError::ParsingFailed("Unterminated string".into()))
Err(ParserError::ParsingFailed("Unterminated string".into()))
}
}
fn parse_unicode_escape(&mut self) -> ToolParserResult<char> {
fn parse_unicode_escape(&mut self) -> ParserResult<char> {
let mut hex = String::new();
for _ in 0..4 {
if let Some(ch) = self.peek() {
@@ -327,17 +352,17 @@ impl<'a> Parser<'a> {
u32::from_str_radix(&hex, 16)
.ok()
.and_then(char::from_u32)
.ok_or_else(|| ToolParserError::ParsingFailed("Invalid unicode escape".into()))
.ok_or_else(|| ParserError::ParsingFailed("Invalid unicode escape".into()))
} else if self.allow_incomplete {
Ok('\u{FFFD}') // Replacement character
} else {
Err(ToolParserError::ParsingFailed(
Err(ParserError::ParsingFailed(
"Incomplete unicode escape".into(),
))
}
}
fn parse_number(&mut self) -> ToolParserResult<Value> {
fn parse_number(&mut self) -> ParserResult<Value> {
let mut number = String::new();
// Handle negative sign
@@ -410,11 +435,11 @@ impl<'a> Parser<'a> {
} else if self.allow_incomplete {
Ok(Value::Number(serde_json::Number::from(0)))
} else {
Err(ToolParserError::ParsingFailed("Invalid number".into()))
Err(ParserError::ParsingFailed("Invalid number".into()))
}
}
fn parse_bool(&mut self) -> ToolParserResult<Value> {
fn parse_bool(&mut self) -> ParserResult<Value> {
let mut word = String::new();
// Peek at upcoming characters to validate it looks like a boolean
@@ -435,7 +460,7 @@ impl<'a> Parser<'a> {
|| (self.allow_incomplete && ("true".starts_with(&word) || "false".starts_with(&word)));
if !is_valid {
return Err(ToolParserError::ParsingFailed("Invalid boolean".into()));
return Err(ParserError::ParsingFailed("Invalid boolean".into()));
}
// Now actually consume the characters
@@ -458,14 +483,14 @@ impl<'a> Parser<'a> {
} else if "false".starts_with(partial) {
Ok(Value::Bool(false))
} else {
Err(ToolParserError::ParsingFailed("Invalid boolean".into()))
Err(ParserError::ParsingFailed("Invalid boolean".into()))
}
}
_ => Err(ToolParserError::ParsingFailed("Invalid boolean".into())),
_ => Err(ParserError::ParsingFailed("Invalid boolean".into())),
}
}
fn parse_null(&mut self) -> ToolParserResult<Value> {
fn parse_null(&mut self) -> ParserResult<Value> {
let mut word = String::new();
// Peek at upcoming characters to validate it looks like "null"
@@ -484,7 +509,7 @@ impl<'a> Parser<'a> {
let is_valid = word == "null" || (self.allow_incomplete && "null".starts_with(&word));
if !is_valid {
return Err(ToolParserError::ParsingFailed("Invalid null".into()));
return Err(ParserError::ParsingFailed("Invalid null".into()));
}
// Now actually consume the characters
@@ -501,7 +526,7 @@ impl<'a> Parser<'a> {
if word == "null" || (self.allow_incomplete && "null".starts_with(&word)) {
Ok(Value::Null)
} else {
Err(ToolParserError::ParsingFailed("Invalid null".into()))
Err(ParserError::ParsingFailed("Invalid null".into()))
}
}
}