36
//!
37
//! One spelling difference is folded away before matching: DSML markers are
38
//! built from fullwidth bars, and some stacks hand them back with each bar
39
-//! doubled (`<||DSML||invoke …>`). See [`fold_dsml_bars`].
39
+//! doubled (`<||DSML||invoke …>`). See [`fold_dsml_bars`]. DSML envelopes
40
+//! come as `tool_calls`, `function_calls` or bare `calls`, and a closer may put
41
+//! its slash before the bar (`</|DSML|invoke>`); see [`fold_dsml_closers`].
42
+//!
43
+//! Some models wrap the call in `<tool_call>` or `<minimax:tool_call>` and
44
+//! write `<invoke name="…">` with parameters as `<parameter name="k">v</parameter>`
45
+//! or `<k>v</k>`. The tokenizer's `]<]minimax[>[` special token also leaks
46
+//! into content, between tags and in prose; it is stripped before matching.
47
48
use std::collections::HashMap;
49
52
const XML_OPEN_TAG: &str = "<tool_call>";
53
const XML_CLOSE_TAG: &str = "</tool_call>";
54
55
+const NAMESPACED_OPEN_TAG: &str = "<minimax:tool_call>";
56
+const NAMESPACED_CLOSE_TAG: &str = "</minimax:tool_call>";
57
+/// A tokenizer special token that leaks into content. Never meant to
58
+/// be visible, so it is removed wherever it appears.
59
+const LEAKED_SENTINEL: &str = "]<]minimax[>[";
60
+const INVOKE_OPEN: &str = "<invoke";
61
+const INVOKE_CLOSE: &str = "</invoke>";
62
+const PARAMETER_OPEN: &str = "<parameter";
63
+const PARAMETER_CLOSE: &str = "</parameter>";
64
+
65
const K3_TOOLS_OPEN: &str = "<|open|>tools<|sep|>";
66
const K3_TOOLS_CLOSE: &str = "<|close|>tools<|sep|>";
67
const K3_CALL_OPEN: &str = "<|open|>call";
83
84
const DSML_TOOL_CALLS_OPEN: &str = "<|DSML|tool_calls>";
85
const DSML_TOOL_CALLS_CLOSE: &str = "<|DSML|/tool_calls>";
86
+// The published DSML template names the envelope `function_calls`; one stack
87
+// was seen writing bare `calls`.
88
+const DSML_FUNCTION_CALLS_OPEN: &str = "<|DSML|function_calls>";
89
+const DSML_FUNCTION_CALLS_CLOSE: &str = "<|DSML|/function_calls>";
90
+const DSML_CALLS_OPEN: &str = "<|DSML|calls>";
91
+const DSML_CALLS_CLOSE: &str = "<|DSML|/calls>";
92
+/// Every DSML envelope spelling as `(open, close)`.
93
+const DSML_ENVELOPES: &[(&str, &str)] = &[
94
+ (DSML_TOOL_CALLS_OPEN, DSML_TOOL_CALLS_CLOSE),
95
+ (DSML_FUNCTION_CALLS_OPEN, DSML_FUNCTION_CALLS_CLOSE),
96
+ (DSML_CALLS_OPEN, DSML_CALLS_CLOSE),
97
+];
98
+/// A closer with the slash before the bar, and the canonical spelling the
99
+/// scanner matches instead.
100
+const DSML_SLASH_FIRST_CLOSER: &str = "</|DSML|";
101
+const DSML_CANONICAL_CLOSER: &str = "<|DSML|/";
102
const DSML_INVOKE_OPEN: &str = "<|DSML|invoke";
103
const DSML_INVOKE_CLOSE: &str = "<|DSML|/invoke>";
104
const DSML_PARAMETER_OPEN: &str = "<|DSML|parameter";
106
107
const START_MARKERS: &[&str] = &[
108
XML_OPEN_TAG,
109
+ NAMESPACED_OPEN_TAG,
110
K3_TOOLS_OPEN,
111
K3_RESPONSE_OPEN,
112
K3_THINK_OPEN,
113
DSML_TOOL_CALLS_OPEN,
114
+ DSML_FUNCTION_CALLS_OPEN,
115
+ DSML_CALLS_OPEN,
116
];
117
118
/// One tool call recovered from inline text.
407
Some(args)
408
}
409
410
+/// Parse an `<invoke>`-shaped `<tool_call>` body of one or more
411
+/// `<invoke name="$TOOL">…</invoke>` blocks. A parameter is either
412
+/// `<parameter name="k">v</parameter>` or `<k>v</k>`.
413
+fn parse_invoke_block(inner: &str, tools: &[ToolSpec]) -> Option<Vec<Recovered>> {
414
+ let mut calls = Vec::new();
415
+ let mut rest = inner;
416
+
417
+ while !rest.trim().is_empty() {
418
+ rest = rest.trim_start();
419
+ let header = rest.strip_prefix(INVOKE_OPEN)?;
420
+ let (attr_text, after_header) = split_once(header, ">")?;
421
+ let attrs = parse_k3_attributes(attr_text)?;
422
+ let name = attrs.get("name")?;
423
+ if name.is_empty() || !offered_tool(tools, name) {
424
+ return None;
425
+ }
426
+
427
+ let (call_inner, after_call) = split_once(after_header, INVOKE_CLOSE)?;
428
+ let arguments = parse_invoke_parameters(call_inner, tools, name)?;
429
+ calls.push(Recovered {
430
+ name: name.to_string(),
431
+ arguments: serde_json::Value::Object(arguments).to_string(),
432
+ });
433
+ rest = after_call;
434
+ }
435
+
436
+ if calls.is_empty() { None } else { Some(calls) }
437
+}
438
+
439
+fn parse_invoke_parameters(
440
+ mut rest: &str,
441
+ tools: &[ToolSpec],
442
+ tool_name: &str,
443
+) -> Option<serde_json::Map<String, serde_json::Value>> {
444
+ let mut args = serde_json::Map::new();
445
+
446
+ while !rest.trim().is_empty() {
447
+ rest = rest.trim_start();
448
+ let (key, raw_value, after) = if let Some(header) = rest.strip_prefix(PARAMETER_OPEN) {
449
+ let (attr_text, after_header) = split_once(header, ">")?;
450
+ let attrs = parse_k3_attributes(attr_text)?;
451
+ let key = attrs.get("name")?.trim().to_string();
452
+ let (raw_value, after) = split_once(after_header, PARAMETER_CLOSE)?;
453
+ (key, raw_value, after)
454
+ } else {
455
+ let header = rest.strip_prefix('<')?;
456
+ let (key, after_header) = split_once(header, ">")?;
457
+ if !key
458
+ .chars()
459
+ .all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '-')
460
+ {
461
+ return None;
462
+ }
463
+ let (raw_value, after) = split_once(after_header, &format!("</{key}>"))?;
464
+ (key.to_string(), raw_value, after)
465
+ };
466
+ if key.is_empty() {
467
+ return None;
468
+ }
469
+
470
+ let declared = declared_type(tools, tool_name, &key);
471
+ let value = match declared.as_deref() {
472
+ Some("array" | "object") => serde_json::from_str(raw_value).ok()?,
473
+ other => coerce(raw_value, other),
474
+ };
475
+ args.insert(key, value);
476
+ rest = after;
477
+ }
478
+
479
+ Some(args)
480
+}
481
+
482
fn split_once<'a>(s: &'a str, delim: &str) -> Option<(&'a str, &'a str)> {
483
let idx = s.find(delim)?;
484
Some((&s[..idx], &s[idx + delim.len()..]))
566
/// be empty if the chunk was absorbed into a tag that hasn't closed yet.
567
pub fn push(&mut self, chunk: &str) -> Vec<ScanEvent> {
568
self.pending.push_str(chunk);
569
+ // Like the bar fold, this runs on the whole buffer so a sentinel split
570
+ // across deltas is seen joined; `drain` holds back a half sentinel.
571
+ if self.pending.contains(LEAKED_SENTINEL) {
572
+ self.pending = self.pending.replace(LEAKED_SENTINEL, "");
573
+ }
574
// The fold runs on the whole buffer, not the chunk: a doubled bar split
575
// across two deltas is only visible once they are joined, and a marker
576
// is never emitted before then because a partial prefix is held back.
581
if let Some(folded) = fold_dsml_bars(&self.pending) {
582
self.pending = folded;
583
}
584
+ if let Some(folded) = fold_dsml_closers(&self.pending) {
585
+ self.pending = folded;
586
+ }
587
self.drain()
588
}
589
601
if rest
602
.strip_prefix(K3_TOOLS_OPEN)
603
.is_some_and(looks_like_k3_call)
488
- || rest
489
- .strip_prefix(XML_OPEN_TAG)
490
- .is_some_and(looks_like_xml_call)
491
- || rest
492
- .strip_prefix(DSML_TOOL_CALLS_OPEN)
493
- .is_some_and(looks_like_dsml_call)
604
+ || [XML_OPEN_TAG, NAMESPACED_OPEN_TAG]
605
+ .iter()
606
+ .any(|open| rest.strip_prefix(open).is_some_and(looks_like_xml_call))
607
+ || DSML_ENVELOPES
608
+ .iter()
609
+ .any(|(open, _)| rest.strip_prefix(open).is_some_and(looks_like_dsml_call))
610
{
611
ScanEvent::Malformed(rest)
612
} else {
619
let mut events = Vec::new();
620
loop {
621
let Some((open_idx, marker)) = find_next_marker(&self.pending) else {
506
- let keep = max_partial_prefix_len(&self.pending, START_MARKERS);
622
+ let keep = max_partial_prefix_len(&self.pending, START_MARKERS)
623
+ .max(partial_prefix_len(&self.pending, LEAKED_SENTINEL));
624
let flush_len = self.pending.len() - keep;
625
if flush_len > 0 {
626
events.push(ScanEvent::Text(
637
}
638
639
match marker {
523
- XML_OPEN_TAG => {
524
- let Some(event) = self.drain_xml() else {
525
- break;
640
+ XML_OPEN_TAG | NAMESPACED_OPEN_TAG => {
641
+ let close = if marker == XML_OPEN_TAG {
642
+ XML_CLOSE_TAG
643
+ } else {
644
+ NAMESPACED_CLOSE_TAG
645
};
527
- events.push(event);
646
+ match self.drain_xml(marker, close) {
647
+ Some(mut recovered) => events.append(&mut recovered),
648
+ None => break,
649
+ }
650
}
651
K3_TOOLS_OPEN => match self.drain_k3_tools() {
652
Some(mut recovered) => events.append(&mut recovered),
653
None => break,
654
},
533
- DSML_TOOL_CALLS_OPEN => match self.drain_dsml() {
534
- Some(mut recovered) => events.append(&mut recovered),
535
- None => break,
536
- },
655
+ DSML_TOOL_CALLS_OPEN | DSML_FUNCTION_CALLS_OPEN | DSML_CALLS_OPEN => {
656
+ let close = DSML_ENVELOPES
657
+ .iter()
658
+ .find_map(|(open, close)| (*open == marker).then_some(*close))
659
+ .expect("every DSML open marker has a close");
660
+ match self.drain_dsml(marker, close) {
661
+ Some(mut recovered) => events.append(&mut recovered),
662
+ None => break,
663
+ }
664
+ }
665
K3_RESPONSE_OPEN => {
666
let Some(mut unwrapped) =
667
self.drain_wrapped_text(K3_RESPONSE_OPEN, K3_RESPONSE_CLOSE, true)
684
events
685
}
686
559
- fn drain_xml(&mut self) -> Option<ScanEvent> {
560
- let close_rel = self.pending[XML_OPEN_TAG.len()..].find(XML_CLOSE_TAG)?;
561
- let close_idx = XML_OPEN_TAG.len() + close_rel;
562
- let block: String = self
563
- .pending
564
- .drain(..close_idx + XML_CLOSE_TAG.len())
565
- .collect();
566
- let inner = &block[XML_OPEN_TAG.len()..block.len() - XML_CLOSE_TAG.len()];
567
- Some(match parse_xml_block(inner, self.tools) {
687
+ fn drain_xml(&mut self, open: &str, close: &str) -> Option<Vec<ScanEvent>> {
688
+ let close_rel = self.pending[open.len()..].find(close)?;
689
+ let close_idx = open.len() + close_rel;
690
+ let block: String = self.pending.drain(..close_idx + close.len()).collect();
691
+ let inner = &block[open.len()..block.len() - close.len()];
692
+ if inner.trim_start().starts_with(INVOKE_OPEN) {
693
+ return Some(match parse_invoke_block(inner, self.tools) {
694
+ Some(calls) => calls.into_iter().map(ScanEvent::ToolCall).collect(),
695
+ None => vec![ScanEvent::Malformed(block)],
696
+ });
697
+ }
698
+ Some(vec![match parse_xml_block(inner, self.tools) {
699
Some(call) => ScanEvent::ToolCall(call),
700
None if looks_like_xml_call(inner) => ScanEvent::Malformed(block),
701
None => ScanEvent::Text(block),
571
- })
702
+ }])
703
}
704
705
fn drain_k3_tools(&mut self) -> Option<Vec<ScanEvent>> {
718
})
719
}
720
590
- fn drain_dsml(&mut self) -> Option<Vec<ScanEvent>> {
591
- let close_rel = self.pending[DSML_TOOL_CALLS_OPEN.len()..].find(DSML_TOOL_CALLS_CLOSE)?;
592
- let close_idx = DSML_TOOL_CALLS_OPEN.len() + close_rel;
593
- let block: String = self
594
- .pending
595
- .drain(..close_idx + DSML_TOOL_CALLS_CLOSE.len())
596
- .collect();
597
- let inner = &block[DSML_TOOL_CALLS_OPEN.len()..block.len() - DSML_TOOL_CALLS_CLOSE.len()];
721
+ fn drain_dsml(&mut self, open: &str, close: &str) -> Option<Vec<ScanEvent>> {
722
+ let close_rel = self.pending[open.len()..].find(close)?;
723
+ let close_idx = open.len() + close_rel;
724
+ let block: String = self.pending.drain(..close_idx + close.len()).collect();
725
+ let inner = &block[open.len()..block.len() - close.len()];
726
727
Some(match parse_dsml_block(inner, self.tools) {
728
Some(calls) => calls.into_iter().map(ScanEvent::ToolCall).collect(),
865
bytes / FULLWIDTH_BAR.len_utf8()
866
}
867
868
+/// Some stacks write a DSML closer with the slash before the bar
869
+/// (`</|DSML|invoke>`), as the published template does, instead of after it.
870
+/// Rewritten to the canonical `<|DSML|/invoke>` once bars are folded. Returns
871
+/// `None` when nothing needs rewriting.
872
+fn fold_dsml_closers(text: &str) -> Option<String> {
873
+ text.contains(DSML_SLASH_FIRST_CLOSER)
874
+ .then(|| text.replace(DSML_SLASH_FIRST_CLOSER, DSML_CANONICAL_CLOSER))
875
+}
876
+
877
fn scan_complete_text(text: &str, tools: &[ToolSpec]) -> Vec<ScanEvent> {
878
let mut scanner = StreamScanner::new(tools);
879
let mut events = scanner.push(text);
1806
assert!(calls.is_empty());
1807
assert_eq!(malformed, 0);
1808
}
1809
+
1810
+ const ISSUE_246: &str = "Now let me run the lint/typecheck:]<]minimax[>[<tool_call>\n\
1811
+ ]<]minimax[>[<invoke name=\"run_command\">]<]minimax[>[<command>cd app && npm run lint | tail -60]<]minimax[>[</command>\
1812
+ ]<]minimax[>[<timeout>120000]<]minimax[>[</timeout>]<]minimax[>[</invoke>\n]<]minimax[>[";
1813
+
1814
+ fn args_of(call: &Recovered) -> serde_json::Value {
1815
+ serde_json::from_str(&call.arguments).unwrap()
1816
+ }
1817
+
1818
+ fn feed(tools: &[ToolSpec], chunks: &[&str]) -> (String, Vec<Recovered>, usize) {
1819
+ let mut scanner = StreamScanner::new(tools);
1820
+ let (mut text, mut calls, mut malformed) = (String::new(), Vec::new(), 0);
1821
+ let mut handle = |event| match event {
1822
+ ScanEvent::Text(t) => text.push_str(&t),
1823
+ ScanEvent::ToolCall(c) => calls.push(c),
1824
+ ScanEvent::Malformed(_) => malformed += 1,
1825
+ };
1826
+ for chunk in chunks {
1827
+ scanner.push(chunk).into_iter().for_each(&mut handle);
1828
+ }
1829
+ scanner.finish().into_iter().for_each(&mut handle);
1830
+ (text, calls, malformed)
1831
+ }
1832
+
1833
+ #[test]
1834
+ fn an_unclosed_invoke_call_with_sentinels_is_dropped() {
1835
+ let tools = vec![run_command_spec()];
1836
+ let extracted = extract(ISSUE_246, &tools);
1837
+ assert_eq!(extracted.text, "Now let me run the lint/typecheck:");
1838
+ assert_eq!(extracted.malformed, 1);
1839
+ assert!(extracted.calls.is_empty());
1840
+ assert!(!extracted.text.contains("minimax"));
1841
+ }
1842
+
1843
+ #[test]
1844
+ fn a_closed_invoke_call_with_sentinels_is_recovered() {
1845
+ let tools = vec![run_command_spec()];
1846
+ let text = format!("{ISSUE_246}</tool_call>");
1847
+ let extracted = extract(&text, &tools);
1848
+ assert_eq!(extracted.malformed, 0);
1849
+ assert_eq!(extracted.calls.len(), 1);
1850
+ assert_eq!(extracted.calls[0].name, "run_command");
1851
+ assert_eq!(
1852
+ args_of(&extracted.calls[0])["command"],
1853
+ "cd app && npm run lint | tail -60"
1854
+ );
1855
+ }
1856
+
1857
+ #[test]
1858
+ fn a_namespaced_envelope_is_recovered_with_typed_parameters() {
1859
+ let tools = vec![run_command_spec(), command_output_spec()];
1860
+ let extracted = extract(
1861
+ "<minimax:tool_call><invoke name=\"run_command\"><parameter name=\"command\">echo hi</parameter></invoke>\
1862
+ <invoke name=\"command_output\"><parameter name=\"task_id\">3</parameter></invoke></minimax:tool_call>",
1863
+ &tools,
1864
+ );
1865
+ assert_eq!(extracted.malformed, 0);
1866
+ assert_eq!(extracted.calls.len(), 2);
1867
+ assert_eq!(args_of(&extracted.calls[0])["command"], "echo hi");
1868
+ assert_eq!(args_of(&extracted.calls[1])["task_id"], 3);
1869
+ }
1870
+
1871
+ #[test]
1872
+ fn a_sentinel_split_across_pushes_is_removed() {
1873
+ let tools = vec![run_command_spec()];
1874
+ let (text, calls, malformed) = feed(&tools, &["a]<]mini", "max[>[b"]);
1875
+ assert_eq!(text, "ab");
1876
+ assert!(calls.is_empty());
1877
+ assert_eq!(malformed, 0);
1878
+ }
1879
+
1880
+ #[test]
1881
+ fn an_unknown_tool_in_an_invoke_block_is_malformed() {
1882
+ let tools = vec![run_command_spec()];
1883
+ let extracted = extract(
1884
+ "<tool_call><invoke name=\"delete_everything\"><x>1</x></invoke></tool_call>",
1885
+ &tools,
1886
+ );
1887
+ assert_eq!(extracted.malformed, 1);
1888
+ assert!(extracted.calls.is_empty());
1889
+ assert_eq!(extracted.text, "");
1890
+ }
1891
+
1892
+ const ISSUE_236_BODY: &str = "<||DSML||invoke name=\"run_command\">\
1893
+ <||DSML||parameter name=\"command\" string=\"true\">echo hi</||DSML||parameter>\
1894
+ </||DSML||invoke>";
1895
+
1896
+ #[test]
1897
+ fn doubled_bar_dsml_with_slash_first_closers_is_recovered_for_every_envelope() {
1898
+ let tools = vec![run_command_spec()];
1899
+ for name in ["function_calls", "calls", "tool_calls"] {
1900
+ let text =
1901
+ format!("before <||DSML||{name}>{ISSUE_236_BODY}</||DSML||{name}> after");
1902
+ let extracted = extract(&text, &tools);
1903
+ assert_eq!(extracted.text, "before after", "{name}");
1904
+ assert_eq!(extracted.malformed, 0, "{name}");
1905
+ assert_eq!(extracted.calls.len(), 1, "{name}");
1906
+ assert_eq!(args_of(&extracted.calls[0])["command"], "echo hi");
1907
+ }
1908
+ }
1909
+
1910
+ #[test]
1911
+ fn a_dsml_function_calls_open_marker_split_mid_name_is_recovered() {
1912
+ let tools = vec![run_command_spec()];
1913
+ let (text, calls, malformed) = feed(
1914
+ &tools,
1915
+ &[
1916
+ "pre <|DSML|func",
1917
+ "tion_calls><|DSML|invoke name=\"run_command\"><|DSML|parameter name=\"command\" string=\"true\">echo hi</|DSML|parameter></|DSML|invoke></|DSML|function_calls>",
1918
+ " post",
1919
+ ],
1920
+ );
1921
+ assert_eq!(text, "pre post");
1922
+ assert_eq!(malformed, 0);
1923
+ assert_eq!(calls.len(), 1);
1924
+ }
1925
+
1926
+ #[test]
1927
+ fn a_dsml_function_calls_marker_mentioned_in_prose_stays_text() {
1928
+ let tools = vec![run_command_spec()];
1929
+ for text in [
1930
+ "The marker is `<|DSML|function_calls>` and it never closes here.",
1931
+ "Quote `<|DSML|function_calls>` then `<|DSML|/function_calls>` in docs.",
1932
+ ] {
1933
+ assert_eq!(
1934
+ extract(text, &tools),
1935
+ Extracted {
1936
+ text: text.to_string(),
1937
+ ..Extracted::default()
1938
+ },
1939
+ "{text}"
1940
+ );
1941
+ }
1942
+ }
1943
}