Skip to main content

trace_stream/
dsml.rs

1// Copyright (c) 2026 Enzo Lombardi
2// SPDX-License-Identifier: MIT
3
4//! Streaming DSML tool-call parser.
5//!
6//! The model streams raw text tokens. This parser recognizes completed DSML
7//! tool stanzas (`<|DSML|tool_calls>` ... `</|DSML|tool_calls|>`) and keeps
8//! a copy of the raw stanza for diagnostics. Inner tags tolerate the one
9//! observed typo (a dropped leading `|`, e.g. `<DSML|invoke ...>`), matching
10//! the tolerance the stanza opener already had; beyond that the parser stays
11//! strict, so the actual tool parser stays small and predictable.
12//!
13//! Port of the `agent_dsml_*` family from `ds4_agent.c`.
14
15const DSML_START: &[u8] = "<|DSML|tool_calls>".as_bytes();
16const SSML_START: &[u8] = "<|SSML|tool_calls>".as_bytes();
17/// The same openers with a trailing `|` before `>`. Closing tags have always
18/// tolerated that bar; post-update weights emit it on the opener too, and
19/// without these forms the stanza never opens and the model only sees the
20/// downstream "DSML markup outside a valid `tool_calls` block" error.
21const DSML_START_BAR: &[u8] = "<|DSML|tool_calls|>".as_bytes();
22const SSML_START_BAR: &[u8] = "<|SSML|tool_calls|>".as_bytes();
23/// Cheap scan filter used to locate candidate closing tags: any `</` byte
24/// pair, not just a validated close marker. Real validation happens in
25/// [`close_tag_at`], which requires a full [`tag_prefix_len`] match against
26/// the accepted marker/name spellings — so a bare `</` inside a parameter
27/// value (e.g. HTML written through a `write` or `edit` call) never
28/// terminates the parameter on its own.
29const CLOSE_SCAN_HEAD: &[u8] = "</".as_bytes();
30const DSML_BAR: &[u8] = "|".as_bytes();
31
32/// Marker names accepted inside a tag: `<|NAME|invoke ...>`.
33///
34/// `DSML` is canonical and the only form the system prompt teaches. `SSML` is
35/// an alias for a misspelling the model actually emits: `|DSML|` is a
36/// dedicated vocab token, but plank composes the tools prompt as an ordinary
37/// system message, so the marker arrives as ordinary BPE pieces and the model
38/// spells it back out — where the far more common pretraining string "SSML"
39/// occasionally wins the "D". Without the alias the stanza parses as nothing,
40/// prints raw, and the turn ends with no tool error for the model to retry
41/// from. The prompt tells the model SSML is unsupported so this stays a
42/// recovery path rather than a second syntax.
43pub(crate) const MARKER_NAMES: [&str; 2] = ["DSML", "SSML"];
44
45/// Matches an opening or closing tag prefix for `name` under any accepted
46/// marker, returning the matched length.
47///
48/// Both the canonical `<|NAME|tag` and the dropped-leading-bar `<NAME|tag`
49/// typo are accepted, mirroring the tolerance `dsml_start_match` has always
50/// had on the stanza opener. The two forms differ in length, so the matched
51/// length is taken from the form that actually matched.
52pub(crate) fn tag_prefix_len(s: &[u8], closing: bool, name: &str) -> Option<usize> {
53    MARKER_NAMES.iter().find_map(|marker| {
54        tag_prefix_forms(marker, closing, name)
55            .into_iter()
56            .find_map(|form| segments_prefix_of(&form, s))
57    })
58}
59
60/// True when `s` is a (possibly incomplete) prefix of a tag opener for `name`
61/// under any accepted marker, in either the canonical or dropped-bar form.
62pub(crate) fn tag_prefix_partial(s: &[u8], closing: bool, name: &str) -> bool {
63    MARKER_NAMES.iter().any(|marker| {
64        tag_prefix_forms(marker, closing, name)
65            .iter()
66            .any(|form| is_prefix_of_segments(s, form))
67    })
68}
69
70/// The accepted spellings of a tag prefix, as segment lists: canonical first,
71/// then the dropped-leading-bar typo the model actually emits.
72///
73/// Segments rather than assembled `String`s because this sits on the hot path:
74/// [`find_close_tag_any`] re-scans the accumulated parameter value on every
75/// `feed`, and `CLOSE_SCAN_HEAD` is a bare `</`, which occurs on nearly every
76/// line of HTML written through a `write` or `edit` parameter. Building the
77/// spellings with `format!` cost four transient allocations per candidate,
78/// which is order 10^6 for a few hundred lines of markup.
79fn tag_prefix_forms<'a>(marker: &'a str, closing: bool, name: &'a str) -> [[&'a [u8]; 6]; 2] {
80    let slash: &[u8] = if closing { b"/" } else { b"" };
81    let (marker, name) = (marker.as_bytes(), name.as_bytes());
82    [
83        [b"<", slash, DSML_BAR, marker, DSML_BAR, name],
84        [b"<", slash, marker, DSML_BAR, name, b""],
85    ]
86}
87
88/// Length of the concatenated `segments` when they are a prefix of `s`.
89fn segments_prefix_of(segments: &[&[u8]], s: &[u8]) -> Option<usize> {
90    let mut at = 0;
91    for seg in segments {
92        if !s[at..].starts_with(seg) {
93            return None;
94        }
95        at += seg.len();
96    }
97    Some(at)
98}
99
100/// True when `s` is a prefix of the concatenated `segments`, `s` possibly
101/// stopping part-way through one of them.
102fn is_prefix_of_segments(s: &[u8], segments: &[&[u8]]) -> bool {
103    let mut rest = s;
104    for seg in segments {
105        if rest.len() < seg.len() {
106            return seg.starts_with(rest);
107        }
108        if !rest.starts_with(seg) {
109            return false;
110        }
111        rest = &rest[seg.len()..];
112    }
113    rest.is_empty()
114}
115
116/// Byte offset of the earliest complete tool-call stanza opening in `s`, if any.
117///
118/// Port of the C server's `find_any_tool_start`: the wrapper opener under any
119/// accepted marker, its dropped-leading-bar typo, and the bare `<tool_calls>`
120/// the model sometimes emits. Deliberately *not* the bare `invoke` opener the
121/// streaming detector also accepts — this feeds mid-generation recovery, where
122/// acting on a weaker signal costs a forced injection.
123///
124/// Matching is on accumulated text, so how the marker was tokenized does not
125/// matter; an incomplete opening does not match, and the caller is expected to
126/// re-scan from far enough back that one split across tokens is still seen.
127#[must_use]
128pub fn find_tool_start(s: &str) -> Option<usize> {
129    let mut forms: Vec<String> = vec!["<tool_calls>".to_owned()];
130    for m in MARKER_NAMES {
131        forms.push(format!("<|{m}|tool_calls>"));
132        forms.push(format!("<|{m}|tool_calls|>"));
133        forms.push(format!("<{m}|tool_calls>"));
134        forms.push(format!("<{m}|tool_calls|>"));
135    }
136    forms.iter().filter_map(|f| s.find(f.as_str())).min()
137}
138
139/// Bytes held back when re-scanning a stream for [`find_tool_start`]: longer
140/// than the longest opening, so one split across future tokens is still seen
141/// from its first byte.
142pub const TOOL_START_SCAN_HOLD: usize = 80;
143
144/// One named argument of a parsed tool call.
145#[derive(Debug, Clone, PartialEq, Eq)]
146pub struct ToolArg {
147    /// Argument name from the `name="..."` attribute.
148    pub name: String,
149    /// Raw argument value (bytes between the parameter tags).
150    pub value: String,
151    /// True when the parameter carried `string="true"`.
152    pub is_string: bool,
153}
154
155/// A parsed tool invocation: tool name plus its arguments in stream order.
156#[derive(Debug, Clone, Default, PartialEq, Eq)]
157pub struct ToolCall {
158    /// Tool name from the invoke tag's `name="..."` attribute.
159    pub name: String,
160    /// Arguments in the order they were streamed.
161    pub args: Vec<ToolArg>,
162}
163
164impl ToolCall {
165    /// Returns the value of the named argument, if present.
166    pub fn arg_value(&self, name: impl AsRef<str>) -> Option<&str> {
167        let name = name.as_ref();
168        self.args
169            .iter()
170            .find(|a| a.name == name)
171            .map(|a| a.value.as_str())
172    }
173}
174
175/// Parser progress; terminal states are `Done` and `Error`.
176#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
177pub enum DsmlState {
178    /// Scanning free text for the opening `<|DSML|tool_calls>` marker.
179    #[default]
180    Search,
181    /// Between tags: expecting invoke/parameter open tags or close tags.
182    Structural,
183    /// Accumulating a parameter value until its close tag arrives.
184    ParamValue,
185    /// A full `tool_calls` stanza was parsed.
186    Done,
187    /// The stanza was malformed; see [`DsmlParser::error`].
188    Error,
189}
190
191/// Incremental parser for one DSML tool-call stanza.
192///
193/// Feed streamed bytes with [`feed`](Self::feed); it can be called after every
194/// byte. Incomplete input leaves state unchanged until enough bytes arrive,
195/// while malformed completed input switches to [`DsmlState::Error`] so the
196/// model gets a retryable tool error.
197#[derive(Debug, Default)]
198pub struct DsmlParser {
199    state: DsmlState,
200    search_tail: Vec<u8>,
201    raw: Vec<u8>,
202    parse_pos: usize,
203    current: Option<PendingCall>,
204    param_name: Option<String>,
205    param_is_string: bool,
206    param_value_start: usize,
207    /// Element name when the open parameter used the shorthand form
208    /// (`<|DSML|command …>` instead of `<|DSML|parameter name="command" …>`),
209    /// which widens the accepted terminators for *this value only*. `None` for
210    /// canonical parameters, so their strict `parameter`-only terminator — the
211    /// thing that keeps a `</` inside a `write` payload from ending the value —
212    /// is never relaxed.
213    param_elem: Option<String>,
214    /// True while the raw tail looks like a partial parameter close tag, so
215    /// online rendering can hide it before the full tag arrives.
216    param_close_prefix: bool,
217    calls: Vec<ToolCall>,
218    error: String,
219}
220
221#[derive(Debug, Default)]
222struct PendingCall {
223    name: String,
224    args: Vec<ToolArg>,
225}
226
227/// True when a name is an unsubstituted placeholder copied from the tools
228/// prompt (`$TOOL_NAME`, `$PARAMETER_NAME`, `$PARAMETER_VALUE`).
229///
230/// Exactly those three: a tool or parameter genuinely named `$path` is a
231/// different mistake and must not be told it copied the prompt.
232fn is_prompt_placeholder(name: &str) -> bool {
233    matches!(name, "$TOOL_NAME" | "$PARAMETER_NAME" | "$PARAMETER_VALUE")
234}
235
236impl DsmlParser {
237    /// Creates a parser in the `Search` state.
238    #[must_use]
239    pub fn new() -> Self {
240        Self::default()
241    }
242
243    /// Current parser state.
244    #[must_use]
245    pub fn state(&self) -> DsmlState {
246        self.state
247    }
248
249    /// Tool calls completed so far, in stream order.
250    #[must_use]
251    pub fn calls(&self) -> &[ToolCall] {
252        &self.calls
253    }
254
255    /// Error message; empty unless the state is [`DsmlState::Error`].
256    #[must_use]
257    pub fn error(&self) -> &str {
258        &self.error
259    }
260
261    /// Snapshot of the invoke currently being parsed (name plus the
262    /// arguments whose close tags have arrived), for mid-stream preflight.
263    #[must_use]
264    pub fn pending_call(&self) -> Option<ToolCall> {
265        self.current.as_ref().map(|c| ToolCall {
266            name: c.name.clone(),
267            args: c.args.clone(),
268        })
269    }
270
271    /// Raw bytes of the stanza accumulated so far, for diagnostics.
272    #[must_use]
273    pub fn raw(&self) -> &[u8] {
274        &self.raw
275    }
276
277    /// True while the raw tail is a partial parameter close tag.
278    #[must_use]
279    pub fn param_close_prefix(&self) -> bool {
280        self.param_close_prefix
281    }
282
283    /// Resets the parser to a fresh `Search` state, discarding all results.
284    pub fn reset(&mut self) {
285        *self = Self::default();
286    }
287
288    /// Feeds streamed bytes; no-op once the parser is `Done` or `Error`.
289    pub fn feed(&mut self, s: impl AsRef<[u8]>) {
290        let s = s.as_ref();
291        if matches!(self.state, DsmlState::Done | DsmlState::Error) {
292            return;
293        }
294        for &c in s {
295            if self.state == DsmlState::Search {
296                if self.search_tail.len() == 64 {
297                    self.search_tail.remove(0);
298                }
299                self.search_tail.push(c);
300                if [DSML_START, SSML_START, DSML_START_BAR, SSML_START_BAR]
301                    .iter()
302                    .any(|f| self.search_tail.ends_with(f))
303                {
304                    self.start();
305                }
306                continue;
307            }
308
309            self.raw.push(c);
310            self.parse();
311            if self.state == DsmlState::ParamValue {
312                self.update_param_close_prefix();
313            } else {
314                self.param_close_prefix = false;
315            }
316        }
317    }
318
319    fn start(&mut self) {
320        self.state = DsmlState::Structural;
321        self.search_tail.clear();
322        self.raw.extend_from_slice(DSML_START);
323        self.parse_pos = DSML_START.len();
324    }
325
326    fn set_error(&mut self, msg: impl Into<String>) {
327        self.state = DsmlState::Error;
328        self.error = msg.into();
329    }
330
331    fn push_current(&mut self) {
332        if let Some(call) = self.current.take() {
333            self.calls.push(ToolCall {
334                name: call.name,
335                args: call.args,
336            });
337        }
338    }
339
340    /// Parses as much of the accumulated buffer as possible.
341    fn parse(&mut self) {
342        loop {
343            match self.state {
344                DsmlState::ParamValue => {
345                    // A shorthand parameter is closed inconsistently: the
346                    // recorded repro ends `<|DSML|command …>` with
347                    // `</|DSML|invoke>`, not `</|DSML|command>`. Accept either,
348                    // plus `parameter`, and take whichever lands first.
349                    let mut names: Vec<&str> = vec!["parameter"];
350                    if let Some(elem) = self.param_elem.as_deref() {
351                        names.push(elem);
352                        names.push("invoke");
353                    }
354                    let Some((end, tag_len)) =
355                        find_close_tag_any(&self.raw[self.param_value_start..], &names)
356                    else {
357                        return;
358                    };
359                    let value_bytes =
360                        &self.raw[self.param_value_start..self.param_value_start + end];
361                    let arg = ToolArg {
362                        name: self.param_name.take().unwrap_or_default(),
363                        value: String::from_utf8_lossy(value_bytes).into_owned(),
364                        is_string: self.param_is_string,
365                    };
366                    self.current
367                        .get_or_insert_with(Default::default)
368                        .args
369                        .push(arg);
370                    self.param_close_prefix = false;
371                    self.param_elem = None;
372                    self.parse_pos = self.param_value_start + end + tag_len;
373                    self.state = DsmlState::Structural;
374                }
375                DsmlState::Structural => {
376                    while self.parse_pos < self.raw.len()
377                        && self.raw[self.parse_pos].is_ascii_whitespace()
378                    {
379                        self.parse_pos += 1;
380                    }
381                    if self.parse_pos >= self.raw.len() {
382                        return;
383                    }
384
385                    let rest = &self.raw[self.parse_pos..];
386                    if let Some(close_len) = close_tag_at(rest, "tool_calls") {
387                        self.push_current();
388                        self.parse_pos += close_len;
389                        self.state = DsmlState::Done;
390                        return;
391                    }
392                    if let Some(close_len) = close_tag_at(rest, "invoke") {
393                        self.push_current();
394                        self.parse_pos += close_len;
395                        continue;
396                    }
397
398                    let Some(gt) = rest.iter().position(|&b| b == b'>') else {
399                        return;
400                    };
401                    let tag_len = gt + 1;
402                    let tag = String::from_utf8_lossy(&rest[..tag_len]).into_owned();
403
404                    if !self.open_tag(&tag, tag_len) {
405                        return;
406                    }
407                }
408                _ => return,
409            }
410        }
411    }
412
413    /// Handles one opening tag inside a stanza. Returns false when parsing must
414    /// stop, having set the error.
415    ///
416    /// The four accepted spellings, in precedence order: a canonical `invoke`,
417    /// a canonical `parameter`, and the two bare-element shorthands — parameter
418    /// first, since it is the one that applies while an invoke is open (see
419    /// [`Self::shorthand_invoke_name`] for why the open invoke is the whole
420    /// distinction).
421    fn open_tag(&mut self, tag: &str, tag_len: usize) -> bool {
422        // A repeated wrapper opener is the model restating itself, not a second
423        // stanza — `repro-1785770781.md` does it 37 times. The stanza is already
424        // open, so consuming the tag and moving on is idempotent, and strictly
425        // better than the alternatives: erroring costs the turn, and treating it
426        // as a name invents a `tool_calls` call that swallows the parameters.
427        if open_tag_is(tag, "tool_calls") {
428            self.parse_pos += tag_len;
429            return true;
430        }
431        if open_tag_is(tag, "invoke") {
432            let Some(name) = parse_attr(tag, "name") else {
433                self.set_error("tool invoke without name");
434                return false;
435            };
436            if is_prompt_placeholder(&name) {
437                self.set_error(format!(
438                    "tool name is the prompt's placeholder {name}, not a real tool; substitute the actual tool name"
439                ));
440                return false;
441            }
442            self.open_invoke(name, tag_len);
443        } else if open_tag_is(tag, "parameter") {
444            let Some(name) = parse_attr(tag, "name") else {
445                self.set_error("tool parameter without name");
446                return false;
447            };
448            if is_prompt_placeholder(&name) {
449                self.set_error(format!(
450                    "parameter name is the prompt's placeholder {name}, not a real parameter; substitute the actual parameter name"
451                ));
452                return false;
453            }
454            self.param_elem = None;
455            self.open_param(name, tag, tag_len);
456        } else if let Some(elem) = self.shorthand_param_name(tag) {
457            self.param_elem = Some(elem.clone());
458            self.open_param(elem, tag, tag_len);
459        } else if let Some(elem) = self.shorthand_invoke_name(tag) {
460            self.open_invoke(elem, tag_len);
461        } else {
462            let shown: String = tag.chars().take(80).collect();
463            self.set_error(format!("unexpected DSML tag: {shown}"));
464            return false;
465        }
466        true
467    }
468
469    /// Opens a tool call, however its name was spelled.
470    fn open_invoke(&mut self, name: String, tag_len: usize) {
471        self.current = Some(PendingCall {
472            name,
473            args: Vec::new(),
474        });
475        self.parse_pos += tag_len;
476    }
477
478    /// Enters `ParamValue` for a parameter named `name` opened by `tag`.
479    fn open_param(&mut self, name: String, tag: &str, tag_len: usize) {
480        self.param_name = Some(name);
481        self.param_is_string = parse_attr(tag, "string").as_deref() == Some("true");
482        self.parse_pos += tag_len;
483        self.param_value_start = self.parse_pos;
484        self.param_close_prefix = false;
485        self.state = DsmlState::ParamValue;
486    }
487
488    /// The parameter name for the shorthand form the post-update weights emit:
489    /// the parameter name written as the element name,
490    /// `<|DSML|command string="true">ls</|DSML|invoke>` in place of
491    /// `<|DSML|parameter name="command" string="true">ls</|DSML|parameter>`.
492    ///
493    /// Deliberately narrow, because accepting it means *running* a tool call
494    /// that was not written the way the prompt teaches. All of these must hold:
495    /// an invoke is already open, the tag carries the DSML marker, its element
496    /// name is a plain identifier, and it has no `name` attribute — a tag with
497    /// one is some other malformation and still errors. Rejecting instead is
498    /// not free: the recorded repro shows the model unable to find its way back
499    /// from the error, re-emitting the same shape and then breaking the think
500    /// gate, so the turn is lost either way.
501    fn shorthand_param_name(&self, tag: &str) -> Option<String> {
502        if self.current.is_none() || parse_attr(tag, "name").is_some() {
503            return None;
504        }
505        let elem = element_name(tag)?;
506        (!is_prompt_placeholder(&elem) && !Self::STRUCTURAL_ELEMS.contains(&elem.as_str()))
507            .then_some(elem)
508    }
509
510    /// Structural element names, which can never be a tool or a parameter name.
511    ///
512    /// Without this the shorthands read the model restating a wrapper tag as a
513    /// name: `repro-1785770781.md` opens 37 stanzas with `<|DSML|tool_calls>`
514    /// twice in a row, and the second one parsed as a *successful* call named
515    /// `tool_calls` that swallowed the real parameters. Erroring would be better
516    /// than that; skipping it, as [`Self::open_tag`] now does, is better still.
517    const STRUCTURAL_ELEMS: [&'static str; 3] = ["tool_calls", "invoke", "parameter"];
518
519    /// The same shorthand one level up: the *tool* name written as the element
520    /// name, `<|DSML|edit>…</|DSML|invoke>` in place of
521    /// `<|DSML|invoke name="edit">…</|DSML|invoke>`.
522    ///
523    /// Which of the two shorthands a bare element is depends on one thing:
524    /// whether an invoke is open. Before one, the model is naming the tool it
525    /// wants; inside one, a parameter. That is why this is checked *after*
526    /// [`Self::shorthand_param_name`] — the parameter reading wins whenever
527    /// both could apply, which is exactly when an invoke is already open.
528    ///
529    /// Same narrowness as the parameter form, and for the same reason:
530    /// accepting it means *running* a call written the way the prompt does not
531    /// teach. The tag must carry the DSML marker, its element name must be a
532    /// plain identifier, and it must have no `name` attribute — a tag with one
533    /// is some other malformation and still errors. An element name that is not
534    /// a real tool reaches dispatch and fails there by name, which is a clear
535    /// error the model can act on; rejecting the stanza outright is not free.
536    /// `repro-1785754509.md` is the recorded cost: five rejections of this
537    /// shape, no recovery, and the model finally breaking the think gate while
538    /// trying to restate the syntax back to itself.
539    fn shorthand_invoke_name(&self, tag: &str) -> Option<String> {
540        if self.current.is_some() || parse_attr(tag, "name").is_some() {
541            return None;
542        }
543        let elem = element_name(tag)?;
544        (!is_prompt_placeholder(&elem) && !Self::STRUCTURAL_ELEMS.contains(&elem.as_str()))
545            .then_some(elem)
546    }
547
548    /// Tracks whether the raw tail is a partial parameter close tag, so the
549    /// terminal renderer can hide it without waiting for the whole parameter.
550    fn update_param_close_prefix(&mut self) {
551        self.param_close_prefix = false;
552        if self.state != DsmlState::ParamValue || self.raw.len() <= self.param_value_start {
553            return;
554        }
555        let value = &self.raw[self.param_value_start..];
556        let Some(lt) = value.iter().rposition(|&b| b == b'<') else {
557            return;
558        };
559        let tail = &value[lt..];
560        if tail.len() > 64 || tag_prefix_len(tail, true, "").is_none() {
561            return;
562        }
563        let mut complete = false;
564        self.param_close_prefix = parameter_close_tail(tail, &mut complete) && !complete;
565    }
566}
567
568/// The element name of a DSML-marked opening tag, e.g. `command` for
569/// `<|DSML|command string="true">`.
570///
571/// Used only to sharpen the "unexpected DSML tag" error. Post-update weights
572/// write the *parameter name* as the element name; echoing the tag back taught
573/// the model nothing, and the recorded repro shows it guessing at the marker
574/// spelling for three turns and then emitting DSML inside `<think>`. Naming the
575/// rewrite gives it something to act on.
576pub(crate) fn element_name(tag: &str) -> Option<String> {
577    let len = tag_prefix_len(tag.as_bytes(), false, "")?;
578    let name: String = tag[len..]
579        .chars()
580        .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
581        .collect();
582    (!name.is_empty()).then_some(name)
583}
584
585/// Checks whether `tag` is an opening DSML tag with the given element name.
586fn open_tag_is(tag: &str, name: &str) -> bool {
587    let Some(len) = tag_prefix_len(tag.as_bytes(), false, name) else {
588        return false;
589    };
590    tag.as_bytes()
591        .get(len)
592        .is_some_and(|&c| c == b'>' || c.is_ascii_whitespace())
593}
594
595/// Recognizes a DSML closing tag at the start of `s`, returning its length.
596///
597/// Accepts the few harmless closing-tag variants the model has been observed
598/// to emit (whitespace and an optional trailing `|` before `>`). Opening tags
599/// stay strict so accidental prose does not become a tool call.
600fn close_tag_at(s: &[u8], name: &str) -> Option<usize> {
601    let mut i = tag_prefix_len(s, true, name)?;
602    while i < s.len() && s[i].is_ascii_whitespace() {
603        i += 1;
604    }
605    if s[i..].starts_with(DSML_BAR) {
606        i += DSML_BAR.len();
607    }
608    while i < s.len() && s[i].is_ascii_whitespace() {
609        i += 1;
610    }
611    if s.get(i) != Some(&b'>') {
612        return None;
613    }
614    Some(i + 1)
615}
616
617/// Finds the earliest DSML closing tag for any of `names`; returns
618/// (offset, tag length).
619///
620/// Scanning position-first rather than name-first matters: the winner must be
621/// the tag that appears earliest in the value, not the one whose name happens
622/// to come first in the list.
623fn find_close_tag_any(s: &[u8], names: &[&str]) -> Option<(usize, usize)> {
624    let mut from = 0;
625    while let Some(pos) = find_bytes(&s[from..], CLOSE_SCAN_HEAD) {
626        let at = from + pos;
627        if let Some(tag_len) = names.iter().find_map(|n| close_tag_at(&s[at..], n)) {
628            return Some((at, tag_len));
629        }
630        from = at + 1;
631    }
632    None
633}
634
635fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
636    haystack.windows(needle.len()).position(|w| w == needle)
637}
638
639/// Recognizes a streamed parameter close tag prefix.
640///
641/// Full close detection is handled by [`close_tag_at`]; this exists for online
642/// behavior: terminal rendering must hide partial close tags without waiting
643/// for the whole parameter to finish. Sets `complete` when the tail is a full
644/// close tag ending exactly at the last byte.
645fn parameter_close_tail(tail: &[u8], complete: &mut bool) -> bool {
646    *complete = false;
647    if tag_prefix_partial(tail, true, "parameter") {
648        return true;
649    }
650    let Some(mut i) = tag_prefix_len(tail, true, "parameter") else {
651        return false;
652    };
653    while i < tail.len() && tail[i].is_ascii_whitespace() {
654        i += 1;
655    }
656    if i < tail.len() && tail.len() - i <= DSML_BAR.len() && DSML_BAR.starts_with(&tail[i..]) {
657        return true;
658    }
659    if tail[i..].starts_with(DSML_BAR) {
660        i += DSML_BAR.len();
661    }
662    while i < tail.len() {
663        if tail[i] == b'>' {
664            *complete = i == tail.len() - 1;
665            return *complete;
666        }
667        if !tail[i].is_ascii_whitespace() {
668            return false;
669        }
670        i += 1;
671    }
672    true
673}
674
675/// Extracts a `name="value"` attribute from a tag, if present.
676fn parse_attr(tag: &str, name: &str) -> Option<String> {
677    let pat = format!("{name}=\"");
678    let start = tag.find(&pat)? + pat.len();
679    let end = tag[start..].find('"')? + start;
680    Some(tag[start..end].to_string())
681}
682
683#[cfg(test)]
684mod tests {
685    use super::*;
686
687    const STANZA: &str = concat!(
688        "<|DSML|tool_calls>",
689        "<|DSML|invoke name=\"read_file\">",
690        "<|DSML|parameter name=\"path\" string=\"true\">src/main.rs</|DSML|parameter|>",
691        "<|DSML|parameter name=\"offset\">42</|DSML|parameter|>",
692        "</|DSML|invoke|>",
693        "</|DSML|tool_calls|>",
694    );
695
696    fn feed_all(p: &mut DsmlParser, s: &str) {
697        p.feed(s.as_bytes());
698    }
699
700    fn feed_bytewise(p: &mut DsmlParser, s: &str) {
701        for b in s.as_bytes() {
702            p.feed([*b]);
703        }
704    }
705
706    /// Post-update weights close the stanza opener with `|>` rather than `>`.
707    /// Before this was accepted the stanza never opened at all, and the model
708    /// saw only "DSML markup outside a valid `tool_calls` block" — with no way to
709    /// tell which part of its syntax was rejected — turn after turn.
710    #[test]
711    fn opener_tolerates_trailing_bar() {
712        let stanza = STANZA.replacen("<|DSML|tool_calls>", "<|DSML|tool_calls|>", 1);
713        for feed in [feed_all as fn(&mut DsmlParser, &str), feed_bytewise] {
714            let mut p = super::DsmlParser::new();
715            feed(&mut p, &stanza);
716            assert_eq!(p.state(), super::DsmlState::Done);
717            assert_eq!(p.calls().len(), 1);
718            assert_eq!(p.calls()[0].name, "read_file");
719            assert_eq!(p.calls()[0].arg_value("path"), Some("src/main.rs"));
720        }
721        assert_eq!(super::find_tool_start(&stanza), Some(0));
722    }
723
724    /// Post-update weights write the parameter name as the element name, and
725    /// close it with `</|DSML|invoke>`. Verbatim from the recorded repro, in
726    /// which rejecting it cost the whole turn.
727    #[test]
728    fn shorthand_parameter_element_is_executed() {
729        let stanza = concat!(
730            "<|DSML|tool_calls|>",
731            "<|DSML|invoke name=\"bash\">",
732            "<|DSML|command string=\"true\">cd /tmp && ls</|DSML|invoke>",
733            "</|DSML|invoke>",
734            "</|DSML|tool_calls|>",
735        );
736        for feed in [feed_all as fn(&mut DsmlParser, &str), feed_bytewise] {
737            let mut p = super::DsmlParser::new();
738            feed(&mut p, stanza);
739            assert_eq!(p.state(), super::DsmlState::Done, "{}", p.error());
740            assert_eq!(p.calls().len(), 1);
741            assert_eq!(p.calls()[0].name, "bash");
742            assert_eq!(p.calls()[0].arg_value("command"), Some("cd /tmp && ls"));
743            assert!(p.calls()[0].args[0].is_string);
744        }
745    }
746
747    /// The same shorthand one level up: the *tool* name written as the element
748    /// name, with no `invoke` wrapper. Verbatim from `repro-1785754509.md`,
749    /// where the model emitted this shape five times for `write` and `edit`,
750    /// never recovered from the rejection, and finally broke the think gate
751    /// trying to restate the syntax.
752    #[test]
753    fn shorthand_invoke_element_is_executed() {
754        let stanza = concat!(
755            "<|DSML|tool_calls>",
756            "<|DSML|edit>",
757            "<|DSML|parameter name=\"path\" string=\"true\">/tmp/a.rs</|DSML|parameter>",
758            "<|DSML|parameter name=\"old\" string=\"true\">one</|DSML|parameter>",
759            "<|DSML|parameter name=\"new\" string=\"true\">two</|DSML|parameter>",
760            "</|DSML|invoke>",
761            "</|DSML|tool_calls>",
762        );
763        for feed in [feed_all as fn(&mut DsmlParser, &str), feed_bytewise] {
764            let mut p = super::DsmlParser::new();
765            feed(&mut p, stanza);
766            assert_eq!(p.state(), super::DsmlState::Done, "{}", p.error());
767            assert_eq!(p.calls().len(), 1);
768            assert_eq!(p.calls()[0].name, "edit");
769            assert_eq!(p.calls()[0].arg_value("path"), Some("/tmp/a.rs"));
770            assert_eq!(p.calls()[0].arg_value("old"), Some("one"));
771            assert_eq!(p.calls()[0].arg_value("new"), Some("two"));
772        }
773    }
774
775    /// A repeated `<|DSML|tool_calls>` opener is the model restating itself and
776    /// must not become a call. Verbatim from `repro-1785770781.md`, which opens
777    /// 37 stanzas that way; the invoke shorthand read the second one as a tool
778    /// named `tool_calls` and reported Done, so a nonsense call reached dispatch
779    /// carrying the real parameters.
780    #[test]
781    fn a_repeated_wrapper_opener_is_skipped_not_named() {
782        // The benign case: the repeat is absorbed and the real call is intact.
783        let stanza = concat!(
784            "<|DSML|tool_calls>",
785            "<|DSML|tool_calls>",
786            "<|DSML|invoke name=\"bash\">",
787            "<|DSML|parameter name=\"command\" string=\"true\">ls -la</|DSML|parameter>",
788            "</|DSML|invoke>",
789            "</|DSML|tool_calls>",
790        );
791        for feed in [feed_all as fn(&mut DsmlParser, &str), feed_bytewise] {
792            let mut p = super::DsmlParser::new();
793            feed(&mut p, stanza);
794            assert_eq!(p.state(), super::DsmlState::Done, "{}", p.error());
795            assert_eq!(p.calls().len(), 1, "{:?}", p.calls());
796            assert_eq!(p.calls()[0].name, "bash");
797            assert_eq!(p.calls()[0].arg_value("command"), Some("ls -la"));
798        }
799
800        // The dangerous case, and the one that regressed: with no real invoke
801        // after the repeat, the shorthand read `tool_calls` as the tool name and
802        // reported Done, so a call literally named `tool_calls` — carrying the
803        // parameters meant for the real tool — reached dispatch. Whatever else
804        // this shape yields, that name must never appear.
805        let no_invoke = concat!(
806            "<|DSML|tool_calls>",
807            "<|DSML|tool_calls>",
808            "<|DSML|parameter name=\"command\" string=\"true\">ls</|DSML|parameter>",
809            "</|DSML|invoke>",
810            "</|DSML|tool_calls>",
811        );
812        for feed in [feed_all as fn(&mut DsmlParser, &str), feed_bytewise] {
813            let mut p = super::DsmlParser::new();
814            feed(&mut p, no_invoke);
815            assert!(
816                p.calls().iter().all(|c| c.name != "tool_calls"),
817                "the wrapper name must never become a tool: {:?}",
818                p.calls()
819            );
820        }
821    }
822
823    /// The full malformed stanza from `repro-1785770781.md`, verbatim: a repeated
824    /// wrapper opener followed by `<|DSML|tool ATTR="...">`, where the model
825    /// folded the element name and the attribute name together.
826    ///
827    /// The repeated opener is tolerated, but the rest is *not* guessed at. The
828    /// shape is self-inconsistent — `path="/tmp/lib.rs"` puts the value in the
829    /// attribute while `old="true"` uses it as the `string=` flag with the value
830    /// as element text — so any reading would be wrong half the time, and this is
831    /// an `edit`. Erroring is the correct outcome; fabricating a call is not.
832    #[test]
833    fn the_folded_attribute_shape_errors_without_fabricating_a_call() {
834        let stanza = concat!(
835            "<|DSML|tool_calls>\n",
836            "<|DSML|tool_calls>\n",
837            "<|DSML|tool name=\"edit\">\n",
838            "<|DSML|tool path=\"/tmp/lib.rs\">\n",
839            "<|DSML|tool old=\"true\">OLD TEXT</|DSML|tool>\n",
840            "</|DSML|invoke>\n",
841            "</|DSML|tool_calls>",
842        );
843        for feed in [feed_all as fn(&mut DsmlParser, &str), feed_bytewise] {
844            let mut p = super::DsmlParser::new();
845            feed(&mut p, stanza);
846            assert_eq!(p.state(), super::DsmlState::Error, "{:?}", p.calls());
847            assert!(
848                p.error().starts_with("unexpected DSML tag:"),
849                "{}",
850                p.error()
851            );
852            // Nothing executable may survive under the wrapper's name, and no
853            // half-built `edit` carrying a bogus `old`.
854            assert!(
855                p.calls().iter().all(|c| c.name != "tool_calls"),
856                "{:?}",
857                p.calls()
858            );
859            assert!(
860                p.calls().iter().all(|c| c.arg_value("old") != Some("true")),
861                "`old=\"true\"` is a string flag, never the old text: {:?}",
862                p.calls()
863            );
864        }
865    }
866
867    /// The shorthands must never read a structural element as a name, whichever
868    /// level they are at. `tool_calls` is the wrapper; a bare `invoke` or
869    /// `parameter` is a missing-name error, which is its own clearer message.
870    #[test]
871    fn structural_elements_are_never_names() {
872        // At invoke level: no `tool_calls` call is fabricated.
873        let mut p = super::DsmlParser::new();
874        feed_all(
875            &mut p,
876            "<|DSML|tool_calls><|DSML|tool_calls>\
877             <|DSML|parameter name=\"command\">ls</|DSML|parameter>",
878        );
879        assert!(
880            p.calls().iter().all(|c| c.name != "tool_calls"),
881            "{:?}",
882            p.calls()
883        );
884
885        // At parameter level: an invoke is open, and a repeated wrapper tag
886        // still must not become a parameter called `tool_calls`.
887        let mut p = super::DsmlParser::new();
888        feed_all(
889            &mut p,
890            "<|DSML|tool_calls><|DSML|invoke name=\"bash\"><|DSML|tool_calls>\
891             <|DSML|parameter name=\"command\" string=\"true\">ls</|DSML|parameter>\
892             </|DSML|invoke></|DSML|tool_calls>",
893        );
894        assert_eq!(p.state(), super::DsmlState::Done, "{}", p.error());
895        assert_eq!(p.calls().len(), 1);
896        assert_eq!(p.calls()[0].name, "bash");
897        assert!(
898            p.calls()[0].args.iter().all(|a| a.name != "tool_calls"),
899            "{:?}",
900            p.calls()[0].args
901        );
902        assert_eq!(p.calls()[0].arg_value("command"), Some("ls"));
903
904        // A bare `invoke` / `parameter` keeps its own missing-name error rather
905        // than being silently accepted as a shorthand name.
906        for (text, want) in [
907            (
908                "<|DSML|tool_calls><|DSML|invoke>",
909                "tool invoke without name",
910            ),
911            (
912                "<|DSML|tool_calls><|DSML|invoke name=\"bash\"><|DSML|parameter>",
913                "tool parameter without name",
914            ),
915        ] {
916            let mut p = super::DsmlParser::new();
917            feed_all(&mut p, text);
918            assert_eq!(p.state(), super::DsmlState::Error, "{text}");
919            assert_eq!(p.error(), want, "{text}");
920        }
921    }
922
923    /// Which shorthand a bare element is depends only on whether an invoke is
924    /// open: before one it names the tool, inside one it names a parameter.
925    /// Both spellings in a single stanza must land in the right slots.
926    #[test]
927    fn bare_elements_are_tool_then_parameter_names() {
928        let mut p = super::DsmlParser::new();
929        feed_all(
930            &mut p,
931            "<|DSML|tool_calls>\
932             <|DSML|bash>\
933             <|DSML|command string=\"true\">ls -la</|DSML|command>\
934             </|DSML|invoke>\
935             </|DSML|tool_calls>",
936        );
937        assert_eq!(p.state(), super::DsmlState::Done, "{}", p.error());
938        assert_eq!(p.calls().len(), 1);
939        assert_eq!(p.calls()[0].name, "bash");
940        assert_eq!(p.calls()[0].arg_value("command"), Some("ls -la"));
941    }
942
943    /// The self-consistent spelling of the shorthand closes with its own
944    /// element name and a single `</|DSML|invoke>`.
945    #[test]
946    fn shorthand_parameter_closed_by_its_own_element() {
947        let mut p = super::DsmlParser::new();
948        feed_all(
949            &mut p,
950            "<|DSML|tool_calls>\
951             <|DSML|invoke name=\"read\">\
952             <|DSML|path string=\"true\">src/main.rs</|DSML|path>\
953             </|DSML|invoke>\
954             </|DSML|tool_calls>",
955        );
956        assert_eq!(p.state(), super::DsmlState::Done, "{}", p.error());
957        assert_eq!(p.calls()[0].arg_value("path"), Some("src/main.rs"));
958    }
959
960    /// The tolerance must not reach canonical parameters: a `write` payload
961    /// that itself contains `</|DSML|invoke>` (this repo's own sources and
962    /// docs do) still runs to its real `</|DSML|parameter>` terminator.
963    #[test]
964    fn canonical_parameter_value_is_not_truncated_by_a_foreign_close_tag() {
965        let content = "docs mentioning </|DSML|invoke> and </|DSML|command> inline";
966        let mut p = super::DsmlParser::new();
967        feed_all(
968            &mut p,
969            &format!(
970                "<|DSML|tool_calls>\
971                 <|DSML|invoke name=\"write\">\
972                 <|DSML|parameter name=\"content\" string=\"true\">{content}</|DSML|parameter|>\
973                 </|DSML|invoke|>\
974                 </|DSML|tool_calls|>"
975            ),
976        );
977        assert_eq!(p.state(), super::DsmlState::Done, "{}", p.error());
978        assert_eq!(p.calls()[0].arg_value("content"), Some(content));
979    }
980
981    /// A tag carrying `name=` is a different malformation, not the shorthand,
982    /// so it must not be silently turned into a parameter and run.
983    #[test]
984    fn unknown_element_with_a_name_attribute_still_errors() {
985        let mut p = super::DsmlParser::new();
986        feed_all(
987            &mut p,
988            "<|DSML|tool_calls><|DSML|invoke name=\"bash\">\
989             <|DSML|argument name=\"command\">ls</|DSML|argument>",
990        );
991        assert_eq!(p.state(), super::DsmlState::Error);
992        assert!(
993            p.error().starts_with("unexpected DSML tag:"),
994            "{}",
995            p.error()
996        );
997    }
998
999    /// The hint is only meaningful inside an open invoke; stray markup before
1000    /// one keeps the plain echo rather than inventing a parameter name.
1001    #[test]
1002    fn unexpected_tag_outside_an_invoke_keeps_the_plain_error() {
1003        let mut p = super::DsmlParser::new();
1004        feed_all(&mut p, "<|DSML|tool_calls><b>");
1005        assert_eq!(p.state(), super::DsmlState::Error);
1006        assert_eq!(p.error(), "unexpected DSML tag: <b>");
1007    }
1008
1009    // The model copies TOOLS_PROMPT_INTRO verbatim (4 recorded occurrences).
1010    // Telling it "not allowed inside <think>" sends it fixing placement when the
1011    // real mistake is that it never substituted anything.
1012    #[test]
1013    fn placeholder_tool_name_is_named_as_such() {
1014        let mut p = super::DsmlParser::new();
1015        p.feed("<|DSML|tool_calls><|DSML|invoke name=\"$TOOL_NAME\">".as_bytes());
1016        assert_eq!(p.state(), super::DsmlState::Error);
1017        assert_eq!(
1018            p.error(),
1019            "tool name is the prompt's placeholder $TOOL_NAME, not a real tool; substitute the actual tool name"
1020        );
1021    }
1022
1023    #[test]
1024    fn placeholder_parameter_name_is_named_as_such() {
1025        let mut p = super::DsmlParser::new();
1026        p.feed(
1027            "<|DSML|tool_calls><|DSML|invoke name=\"bash\">\
1028             <|DSML|parameter name=\"$PARAMETER_NAME\" string=\"true\">x"
1029                .as_bytes(),
1030        );
1031        assert_eq!(p.state(), super::DsmlState::Error);
1032        assert_eq!(
1033            p.error(),
1034            "parameter name is the prompt's placeholder $PARAMETER_NAME, not a real parameter; substitute the actual parameter name"
1035        );
1036    }
1037
1038    // A name with a dollar sign in it is untouched: only the three literal
1039    // placeholders from the tools prompt count, so a tool named `$path` is not
1040    // told it copied the prompt.
1041    #[test]
1042    fn dollar_inside_a_name_is_not_a_placeholder() {
1043        for name in ["we$rd", "$path"] {
1044            let mut p = super::DsmlParser::new();
1045            p.feed(
1046                format!(
1047                    "<|DSML|tool_calls><|DSML|invoke name=\"{name}\">\
1048                     </|DSML|invoke|></|DSML|tool_calls|>"
1049                )
1050                .as_bytes(),
1051            );
1052            assert_eq!(p.state(), super::DsmlState::Done, "error: {}", p.error());
1053            assert_eq!(p.calls()[0].name, name);
1054        }
1055    }
1056
1057    /// The SSML alias (see [`MARKER_NAMES`]) must parse identically to the
1058    /// canonical spelling, including when only some tags drifted, and `raw()`
1059    /// must stay usable for the diagnostics that quote it.
1060    #[test]
1061    fn ssml_alias_parses_like_dsml() {
1062        let ssml = STANZA.replace("DSML", "SSML");
1063        let mixed = STANZA.replacen("DSML", "SSML", 2);
1064        for text in [ssml.as_str(), mixed.as_str()] {
1065            for mut p in [DsmlParser::new(), DsmlParser::new()] {
1066                feed_all(&mut p, text);
1067                assert_eq!(p.state(), DsmlState::Done, "{text:?}");
1068                assert_eq!(p.calls().len(), 1);
1069                assert_eq!(p.calls()[0].name, "read_file");
1070                assert_eq!(p.calls()[0].arg_value("path"), Some("src/main.rs"));
1071                assert_eq!(p.calls()[0].arg_value("offset"), Some("42"));
1072                assert!(!p.raw().is_empty());
1073            }
1074            let mut p = DsmlParser::new();
1075            feed_bytewise(&mut p, text);
1076            assert_eq!(p.state(), DsmlState::Done, "bytewise {text:?}");
1077            assert_eq!(p.calls()[0].arg_value("path"), Some("src/main.rs"));
1078        }
1079    }
1080
1081    /// Only the one observed misspelling is an alias; other marker names stay
1082    /// unrecognized so prose cannot open a stanza.
1083    #[test]
1084    fn other_marker_names_do_not_open_a_stanza() {
1085        let mut p = DsmlParser::new();
1086        feed_all(&mut p, &STANZA.replace("DSML", "XSML"));
1087        assert_eq!(p.state(), DsmlState::Search);
1088        assert!(p.calls().is_empty());
1089    }
1090
1091    #[test]
1092    fn parses_full_stanza() {
1093        let mut p = DsmlParser::new();
1094        feed_all(&mut p, STANZA);
1095        assert_eq!(p.state(), DsmlState::Done);
1096        assert_eq!(p.calls().len(), 1);
1097        let call = &p.calls()[0];
1098        assert_eq!(call.name, "read_file");
1099        assert_eq!(call.arg_value("path"), Some("src/main.rs"));
1100        assert_eq!(call.arg_value("offset"), Some("42"));
1101        assert_eq!(call.arg_value("missing"), None);
1102        assert!(call.args[0].is_string);
1103        assert!(!call.args[1].is_string);
1104    }
1105
1106    #[test]
1107    fn parses_bytewise_identically() {
1108        let mut p = DsmlParser::new();
1109        feed_bytewise(&mut p, STANZA);
1110        assert_eq!(p.state(), DsmlState::Done);
1111        assert_eq!(p.calls().len(), 1);
1112        assert_eq!(p.calls()[0].arg_value("path"), Some("src/main.rs"));
1113    }
1114
1115    // `find_tool_start` reports the *earliest* opening under any accepted
1116    // form, so recovery reacts to the first one the model wrote.
1117    #[test]
1118    fn find_tool_start_matches_every_accepted_wrapper_form() {
1119        for form in [
1120            "<|DSML|tool_calls>",
1121            "<DSML|tool_calls>",
1122            "<|SSML|tool_calls>",
1123            "<tool_calls>",
1124        ] {
1125            let text = format!("prose {form} rest");
1126            assert_eq!(
1127                super::find_tool_start(&text),
1128                Some("prose ".len()),
1129                "{form}"
1130            );
1131        }
1132    }
1133
1134    // Incomplete openings and the bare invoke opener are deliberately not
1135    // matched: acting on a weaker signal costs a forced injection.
1136    #[test]
1137    fn find_tool_start_ignores_partial_and_bare_invoke() {
1138        assert_eq!(super::find_tool_start("<"), None);
1139        assert_eq!(super::find_tool_start("<|DSML|tool_call"), None);
1140        assert_eq!(super::find_tool_start("<|DSML|invoke name=\"a\">"), None);
1141    }
1142
1143    #[test]
1144    fn skips_leading_prose_before_marker() {
1145        let mut p = DsmlParser::new();
1146        feed_all(&mut p, "Some thinking text first. ");
1147        assert_eq!(p.state(), DsmlState::Search);
1148        feed_all(&mut p, STANZA);
1149        assert_eq!(p.state(), DsmlState::Done);
1150    }
1151
1152    #[test]
1153    fn incomplete_input_stays_pending() {
1154        let mut p = DsmlParser::new();
1155        feed_all(
1156            &mut p,
1157            "<|DSML|tool_calls><|DSML|invoke name=\"bash\"><|DSML|parameter name=\"command\">ls -la",
1158        );
1159        assert_eq!(p.state(), DsmlState::ParamValue);
1160        assert!(p.calls().is_empty());
1161    }
1162
1163    #[test]
1164    fn close_tag_variants_accepted() {
1165        // Whitespace and missing trailing bar in close tags are tolerated.
1166        let s = concat!(
1167            "<|DSML|tool_calls>",
1168            "<|DSML|invoke name=\"t\">",
1169            "<|DSML|parameter name=\"a\">v</|DSML|parameter >",
1170            "</|DSML|invoke | >",
1171            "</|DSML|tool_calls>",
1172        );
1173        let mut p = DsmlParser::new();
1174        feed_all(&mut p, s);
1175        assert_eq!(p.state(), DsmlState::Done);
1176        assert_eq!(p.calls()[0].arg_value("a"), Some("v"));
1177    }
1178
1179    /// A literal `</` in a parameter value (e.g. HTML written through a
1180    /// `write` call's `content` param) must not terminate the parameter: the
1181    /// cheap `</` scan in `find_close_tag_any` is only a candidate filter, and
1182    /// `close_tag_at` requires the full `</|DSML|parameter` prefix (or its
1183    /// dropped-bar variant) before accepting a close. This pins the safety
1184    /// that let `CLOSE_SCAN_HEAD` widen from `"</|"` to `"</"`.
1185    #[test]
1186    fn literal_close_bytes_in_param_value_do_not_terminate_it() {
1187        // Includes a bare `</parameter>` (no DSML marker) so a validator
1188        // that dropped the marker check would truncate the value here.
1189        let html = "<div>hi</div></p> see </parameter> too";
1190        let s = format!(
1191            concat!(
1192                "<|DSML|tool_calls>",
1193                "<|DSML|invoke name=\"write\">",
1194                "<|DSML|parameter name=\"content\" string=\"true\">{html}</|DSML|parameter|>",
1195                "</|DSML|invoke|>",
1196                "</|DSML|tool_calls|>",
1197            ),
1198            html = html
1199        );
1200        let mut p = DsmlParser::new();
1201        feed_all(&mut p, &s);
1202        assert_eq!(p.state(), DsmlState::Done);
1203        assert_eq!(p.calls().len(), 1);
1204        assert_eq!(p.calls()[0].arg_value("content"), Some(html));
1205    }
1206
1207    #[test]
1208    fn multiple_invokes() {
1209        let s = concat!(
1210            "<|DSML|tool_calls>",
1211            "<|DSML|invoke name=\"a\"></|DSML|invoke|>",
1212            "<|DSML|invoke name=\"b\"></|DSML|invoke|>",
1213            "</|DSML|tool_calls|>",
1214        );
1215        let mut p = DsmlParser::new();
1216        feed_all(&mut p, s);
1217        assert_eq!(p.state(), DsmlState::Done);
1218        let names: Vec<_> = p.calls().iter().map(|c| c.name.as_str()).collect();
1219        assert_eq!(names, ["a", "b"]);
1220    }
1221
1222    #[test]
1223    fn invoke_without_name_errors() {
1224        let mut p = DsmlParser::new();
1225        feed_all(&mut p, "<|DSML|tool_calls><|DSML|invoke>");
1226        assert_eq!(p.state(), DsmlState::Error);
1227        assert_eq!(p.error(), "tool invoke without name");
1228    }
1229
1230    #[test]
1231    fn unexpected_tag_errors() {
1232        let mut p = DsmlParser::new();
1233        feed_all(&mut p, "<|DSML|tool_calls><b>");
1234        assert_eq!(p.state(), DsmlState::Error);
1235        assert!(p.error().starts_with("unexpected DSML tag:"));
1236    }
1237
1238    #[test]
1239    fn param_value_may_contain_angle_brackets() {
1240        let s = concat!(
1241            "<|DSML|tool_calls>",
1242            "<|DSML|invoke name=\"write\">",
1243            "<|DSML|parameter name=\"content\">if a < b { x > y }</|DSML|parameter|>",
1244            "</|DSML|invoke|>",
1245            "</|DSML|tool_calls|>",
1246        );
1247        let mut p = DsmlParser::new();
1248        feed_all(&mut p, s);
1249        assert_eq!(p.state(), DsmlState::Done);
1250        assert_eq!(
1251            p.calls()[0].arg_value("content"),
1252            Some("if a < b { x > y }")
1253        );
1254    }
1255
1256    #[test]
1257    fn param_close_prefix_tracks_partial_close_tag() {
1258        let mut p = DsmlParser::new();
1259        feed_all(
1260            &mut p,
1261            "<|DSML|tool_calls><|DSML|invoke name=\"t\"><|DSML|parameter name=\"a\">v",
1262        );
1263        assert!(!p.param_close_prefix());
1264        feed_all(&mut p, "</|DSML|parameter");
1265        assert!(p.param_close_prefix());
1266        feed_all(&mut p, "|>");
1267        assert!(!p.param_close_prefix());
1268        assert_eq!(p.state(), DsmlState::Structural);
1269    }
1270
1271    #[test]
1272    fn reset_returns_to_search() {
1273        let mut p = DsmlParser::new();
1274        feed_all(&mut p, STANZA);
1275        p.reset();
1276        assert_eq!(p.state(), DsmlState::Search);
1277        assert!(p.calls().is_empty());
1278        feed_all(&mut p, STANZA);
1279        assert_eq!(p.state(), DsmlState::Done);
1280    }
1281
1282    #[test]
1283    fn ignores_input_after_done() {
1284        let mut p = DsmlParser::new();
1285        feed_all(&mut p, STANZA);
1286        feed_all(&mut p, "trailing garbage <b>");
1287        assert_eq!(p.state(), DsmlState::Done);
1288        assert_eq!(p.calls().len(), 1);
1289    }
1290
1291    // The model drops the leading fullwidth bar on inner tags (~35 recorded
1292    // occurrences). The opener matcher already tolerates it; without the same
1293    // tolerance here the stanza opens and dies on its first inner tag, and the
1294    // model reads "unexpected DSML tag" as a claim that its `|` was wrong.
1295    #[test]
1296    fn inner_tags_tolerate_the_dropped_leading_bar() {
1297        let mut p = super::DsmlParser::new();
1298        p.feed(
1299            "<|DSML|tool_calls><DSML|invoke name=\"bash\">\
1300             <DSML|parameter name=\"command\" string=\"true\">ls</DSML|parameter|>\
1301             </DSML|invoke|></|DSML|tool_calls|>"
1302                .as_bytes(),
1303        );
1304        assert_eq!(p.state(), super::DsmlState::Done, "error: {}", p.error());
1305        let calls = p.calls();
1306        assert_eq!(calls.len(), 1);
1307        assert_eq!(calls[0].name, "bash");
1308        assert_eq!(calls[0].arg_value("command"), Some("ls"));
1309    }
1310
1311    // The canonical form must keep parsing identically.
1312    #[test]
1313    fn canonical_inner_tags_still_parse() {
1314        let mut p = super::DsmlParser::new();
1315        p.feed(
1316            "<|DSML|tool_calls><|DSML|invoke name=\"bash\">\
1317             <|DSML|parameter name=\"command\" string=\"true\">ls</|DSML|parameter|>\
1318             </|DSML|invoke|></|DSML|tool_calls|>"
1319                .as_bytes(),
1320        );
1321        assert_eq!(p.state(), super::DsmlState::Done, "error: {}", p.error());
1322        assert_eq!(p.calls()[0].arg_value("command"), Some("ls"));
1323    }
1324}