1use crate::codec::Delims;
4
5#[derive(Clone)]
6pub enum ToolFragmentMode {
7 Json,
8 Delimited(Delims),
9 None,
10}
11
12#[derive(Debug, PartialEq)]
13pub struct ToolDelta {
14 pub opening: bool,
15 pub name: String,
16 pub arguments: String,
17}
18
19pub struct ToolCallStream {
20 mode: ToolFragmentMode,
21 body: String,
22 opened: bool,
23 complete: bool,
24 sent: String,
25 args_start: Option<usize>,
26 depth: usize,
27 scan_pos: usize,
28 scan_in_str: bool,
29 scan_esc: bool,
30 arg_cursor: Option<usize>,
31 sent_keys: Vec<String>,
32 schemas: Option<std::sync::Arc<crate::codec::ToolSchemas>>,
33}
34
35impl ToolCallStream {
36 pub fn new(mode: ToolFragmentMode) -> Self {
37 Self {
38 mode,
39 body: String::new(),
40 opened: false,
41 complete: false,
42 sent: String::new(),
43 args_start: None,
44 depth: 0,
45 scan_pos: 0,
46 scan_in_str: false,
47 scan_esc: false,
48 arg_cursor: None,
49 sent_keys: Vec::new(),
50 schemas: None,
51 }
52 }
53
54 pub fn with_schemas(mut self, schemas: Option<std::sync::Arc<crate::codec::ToolSchemas>>) -> Self {
55 self.schemas = schemas;
56 self
57 }
58
59 pub fn opened(&self) -> bool {
60 self.opened
61 }
62
63 pub fn reset(&mut self) {
64 let mode = self.mode.clone();
65 let schemas = self.schemas.take();
66 *self = Self::new(mode).with_schemas(schemas);
67 }
68
69 pub fn feed(&mut self, text: &str) -> Vec<ToolDelta> {
70 self.body.push_str(text);
71 let mut out = Vec::new();
72 if self.complete {
73 return out;
74 }
75 match self.mode.clone() {
76 ToolFragmentMode::Json => self.step_json(&mut out),
77 ToolFragmentMode::Delimited(d) => self.step_delimited(&d, &mut out),
78 ToolFragmentMode::None => {}
79 }
80 out
81 }
82
83 pub fn finish(&mut self, arguments: &str) -> Vec<ToolDelta> {
84 let mut out = Vec::new();
85 if self.complete {
86 return out;
87 }
88 match self.mode.clone() {
89 ToolFragmentMode::Json => self.step_json(&mut out),
90 ToolFragmentMode::Delimited(d) => self.step_delimited(&d, &mut out),
91 ToolFragmentMode::None => {}
92 }
93 if !self.opened || self.complete {
94 return out;
95 }
96 if self.sent.is_empty() {
97 self.emit(&mut out, arguments.to_string());
98 } else {
99 match self.mode {
100 ToolFragmentMode::Json => {
101 self.emit(&mut out, "}".repeat(self.depth));
102 }
103 _ => {
104 for (key, value) in missing_members(arguments, &self.sent_keys) {
105 self.emit(
106 &mut out,
107 format!(",{}:{}", serde_json::Value::String(key), value),
108 );
109 }
110 self.emit(&mut out, "}".to_string());
111 }
112 }
113 }
114 self.complete = true;
115 out
116 }
117
118 fn step_json(&mut self, out: &mut Vec<ToolDelta>) {
119 if !self.opened {
120 match json_member_string(&self.body, "name") {
121 Some(name) if !name.is_empty() => self.open(out, name),
122 _ => return,
123 }
124 }
125 if self.args_start.is_none() {
126 let v = json_member_value_start(&self.body, "arguments")
127 .or_else(|| json_member_value_start(&self.body, "parameters"));
128 let Some(v) = v else { return };
129 let first = self.body.as_bytes()[v];
130 if first != b'{' && first != b'[' {
131 self.mode = ToolFragmentMode::None;
132 return;
133 }
134 self.args_start = Some(v);
135 self.scan_pos = v;
136 }
137 let start = self.args_start.expect("located above");
138 let end = self.scan_json_value();
139 let already = start + self.sent.len();
140 if end > already {
141 let fragment = self.body[already..end].to_string();
142 self.emit(out, fragment);
143 }
144 }
145
146 fn scan_json_value(&mut self) -> usize {
147 let b = self.body.as_bytes();
148 while self.scan_pos < b.len() {
149 let c = b[self.scan_pos];
150 self.scan_pos += 1;
151 if self.scan_in_str {
152 if self.scan_esc {
153 self.scan_esc = false;
154 } else if c == b'\\' {
155 self.scan_esc = true;
156 } else if c == b'"' {
157 self.scan_in_str = false;
158 }
159 continue;
160 }
161 match c {
162 b'"' => self.scan_in_str = true,
163 b'{' | b'[' => self.depth += 1,
164 b'}' | b']' => {
165 self.depth -= 1;
166 if self.depth == 0 {
167 self.complete = true;
168 return self.scan_pos;
169 }
170 }
171 _ => {}
172 }
173 }
174 if self.scan_esc {
175 b.len() - 1
176 } else {
177 b.len()
178 }
179 }
180
181 fn step_delimited(&mut self, d: &Delims, out: &mut Vec<ToolDelta>) {
182 if !self.opened {
183 match d.peek_name(&self.body) {
184 Some(name) => self.open(out, name),
185 None => return,
186 }
187 }
188 let mut cursor = self.arg_cursor;
189 let closed = d.closed_args(&self.body, &mut cursor, self.schemas.as_deref());
190 self.arg_cursor = cursor;
191 for (key, value) in closed {
192 if self.sent_keys.contains(&key) {
193 continue;
194 }
195 let sep = if self.sent_keys.is_empty() { "{" } else { "," };
196 let fragment = format!("{sep}{}:{}", serde_json::Value::String(key.clone()), value);
197 self.sent_keys.push(key);
198 self.emit(out, fragment);
199 }
200 }
201
202 fn open(&mut self, out: &mut Vec<ToolDelta>, name: String) {
203 out.push(ToolDelta {
204 opening: true,
205 name,
206 arguments: String::new(),
207 });
208 self.opened = true;
209 }
210
211 fn emit(&mut self, out: &mut Vec<ToolDelta>, fragment: String) {
212 if fragment.is_empty() {
213 return;
214 }
215 self.sent.push_str(&fragment);
216 if let Some(last) = out.last_mut() {
217 if last.opening && last.arguments.is_empty() {
218 last.arguments = fragment;
219 return;
220 }
221 }
222 out.push(ToolDelta {
223 opening: false,
224 name: String::new(),
225 arguments: fragment,
226 });
227 }
228}
229
230fn missing_members(arguments: &str, sent: &[String]) -> Vec<(String, serde_json::Value)> {
231 let Ok(serde_json::Value::Object(map)) = serde_json::from_str::<serde_json::Value>(arguments)
232 else {
233 return Vec::new();
234 };
235 map.into_iter()
236 .filter(|(k, _)| !sent.iter().any(|s| s == k))
237 .collect()
238}
239
240fn json_member_value_start(body: &str, key: &str) -> Option<usize> {
241 let quoted = format!("\"{key}\"");
242 let b = body.as_bytes();
243 let (mut depth, mut in_str, mut esc, mut expect_key) = (0i32, false, false, false);
244 let mut found = None;
245 for (i, &c) in b.iter().enumerate() {
246 if in_str {
247 if esc {
248 esc = false;
249 } else if c == b'\\' {
250 esc = true;
251 } else if c == b'"' {
252 in_str = false;
253 }
254 continue;
255 }
256 match c {
257 b'"' => {
258 if depth == 1 && expect_key && body[i..].starts_with("ed) {
259 found = Some(i);
260 break;
261 }
262 in_str = true;
263 }
264 b'{' => {
265 depth += 1;
266 expect_key = depth == 1;
267 }
268 b'[' => {
269 depth += 1;
270 expect_key = false;
271 }
272 b'}' | b']' => {
273 depth -= 1;
274 expect_key = false;
275 }
276 b',' => expect_key = depth == 1,
277 b':' => expect_key = false,
278 _ => {}
279 }
280 }
281 let k = found?;
282 let colon = k + body[k..].find(':')?;
283 body[colon + 1..]
284 .find(|c: char| !c.is_whitespace())
285 .map(|off| colon + 1 + off)
286}
287
288fn json_member_string(body: &str, key: &str) -> Option<String> {
289 let p = json_member_value_start(body, key)?;
290 if *body.as_bytes().get(p)? != b'"' {
291 return None;
292 }
293 let mut esc = false;
294 for (off, c) in body[p + 1..].char_indices() {
295 if esc {
296 esc = false;
297 } else if c == '\\' {
298 esc = true;
299 } else if c == '"' {
300 return serde_json::from_str::<String>(&body[p..p + off + 2]).ok();
301 }
302 }
303 None
304}
305
306#[cfg(test)]
307mod tests {
308 use super::*;
309
310 fn drive(mode: ToolFragmentMode, chunks: &[&str], parsed_args: &str) -> (Vec<String>, String) {
311 let mut s = ToolCallStream::new(mode);
312 let mut names = Vec::new();
313 let mut args = String::new();
314 let mut open_seen = false;
315 let mut take = |deltas: Vec<ToolDelta>, names: &mut Vec<String>, args: &mut String| {
316 for d in deltas {
317 if d.opening {
318 open_seen = true;
319 names.push(d.name);
320 } else {
321 assert!(open_seen, "an arguments fragment preceded the name");
322 }
323 args.push_str(&d.arguments);
324 }
325 };
326 for c in chunks {
327 take(s.feed(c), &mut names, &mut args);
328 }
329 take(s.finish(parsed_args), &mut names, &mut args);
330 (names, args)
331 }
332
333 #[test]
334 fn json_streams_raw_argument_bytes_after_the_name() {
335 let (names, args) = drive(
336 ToolFragmentMode::Json,
337 &[
338 "{\"name\": \"get_wea",
339 "ther\", \"arguments\": {\"city\": \"Tok",
340 "yo\", \"unit\": \"c\"}}",
341 ],
342 "{\"city\":\"Tokyo\",\"unit\":\"c\"}",
343 );
344 assert_eq!(names, vec!["get_weather"]);
345 assert_eq!(args, "{\"city\": \"Tokyo\", \"unit\": \"c\"}");
346 }
347
348 #[test]
349 fn json_opens_before_any_argument_byte() {
350 let mut s = ToolCallStream::new(ToolFragmentMode::Json);
351 let d = s.feed("{\"name\": \"ping\", \"argu");
352 assert_eq!(d.len(), 1);
353 assert!(d[0].opening && d[0].name == "ping" && d[0].arguments.is_empty());
354 assert!(s.opened());
355 }
356
357 #[test]
358 fn json_holds_back_a_partial_escape() {
359 let mut s = ToolCallStream::new(ToolFragmentMode::Json);
360 let mut sent = String::new();
361 for d in s.feed("{\"name\":\"f\",\"arguments\":{\"s\":\"a\\") {
362 assert!(!d.arguments.ends_with('\\'), "fragment ended mid-escape");
363 sent.push_str(&d.arguments);
364 }
365 for d in s.feed("n b\"}}") {
366 sent.push_str(&d.arguments);
367 }
368 assert_eq!(sent, "{\"s\":\"a\\n b\"}");
369 }
370
371 #[test]
372 fn json_closes_a_call_the_parser_salvaged_from_a_cut_body() {
373 let (_, args) = drive(
374 ToolFragmentMode::Json,
375 &["{\"name\":\"f\",\"arguments\":{\"a\":{\"b\":1"],
376 "{\"a\":{\"b\":1}}",
377 );
378 assert_eq!(args, "{\"a\":{\"b\":1}}");
379 assert!(serde_json::from_str::<serde_json::Value>(&args).is_ok());
380 }
381
382 #[test]
383 fn json_argument_less_call_still_gets_an_object() {
384 let (names, args) = drive(ToolFragmentMode::Json, &["{\"name\":\"ping\"}"], "{}");
385 assert_eq!(names, vec!["ping"]);
386 assert_eq!(args, "{}");
387 }
388
389 #[test]
390 fn json_string_valued_arguments_stay_one_shot() {
391 let mut s = ToolCallStream::new(ToolFragmentMode::Json);
392 let deltas = s.feed("{\"name\":\"f\",\"arguments\":\"{\\\"a\\\":1}\"}");
393 assert!(s.opened());
394 assert!(deltas.iter().all(|d| d.arguments.is_empty()));
395 let rest = s.finish("{\"a\":1}");
396 assert_eq!(rest.len(), 1);
397 assert_eq!(rest[0].arguments, "{\"a\":1}");
398 }
399
400 fn qwen_wire() -> ToolFragmentMode {
401 ToolFragmentMode::Delimited(Delims::for_test(
402 "<function=",
403 ">\n",
404 "<parameter=",
405 ">\n",
406 "\n</parameter>\n",
407 ))
408 }
409
410 #[test]
411 fn delimited_streams_one_closed_argument_at_a_time() {
412 let mut s = ToolCallStream::new(qwen_wire());
413 let mut args = String::new();
414 let d = s.feed("<function=get_weather>\n<parameter=city>\nTok");
415 assert_eq!(d.len(), 1);
416 assert!(d[0].opening && d[0].name == "get_weather" && d[0].arguments.is_empty());
417 for d in s.feed("yo\n</parameter>\n") {
418 args.push_str(&d.arguments);
419 }
420 assert_eq!(
421 args, "{\"city\":\"Tokyo\"",
422 "a closed value streams, the object stays open"
423 );
424 for d in s.feed("<parameter=days>\n3\n</parameter>\n") {
425 args.push_str(&d.arguments);
426 }
427 assert_eq!(
428 args, "{\"city\":\"Tokyo\",\"days\":3",
429 "typed on the way out"
430 );
431 for d in s.finish("{\"city\":\"Tokyo\",\"days\":3}") {
432 args.push_str(&d.arguments);
433 }
434 assert_eq!(args, "{\"city\":\"Tokyo\",\"days\":3}");
435 let v: serde_json::Value =
436 serde_json::from_str(&args).expect("fragments reassemble to JSON");
437 assert_eq!(v["days"], 3);
438 }
439
440 #[test]
441 fn delimited_types_streamed_values_from_the_declared_schema() {
442 let schemas = crate::codec::ToolSchemas::of_values(&[serde_json::json!({
443 "type": "function",
444 "function": {"name": "serve", "parameters": {"type": "object", "properties": {
445 "port": {"type": "string"}, "workers": {"type": "integer"}
446 }}}
447 })]);
448 let mut s = ToolCallStream::new(qwen_wire()).with_schemas(schemas);
449 for round in 0..2 {
450 let mut args = String::new();
451 for chunk in [
452 "<function=serve>\n<parameter=port>\n8080\n</parameter>\n",
453 "<parameter=workers>\n4\n</parameter>\n</function>",
454 ] {
455 for d in s.feed(chunk) {
456 args.push_str(&d.arguments);
457 }
458 }
459 assert_eq!(args, "{\"port\":\"8080\",\"workers\":4", "round {round}");
460 s.reset();
461 }
462 let mut blind = ToolCallStream::new(qwen_wire());
463 let args: String = blind
464 .feed("<function=serve>\n<parameter=port>\n8080\n</parameter>\n")
465 .into_iter()
466 .map(|d| d.arguments)
467 .collect();
468 assert_eq!(args, "{\"port\":8080");
469 }
470
471 #[test]
472 fn delimited_zero_argument_call_gets_an_object() {
473 let mut s = ToolCallStream::new(qwen_wire());
474 let _ = s.feed("<function=ping>\n");
475 let rest = s.finish("{}");
476 assert_eq!(rest.len(), 1);
477 assert_eq!(rest[0].arguments, "{}");
478 }
479
480 #[test]
481 fn delimited_hands_over_a_value_only_the_final_parse_recovered() {
482 let mut s = ToolCallStream::new(qwen_wire());
483 let mut args = String::new();
484 for chunk in [
485 "<function=f>\n<parameter=a>\n1\n</parameter>\n",
486 "<parameter=b>\n22",
487 ] {
488 for d in s.feed(chunk) {
489 args.push_str(&d.arguments);
490 }
491 }
492 assert_eq!(args, "{\"a\":1");
493 for d in s.finish("{\"a\":1,\"b\":22}") {
494 args.push_str(&d.arguments);
495 }
496 assert_eq!(
497 args, "{\"a\":1,\"b\":22}",
498 "the value only the final parse recovered is handed over at the close"
499 );
500 assert!(serde_json::from_str::<serde_json::Value>(&args).is_ok());
501 }
502
503 #[test]
504 fn json_name_is_decoded_by_the_json_parser() {
505 let mut s = ToolCallStream::new(ToolFragmentMode::Json);
506 let d = s.feed(r#"{"name":"get\u005fweather","arguments":{}}"#);
507 assert_eq!(d[0].name, "get_weather");
508 }
509
510 #[test]
511 fn json_name_comes_from_a_key_not_a_value() {
512 let mut s = ToolCallStream::new(ToolFragmentMode::Json);
513 let d = s.feed(r#"{"note":"name","alias":"wrong","name":"right","arguments":{}}"#);
514 assert_eq!(d[0].name, "right");
515 }
516
517 #[test]
518 fn none_mode_never_fragments() {
519 let mut s = ToolCallStream::new(ToolFragmentMode::None);
520 assert!(s.feed("call:f{x:1}").is_empty());
521 assert!(s.finish("{\"x\":1}").is_empty());
522 assert!(
523 !s.opened(),
524 "an unfragmented call is emitted whole by the route"
525 );
526 }
527
528 #[test]
529 fn a_reset_stream_announces_the_next_call() {
530 let mut s = ToolCallStream::new(ToolFragmentMode::Json);
531 let _ = s.feed("{\"name\":\"a\",\"arguments\":{\"x\":1}}");
532 let _ = s.finish("{\"x\":1}");
533 s.reset();
534 assert!(!s.opened());
535 let d = s.feed("{\"name\":\"b\",\"arguments\":{\"y\":2}}");
536 assert!(d[0].opening && d[0].name == "b");
537 }
538}