[Rust Frontend] Make Granite4 string argument scanning incremental (#46507)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jun 26, 2026 at 11:54 UTC
e312c5cb25427e76fc3830ab14e7b6bc0963a55c
4 files changed
+256
-33
rust/src/parser/Cargo.toml
+5
@@ -74,5 +74,10 @@ name = "gemma4"
74
harness = false
75
required-features = ["test-util"]
76
77
+[[bench]]
78
+name = "granite4"
79
+harness = false
80
+required-features = ["test-util"]
81
+
82
[lints]
83
workspace = true
rust/src/parser/benches/granite4.rs
new
+75
@@ -0,0 +1,75 @@
1
+use std::time::Duration;
2
+
3
+use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
4
+use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
5
+use vllm_parser::tool::{Granite4ToolParser, Tool, ToolParser};
6
+
7
+mod utils;
8
+use utils::feed_parser;
9
+
10
+const CHUNK_CHARS: usize = 7;
11
+const LONG_ARGUMENT_BYTES: usize = 64 * 1024;
12
+
13
+fn string_args_fixture() -> String {
14
+ let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(LONG_ARGUMENT_BYTES));
15
+ let encoded_arguments = serde_json::to_string(&arguments).unwrap();
16
+ format!(r#"<tool_call>{{"name":"f","arguments":{encoded_arguments}}}</tool_call>"#)
17
+}
18
+
19
+fn object_args_fixture() -> String {
20
+ format!(
21
+ r#"<tool_call>{{"name":"f","arguments":{{"data":"{}"}}}}</tool_call>"#,
22
+ "x".repeat(LONG_ARGUMENT_BYTES)
23
+ )
24
+}
25
+
26
+fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
27
+ Granite4ToolParser::create(tools).expect("Granite4 parser should initialize")
28
+}
29
+
30
+fn run_stream_group(c: &mut Criterion, name: &str, tools: &[Tool], text: &str) {
31
+ let chunks = split_by_chars(text, CHUNK_CHARS);
32
+
33
+ let mut group = c.benchmark_group(name);
34
+ group.sample_size(50);
35
+ group.warm_up_time(Duration::from_millis(300));
36
+ group.measurement_time(Duration::from_secs(2));
37
+ group.throughput(Throughput::Bytes(text.len() as u64));
38
+
39
+ group.bench_function("reuse_parser", |b| {
40
+ let mut parser = parser(tools);
41
+ b.iter(|| {
42
+ let result = feed_parser(&mut *parser, black_box(&chunks));
43
+ debug_assert_eq!(result.0, "");
44
+ debug_assert_eq!(result.1, 1);
45
+ black_box(result);
46
+ })
47
+ });
48
+
49
+ group.bench_function("create_parser", |b| {
50
+ b.iter_batched(
51
+ || parser(tools),
52
+ |mut parser| {
53
+ let result = feed_parser(&mut *parser, black_box(&chunks));
54
+ debug_assert_eq!(result.0, "");
55
+ debug_assert_eq!(result.1, 1);
56
+ black_box(result);
57
+ },
58
+ BatchSize::SmallInput,
59
+ )
60
+ });
61
+
62
+ group.finish();
63
+}
64
+
65
+fn bench_granite4(c: &mut Criterion) {
66
+ let tools = test_tools();
67
+ let string_args = string_args_fixture();
68
+ let object_args = object_args_fixture();
69
+
70
+ run_stream_group(c, "granite4/long_string_arguments", &tools, &string_args);
71
+ run_stream_group(c, "granite4/long_object_arguments", &tools, &object_args);
72
+}
73
+
74
+criterion_group!(benches, bench_granite4);
75
+criterion_main!(benches);
rust/src/parser/src/tool/json/granite4.rs
+59
-16
@@ -9,7 +9,8 @@ use super::{
9
tool_call_header_event,
10
};
11
use crate::tool::utils::{
12
- JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object,
12
+ JsonObjectScanState, JsonStringScanState, decode_json_str, parse_buffered_event, safe_text_len,
13
+ take_json_object, take_json_string,
14
};
15
use crate::tool::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
16
@@ -22,14 +23,20 @@ enum Granite4Mode {
23
Header,
24
/// Parsing the arguments value:
25
/// `None` until the first byte decides object vs string;
25
- /// `Some` while streaming an object value.
26
+ /// `Some` while streaming the selected value shape.
27
Args {
27
- json_scan: Option<JsonObjectScanState>,
28
+ args_scan: Option<Granite4ArgsScan>,
29
},
30
/// Arguments done; consume the object's closing `}` and `</tool_call>`.
31
Close,
32
}
33
34
+#[derive(Debug, Clone, PartialEq, Eq)]
35
+enum Granite4ArgsScan {
36
+ Object(JsonObjectScanState),
37
+ String(JsonStringScanState),
38
+}
39
+
40
#[derive(Debug, Clone, PartialEq, Eq)]
41
enum Granite4Event {
42
Text {
@@ -92,7 +99,7 @@ impl Granite4ToolParser {
99
let tool_index = self.emitted_tool_count;
100
self.emitted_tool_count += 1;
101
self.active_tool_index = Some(tool_index);
95
- self.mode = Granite4Mode::Args { json_scan: None };
102
+ self.mode = Granite4Mode::Args { args_scan: None };
103
output.push_call(ToolCallDelta {
104
tool_index,
105
name: Some(function_name),
@@ -187,7 +194,7 @@ fn parse_next_granite4_event(
194
match mode {
195
Granite4Mode::Text => text_event(input),
196
Granite4Mode::Header => header_event(input),
190
- Granite4Mode::Args { json_scan } => args_event(input, json_scan),
197
+ Granite4Mode::Args { args_scan } => args_event(input, args_scan),
198
Granite4Mode::Close => close_event(input),
199
}
200
}
@@ -237,14 +244,19 @@ fn header_event(input: &mut JsonToolInput<'_>) -> ModalResult<Granite4Event> {
244
/// once seen whole and unescaped.
245
fn args_event(
246
input: &mut JsonToolInput<'_>,
240
- json_scan: &mut Option<JsonObjectScanState>,
247
+ args_scan: &mut Option<Granite4ArgsScan>,
248
) -> ModalResult<Granite4Event> {
242
- if let Some(scan) = json_scan {
243
- let len = take_json_object(input, scan)?;
244
- return Ok(Granite4Event::ObjectArgsDelta {
245
- len,
246
- complete: scan.complete(),
247
- });
249
+ if let Some(scan) = args_scan {
250
+ return match scan {
251
+ Granite4ArgsScan::Object(scan) => {
252
+ let len = take_json_object(input, scan)?;
253
+ Ok(Granite4Event::ObjectArgsDelta {
254
+ len,
255
+ complete: scan.complete(),
256
+ })
257
+ }
258
+ Granite4ArgsScan::String(scan) => string_args_event(input, scan),
259
+ };
260
}
261
262
match peek(any).parse_next(input)? {
@@ -252,12 +264,16 @@ fn args_event(
264
let mut scan = JsonObjectScanState::default();
265
let len = take_json_object(input, &mut scan)?;
266
let complete = scan.complete();
255
- *json_scan = Some(scan);
267
+ *args_scan = Some(Granite4ArgsScan::Object(scan));
268
Ok(Granite4Event::ObjectArgsDelta { len, complete })
269
}
258
- '"' => Ok(Granite4Event::StringArgs {
259
- decoded: json_str(input)?,
260
- }),
270
+ '"' => {
271
+ *args_scan = Some(Granite4ArgsScan::String(JsonStringScanState::default()));
272
+ let Some(Granite4ArgsScan::String(scan)) = args_scan else {
273
+ unreachable!("Granite4 string scan state was just initialized")
274
+ };
275
+ string_args_event(input, scan)
276
+ }
277
_ => {
278
let mut error = ContextError::new();
279
error.push(StrContext::Label("Granite4 arguments"));
@@ -266,6 +282,17 @@ fn args_event(
282
}
283
}
284
285
+fn string_args_event(
286
+ input: &mut JsonToolInput<'_>,
287
+ scan: &mut JsonStringScanState,
288
+) -> ModalResult<Granite4Event> {
289
+ let text = **input;
290
+ let len = take_json_string(input, scan)?;
291
+ Ok(Granite4Event::StringArgs {
292
+ decoded: decode_json_str(&text[..len])?,
293
+ })
294
+}
295
+
296
/// Parse the tool-call object's closing `}` and the `</tool_call>` end marker.
297
fn close_event(input: &mut JsonToolInput<'_>) -> ModalResult<Granite4Event> {
298
seq!(_: ws0, _: literal("}"), _: ws0, _: literal(TOOL_CALL_END))
@@ -420,6 +447,22 @@ mod tests {
447
assert_eq!(output.calls()[0].arguments, r#"{"a":1}"#);
448
}
449
450
+ #[test]
451
+ fn granite4_long_string_args_stream_without_reparse() {
452
+ let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(64 * 1024));
453
+ let encoded_arguments = serde_json::to_string(&arguments).unwrap();
454
+ let input =
455
+ format!(r#"<tool_call>{{"name":"f","arguments":{encoded_arguments}}}</tool_call>"#);
456
+ let chunks = split_by_chars(&input, 7);
457
+ let mut parser = Granite4ToolParser::new(&test_tools());
458
+
459
+ let output = collect_stream(&mut parser, &chunks);
460
+
461
+ assert_eq!(output.calls().len(), 1);
462
+ assert_eq!(output.calls()[0].name.as_deref(), Some("f"));
463
+ assert_eq!(output.calls()[0].arguments, arguments);
464
+ }
465
+
466
#[test]
467
fn granite4_streaming_handles_marker_and_json_whitespace() {
468
// Granite spaces the markers (`<tool_call> {…} </tool_call>`) and the JSON
rust/src/parser/src/utils.rs
+117
-17
@@ -290,8 +290,22 @@ pub fn take_json_object(
290
Ok(text.len())
291
}
292
293
-/// Parse a JSON string literal.
294
-pub fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
293
+/// Streaming lexical state for a JSON string literal.
294
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
295
+pub struct JsonStringScanState {
296
+ scanned_len: usize,
297
+ escape: bool,
298
+}
299
+
300
+/// Parse a raw JSON string literal, resuming from the last scanned byte.
301
+///
302
+/// The returned length covers the quoted JSON string. This only scans for the
303
+/// string boundary; callers that need the decoded value should pass the raw
304
+/// slice to [`decode_json_str`].
305
+pub fn take_json_string(
306
+ input: &mut Partial<&str>,
307
+ state: &mut JsonStringScanState,
308
+) -> ModalResult<usize> {
309
let text = **input;
310
if text.is_empty() {
311
return incomplete();
@@ -305,37 +319,58 @@ pub fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
319
));
320
}
321
308
- let mut escape = false;
309
- let mut index = 1;
322
+ let mut index = if state.scanned_len == 0 {
323
+ 1
324
+ } else if state.scanned_len <= bytes.len() {
325
+ state.scanned_len
326
+ } else {
327
+ return incomplete();
328
+ };
329
+
330
while index < bytes.len() {
331
let byte = bytes[index];
332
index += 1;
333
314
- if escape {
315
- escape = false;
334
+ if state.escape {
335
+ state.escape = false;
336
continue;
337
}
338
339
match byte {
320
- b'\\' => escape = true,
340
+ b'\\' => state.escape = true,
341
b'"' => {
322
- let raw = &text[..index];
323
- let value = serde_json::from_str::<String>(raw).map_err(|_| {
324
- json_scan_error(
325
- "JSON string",
326
- StrContextValue::Description("valid JSON string"),
327
- )
328
- })?;
342
input.next_slice(index);
330
- return Ok(value);
343
+ return Ok(index);
344
}
345
_ => {}
346
}
347
}
348
349
+ state.scanned_len = text.len();
350
incomplete()
351
}
352
353
+/// Parse a JSON string literal.
354
+pub fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
355
+ let text = **input;
356
+ let checkpoint = input.checkpoint();
357
+ let mut state = JsonStringScanState::default();
358
+ let len = take_json_string(input, &mut state)?;
359
+ decode_json_str(&text[..len]).inspect_err(|_| {
360
+ input.reset(&checkpoint);
361
+ })
362
+}
363
+
364
+/// Decode a complete JSON string literal.
365
+pub fn decode_json_str(raw: &str) -> ModalResult<String> {
366
+ serde_json::from_str::<String>(raw).map_err(|_| {
367
+ json_scan_error(
368
+ "JSON string",
369
+ StrContextValue::Description("valid JSON string"),
370
+ )
371
+ })
372
+}
373
+
374
fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode<ContextError> {
375
let mut error = ContextError::new();
376
error.push(StrContext::Label(label));
@@ -388,8 +423,8 @@ mod tests {
423
use winnow::stream::{Offset, Partial, Stream};
424
425
use super::{
391
- JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len,
392
- safe_text_len_mul, take_json_object, take_until_marker,
426
+ JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, partial_prefix_len,
427
+ safe_text_len, safe_text_len_mul, take_json_object, take_json_string, take_until_marker,
428
};
429
430
#[test]
@@ -714,6 +749,71 @@ mod tests {
749
.assert_eq(&error.to_string());
750
}
751
752
+ #[test]
753
+ fn take_json_string_consumes_complete_string() {
754
+ let mut state = JsonStringScanState::default();
755
+ let mut input = Partial::new(r#""say_\"hi\u0021" rest"#);
756
+ let checkpoint = input.checkpoint();
757
+
758
+ let len = take_json_string(&mut input, &mut state).unwrap();
759
+
760
+ assert_eq!(len, r#""say_\"hi\u0021""#.len());
761
+ assert_eq!(input.offset_from(&checkpoint), len);
762
+ assert_eq!(*input, " rest");
763
+ }
764
+
765
+ #[test]
766
+ fn take_json_string_resumes_after_incomplete_input() {
767
+ let mut state = JsonStringScanState::default();
768
+ let mut input = Partial::new(r#""{\"data\":\"partial"#);
769
+ let checkpoint = input.checkpoint();
770
+
771
+ let error = take_json_string(&mut input, &mut state).unwrap_err();
772
+
773
+ assert!(matches!(error, ErrMode::Incomplete(_)));
774
+ assert_eq!(input.offset_from(&checkpoint), 0);
775
+ assert_eq!(state.scanned_len, r#""{\"data\":\"partial"#.len());
776
+
777
+ let mut input = Partial::new(r#""{\"data\":\"partial string\"}" tail"#);
778
+ let len = take_json_string(&mut input, &mut state).unwrap();
779
+
780
+ assert_eq!(len, r#""{\"data\":\"partial string\"}""#.len());
781
+ assert_eq!(*input, " tail");
782
+ }
783
+
784
+ #[test]
785
+ fn take_json_string_tracks_escape_across_chunks() {
786
+ let mut state = JsonStringScanState::default();
787
+ let mut input = Partial::new(r#""abc\"#);
788
+
789
+ let error = take_json_string(&mut input, &mut state).unwrap_err();
790
+
791
+ assert!(matches!(error, ErrMode::Incomplete(_)));
792
+ assert!(state.escape);
793
+
794
+ let mut input = Partial::new(r#""abc\"def" tail"#);
795
+ let len = take_json_string(&mut input, &mut state).unwrap();
796
+
797
+ assert_eq!(len, r#""abc\"def""#.len());
798
+ assert_eq!(*input, " tail");
799
+ }
800
+
801
+ #[test]
802
+ fn take_json_string_rejects_non_string_start() {
803
+ let mut state = JsonStringScanState::default();
804
+ let mut input = Partial::new("42");
805
+
806
+ let error = take_json_string(&mut input, &mut state).unwrap_err();
807
+
808
+ let ErrMode::Cut(error) = error else {
809
+ panic!("expected cut error");
810
+ };
811
+ expect![[r#"
812
+ invalid JSON string
813
+ expected `"`"#]]
814
+ .assert_eq(&error.to_string());
815
+ }
816
+
817
#[test]
818
fn json_str_decodes_escaped_content() {
819
let mut input = Partial::new(r#""say_\"hi\u0021" rest"#);