@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e1a5fc406b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Rust Frontend][Perf] O(n) argument scan in tool parser (#45826)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 18, 2026 at 09:42 UTC e1a5fc406b248925427a84beefbd123924f0c9e5
8 files changed +408 -72
rust/src/tool-parser/benches/qwen3_coder.rs
+23
@@ -10,6 +10,7 @@ use utils::{feed_external_parser, feed_parser, openai_tools};
10
11 const CHUNK_CHARS: usize = 7;
12 const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
13 +const LONG_TOOL_BODY_REPEATS: usize = 8192;
14
15 fn mixed_fixture() -> String {
16 concat!(
@@ -39,6 +40,17 @@ fn long_normal_text_fixture() -> String {
40 line.repeat(LONG_NORMAL_TEXT_REPEATS)
41 }
42
43 +fn long_tool_call_fixture() -> String {
44 + let location = "x".repeat(LONG_TOOL_BODY_REPEATS);
45 + format!(
46 + "<tool_call>\n\
47 + <function=get_weather>\n\
48 + <parameter=location>{location}</parameter>\n\
49 + </function>\n\
50 + </tool_call>"
51 + )
52 +}
53 +
54 fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
55 Qwen3CoderToolParser::create(tools).expect("Qwen Coder parser should initialize")
56 }
@@ -112,6 +124,7 @@ fn bench_qwen3_coder(c: &mut Criterion) {
124 let tools = test_tools();
125 let mixed_text = mixed_fixture();
126 let long_normal_text = long_normal_text_fixture();
127 + let long_tool_call = long_tool_call_fixture();
128
129 run_stream_group(
130 c,
@@ -132,6 +145,16 @@ fn bench_qwen3_coder(c: &mut Criterion) {
145 &long_normal_text,
146 0,
147 );
148 +
149 + run_stream_group(
150 + c,
151 + "qwen3_coder/long_tool_call_body",
152 + &tools,
153 + &long_tool_call,
154 + CHUNK_CHARS,
155 + "",
156 + 1,
157 + );
158 }
159
160 criterion_group!(benches, bench_qwen3_coder);
rust/src/tool-parser/src/deepseek_dsml/mod.rs
+28 -14
@@ -5,7 +5,7 @@ use winnow::stream::Partial;
5 use winnow::token::{literal, rest, take_until};
6
7 use super::parameters::ToolSchemas;
8 -use super::utils::{parse_buffered_event, safe_text_len};
8 +use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
9 use super::{Result, ToolCallDelta, ToolParserOutput};
10 use crate::Tool;
11
@@ -39,10 +39,10 @@ impl DsmlTokens {
39 };
40 }
41
42 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42 +#[derive(Debug, Clone, PartialEq, Eq)]
43 enum DsmlMode {
44 Text,
45 - ToolBlock,
45 + ToolBlock { invoke_end_scan: MarkerScanState },
46 Done,
47 }
48
@@ -94,7 +94,11 @@ impl DeepSeekDsmlToolParser {
94 DsmlEvent::Text { len: consumed_len } => {
95 output.normal_text.push_str(&self.buffer[..consumed_len]);
96 }
97 - DsmlEvent::ToolCallsStart => self.mode = DsmlMode::ToolBlock,
97 + DsmlEvent::ToolCallsStart => {
98 + self.mode = DsmlMode::ToolBlock {
99 + invoke_end_scan: MarkerScanState::default(),
100 + };
101 + }
102 DsmlEvent::Invoke { name, raw_params } => {
103 let mut arguments = serde_json::Map::with_capacity(raw_params.len());
104 for param in raw_params {
@@ -140,7 +144,7 @@ impl DeepSeekDsmlToolParser {
144 self.buffer.push_str(chunk);
145
146 while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
143 - parse_next_dsml_event(input, self.mode, self.tokens)
147 + parse_next_dsml_event(input, &mut self.mode, self.tokens)
148 })? {
149 self.apply_event(event, output)?;
150 self.buffer.drain(..consumed_len);
@@ -154,7 +158,7 @@ impl DeepSeekDsmlToolParser {
158 match self.mode {
159 DsmlMode::Text => output.normal_text.push_str(&self.buffer),
160 DsmlMode::Done => {}
157 - DsmlMode::ToolBlock => {
161 + DsmlMode::ToolBlock { .. } => {
162 return Err(parsing_failed!("incomplete DeepSeek DSML tool call"));
163 }
164 }
@@ -166,12 +170,14 @@ impl DeepSeekDsmlToolParser {
170 /// Parse a DSML event for the current parser mode.
171 fn parse_next_dsml_event(
172 input: &mut DsmlInput<'_>,
169 - mode: DsmlMode,
173 + mode: &mut DsmlMode,
174 tokens: DsmlTokens,
175 ) -> ModalResult<DsmlEvent> {
176 match mode {
177 DsmlMode::Text => parse_text_event(input, tokens),
174 - DsmlMode::ToolBlock => parse_tool_block_event(input, tokens),
178 + DsmlMode::ToolBlock { invoke_end_scan } => {
179 + parse_tool_block_event(input, tokens, invoke_end_scan)
180 + }
181 DsmlMode::Done => ignored_rest_event(input),
182 }
183 }
@@ -186,11 +192,16 @@ fn parse_text_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResul
192 }
193
194 /// Parse a tool-block DSML event.
189 -fn parse_tool_block_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResult<DsmlEvent> {
195 +fn parse_tool_block_event(
196 + input: &mut DsmlInput<'_>,
197 + tokens: DsmlTokens,
198 + invoke_end_scan: &mut MarkerScanState,
199 +) -> ModalResult<DsmlEvent> {
200 ws0.void().parse_next(input)?;
191 - alt((invoke_event, |input: &mut DsmlInput<'_>| {
192 - tool_calls_end_event(input, tokens)
193 - }))
201 + alt((
202 + |input: &mut DsmlInput<'_>| invoke_event(input, invoke_end_scan),
203 + |input: &mut DsmlInput<'_>| tool_calls_end_event(input, tokens),
204 + ))
205 .parse_next(input)
206 }
207
@@ -217,14 +228,17 @@ fn safe_text_event(input: &mut DsmlInput<'_>, tokens: DsmlTokens) -> ModalResult
228 }
229
230 /// Parse a DSML invoke block.
220 -fn invoke_event(input: &mut DsmlInput<'_>) -> ModalResult<DsmlEvent> {
231 +fn invoke_event(
232 + input: &mut DsmlInput<'_>,
233 + invoke_end_scan: &mut MarkerScanState,
234 +) -> ModalResult<DsmlEvent> {
235 let (name, body) = seq!(
236 _: literal(INVOKE_START),
237 _: ws1,
238 dsml_name_attr,
239 _: ws0,
240 _: ">",
227 - take_until(0.., INVOKE_END),
241 + take_until_marker(INVOKE_END, invoke_end_scan),
242 _: literal(INVOKE_END),
243 )
244 .parse_next(input)?;
rust/src/tool-parser/src/glm_xml/mod.rs
+22 -10
@@ -5,7 +5,7 @@ use winnow::stream::Partial;
5 use winnow::token::{literal, rest, take_until, take_while};
6
7 use super::parameters::ToolSchemas;
8 -use super::utils::{parse_buffered_event, safe_text_len};
8 +use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
9 use super::{Result, ToolCallDelta, ToolParserOutput};
10 use crate::Tool;
11
@@ -24,10 +24,10 @@ const ARG_VALUE_END: &str = "</arg_value>";
24
25 type GlmInput<'i> = Partial<&'i str>;
26
27 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27 +#[derive(Debug, Clone, PartialEq, Eq)]
28 enum GlmMode {
29 Text,
30 - ToolCall,
30 + ToolCall { tool_call_end_scan: MarkerScanState },
31 AfterToolCall,
32 }
33
@@ -81,7 +81,11 @@ impl GlmXmlToolParser {
81 GlmEvent::Text { len: consumed_len } => {
82 output.normal_text.push_str(&self.buffer[..consumed_len]);
83 }
84 - GlmEvent::ToolCallStart => self.mode = GlmMode::ToolCall,
84 + GlmEvent::ToolCallStart => {
85 + self.mode = GlmMode::ToolCall {
86 + tool_call_end_scan: MarkerScanState::default(),
87 + };
88 + }
89 GlmEvent::ToolCall { name, raw_params } => {
90 self.mode = GlmMode::AfterToolCall;
91 let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
@@ -110,7 +114,7 @@ impl GlmXmlToolParser {
114 self.buffer.push_str(chunk);
115
116 while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
113 - parse_next_glm_event(input, self.mode, self.separator)
117 + parse_next_glm_event(input, &mut self.mode, self.separator)
118 })? {
119 self.apply_event(event, output)?;
120 self.buffer.drain(..consumed_len);
@@ -124,7 +128,9 @@ impl GlmXmlToolParser {
128 if !self.buffer.is_empty() {
129 match self.mode {
130 GlmMode::Text => output.normal_text.push_str(&self.buffer),
127 - GlmMode::ToolCall => return Err(parsing_failed!("incomplete GLM MoE tool call")),
131 + GlmMode::ToolCall { .. } => {
132 + return Err(parsing_failed!("incomplete GLM MoE tool call"));
133 + }
134 GlmMode::AfterToolCall => {}
135 }
136 }
@@ -136,12 +142,14 @@ impl GlmXmlToolParser {
142 /// Parse a GLM event for the current parser mode.
143 fn parse_next_glm_event(
144 input: &mut GlmInput<'_>,
139 - mode: GlmMode,
145 + mode: &mut GlmMode,
146 separator: Separator,
147 ) -> ModalResult<GlmEvent> {
148 match mode {
149 GlmMode::Text => parse_text_event(input),
144 - GlmMode::ToolCall => tool_call_event(input, separator),
150 + GlmMode::ToolCall { tool_call_end_scan } => {
151 + tool_call_event(input, separator, tool_call_end_scan)
152 + }
153 GlmMode::AfterToolCall => after_tool_call_event(input),
154 }
155 }
@@ -173,9 +181,13 @@ fn ignored_rest_event(input: &mut GlmInput<'_>) -> ModalResult<GlmEvent> {
181 }
182
183 /// Parse a complete GLM tool call.
176 -fn tool_call_event(input: &mut GlmInput<'_>, separator: Separator) -> ModalResult<GlmEvent> {
184 +fn tool_call_event(
185 + input: &mut GlmInput<'_>,
186 + separator: Separator,
187 + tool_call_end_scan: &mut MarkerScanState,
188 +) -> ModalResult<GlmEvent> {
189 let (body,) = seq!(
178 - take_until(0.., TOOL_CALL_END),
190 + take_until_marker(TOOL_CALL_END, tool_call_end_scan),
191 _: literal(TOOL_CALL_END),
192 )
193 .parse_next(input)?;
rust/src/tool-parser/src/hy_v3.rs
+30 -12
@@ -5,7 +5,7 @@ use winnow::stream::Partial;
5 use winnow::token::{literal, rest, take_until};
6
7 use super::parameters::ToolSchemas;
8 -use super::utils::{parse_buffered_event, safe_text_len};
8 +use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
9 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
10 use crate::Tool;
11
@@ -21,10 +21,10 @@ const ARG_VALUE_END: &str = "</arg_value>";
21
22 type HyV3Input<'i> = Partial<&'i str>;
23
24 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24 +#[derive(Debug, Clone, PartialEq, Eq)]
25 enum HyV3Mode {
26 Text,
27 - ToolBlock,
27 + ToolBlock { tool_call_end_scan: MarkerScanState },
28 Done,
29 }
30
@@ -81,7 +81,11 @@ impl HyV3ToolParser {
81 HyV3Event::Text { len: consumed_len } => {
82 output.normal_text.push_str(&self.buffer[..consumed_len]);
83 }
84 - HyV3Event::ToolBlockStart => self.mode = HyV3Mode::ToolBlock,
84 + HyV3Event::ToolBlockStart => {
85 + self.mode = HyV3Mode::ToolBlock {
86 + tool_call_end_scan: MarkerScanState::default(),
87 + };
88 + }
89 HyV3Event::ToolCall { name, raw_params } => {
90 let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
91 let arguments = serde_json::to_string(&arguments)
@@ -113,7 +117,7 @@ impl ToolParser for HyV3ToolParser {
117 self.buffer.push_str(chunk);
118
119 while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
116 - parse_next_hy_v3_event(input, self.mode)
120 + parse_next_hy_v3_event(input, &mut self.mode)
121 })? {
122 self.apply_event(event, output)?;
123 self.buffer.drain(..consumed_len);
@@ -126,7 +130,7 @@ impl ToolParser for HyV3ToolParser {
130 let mut output = ToolParserOutput::default();
131 match self.mode {
132 HyV3Mode::Text => output.normal_text.push_str(&self.buffer),
129 - HyV3Mode::ToolBlock => return Err(parsing_failed!("incomplete HY3 tool call")),
133 + HyV3Mode::ToolBlock { .. } => return Err(parsing_failed!("incomplete HY3 tool call")),
134 HyV3Mode::Done => {}
135 }
136 let _ = self.reset();
@@ -141,10 +145,15 @@ impl ToolParser for HyV3ToolParser {
145 }
146
147 /// Parse a HY3 event for the current parser mode.
144 -fn parse_next_hy_v3_event(input: &mut HyV3Input<'_>, mode: HyV3Mode) -> ModalResult<HyV3Event> {
148 +fn parse_next_hy_v3_event(
149 + input: &mut HyV3Input<'_>,
150 + mode: &mut HyV3Mode,
151 +) -> ModalResult<HyV3Event> {
152 match mode {
153 HyV3Mode::Text => parse_text_event(input),
147 - HyV3Mode::ToolBlock => parse_tool_block_event(input),
154 + HyV3Mode::ToolBlock { tool_call_end_scan } => {
155 + parse_tool_block_event(input, tool_call_end_scan)
156 + }
157 HyV3Mode::Done => ignored_rest_event(input),
158 }
159 }
@@ -165,8 +174,14 @@ fn safe_text_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
174 }
175
176 /// Parse one event inside a HY3 tool block.
168 -fn parse_tool_block_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
169 - alt((tool_block_end_event, tool_call_event)).parse_next(input)
177 +fn parse_tool_block_event(
178 + input: &mut HyV3Input<'_>,
179 + tool_call_end_scan: &mut MarkerScanState,
180 +) -> ModalResult<HyV3Event> {
181 + alt((tool_block_end_event, |input: &mut HyV3Input<'_>| {
182 + tool_call_event(input, tool_call_end_scan)
183 + }))
184 + .parse_next(input)
185 }
186
187 /// Parse a HY3 tool-block end marker.
@@ -175,13 +190,16 @@ fn tool_block_end_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
190 }
191
192 /// Parse a complete HY3 tool-call block.
178 -fn tool_call_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
193 +fn tool_call_event(
194 + input: &mut HyV3Input<'_>,
195 + tool_call_end_scan: &mut MarkerScanState,
196 +) -> ModalResult<HyV3Event> {
197 let (name, body) = seq!(
198 _: ws0,
199 _: literal(TOOL_CALL_START),
200 take_until(0.., TOOL_SEP),
201 _: literal(TOOL_SEP),
184 - take_until(0.., TOOL_CALL_END),
202 + take_until_marker(TOOL_CALL_END, tool_call_end_scan),
203 _: literal(TOOL_CALL_END),
204 )
205 .parse_next(input)?;
rust/src/tool-parser/src/minimax_m2.rs
+27 -12
@@ -5,7 +5,7 @@ use winnow::stream::Partial;
5 use winnow::token::{literal, rest, take_until};
6
7 use super::parameters::ToolSchemas;
8 -use super::utils::{parse_buffered_event, safe_text_len};
8 +use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
9 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
10 use crate::Tool;
11
@@ -18,10 +18,10 @@ const PARAMETER_END: &str = "</parameter>";
18
19 type MinimaxM2Input<'i> = Partial<&'i str>;
20
21 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21 +#[derive(Debug, Clone, PartialEq, Eq)]
22 enum MinimaxM2Mode {
23 Text,
24 - ToolBlock,
24 + ToolBlock { invoke_end_scan: MarkerScanState },
25 Done,
26 }
27
@@ -74,7 +74,11 @@ impl MinimaxM2ToolParser {
74 MinimaxM2Event::Text { len: consumed_len } => {
75 output.normal_text.push_str(&self.buffer[..consumed_len]);
76 }
77 - MinimaxM2Event::ToolBlockStart => self.mode = MinimaxM2Mode::ToolBlock,
77 + MinimaxM2Event::ToolBlockStart => {
78 + self.mode = MinimaxM2Mode::ToolBlock {
79 + invoke_end_scan: MarkerScanState::default(),
80 + };
81 + }
82 MinimaxM2Event::Invoke { name, raw_params } => {
83 let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
84 let arguments = serde_json::to_string(&arguments)
@@ -112,7 +116,7 @@ impl ToolParser for MinimaxM2ToolParser {
116 self.buffer.push_str(chunk);
117
118 while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
115 - parse_next_minimax_m2_event(input, self.mode)
119 + parse_next_minimax_m2_event(input, &mut self.mode)
120 })? {
121 self.apply_event(event, output)?;
122 self.buffer.drain(..consumed_len);
@@ -127,7 +131,7 @@ impl ToolParser for MinimaxM2ToolParser {
131 MinimaxM2Mode::Text => {
132 output.normal_text.push_str(&self.buffer);
133 }
130 - MinimaxM2Mode::ToolBlock => {
134 + MinimaxM2Mode::ToolBlock { .. } => {
135 return Err(parsing_failed!("incomplete MiniMax M2 tool call"));
136 }
137 MinimaxM2Mode::Done => {}
@@ -144,11 +148,13 @@ impl ToolParser for MinimaxM2ToolParser {
148 /// Parse a MiniMax M2 event for the current parser mode.
149 fn parse_next_minimax_m2_event(
150 input: &mut MinimaxM2Input<'_>,
147 - mode: MinimaxM2Mode,
151 + mode: &mut MinimaxM2Mode,
152 ) -> ModalResult<MinimaxM2Event> {
153 match mode {
154 MinimaxM2Mode::Text => parse_text_event(input),
151 - MinimaxM2Mode::ToolBlock => parse_tool_block_event(input),
155 + MinimaxM2Mode::ToolBlock { invoke_end_scan } => {
156 + parse_tool_block_event(input, invoke_end_scan)
157 + }
158 MinimaxM2Mode::Done => ignored_rest_event(input),
159 }
160 }
@@ -169,8 +175,14 @@ fn safe_text_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event
175 }
176
177 /// Parse one event inside a MiniMax M2 tool block.
172 -fn parse_tool_block_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
173 - alt((tool_block_end_event, invoke_event)).parse_next(input)
178 +fn parse_tool_block_event(
179 + input: &mut MinimaxM2Input<'_>,
180 + invoke_end_scan: &mut MarkerScanState,
181 +) -> ModalResult<MinimaxM2Event> {
182 + alt((tool_block_end_event, |input: &mut MinimaxM2Input<'_>| {
183 + invoke_event(input, invoke_end_scan)
184 + }))
185 + .parse_next(input)
186 }
187
188 /// Parse a MiniMax M2 tool-block end marker.
@@ -181,14 +193,17 @@ fn tool_block_end_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2
193 }
194
195 /// Parse a complete MiniMax M2 invoke block.
184 -fn invoke_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
196 +fn invoke_event(
197 + input: &mut MinimaxM2Input<'_>,
198 + invoke_end_scan: &mut MarkerScanState,
199 +) -> ModalResult<MinimaxM2Event> {
200 let (name, body) = seq!(
201 _: ws0,
202 _: literal(INVOKE_START),
203 _: (ws1, literal("name=")),
204 partial_attr_value,
205 _: literal(">"),
191 - take_until(0.., INVOKE_END),
206 + take_until_marker(INVOKE_END, invoke_end_scan),
207 _: literal(INVOKE_END),
208 )
209 .parse_next(input)?;
rust/src/tool-parser/src/minimax_m3.rs
+27 -12
@@ -6,7 +6,7 @@ use winnow::stream::Partial;
6 use winnow::token::{literal, rest, take_until};
7
8 use super::parameters::{ParamElement, ParamInput, ToolSchemas};
9 -use super::utils::{parse_buffered_event, safe_text_len};
9 +use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
10 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
11 use crate::Tool;
12
@@ -21,10 +21,10 @@ const MIXED_TEXT_FIELD: &str = "$text";
21
22 type MinimaxM3Input<'i> = Partial<&'i str>;
23
24 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24 +#[derive(Debug, Clone, PartialEq, Eq)]
25 enum MinimaxM3Mode {
26 Text,
27 - ToolBlock,
27 + ToolBlock { invoke_end_scan: MarkerScanState },
28 Done,
29 }
30
@@ -109,7 +109,11 @@ impl MinimaxM3ToolParser {
109 MinimaxM3Event::Text { len: consumed_len } => {
110 output.normal_text.push_str(&self.buffer[..consumed_len]);
111 }
112 - MinimaxM3Event::ToolBlockStart => self.mode = MinimaxM3Mode::ToolBlock,
112 + MinimaxM3Event::ToolBlockStart => {
113 + self.mode = MinimaxM3Mode::ToolBlock {
114 + invoke_end_scan: MarkerScanState::default(),
115 + };
116 + }
117 MinimaxM3Event::Invoke { name, params } => {
118 let arguments = self.tool_parameters.convert_params_with_schema(&name, params);
119 let arguments = serde_json::to_string(&arguments)
@@ -141,7 +145,7 @@ impl ToolParser for MinimaxM3ToolParser {
145 self.buffer.push_str(chunk);
146
147 while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
144 - parse_next_minimax_m3_event(input, self.mode)
148 + parse_next_minimax_m3_event(input, &mut self.mode)
149 })? {
150 self.apply_event(event, output)?;
151 self.buffer.drain(..consumed_len);
@@ -156,7 +160,7 @@ impl ToolParser for MinimaxM3ToolParser {
160 MinimaxM3Mode::Text => {
161 output.normal_text.push_str(&self.buffer);
162 }
159 - MinimaxM3Mode::ToolBlock => {
163 + MinimaxM3Mode::ToolBlock { .. } => {
164 if !self.buffer.trim_start().is_empty() {
165 return Err(parsing_failed!("incomplete MiniMax M3 tool call"));
166 }
@@ -177,11 +181,13 @@ impl ToolParser for MinimaxM3ToolParser {
181 /// Parse a MiniMax M3 event for the current parser mode.
182 fn parse_next_minimax_m3_event(
183 input: &mut MinimaxM3Input<'_>,
180 - mode: MinimaxM3Mode,
184 + mode: &mut MinimaxM3Mode,
185 ) -> ModalResult<MinimaxM3Event> {
186 match mode {
187 MinimaxM3Mode::Text => parse_text_event(input),
184 - MinimaxM3Mode::ToolBlock => parse_tool_block_event(input),
188 + MinimaxM3Mode::ToolBlock { invoke_end_scan } => {
189 + parse_tool_block_event(input, invoke_end_scan)
190 + }
191 MinimaxM3Mode::Done => ignored_rest_event(input),
192 }
193 }
@@ -202,8 +208,14 @@ fn safe_text_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event
208 }
209
210 /// Parse one event inside a MiniMax M3 tool block.
205 -fn parse_tool_block_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
206 - alt((tool_block_end_event, invoke_event)).parse_next(input)
211 +fn parse_tool_block_event(
212 + input: &mut MinimaxM3Input<'_>,
213 + invoke_end_scan: &mut MarkerScanState,
214 +) -> ModalResult<MinimaxM3Event> {
215 + alt((tool_block_end_event, |input: &mut MinimaxM3Input<'_>| {
216 + invoke_event(input, invoke_end_scan)
217 + }))
218 + .parse_next(input)
219 }
220
221 /// Parse a MiniMax M3 tool-block end marker.
@@ -214,14 +226,17 @@ fn tool_block_end_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3
226 }
227
228 /// Parse a complete MiniMax M3 invoke block.
217 -fn invoke_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
229 +fn invoke_event(
230 + input: &mut MinimaxM3Input<'_>,
231 + invoke_end_scan: &mut MarkerScanState,
232 +) -> ModalResult<MinimaxM3Event> {
233 let (name, body) = seq!(
234 _: ws0,
235 _: literal(INVOKE_START),
236 _: (ws1, literal("name=")),
237 partial_attr_value,
238 _: literal(">"),
224 - take_until(0.., INVOKE_END),
239 + take_until_marker(INVOKE_END, invoke_end_scan),
240 _: literal(INVOKE_END),
241 )
242 .parse_next(input)?;
rust/src/tool-parser/src/qwen_coder.rs
+82 -11
@@ -5,7 +5,7 @@ use winnow::stream::Partial;
5 use winnow::token::{literal, take_until};
6
7 use super::parameters::ToolSchemas;
8 -use super::utils::{parse_buffered_event, safe_text_len};
8 +use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
9 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
10 use crate::Tool;
11
@@ -18,10 +18,10 @@ const PARAMETER_END: &str = "</parameter>";
18
19 type QwenCoderInput<'i> = Partial<&'i str>;
20
21 -#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21 +#[derive(Debug, Clone, PartialEq, Eq)]
22 enum QwenCoderMode {
23 Text,
24 - ToolCall,
24 + ToolCall { end_marker_scan: MarkerScanState },
25 }
26
27 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -76,7 +76,11 @@ impl Qwen3CoderToolParser {
76 QwenCoderEvent::Text { len: consumed_len } => {
77 output.normal_text.push_str(&self.buffer[..consumed_len]);
78 }
79 - QwenCoderEvent::ToolCallStart => self.mode = QwenCoderMode::ToolCall,
79 + QwenCoderEvent::ToolCallStart => {
80 + self.mode = QwenCoderMode::ToolCall {
81 + end_marker_scan: MarkerScanState::default(),
82 + };
83 + }
84 QwenCoderEvent::ToolCall { name, raw_params } => {
85 self.mode = QwenCoderMode::Text;
86 let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
@@ -113,7 +117,7 @@ impl ToolParser for Qwen3CoderToolParser {
117 self.buffer.push_str(chunk);
118
119 while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
116 - parse_next_qwen_coder_event(input, self.mode)
120 + parse_next_qwen_coder_event(input, &mut self.mode)
121 })? {
122 self.apply_event(event, output)?;
123 self.buffer.drain(..consumed_len);
@@ -125,7 +129,9 @@ impl ToolParser for Qwen3CoderToolParser {
129 fn finish(&mut self) -> Result<ToolParserOutput> {
130 let mut output = ToolParserOutput::default();
131 if !self.buffer.is_empty() {
128 - if self.mode == QwenCoderMode::ToolCall || self.buffer.starts_with(TOOL_CALL_START) {
132 + if matches!(self.mode, QwenCoderMode::ToolCall { .. })
133 + || self.buffer.starts_with(TOOL_CALL_START)
134 + {
135 return Err(parsing_failed!("incomplete Qwen Coder tool call"));
136 }
137 output.normal_text.push_str(&self.buffer);
@@ -142,11 +148,11 @@ impl ToolParser for Qwen3CoderToolParser {
148 /// Parse a Qwen Coder event for the current parser mode.
149 fn parse_next_qwen_coder_event(
150 input: &mut QwenCoderInput<'_>,
145 - mode: QwenCoderMode,
151 + mode: &mut QwenCoderMode,
152 ) -> ModalResult<QwenCoderEvent> {
153 match mode {
154 QwenCoderMode::Text => parse_text_event(input),
149 - QwenCoderMode::ToolCall => tool_call_event(input),
155 + QwenCoderMode::ToolCall { end_marker_scan } => tool_call_event(input, end_marker_scan),
156 }
157 }
158
@@ -166,10 +172,13 @@ fn safe_text_event(input: &mut QwenCoderInput<'_>) -> ModalResult<QwenCoderEvent
172 }
173
174 /// Parse a complete Qwen Coder tool call.
169 -fn tool_call_event(input: &mut QwenCoderInput<'_>) -> ModalResult<QwenCoderEvent> {
175 +fn tool_call_event(
176 + input: &mut QwenCoderInput<'_>,
177 + end_marker_scan: &mut MarkerScanState,
178 +) -> ModalResult<QwenCoderEvent> {
179 let (body,) = seq!(
180 _: ws0,
172 - take_until(0.., TOOL_CALL_END),
181 + take_until_marker(TOOL_CALL_END, end_marker_scan),
182 _: literal(TOOL_CALL_END),
183 )
184 .parse_next(input)?;
@@ -228,8 +237,8 @@ mod tests {
237 use thiserror_ext::AsReport;
238
239 use super::{Qwen3CoderToolParser, ToolParser};
231 - use crate::ToolParserTestExt as _;
240 use crate::test_utils::{collect_stream, split_by_chars, test_tools};
241 + use crate::{ToolParserOutput, ToolParserTestExt as _};
242
243 fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
244 let params = params
@@ -574,6 +583,68 @@ mod tests {
583 );
584 }
585
586 + #[test]
587 + fn qwen_coder_streaming_handles_end_token_split_across_chunks() {
588 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
589 + let output = parser
590 + .parse_chunk(
591 + "<tool_call>\n\
592 + <function=get_weather>\n\
593 + <parameter=location>SF</parameter>\n\
594 + </function>\n\
595 + </tool",
596 + )
597 + .unwrap();
598 +
599 + assert!(output.normal_text.is_empty());
600 + assert!(output.calls.is_empty());
601 +
602 + let mut output = output;
603 + output.append(parser.parse_chunk("_call>").unwrap());
604 + output.append(parser.finish().unwrap());
605 + let output = output.coalesce_calls();
606 +
607 + assert!(output.normal_text.is_empty());
608 + assert_eq!(output.calls.len(), 1);
609 + assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
610 + assert_eq!(
611 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
612 + json!({ "location": "SF" })
613 + );
614 + }
615 +
616 + #[test]
617 + fn qwen_coder_streaming_buffers_long_body_until_end_marker() {
618 + let long_location = format!("SF-{}", "x".repeat(8192));
619 + let text = build_tool_call("get_weather", &[("location", &long_location)]);
620 + let split_at = text.len() - "_call>".len();
621 + let (body_with_partial_end, end_suffix) = text.split_at(split_at);
622 + let chunks = split_by_chars(body_with_partial_end, 31);
623 + let mut parser = Qwen3CoderToolParser::new(&test_tools());
624 + let mut output = ToolParserOutput::default();
625 +
626 + assert_eq!(end_suffix, "_call>");
627 +
628 + for chunk in chunks {
629 + let chunk_output = parser.parse_chunk(chunk).unwrap();
630 + assert!(chunk_output.normal_text.is_empty());
631 + assert!(chunk_output.calls.is_empty());
632 + output.append(chunk_output);
633 + }
634 +
635 + output.append(parser.parse_chunk(end_suffix).unwrap());
636 + output.append(parser.finish().unwrap());
637 + let output = output.coalesce_calls();
638 +
639 + assert!(output.normal_text.is_empty());
640 + assert_eq!(output.calls.len(), 1);
641 + assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
642 + assert_eq!(
643 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
644 + json!({ "location": long_location })
645 + );
646 + }
647 +
648 #[test]
649 fn qwen_coder_streaming_does_not_emit_incomplete_tool_call() {
650 let mut parser = Qwen3CoderToolParser::new(&test_tools());
rust/src/tool-parser/src/utils.rs
+169 -1
@@ -1,5 +1,6 @@
1 //! Shared helpers for tool parsers.
2
3 +use winnow::Parser;
4 use winnow::error::{ContextError, ErrMode, ModalResult, Needed, StrContext, StrContextValue};
5 use winnow::stream::{Offset, Partial, Stream};
6
@@ -66,6 +67,74 @@ pub(super) fn safe_text_len(input: &mut Partial<&str>, marker: &str) -> ModalRes
67 Ok(emit_len)
68 }
69
70 +/// Streaming scan state for a buffered marker search [`take_until_marker`],
71 +/// so that we don't have to rescan the whole buffered prefix when resuming.
72 +#[derive(Debug, Clone, Default, PartialEq, Eq)]
73 +pub(super) struct MarkerScanState {
74 + scan_start: usize,
75 +}
76 +
77 +impl MarkerScanState {
78 + pub(super) fn reset(&mut self) {
79 + self.scan_start = 0;
80 + }
81 +}
82 +
83 +/// Parse text until `marker`, resuming from the last safe scan checkpoint.
84 +///
85 +/// This is the streaming-buffered variant of `winnow::token::take_until(0..,
86 +/// marker)`: it returns the slice before `marker` and leaves `marker` for the
87 +/// caller to consume. On incomplete input, it stores the earliest byte offset
88 +/// that can still match `marker` and returns `Incomplete` without consuming
89 +/// input, so the next parse can avoid rescanning the whole buffered prefix.
90 +///
91 +/// Use this for outer parser states that keep the full buffered input across
92 +/// chunks while waiting for a closing marker. Plain `take_until` is still a
93 +/// better fit for one-shot parsers over a complete body, and for `1..` cases
94 +/// where an empty slice before the marker should be rejected.
95 +pub(super) fn take_until_marker<'i, 'a>(
96 + marker: &'a str,
97 + state: &'a mut MarkerScanState,
98 +) -> impl Parser<Partial<&'i str>, &'i str, ErrMode<ContextError>> + 'a {
99 + move |input: &mut Partial<&'i str>| take_until_marker_(input, marker, state)
100 +}
101 +
102 +fn take_until_marker_<'i>(
103 + input: &mut Partial<&'i str>,
104 + marker: &str,
105 + state: &mut MarkerScanState,
106 +) -> ModalResult<&'i str> {
107 + debug_assert!(!marker.is_empty());
108 +
109 + let text = **input;
110 + if text.is_empty() {
111 + return incomplete();
112 + }
113 +
114 + // Normal updates store a char boundary; this keeps stale or misused state from panicking.
115 + let scan_start = floor_char_boundary(text, state.scan_start);
116 +
117 + if let Some(offset) = text[scan_start..].find(marker) {
118 + let marker_start = scan_start + offset;
119 + let body = &text[..marker_start];
120 + input.next_slice(marker_start);
121 + state.reset();
122 + return Ok(body);
123 + }
124 +
125 + let keep_len = partial_prefix_len(text, marker);
126 + state.scan_start = text.len() - keep_len;
127 + incomplete()
128 +}
129 +
130 +fn floor_char_boundary(text: &str, index: usize) -> usize {
131 + let mut index = index.min(text.len());
132 + while !text.is_char_boundary(index) {
133 + index -= 1;
134 + }
135 + index
136 +}
137 +
138 /// Streaming lexical state for a top-level JSON object.
139 #[derive(Debug, Clone, Default, PartialEq, Eq)]
140 pub(super) struct JsonObjectScanState {
@@ -273,11 +342,13 @@ pub(super) fn incomplete<T>() -> ModalResult<T> {
342 mod tests {
343
344 use expect_test::expect;
345 + use winnow::Parser;
346 use winnow::error::ErrMode;
347 use winnow::stream::{Offset, Partial, Stream};
348
349 use super::{
280 - JsonObjectScanState, json_str, partial_prefix_len, safe_text_len, take_json_object,
350 + JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len,
351 + take_json_object, take_until_marker,
352 };
353
354 #[test]
@@ -335,6 +406,103 @@ mod tests {
406 assert!(matches!(error, ErrMode::Incomplete(_)));
407 }
408
409 + #[test]
410 + fn take_until_marker_stops_before_marker() {
411 + let mut state = MarkerScanState::default();
412 + let mut input = Partial::new("body</end>tail");
413 + let checkpoint = input.checkpoint();
414 +
415 + let body = take_until_marker("</end>", &mut state).parse_next(&mut input).unwrap();
416 +
417 + assert_eq!(body, "body");
418 + assert_eq!(input.offset_from(&checkpoint), "body".len());
419 + assert_eq!(*input, "</end>tail");
420 + assert_eq!(state, MarkerScanState::default());
421 + }
422 +
423 + #[test]
424 + fn take_until_marker_resumes_after_split_marker() {
425 + let mut state = MarkerScanState::default();
426 + let mut input = Partial::new("body</to");
427 +
428 + let error = take_until_marker("</tool_call>", &mut state)
429 + .parse_next(&mut input)
430 + .unwrap_err();
431 +
432 + assert!(matches!(error, ErrMode::Incomplete(_)));
433 + assert_eq!(state.scan_start, "body".len());
434 +
435 + let mut input = Partial::new("body</tool_call>tail");
436 + let body = take_until_marker("</tool_call>", &mut state).parse_next(&mut input).unwrap();
437 +
438 + assert_eq!(body, "body");
439 + assert_eq!(*input, "</tool_call>tail");
440 + assert_eq!(state, MarkerScanState::default());
441 + }
442 +
443 + #[test]
444 + fn take_until_marker_advances_checkpoint_for_long_prefix() {
445 + let mut state = MarkerScanState::default();
446 + let text = format!("{}{}", "x".repeat(1024), "</too");
447 + let mut input = Partial::new(text.as_str());
448 +
449 + let error = take_until_marker("</tool_call>", &mut state)
450 + .parse_next(&mut input)
451 + .unwrap_err();
452 +
453 + assert!(matches!(error, ErrMode::Incomplete(_)));
454 + assert_eq!(state.scan_start, 1024);
455 + }
456 +
457 + #[test]
458 + fn take_until_marker_keeps_unicode_marker_boundaries() {
459 + let marker = "<｜DSML｜function_calls>";
460 + let mut state = MarkerScanState::default();
461 + let mut input = Partial::new("prefix <｜DSML｜fun");
462 +
463 + let error = take_until_marker(marker, &mut state).parse_next(&mut input).unwrap_err();
464 +
465 + assert!(matches!(error, ErrMode::Incomplete(_)));
466 + assert_eq!(state.scan_start, "prefix ".len());
467 + assert!("prefix <｜DSML｜fun".is_char_boundary(state.scan_start));
468 +
469 + let mut input = Partial::new("prefix <｜DSML｜function_calls>tail");
470 + let body = take_until_marker(marker, &mut state).parse_next(&mut input).unwrap();
471 +
472 + assert_eq!(body, "prefix ");
473 + assert_eq!(*input, "<｜DSML｜function_calls>tail");
474 + assert_eq!(state, MarkerScanState::default());
475 + }
476 +
477 + #[test]
478 + fn take_until_marker_floors_stale_checkpoint_to_char_boundary() {
479 + let mut state = MarkerScanState { scan_start: 1 };
480 + let mut input = Partial::new("é</end>");
481 +
482 + let body = take_until_marker("</end>", &mut state).parse_next(&mut input).unwrap();
483 +
484 + assert_eq!(body, "é");
485 + assert_eq!(*input, "</end>");
486 + assert_eq!(state, MarkerScanState::default());
487 + }
488 +
489 + #[test]
490 + fn take_until_marker_handles_overlapping_prefixes() {
491 + let mut state = MarkerScanState::default();
492 + let mut input = Partial::new("xxaba");
493 +
494 + let error = take_until_marker("ababa", &mut state).parse_next(&mut input).unwrap_err();
495 +
496 + assert!(matches!(error, ErrMode::Incomplete(_)));
497 + assert_eq!(state.scan_start, 2);
498 +
499 + let mut input = Partial::new("xxababa!");
500 + let body = take_until_marker("ababa", &mut state).parse_next(&mut input).unwrap();
501 +
502 + assert_eq!(body, "xx");
503 + assert_eq!(*input, "ababa!");
504 + }
505 +
506 #[test]
507 fn take_json_object_consumes_simple_object() {
508 let mut state = JsonObjectScanState::default();