33
//! malformed when what follows the marker looks like the start of a call (see
34
//! [`looks_like_xml_call`] and [`looks_like_k3_call`]); otherwise it stays
35
//! text, as it always did.
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`].
40
41
use std::collections::HashMap;
42
57
const K3_THINK_CLOSE: &str = "<|close|>think<|sep|>";
58
const K3_SEP: &str = "<|sep|>";
59
60
+/// `|` (U+FF5C), the bar every DSML marker above is built from.
61
+const FULLWIDTH_BAR: char = '|';
62
+/// The name those markers wrap, and the only place [`fold_dsml_bars`] folds.
63
+const DSML_MARKER_NAME: &str = "DSML";
64
+
65
const DSML_TOOL_CALLS_OPEN: &str = "<|DSML|tool_calls>";
66
const DSML_TOOL_CALLS_CLOSE: &str = "<|DSML|/tool_calls>";
67
const DSML_INVOKE_OPEN: &str = "<|DSML|invoke";
456
/// be empty if the chunk was absorbed into a tag that hasn't closed yet.
457
pub fn push(&mut self, chunk: &str) -> Vec<ScanEvent> {
458
self.pending.push_str(chunk);
459
+ // The fold runs on the whole buffer, not the chunk: a doubled bar split
460
+ // across two deltas is only visible once they are joined, and a marker
461
+ // is never emitted before then because a partial prefix is held back.
462
+ if let Some(folded) = fold_dsml_bars(&self.pending) {
463
+ self.pending = folded;
464
+ }
465
self.drain()
466
}
467
666
rest.is_empty() || rest.starts_with(DSML_INVOKE_OPEN)
667
}
668
669
+/// Every DSML marker above carries exactly one fullwidth bar on each side of
670
+/// `DSML`. Some serving stacks emit each bar twice (`<||DSML||invoke …>`),
671
+/// which no marker matches, so the block falls through as text and renders raw
672
+/// in the editor.
673
+///
674
+/// Folding the run on both sides of `DSML` down to a single bar, rather than
675
+/// listing a doubled-bar constant next to each one, keeps the marker set at six
676
+/// and covers a run of any length. A bar run elsewhere is never touched, so a
677
+/// fullwidth bar used as prose or as a table rule passes through unchanged. The
678
+/// fold is not conditional on the block parsing, so a doubled-bar marker that
679
+/// turns out not to be a call comes back with single bars — markup either way.
680
+///
681
+/// Returns `None` when there is nothing to fold — in particular for any buffer
682
+/// without a `DSML` marker, which is the common case and costs no allocation.
683
+fn fold_dsml_bars(text: &str) -> Option<String> {
684
+ if !text.contains(DSML_MARKER_NAME) {
685
+ return None;
686
+ }
687
+ let mut folded = String::with_capacity(text.len());
688
+ let mut changed = false;
689
+ let mut rest = text;
690
+ while let Some(idx) = rest.find(DSML_MARKER_NAME) {
691
+ let (before, from_name) = rest.split_at(idx);
692
+ // The run of bars closing the tag being opened, i.e. `<||` in
693
+ // `<||DSML||invoke`. Counted in bars, not bytes: a lone fullwidth bar
694
+ // is three bytes, so a byte count would read as a run of its own.
695
+ let opening_run = bar_run_len(before, false);
696
+ folded.push_str(&before[..before.len() - opening_run * FULLWIDTH_BAR.len_utf8()]);
697
+ if opening_run > 0 {
698
+ changed |= opening_run > 1;
699
+ folded.push(FULLWIDTH_BAR);
700
+ }
701
+ folded.push_str(DSML_MARKER_NAME);
702
+ let after = &from_name[DSML_MARKER_NAME.len()..];
703
+ // The run opening the tag being closed, i.e. `||>` in
704
+ // `<|DSML||/invoke>`.
705
+ let closing_run = bar_run_len(after, true);
706
+ if closing_run > 0 {
707
+ changed |= closing_run > 1;
708
+ folded.push(FULLWIDTH_BAR);
709
+ }
710
+ rest = &after[closing_run * FULLWIDTH_BAR.len_utf8()..];
711
+ }
712
+ folded.push_str(rest);
713
+ changed.then_some(folded)
714
+}
715
+
716
+/// Number of fullwidth bars in the run at the end of `text`, or at its start
717
+/// when `leading` is set.
718
+fn bar_run_len(text: &str, leading: bool) -> usize {
719
+ let bytes = if leading {
720
+ text.len() - text.trim_start_matches(FULLWIDTH_BAR).len()
721
+ } else {
722
+ text.len() - text.trim_end_matches(FULLWIDTH_BAR).len()
723
+ };
724
+ bytes / FULLWIDTH_BAR.len_utf8()
725
+}
726
+
727
fn scan_complete_text(text: &str, tools: &[ToolSpec]) -> Vec<ScanEvent> {
728
let mut scanner = StreamScanner::new(tools);
729
let mut events = scanner.push(text);
1565
))
1566
);
1567
}
1568
+
1569
+ #[test]
1570
+ fn fold_dsml_bars_collapses_runs_glued_to_dsml_only() {
1571
+ assert_eq!(
1572
+ fold_dsml_bars("<||DSML||invoke name=\"run_command\">"),
1573
+ Some("<|DSML|invoke name=\"run_command\">".to_string())
1574
+ );
1575
+ assert_eq!(
1576
+ fold_dsml_bars("<||||DSML|/tool_calls>"),
1577
+ Some("<|DSML|/tool_calls>".to_string())
1578
+ );
1579
+ // Only the bar run adjacent to DSML folds; a doubled bar further along
1580
+ // in the tag name or body is not part of the marker grammar.
1581
+ assert_eq!(
1582
+ fold_dsml_bars("<|DSML||param||eter"),
1583
+ Some("<|DSML|param||eter".to_string())
1584
+ );
1585
+ // A bar run that isn't a marker, and single-bar markup, are left alone.
1586
+ assert_eq!(fold_dsml_bars("a || b <|DSML|invoke>"), None);
1587
+ assert_eq!(fold_dsml_bars("<|DSML|invoke>"), None);
1588
+ assert_eq!(fold_dsml_bars("plain text"), None);
1589
+ }
1590
+
1591
+ #[test]
1592
+ fn a_doubled_bar_dsml_call_is_recovered() {
1593
+ let tools = vec![run_command_spec()];
1594
+ let Extracted {
1595
+ text,
1596
+ calls,
1597
+ malformed,
1598
+ } = extract(
1599
+ "before <||DSML||tool_calls><||DSML||invoke name=\"run_command\">\
1600
+ <||DSML||parameter name=\"command\" string=\"true\">echo hi\
1601
+ <||DSML||/parameter><||DSML||/invoke><||DSML||/tool_calls> after",
1602
+ &tools,
1603
+ );
1604
+
1605
+ assert_eq!(text, "before after");
1606
+ assert_eq!(malformed, 0);
1607
+ assert_eq!(calls.len(), 1);
1608
+ assert_eq!(calls[0].name, "run_command");
1609
+ let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1610
+ assert_eq!(args["command"], "echo hi");
1611
+ }
1612
+
1613
+ #[test]
1614
+ fn a_doubled_bar_split_across_pushes_is_recovered() {
1615
+ let tools = vec![run_command_spec()];
1616
+ let mut scanner = StreamScanner::new(&tools);
1617
+ let mut text = String::new();
1618
+ let mut calls = Vec::new();
1619
+
1620
+ for chunk in [
1621
+ "pre <|",
1622
+ "|DSML||tool_calls><||DSML||invoke name=\"run_com",
1623
+ "mand\"><||DSML||parameter name=\"command\" string=\"true\">echo hi\
1624
+ <||DSML||/parameter><||DSML||/invoke><||DSML||/tool_calls>",
1625
+ " post",
1626
+ ] {
1627
+ for event in scanner.push(chunk) {
1628
+ match event {
1629
+ ScanEvent::Text(t) => text.push_str(&t),
1630
+ ScanEvent::ToolCall(call) => calls.push(call),
1631
+ ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1632
+ }
1633
+ }
1634
+ }
1635
+ if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1636
+ text.push_str(&rest);
1637
+ }
1638
+
1639
+ assert_eq!(text, "pre post");
1640
+ assert_eq!(calls.len(), 1);
1641
+ assert_eq!(calls[0].name, "run_command");
1642
+ let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1643
+ assert_eq!(args["command"], "echo hi");
1644
+ }
1645
+
1646
+ #[test]
1647
+ fn doubled_bar_markup_outside_a_call_returns_folded_as_text() {
1648
+ let tools = vec![run_command_spec()];
1649
+ let Extracted {
1650
+ text,
1651
+ calls,
1652
+ malformed,
1653
+ } = extract("the tag <||DSML||done> is not a call", &tools);
1654
+
1655
+ assert_eq!(text, "the tag <|DSML|done> is not a call");
1656
+ assert!(calls.is_empty());
1657
+ assert_eq!(malformed, 0);
1658
+ }
1659
}