1use crate::dsml::{
18 DsmlParser, DsmlState, MARKER_NAMES, ToolCall, tag_prefix_len, tag_prefix_partial,
19};
20
21pub const IN_THINK_PROHIBITION: &str =
27 "Tool calls are not allowed inside <think></think>; finish thinking before emitting DSML.";
28
29const DSML_START: &[u8] = "<|DSML|tool_calls>".as_bytes();
31const CANONICAL_INVOKE: &[u8] = "<|DSML|invoke".as_bytes();
33const DSML_BAR: &[u8] = "|".as_bytes();
34
35const THINK_OPEN: &[u8] = b"<think>";
36const THINK_CLOSE: &[u8] = b"</think>";
37
38fn logging_enabled_for(opt_out: Option<&std::ffi::OsStr>) -> bool {
42 opt_out.is_none_or(|v| v != "1")
43}
44
45fn tool_error_logging_enabled() -> bool {
51 if cfg!(test) {
52 return false;
53 }
54 logging_enabled_for(std::env::var_os("PLANK_NO_TOOL_ERROR_LOG").as_deref())
55}
56
57fn log_tool_error(reason: &str, raw: &[u8]) {
64 use std::io::Write;
65
66 if !tool_error_logging_enabled() {
67 return;
68 }
69
70 let Some(home) = std::env::var_os("HOME").filter(|h| !h.is_empty()) else {
71 return;
72 };
73 let dir = std::path::PathBuf::from(home).join(".plank");
74 if std::fs::create_dir_all(&dir).is_err() {
75 return;
76 }
77 let secs = std::time::SystemTime::now()
78 .duration_since(std::time::UNIX_EPOCH)
79 .map_or(0, |d| d.as_secs());
80 let snippet = String::from_utf8_lossy(raw);
81 if let Ok(mut f) = std::fs::OpenOptions::new()
82 .create(true)
83 .append(true)
84 .open(dir.join("tool-call-errors.log"))
85 {
86 let record = format!("[{secs}] {reason}\n---\n{snippet}\n===\n");
88 let _ = f.write_all(record.as_bytes());
89 }
90}
91
92pub trait RenderSink {
105 fn visible_text(&mut self, text: &str);
107 fn think_text(&mut self, text: &str);
109 fn tool_text(&mut self, text: &str) {
111 self.visible_text(text);
112 }
113 fn error_text(&mut self, text: &str) {
116 self.tool_text(text);
117 }
118}
119
120impl RenderSink for Box<dyn RenderSink> {
121 fn visible_text(&mut self, text: &str) {
122 (**self).visible_text(text);
123 }
124 fn think_text(&mut self, text: &str) {
125 (**self).think_text(text);
126 }
127 fn tool_text(&mut self, text: &str) {
128 (**self).tool_text(text);
129 }
130 fn error_text(&mut self, text: &str) {
131 (**self).error_text(text);
132 }
133}
134
135impl RenderSink for Box<dyn RenderSink + Send> {
139 fn visible_text(&mut self, text: &str) {
140 (**self).visible_text(text);
141 }
142 fn think_text(&mut self, text: &str) {
143 (**self).think_text(text);
144 }
145 fn tool_text(&mut self, text: &str) {
146 (**self).tool_text(text);
147 }
148 fn error_text(&mut self, text: &str) {
149 (**self).error_text(text);
150 }
151}
152
153#[derive(Debug, Clone, Default)]
160pub struct CollectSink(pub std::sync::Arc<std::sync::Mutex<String>>);
161
162impl CollectSink {
163 #[must_use]
165 pub fn take(&self) -> String {
166 let mut guard = self
167 .0
168 .lock()
169 .unwrap_or_else(std::sync::PoisonError::into_inner);
170 std::mem::take(&mut *guard)
171 }
172
173 fn push(&mut self, text: &str) {
174 self.0
175 .lock()
176 .unwrap_or_else(std::sync::PoisonError::into_inner)
177 .push_str(text);
178 }
179}
180
181impl RenderSink for CollectSink {
182 fn visible_text(&mut self, text: &str) {
183 self.push(text);
184 }
185 fn think_text(&mut self, _text: &str) {}
188 fn tool_text(&mut self, text: &str) {
189 self.push(text);
190 }
191}
192
193#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
195enum ParamKind {
196 #[default]
197 Normal,
198 Path,
199 Content,
200 DiffOld,
201 DiffNew,
202 BashCommand,
203}
204
205fn param_kind_for(tool: &str, param: &str) -> ParamKind {
206 match (tool, param) {
207 ("bash", "command") => ParamKind::BashCommand,
208 ("edit", "old") => ParamKind::DiffOld,
209 ("edit", "new") => ParamKind::DiffNew,
210 (_, "path" | "file" | "filename") => ParamKind::Path,
211 (_, "content" | "text") => ParamKind::Content,
212 _ => ParamKind::Normal,
213 }
214}
215
216fn tool_prefix(name: &str) -> Option<&'static str> {
218 match name {
219 "bash" => Some("$ "),
220 "read" => Some("read "),
221 "write" => Some("write "),
222 "edit" => Some("edit "),
223 "search" => Some("search "),
224 "google_search" => Some("google "),
225 "visit_page" => Some("visit "),
226 name if name.starts_with("mcp_") => Some("mcp "),
227 _ => None,
228 }
229}
230
231fn diff_prefix(kind: ParamKind) -> Option<&'static str> {
232 match kind {
233 ParamKind::DiffOld => Some("- "),
234 ParamKind::DiffNew => Some("+ "),
235 _ => None,
236 }
237}
238
239fn parse_bool_default(s: &str, default: bool) -> bool {
240 if s.is_empty() {
241 return default;
242 }
243 if s.eq_ignore_ascii_case("true") || s.eq_ignore_ascii_case("yes") || s == "1" {
244 return true;
245 }
246 if s.eq_ignore_ascii_case("false") || s.eq_ignore_ascii_case("no") || s == "0" {
247 return false;
248 }
249 default
250}
251
252fn parse_attr(tag: &str, name: &str) -> Option<String> {
254 let pat = format!("{name}=\"");
255 let start = tag.find(&pat)? + pat.len();
256 let end = tag[start..].find('"')? + start;
257 Some(tag[start..end].to_string())
258}
259
260fn parameter_close_tail(tail: &[u8], complete: &mut bool) -> bool {
266 *complete = false;
267 if tag_prefix_partial(tail, true, "parameter") {
268 return true;
269 }
270 let Some(mut i) = tag_prefix_len(tail, true, "parameter") else {
271 return false;
272 };
273 while i < tail.len() && tail[i].is_ascii_whitespace() {
274 i += 1;
275 }
276 if i < tail.len() && tail.len() - i <= DSML_BAR.len() && DSML_BAR.starts_with(&tail[i..]) {
277 return true;
278 }
279 if tail[i..].starts_with(DSML_BAR) {
280 i += DSML_BAR.len();
281 }
282 while i < tail.len() {
283 if tail[i] == b'>' {
284 *complete = i == tail.len() - 1;
285 return *complete;
286 }
287 if !tail[i].is_ascii_whitespace() {
288 return false;
289 }
290 i += 1;
291 }
292 true
293}
294
295fn dsml_start_match(tail: &[u8], complete: &mut bool, implicit_invoke: &mut bool) -> bool {
301 *complete = false;
302 *implicit_invoke = false;
303 let forms = MARKER_NAMES.iter().flat_map(|m| {
308 [
309 (format!("<|{m}|tool_calls>"), false),
310 (format!("<|{m}|tool_calls|>"), false),
311 (format!("<{m}|tool_calls>"), false),
312 (format!("<{m}|tool_calls|>"), false),
313 (format!("<|{m}|invoke"), true),
314 (format!("<{m}|invoke"), true),
315 ]
316 });
317 for (form, implicit) in forms {
318 let form = form.as_bytes();
319 if tail.len() <= form.len() && form[..tail.len()] == *tail {
320 *complete = tail.len() == form.len();
321 *implicit_invoke = implicit;
322 return true;
323 }
324 }
325 false
326}
327
328#[derive(Debug, Default)]
334struct MarkerDetector {
335 tail: Vec<u8>,
336}
337
338impl MarkerDetector {
339 const CAP: usize = 32;
340 fn feed(&mut self, c: u8) -> bool {
341 if self.tail.len() == Self::CAP {
342 self.tail.remove(0);
343 }
344 self.tail.push(c);
345 MARKER_NAMES.iter().any(|m| {
346 [
347 format!("|{m}|"),
348 format!("|{m}|"),
349 format!("<{m}|"),
350 format!("</{m}|"),
351 ]
352 .iter()
353 .any(|n| self.tail.ends_with(n.as_bytes()))
354 })
355 }
356}
357
358const GENERIC_PSEUDO_OPENERS: [&str; 3] = ["<tool_call>", "<function_call>", "<invoke "];
361
362#[derive(Debug, Default)]
372struct PseudoToolDetector {
373 line: Vec<u8>,
375 in_fence: bool,
378 tool_names: Vec<String>,
380 fired: bool,
382}
383
384impl PseudoToolDetector {
385 const CAP: usize = 96;
386
387 fn set_tool_names(&mut self, names: Vec<String>) {
388 self.tool_names = names;
389 }
390
391 fn reset_line(&mut self) {
397 self.line.clear();
398 }
399
400 fn feed(&mut self, c: u8, in_think: bool) -> Option<String> {
417 if c == b'\n' {
418 let hit = self.match_line(in_think);
419 if self.line.starts_with(b"```") {
420 self.in_fence = !self.in_fence;
429 }
430 self.line.clear();
431 return hit;
432 }
433 if self.line.len() < Self::CAP {
434 self.line.push(c);
435 }
436 None
437 }
438
439 fn finish_line(&mut self, in_think: bool) -> Option<String> {
441 let hit = self.match_line(in_think);
442 self.line.clear();
443 hit
444 }
445
446 fn match_line(&mut self, in_think: bool) -> Option<String> {
449 if in_think || self.fired || self.in_fence {
450 return None;
451 }
452 let trimmed = self.line.trim_ascii();
455 if !trimmed.starts_with(b"<") {
456 return None;
457 }
458 let text = std::str::from_utf8(trimmed).ok()?;
459 let matched = GENERIC_PSEUDO_OPENERS.iter().any(|opener| {
460 if opener.ends_with(' ') {
463 text.starts_with(*opener) && is_lone_tag(text)
464 } else {
465 text == *opener
466 }
467 }) || self.tool_names.iter().any(|name| {
468 text.strip_prefix('<')
469 .and_then(|t| t.strip_suffix('>'))
470 .is_some_and(|inner| inner == name)
471 });
472 if !matched {
473 return None;
474 }
475 self.fired = true;
476 Some(text.to_string())
477 }
478}
479
480fn is_lone_tag(text: &str) -> bool {
482 text.ends_with('>') && !text[..text.len() - 1].contains('>')
483}
484
485#[derive(Debug, Default)]
487enum DsmlScan {
488 #[default]
490 Between,
491 Tag(Vec<u8>),
493 Value,
495}
496
497#[allow(clippy::struct_excessive_bools)]
501#[derive(Debug, Default)]
502struct ToolViz {
503 active: bool,
504 tool_announced: bool,
505 param_active: bool,
506 at_line_start: bool,
507 param_kind: ParamKind,
508 tool_name: String,
509 param_name: String,
510 param_end_tail: Vec<u8>,
511 read_style: bool,
512 read_prefix_rendered: bool,
513 read_line_rendered: bool,
514 read_path: String,
515 read_start: String,
516 read_max: String,
517 read_whole: String,
518 code_param_active: bool,
519 write_path: String,
522 write_is_create: bool,
526}
527
528impl ToolViz {
529 const END_TAIL_CAP: usize = 64;
530}
531
532#[derive(Debug, Clone, Copy)]
536pub struct Finished<'a> {
537 pub calls: &'a [ToolCall],
539 pub error: Option<&'a str>,
541 pub dsml_in_think: bool,
543 pub in_think_rejected: bool,
548 pub ended_in_think: bool,
552}
553
554#[allow(clippy::struct_excessive_bools)]
581#[derive(Debug)]
582pub struct StreamRenderer<S> {
583 sink: S,
584 parser: DsmlParser,
585 viz: ToolViz,
586 scan: DsmlScan,
587 in_think: bool,
588 dsml_active: bool,
589 dsml_ignored: bool,
590 pending: Vec<u8>,
592 dsml_start_tail: Vec<u8>,
594 plain_dsml: MarkerDetector,
595 think_dsml: MarkerDetector,
596 pseudo_tool: PseudoToolDetector,
598 dsml_in_think: bool,
599 dsml_in_think_reported: bool,
600 think_close_swallowed: bool,
606 in_think_rejected: bool,
613 pseudo_tool_fired: bool,
618 post_think_gap: bool,
619 stream_error: Option<String>,
622 freeze_on_error: bool,
625 last_output_newline: bool,
626 calls: Vec<ToolCall>,
628 vis_carry: Vec<u8>,
631 think_carry: Vec<u8>,
632 preflight: Preflight,
637 preflight_error: Option<String>,
638 show_tool_calls: bool,
643 show_thinking: bool,
647 thinking_tool_calls: bool,
653 replay: bool,
660}
661
662type PreflightFn = Box<dyn FnMut(&ToolCall) -> Result<(), String>>;
664
665#[derive(Default)]
666struct Preflight(Option<PreflightFn>);
667
668impl std::fmt::Debug for Preflight {
669 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
670 f.write_str(if self.0.is_some() {
671 "Preflight(set)"
672 } else {
673 "Preflight(unset)"
674 })
675 }
676}
677
678const DSML_START_TAIL_CAP: usize = 64;
679
680impl<S: RenderSink> StreamRenderer<S> {
681 pub fn new(sink: S) -> Self {
683 Self {
684 sink,
685 parser: DsmlParser::new(),
686 viz: ToolViz::default(),
687 scan: DsmlScan::Between,
688 in_think: false,
689 freeze_on_error: false,
690 dsml_active: false,
691 dsml_ignored: false,
692 pending: Vec::new(),
693 dsml_start_tail: Vec::new(),
694 plain_dsml: MarkerDetector::default(),
695 think_dsml: MarkerDetector::default(),
696 pseudo_tool: PseudoToolDetector::default(),
697 dsml_in_think: false,
698 dsml_in_think_reported: false,
699 think_close_swallowed: false,
700 in_think_rejected: false,
701 pseudo_tool_fired: false,
702 post_think_gap: false,
703 stream_error: None,
704 last_output_newline: true,
705 calls: Vec::new(),
706 vis_carry: Vec::new(),
707 think_carry: Vec::new(),
708 preflight: Preflight(None),
709 preflight_error: None,
710 show_tool_calls: true,
711 show_thinking: true,
712 thinking_tool_calls: false,
713 replay: false,
714 }
715 }
716
717 pub fn set_replay(&mut self, replay: bool) {
724 self.replay = replay;
725 }
726
727 pub fn set_tool_names(&mut self, names: Vec<String>) {
733 self.pseudo_tool.set_tool_names(names);
734 }
735
736 pub fn set_show_tool_calls(&mut self, show: bool) {
740 self.show_tool_calls = show;
741 }
742
743 pub fn set_show_thinking(&mut self, show: bool) {
748 self.show_thinking = show;
749 }
750
751 pub fn set_freeze_on_error(&mut self, freeze: bool) {
767 self.freeze_on_error = freeze;
768 }
769
770 pub fn set_thinking_tool_calls(&mut self, allow: bool) {
775 self.thinking_tool_calls = allow;
776 }
777
778 pub fn set_preflight(&mut self, f: impl FnMut(&ToolCall) -> Result<(), String> + 'static) {
783 self.preflight = Preflight(Some(Box::new(f)));
784 }
785
786 #[must_use]
788 pub fn preflight_error(&self) -> Option<&str> {
789 self.preflight_error.as_deref()
790 }
791
792 pub fn begin_in_think(&mut self) {
798 self.in_think = true;
799 }
800
801 pub fn push(&mut self, text: impl AsRef<str>) {
803 self.stream_text(text.as_ref().as_bytes(), false);
804 }
805
806 pub fn finish(&mut self) {
811 self.stream_text(b"", true);
812 self.flush_carry();
813 self.flush_pseudo_tool();
814 }
815
816 #[must_use]
823 pub fn finished(&self) -> Finished<'_> {
824 let in_think = (self.in_think_rejected || self.dsml_ignored) && !self.pseudo_tool_fired;
839 let error = self
840 .stream_error
841 .as_deref()
842 .filter(|_| !in_think)
843 .or_else(|| in_think.then_some(IN_THINK_PROHIBITION))
844 .or_else(|| (self.parser.state() == DsmlState::Error).then(|| self.parser.error()))
845 .or_else(|| {
846 matches!(
847 self.parser.state(),
848 DsmlState::Structural | DsmlState::ParamValue
849 )
850 .then_some("incomplete DSML tool call")
851 });
852 Finished {
853 calls: &self.calls,
854 error,
855 dsml_in_think: self.dsml_in_think,
856 in_think_rejected: in_think,
857 ended_in_think: self.in_think,
858 }
859 }
860
861 #[must_use]
868 pub fn wants_greedy_sampling(&self) -> bool {
869 if matches!(self.parser.state(), DsmlState::Error | DsmlState::Done) {
870 return false;
871 }
872 if self.dsml_start_tail.len() > 1 {
875 return true;
876 }
877 if !self.dsml_active {
878 return false;
879 }
880 match self.parser.state() {
881 DsmlState::Structural => true,
882 DsmlState::ParamValue => self.parser.param_close_prefix(),
883 _ => false,
884 }
885 }
886
887 pub fn sink(&self) -> &S {
889 &self.sink
890 }
891
892 pub fn sink_mut(&mut self) -> &mut S {
894 &mut self.sink
895 }
896
897 pub fn into_sink(self) -> S {
899 self.sink
900 }
901
902 fn flush_stream(sink_write: impl FnOnce(&mut S, &str), sink: &mut S, carry: &mut Vec<u8>) {
905 if carry.is_empty() {
906 return;
907 }
908 match std::str::from_utf8(carry) {
909 Ok(s) => {
910 sink_write(sink, s);
911 carry.clear();
912 }
913 Err(e) if e.error_len().is_none() && e.valid_up_to() > 0 => {
914 let tail = carry.split_off(e.valid_up_to());
915 sink_write(sink, std::str::from_utf8(carry).unwrap_or_default());
917 *carry = tail;
918 }
919 Err(e) if e.error_len().is_none() => {}
920 Err(_) => {
921 let s = String::from_utf8_lossy(carry).into_owned();
922 sink_write(sink, &s);
923 carry.clear();
924 }
925 }
926 }
927
928 fn emit_visible_bytes(&mut self, bytes: &[u8]) {
929 if bytes.is_empty() {
930 return;
931 }
932 if self.viz.active && !self.show_tool_calls {
936 return;
937 }
938 self.last_output_newline = bytes.last() == Some(&b'\n');
939 self.vis_carry.extend_from_slice(bytes);
940 let write = if self.viz.active {
941 S::tool_text as fn(&mut S, &str)
942 } else {
943 S::visible_text
944 };
945 Self::flush_stream(write, &mut self.sink, &mut self.vis_carry);
946 }
947
948 fn emit_think_bytes(&mut self, bytes: &[u8]) {
949 if !self.show_thinking {
950 return;
951 }
952 self.think_carry.extend_from_slice(bytes);
953 Self::flush_stream(S::think_text, &mut self.sink, &mut self.think_carry);
954 }
955
956 fn emit_preview_bytes(&mut self, bytes: &[u8]) {
960 if bytes.is_empty() {
961 return;
962 }
963 self.last_output_newline = bytes.last() == Some(&b'\n');
964 self.think_carry.extend_from_slice(bytes);
965 Self::flush_stream(S::think_text, &mut self.sink, &mut self.think_carry);
966 }
967
968 fn viz_preview_puts(&mut self, s: &str) {
969 self.emit_preview_bytes(s.as_bytes());
970 }
971
972 fn viz_is_write_preview(&self) -> bool {
975 self.viz.tool_name == "write" && self.viz.param_kind == ParamKind::Content
976 }
977
978 fn flush_carry(&mut self) {
979 for (write, carry) in [
980 (S::visible_text as fn(&mut S, &str), &mut self.vis_carry),
981 (S::think_text, &mut self.think_carry),
982 ] {
983 if !carry.is_empty() {
984 let s = String::from_utf8_lossy(carry).into_owned();
985 write(&mut self.sink, &s);
986 carry.clear();
987 }
988 }
989 }
990
991 fn write_char(&mut self, c: u8) {
993 if self.in_think {
994 self.emit_think_bytes(&[c]);
995 } else {
996 self.emit_visible_bytes(&[c]);
997 }
998 }
999
1000 fn viz_puts(&mut self, s: &str) {
1001 self.emit_visible_bytes(s.as_bytes());
1002 }
1003
1004 fn viz_error_puts(&mut self, s: &str) {
1006 if s.is_empty() {
1007 return;
1008 }
1009 self.last_output_newline = s.ends_with('\n');
1010 self.sink.error_text(s);
1011 }
1012
1013 fn viz_newline_if_open(&mut self) {
1014 if !self.last_output_newline {
1015 self.viz_puts("\n");
1016 }
1017 }
1018
1019 fn viz_start(&mut self) {
1022 let line_open = !self.last_output_newline;
1023 self.viz = ToolViz {
1024 active: true,
1025 at_line_start: true,
1026 ..ToolViz::default()
1027 };
1028 self.scan = DsmlScan::Between;
1029 if line_open {
1030 self.viz_puts("\n");
1031 }
1032 }
1033
1034 fn viz_line_prefix(&mut self) {
1036 self.viz_newline_if_open();
1037 self.viz_puts("🛠️ ");
1038 self.viz.at_line_start = false;
1039 }
1040
1041 fn viz_tool(&mut self, name: &str) {
1042 if self.viz.tool_announced && self.viz.tool_name == name {
1043 return;
1044 }
1045 if self.viz.tool_announced {
1046 self.viz_newline_if_open();
1047 }
1048 self.viz.tool_name = name.to_string();
1049 self.viz.tool_announced = true;
1050 self.viz.read_style = name == "read";
1051 self.viz_line_prefix();
1052 if self.viz.read_style {
1053 self.viz_puts("Reading ");
1054 self.viz.read_prefix_rendered = true;
1055 return;
1056 }
1057 if let Some(prefix) = tool_prefix(name) {
1058 self.viz_puts(prefix);
1059 } else {
1060 let owned = name.to_string();
1061 self.viz_puts(&owned);
1062 self.viz_puts(" ");
1063 }
1064 }
1065
1066 fn viz_read_value_byte(&mut self, c: u8) {
1067 let field = match self.viz.param_name.as_str() {
1068 "path" => &mut self.viz.read_path,
1069 "start_line" => &mut self.viz.read_start,
1070 "max_lines" => &mut self.viz.read_max,
1071 "whole" => &mut self.viz.read_whole,
1072 _ => return,
1073 };
1074 field.push(c as char);
1075 if self.viz.param_name == "path" && self.viz.read_prefix_rendered {
1076 self.emit_visible_bytes(&[c]);
1077 }
1078 }
1079
1080 fn viz_render_read(&mut self) {
1082 if !self.viz.read_style || self.viz.read_line_rendered {
1083 return;
1084 }
1085 if !self.viz.read_prefix_rendered {
1086 self.viz_line_prefix();
1087 self.viz_puts("Reading ");
1088 let path = if self.viz.read_path.is_empty() {
1089 "<unknown>".to_string()
1090 } else {
1091 self.viz.read_path.clone()
1092 };
1093 self.viz_puts(&path);
1094 } else if self.viz.read_path.is_empty() {
1095 self.viz_puts("<unknown>");
1096 }
1097 let whole = parse_bool_default(&self.viz.read_whole, false);
1098 let range = if whole && (self.viz.read_start.is_empty() || self.viz.read_start == "1") {
1099 " (whole file)".to_string()
1100 } else if whole {
1101 format!(" {}:EOF", self.viz.read_start)
1102 } else {
1103 let start = if self.viz.read_start.is_empty() {
1104 "1"
1105 } else {
1106 &self.viz.read_start
1107 };
1108 let max = if self.viz.read_max.is_empty() {
1109 "500"
1110 } else {
1111 &self.viz.read_max
1112 };
1113 format!(" {start}:{max}")
1114 };
1115 self.viz_puts(&range);
1116 self.viz_puts("...\n");
1117 self.viz.read_line_rendered = true;
1118 }
1119
1120 fn viz_param_is_code_body(&self) -> bool {
1121 match self.viz.tool_name.as_str() {
1122 "write" => self.viz.param_kind == ParamKind::Content,
1123 "edit" => matches!(
1124 self.viz.param_kind,
1125 ParamKind::DiffOld | ParamKind::DiffNew | ParamKind::Content
1126 ),
1127 _ => false,
1128 }
1129 }
1130
1131 fn viz_code_prefix(&mut self) {
1133 if !self.viz.at_line_start {
1134 return;
1135 }
1136 if let Some(prefix) = diff_prefix(self.viz.param_kind) {
1137 self.viz_puts(prefix);
1138 self.viz.at_line_start = false;
1139 }
1140 }
1141
1142 fn viz_code_begin(&mut self) {
1143 self.viz.code_param_active = true;
1144 if matches!(self.viz.param_kind, ParamKind::DiffOld | ParamKind::DiffNew) {
1145 self.viz_code_prefix();
1146 }
1147 }
1148
1149 fn viz_code_end(&mut self) {
1150 if !self.viz.code_param_active {
1151 return;
1152 }
1153 self.viz.code_param_active = false;
1154 self.viz.at_line_start = true;
1155 }
1156
1157 fn viz_code_byte(&mut self, c: u8) {
1158 if self.viz_is_write_preview() {
1161 if self.viz.write_is_create {
1162 self.emit_preview_bytes(&[c]);
1163 }
1164 self.viz.at_line_start = c == b'\n';
1165 return;
1166 }
1167 self.viz_code_prefix();
1168 self.emit_visible_bytes(&[c]);
1169 self.viz.at_line_start = c == b'\n';
1170 }
1171
1172 fn viz_param_begin(&mut self, name: &str) {
1173 self.viz.param_name = name.to_string();
1174 self.viz.param_kind = param_kind_for(&self.viz.tool_name, name);
1175 self.viz.param_active = true;
1176 self.viz.param_end_tail.clear();
1177
1178 if self.viz.read_style {
1179 return;
1180 }
1181 match self.viz.param_kind {
1182 ParamKind::DiffOld | ParamKind::DiffNew => {
1183 self.viz_newline_if_open();
1184 self.viz.at_line_start = true;
1185 self.viz_code_begin();
1186 }
1187 ParamKind::Content => {
1188 self.viz_newline_if_open();
1189 if self.viz.tool_name == "write" {
1190 self.viz.write_is_create = !std::path::Path::new(&self.viz.write_path).exists();
1193 if self.viz.write_is_create && !self.show_tool_calls {
1196 let path = if self.viz.write_path.is_empty() {
1197 "<file>".to_string()
1198 } else {
1199 self.viz.write_path.clone()
1200 };
1201 self.viz_preview_puts(&format!("write {path}\n"));
1202 }
1203 } else {
1204 let label = format!("{name}:\n");
1205 self.viz_puts(&label);
1206 }
1207 self.viz.at_line_start = true;
1208 if self.viz_param_is_code_body() {
1209 self.viz_code_begin();
1210 }
1211 }
1212 ParamKind::BashCommand => {}
1213 ParamKind::Normal | ParamKind::Path => {
1214 if !self.viz.at_line_start {
1215 self.viz_puts(" ");
1216 }
1217 let label = format!("{name}=");
1218 self.viz_puts(&label);
1219 }
1220 }
1221 }
1222
1223 fn viz_param_end(&mut self) {
1224 self.viz.param_end_tail.clear();
1225 if self.viz.code_param_active {
1226 self.viz_code_end();
1227 }
1228 self.viz.param_active = false;
1229 self.viz.param_name.clear();
1230 self.scan = DsmlScan::Between;
1231 }
1232
1233 fn viz_param_raw_byte(&mut self, c: u8) {
1234 if self.viz.read_style {
1235 self.viz_read_value_byte(c);
1236 return;
1237 }
1238 if self.viz.code_param_active {
1239 self.viz_code_byte(c);
1240 return;
1241 }
1242 if matches!(self.viz.param_kind, ParamKind::DiffOld | ParamKind::DiffNew) {
1243 self.viz_code_begin();
1244 self.viz_code_byte(c);
1245 return;
1246 }
1247 if self.viz.tool_name == "write" && self.viz.param_kind == ParamKind::Path {
1249 self.viz.write_path.push(c as char);
1250 }
1251 self.emit_visible_bytes(&[c]);
1252 self.viz.at_line_start = c == b'\n';
1253 }
1254
1255 fn viz_param_value_byte(&mut self, c: u8) {
1260 if !self.viz.param_end_tail.is_empty() || c == b'<' {
1261 if self.viz.param_end_tail.len() == ToolViz::END_TAIL_CAP {
1262 let held = std::mem::take(&mut self.viz.param_end_tail);
1263 for b in held {
1264 self.viz_param_raw_byte(b);
1265 }
1266 if c != b'<' {
1267 self.viz_param_raw_byte(c);
1268 return;
1269 }
1270 }
1271 self.viz.param_end_tail.push(c);
1272 let mut complete = false;
1273 if parameter_close_tail(&self.viz.param_end_tail, &mut complete) {
1274 if complete {
1275 self.viz_param_end();
1276 }
1277 return;
1278 }
1279 let held = std::mem::take(&mut self.viz.param_end_tail);
1280 for b in held {
1281 self.viz_param_raw_byte(b);
1282 }
1283 return;
1284 }
1285 self.viz_param_raw_byte(c);
1286 }
1287
1288 fn viz_invoke_end(&mut self) {
1290 if !self.viz.tool_announced || self.viz.param_active {
1291 return;
1292 }
1293 self.viz_render_read();
1294 self.viz_newline_if_open();
1295 self.viz.read_style = false;
1296 self.viz.read_prefix_rendered = false;
1297 self.viz.read_line_rendered = false;
1298 self.viz.read_path.clear();
1299 self.viz.read_start.clear();
1300 self.viz.read_max.clear();
1301 self.viz.read_whole.clear();
1302 self.viz.tool_announced = false;
1303 }
1304
1305 fn viz_finish(&mut self, status: Option<&str>) {
1306 if !self.viz.active {
1307 return;
1308 }
1309 if self.viz.param_active {
1310 self.viz_param_end();
1311 }
1312 if status.is_none() {
1313 self.viz_render_read();
1314 }
1315 if let Some(status) = status {
1316 self.viz_newline_if_open();
1317 let owned = status.to_string();
1318 self.viz_error_puts(&owned);
1319 }
1320 self.viz_newline_if_open();
1321 self.viz.active = false;
1322 }
1323
1324 fn viz_drop_invalid_dsml(&mut self) {
1329 if !self.viz.active {
1330 return;
1331 }
1332 if self.viz.param_active {
1333 self.viz.param_active = false;
1334 self.viz.param_end_tail.clear();
1335 self.viz.param_name.clear();
1336 }
1337 self.viz_newline_if_open();
1338 }
1339
1340 fn scan_dsml_byte(&mut self, c: u8) {
1344 match &mut self.scan {
1345 DsmlScan::Between => {
1346 if c == b'<' {
1347 self.scan = DsmlScan::Tag(vec![c]);
1348 }
1349 }
1350 DsmlScan::Tag(tag) => {
1351 tag.push(c);
1352 if c == b'>' {
1353 let tag = std::mem::take(tag);
1354 self.scan = DsmlScan::Between;
1355 self.scan_dsml_tag(&tag);
1356 }
1357 }
1358 DsmlScan::Value => self.viz_param_value_byte(c),
1359 }
1360 }
1361
1362 fn scan_dsml_tag(&mut self, tag: &[u8]) {
1363 let tag = String::from_utf8_lossy(tag).into_owned();
1364 let b = tag.as_bytes();
1365 if tag_prefix_len(b, true, "invoke").is_some() {
1366 self.viz_invoke_end();
1367 } else if tag_prefix_len(b, false, "invoke").is_some() {
1368 let name = parse_attr(&tag, "name").unwrap_or_else(|| "tool".to_string());
1369 self.viz_tool(&name);
1370 } else if tag_prefix_len(b, false, "parameter").is_some()
1371 && let Some(name) = parse_attr(&tag, "name")
1372 {
1373 self.viz_param_begin(&name);
1374 self.scan = DsmlScan::Value;
1375 } else if parse_attr(&tag, "name").is_none()
1376 && let Some(elem) = crate::dsml::element_name(&tag)
1377 {
1378 if self.viz.tool_announced {
1385 self.viz_param_begin(&elem);
1386 self.scan = DsmlScan::Value;
1387 } else {
1388 self.viz_tool(&elem);
1389 }
1390 }
1391 }
1393
1394 fn feed_dsml_byte(&mut self, c: u8) {
1395 let was_param = self.parser.state() == DsmlState::ParamValue;
1396 self.parser.feed([c]);
1397 if !self.dsml_ignored {
1398 self.scan_dsml_byte(c);
1399 if was_param && self.parser.state() != DsmlState::ParamValue {
1400 self.preflight_closed_param();
1401 }
1402 }
1403 match self.parser.state() {
1404 DsmlState::Done => {
1405 self.think_close_swallowed = false;
1409 if self.rejects_in_think() {
1410 self.reject_in_think_stanza(
1417 "tool calling is not allowed inside <think></think>",
1418 );
1419 } else {
1420 self.calls = self.parser.calls().to_vec();
1421 self.viz_finish(None);
1422 self.dsml_active = false;
1423 }
1424 }
1425 DsmlState::Error => {
1426 self.resolve_swallowed_think_close();
1430 if self.rejects_in_think() {
1431 self.reject_in_think_stanza("malformed tool call inside <think></think>");
1432 } else {
1433 let err = if self.parser.error().is_empty() {
1434 "parse error"
1435 } else {
1436 self.parser.error()
1437 };
1438 log_tool_error(err, self.parser.raw());
1439 let status = format!("[invalid tool call: {err}]\n");
1440 self.viz_drop_invalid_dsml();
1441 self.viz_finish(Some(&status));
1442 self.dsml_active = false;
1443 }
1444 }
1445 _ => {}
1446 }
1447 }
1448
1449 fn preflight_closed_param(&mut self) {
1454 if self.preflight_error.is_some() {
1455 return;
1456 }
1457 let Preflight(Some(check)) = &mut self.preflight else {
1458 return;
1459 };
1460 let Some(call) = self.parser.pending_call() else {
1461 return;
1462 };
1463 if call.name != "edit" || call.args.last().is_none_or(|a| a.name != "old") {
1464 return;
1465 }
1466 if let Err(err) = check(&call) {
1467 self.preflight_error = Some(format!(
1468 "edit old selector failed before new was generated: {err}"
1469 ));
1470 }
1471 }
1472
1473 fn start_dsml(&mut self) {
1483 self.dsml_active = true;
1484 self.dsml_ignored = self.rejects_in_think();
1485 if self.in_think {
1486 self.dsml_in_think = true;
1487 }
1488 self.dsml_start_tail.clear();
1489 self.post_think_gap = false;
1490 self.parser.feed(DSML_START);
1491 self.scan = DsmlScan::Between;
1492 if !self.dsml_ignored {
1493 self.viz_start();
1494 }
1495 }
1496
1497 fn rejects_in_think(&self) -> bool {
1509 self.in_think && !self.thinking_tool_calls
1510 }
1511
1512 fn reject_in_think_stanza(&mut self, msg: &str) {
1517 self.dsml_ignored = true;
1518 self.viz_drop_invalid_dsml();
1519 self.finish_ignored_dsml(msg);
1520 }
1521
1522 fn think_close_is_control(&self) -> bool {
1530 !self.dsml_active || self.parser.state() != DsmlState::ParamValue
1531 }
1532
1533 fn resolve_swallowed_think_close(&mut self) {
1544 if !std::mem::take(&mut self.think_close_swallowed) || !self.in_think {
1545 return;
1546 }
1547 self.in_think = false;
1548 self.pseudo_tool.reset_line();
1551 self.plain_dsml = MarkerDetector::default();
1552 if self.dsml_active && self.dsml_ignored {
1557 self.dsml_ignored = false;
1558 self.viz_start();
1559 }
1560 }
1561
1562 fn finish_ignored_dsml(&mut self, msg: &str) {
1563 let raw = if self.parser.raw().is_empty() {
1568 self.dsml_start_tail.clone()
1569 } else {
1570 self.parser.raw().to_vec()
1571 };
1572 log_tool_error(msg, &raw);
1573 self.resolve_swallowed_think_close();
1574 self.dsml_in_think = true;
1575 self.in_think_rejected = true;
1576 self.dsml_in_think_reported = true;
1577 self.viz_newline_if_open();
1578 let line = format!("[tool call ignored: {msg}]\n");
1579 self.viz_error_puts(&line);
1580 self.parser.reset();
1581 self.dsml_active = false;
1582 self.dsml_ignored = false;
1583 }
1584
1585 fn malformed_dsml(&mut self, msg: &str) {
1586 if self.stream_error.is_some() {
1587 return;
1588 }
1589 self.stream_error = Some(msg.to_string());
1590 log_tool_error(msg, self.parser.raw());
1591 self.viz_newline_if_open();
1592 let line = format!("[invalid tool call: {msg}]\n");
1593 self.viz_error_puts(&line);
1594 }
1595
1596 fn pseudo_tool_call(&mut self, opener: &str) {
1597 self.pseudo_tool_fired = true;
1598 self.malformed_dsml(&format!(
1599 "{opener} is not a tool call; tools are invoked with the DSML syntax in the system prompt"
1600 ));
1601 }
1602
1603 fn output_frozen(&self) -> bool {
1604 self.freeze_on_error
1605 && (self.stream_error.is_some() || self.parser.state() == DsmlState::Error)
1606 }
1607
1608 fn note_thinking_dsml_byte(&mut self, c: u8) {
1609 if !self.in_think || self.dsml_in_think {
1610 return;
1611 }
1612 if self.think_dsml.feed(c) {
1613 self.dsml_in_think = true;
1614 }
1615 }
1616
1617 fn note_plain_dsml_byte(&mut self, c: u8) {
1618 if self.output_frozen() || self.dsml_active || self.in_think || self.in_think_rejected {
1628 return;
1629 }
1630 if self.plain_dsml.feed(c) {
1631 self.malformed_dsml("DSML markup outside a valid tool_calls block");
1632 }
1633 }
1634
1635 fn note_pseudo_tool_byte(&mut self, c: u8) {
1636 if self.output_frozen() || self.dsml_active || self.replay {
1648 return;
1649 }
1650 if let Some(opener) = self.pseudo_tool.feed(c, self.in_think) {
1651 self.pseudo_tool_call(&opener);
1652 }
1653 }
1654
1655 fn flush_pseudo_tool(&mut self) {
1658 if self.output_frozen() || self.dsml_active || self.replay {
1659 return;
1660 }
1661 if let Some(opener) = self.pseudo_tool.finish_line(self.in_think) {
1662 self.pseudo_tool_call(&opener);
1663 }
1664 }
1665
1666 fn flush_start_tail(&mut self) {
1667 if self.dsml_start_tail.is_empty() {
1668 return;
1669 }
1670 self.post_think_gap = false;
1671 let held = std::mem::take(&mut self.dsml_start_tail);
1672 for b in held {
1673 self.write_char(b);
1674 self.note_plain_dsml_byte(b);
1675 self.note_pseudo_tool_byte(b);
1676 if self.output_frozen() {
1677 break;
1678 }
1679 }
1680 }
1681
1682 fn normal_byte(&mut self, c: u8) {
1687 if self.output_frozen() {
1688 return;
1689 }
1690 self.note_thinking_dsml_byte(c);
1691
1692 if self.post_think_gap && matches!(c, b' ' | b'\t' | b'\r' | b'\n') {
1695 return;
1696 }
1697
1698 if !self.dsml_start_tail.is_empty() || c == b'<' {
1699 if self.dsml_start_tail.len() < DSML_START_TAIL_CAP {
1700 self.dsml_start_tail.push(c);
1701 }
1702 let (mut complete, mut implicit_invoke) = (false, false);
1703 if dsml_start_match(&self.dsml_start_tail, &mut complete, &mut implicit_invoke) {
1704 if complete {
1705 self.start_dsml();
1708 if implicit_invoke {
1709 for &b in CANONICAL_INVOKE {
1710 self.feed_dsml_byte(b);
1711 }
1712 }
1713 }
1714 return;
1715 }
1716 if self.dsml_start_tail.len() > 1 && self.dsml_start_tail.last() == Some(&b'<') {
1719 self.post_think_gap = false;
1720 let held = std::mem::take(&mut self.dsml_start_tail);
1721 for &b in &held[..held.len() - 1] {
1722 self.write_char(b);
1723 self.note_plain_dsml_byte(b);
1724 self.note_pseudo_tool_byte(b);
1725 if self.output_frozen() {
1726 return;
1727 }
1728 }
1729 self.dsml_start_tail.push(b'<');
1730 return;
1731 }
1732 self.flush_start_tail();
1733 return;
1734 }
1735
1736 self.post_think_gap = false;
1737 self.write_char(c);
1738 self.note_plain_dsml_byte(c);
1739 self.note_pseudo_tool_byte(c);
1740 }
1741
1742 fn stream_text(&mut self, text: &[u8], finish: bool) {
1744 let mut buf = std::mem::take(&mut self.pending);
1745 buf.extend_from_slice(text);
1746
1747 let mut i = 0;
1748 while i < buf.len() {
1749 let rem = &buf[i..];
1750 if !self.dsml_active && rem.starts_with(THINK_OPEN) {
1751 self.flush_start_tail();
1752 self.post_think_gap = false;
1753 self.in_think = true;
1754 self.plain_dsml = MarkerDetector::default();
1758 i += THINK_OPEN.len();
1759 continue;
1760 }
1761 if rem.starts_with(THINK_CLOSE) && self.think_close_is_control() {
1762 if self.dsml_active {
1763 self.in_think = false;
1770 if self.dsml_ignored {
1771 self.dsml_ignored = false;
1772 self.viz_start();
1773 }
1774 } else {
1775 self.flush_start_tail();
1776 self.in_think = false;
1777 self.pseudo_tool.reset_line();
1778 self.viz_newline_if_open();
1779 self.emit_visible_bytes(b"\n");
1780 self.post_think_gap = true;
1781 }
1782 self.plain_dsml = MarkerDetector::default();
1783 i += THINK_CLOSE.len();
1784 continue;
1785 }
1786 if rem.starts_with(THINK_CLOSE) {
1787 self.think_close_swallowed = true;
1792 }
1793 if !finish
1794 && rem[0] == b'<'
1795 && self.think_close_is_control()
1796 && (is_partial_prefix(rem, THINK_OPEN) || is_partial_prefix(rem, THINK_CLOSE))
1797 {
1798 self.pending = rem.to_vec();
1799 break;
1800 }
1801
1802 let c = rem[0];
1803 if self.dsml_active {
1804 self.feed_dsml_byte(c);
1805 } else {
1806 self.normal_byte(c);
1809 }
1810 i += 1;
1811 }
1812
1813 if finish {
1814 self.flush_start_tail();
1815 self.post_think_gap = false;
1816 if self.dsml_active {
1817 self.resolve_swallowed_think_close();
1822 if self.rejects_in_think() {
1823 self.reject_in_think_stanza(
1826 "tool calling is not allowed inside <think></think>",
1827 );
1828 } else {
1829 self.viz_finish(Some(if self.preflight_error.is_some() {
1830 "[tool call stopped: edit old selector failed]\n"
1831 } else {
1832 "[tool call interrupted]\n"
1833 }));
1834 self.dsml_active = false;
1835 }
1836 }
1837 }
1846 }
1847}
1848
1849fn is_partial_prefix(bytes: &[u8], prefix: &[u8]) -> bool {
1850 bytes.len() < prefix.len() && prefix[..bytes.len()] == *bytes
1851}
1852
1853#[cfg(test)]
1854mod tests {
1855 use super::*;
1856
1857 #[derive(Debug, Default)]
1858 struct Cap {
1859 visible: String,
1860 think: String,
1861 errors: String,
1862 }
1863
1864 impl RenderSink for Cap {
1865 fn visible_text(&mut self, text: &str) {
1866 self.visible.push_str(text);
1867 }
1868 fn error_text(&mut self, text: &str) {
1869 self.errors.push_str(text);
1870 self.visible.push_str(text);
1871 }
1872 fn think_text(&mut self, text: &str) {
1873 self.think.push_str(text);
1874 }
1875 }
1876
1877 fn run_chunked(text: &str) -> StreamRenderer<Cap> {
1878 let mut sr = StreamRenderer::new(Cap::default());
1879 sr.push(text);
1880 sr.finish();
1881 sr
1882 }
1883
1884 fn pseudo_tool_renderer() -> StreamRenderer<Cap> {
1885 let mut sr = StreamRenderer::new(Cap::default());
1886 sr.set_tool_names(vec!["task".to_string(), "read".to_string()]);
1887 sr
1888 }
1889
1890 #[test]
1894 fn pseudo_tool_block_after_think_is_reported() {
1895 let mut sr = pseudo_tool_renderer();
1896 sr.push("<think>planning</think>");
1897 sr.push("<task>\ntask_context: \"add headers\"\n</task>");
1898 sr.finish();
1899 assert!(
1900 sr.finished().error.is_some(),
1901 "hallucinated call produced no error for the model to correct from"
1902 );
1903 }
1904
1905 #[test]
1908 fn pseudo_tool_block_inside_think_is_ignored() {
1909 let mut sr = pseudo_tool_renderer();
1910 sr.push("<think>\n<task>\n</task>\n</think>");
1911 sr.push("done");
1912 sr.finish();
1913 assert!(sr.finished().error.is_none());
1914 }
1915
1916 #[test]
1923 fn stray_fence_close_leaking_from_think_does_not_disarm_the_detector() {
1924 let mut sr = pseudo_tool_renderer();
1925 sr.push("<think>```rust\nfoo\n</think>```\n<task>\nx\n</task>");
1926 sr.finish();
1927 assert!(
1928 sr.finished().error.is_some(),
1929 "a fence opened inside <think> disarmed the detector for the rest of the stream"
1930 );
1931 }
1932
1933 #[test]
1935 fn pseudo_tool_name_in_prose_or_fence_is_ignored() {
1936 let mut sr = pseudo_tool_renderer();
1937 sr.push("the <task> element is XML, not DSML\n");
1938 sr.push("```xml\n<task>\n</task>\n```\n");
1939 sr.finish();
1940 assert!(sr.finished().error.is_none());
1941 }
1942
1943 #[test]
1948 fn pseudo_tool_error_wins_over_the_in_think_prohibition() {
1949 let mut sr = pseudo_tool_renderer();
1950 sr.push(concat!(
1951 "<think><|DSML|tool_calls><|DSML|invoke name=\"bash\">",
1952 "</|DSML|invoke|></|DSML|tool_calls|></think>\n",
1953 "<task>\nx\n</task>",
1954 ));
1955 sr.finish();
1956 let fin = sr.finished();
1957 assert_eq!(
1958 fin.error,
1959 Some(
1960 "<task> is not a tool call; tools are invoked with the DSML syntax in the system prompt"
1961 ),
1962 "pseudo-tool hit must not be masked by the in-think prohibition"
1963 );
1964 assert!(
1965 !fin.in_think_rejected,
1966 "in-think framing would wrap the answer-region error in the wrong advice"
1967 );
1968 }
1969
1970 #[test]
1973 fn leaked_in_think_stanza_alone_still_reports_the_prohibition() {
1974 let mut sr = pseudo_tool_renderer();
1975 sr.push(concat!(
1976 "<think><|DSML|tool_calls><|DSML|invoke name=\"bash\">",
1977 "</|DSML|invoke|></|DSML|tool_calls|></think>\n",
1978 "all done\n",
1979 ));
1980 sr.finish();
1981 let fin = sr.finished();
1982 assert_eq!(fin.error, Some(IN_THINK_PROHIBITION));
1983 assert!(fin.in_think_rejected);
1984 }
1985
1986 #[test]
1990 fn prose_continuing_after_a_tool_tag_does_not_fire() {
1991 let mut sr = pseudo_tool_renderer();
1992 sr.push("<read> is how you spell it, unlike the others.\nand more text\n");
1993 sr.finish();
1994 assert!(sr.finished().error.is_none(), "{:?}", sr.finished().error);
1995 assert!(
1996 sr.sink().visible.contains("and more text"),
1997 "answer truncated: {:?}",
1998 sr.sink().visible
1999 );
2000 }
2001
2002 #[test]
2004 fn bare_tool_tag_on_its_own_line_fires() {
2005 let mut sr = pseudo_tool_renderer();
2006 sr.push("here goes\n<read>\npath: /tmp/x\n</read>\n");
2007 sr.finish();
2008 assert!(sr.finished().error.is_some());
2009 }
2010
2011 #[test]
2013 fn bare_tool_tag_as_final_line_without_newline_fires() {
2014 let mut sr = pseudo_tool_renderer();
2015 sr.push("here goes\n<task>");
2016 sr.finish();
2017 assert_eq!(
2018 sr.finished().error,
2019 Some(
2020 "<task> is not a tool call; tools are invoked with the DSML syntax in the system prompt"
2021 )
2022 );
2023 }
2024
2025 #[test]
2027 fn pseudo_tool_message_quotes_the_matched_line() {
2028 let mut sr = StreamRenderer::new(Cap::default());
2029 sr.push("<invoke name=\"bash\">\n");
2030 sr.finish();
2031 assert_eq!(
2032 sr.finished().error,
2033 Some(
2034 "<invoke name=\"bash\"> is not a tool call; tools are invoked with the DSML syntax in the system prompt"
2035 )
2036 );
2037 }
2038
2039 #[test]
2041 fn generic_tool_call_wrapper_is_reported() {
2042 let mut sr = StreamRenderer::new(Cap::default());
2043 sr.push("<tool_call>\n{\"name\": \"read\"}\n</tool_call>");
2044 sr.finish();
2045 assert!(sr.finished().error.is_some());
2046 }
2047
2048 #[test]
2051 fn pseudo_tool_block_is_not_reported_during_replay() {
2052 let mut sr = pseudo_tool_renderer();
2053 sr.set_replay(true);
2054 sr.push("<task>\ntask_context: \"add headers\"\n</task>");
2055 sr.finish();
2056 assert!(sr.finished().error.is_none());
2057 }
2058
2059 #[test]
2061 fn real_dsml_call_is_untouched_by_the_pseudo_detector() {
2062 let mut sr = pseudo_tool_renderer();
2063 sr.push(
2064 "<|DSML|tool_calls><|DSML|invoke name=\"read\">\
2065 <|DSML|parameter name=\"path\" string=\"true\">/tmp/x</|DSML|parameter|>\
2066 </|DSML|invoke|></|DSML|tool_calls|>",
2067 );
2068 sr.finish();
2069 let fin = sr.finished();
2070 assert!(fin.error.is_none(), "error: {:?}", fin.error);
2071 assert_eq!(fin.calls.len(), 1);
2072 assert_eq!(fin.calls[0].name, "read");
2073 }
2074
2075 #[test]
2076 fn write_content_previews_dim_even_with_banners_off() {
2077 let stanza = concat!(
2078 "<|DSML|tool_calls>",
2079 "<|DSML|invoke name=\"write\">",
2080 "<|DSML|parameter name=\"path\">src/foo.rs</|DSML|parameter>",
2081 "<|DSML|parameter name=\"content\">fn main() {}\n</|DSML|parameter>",
2082 "</|DSML|invoke>",
2083 "</|DSML|tool_calls>",
2084 );
2085 let mut sr = StreamRenderer::new(Cap::default());
2088 sr.set_show_tool_calls(false);
2089 sr.push(stanza);
2090 sr.finish();
2091 let think = &sr.sink().think;
2092 assert!(think.contains("write src/foo.rs"), "header: {think:?}");
2093 assert!(think.contains("fn main() {}"), "content preview: {think:?}");
2094 assert!(
2095 !sr.sink().visible.contains("fn main()"),
2096 "content not on the visible channel: {:?}",
2097 sr.sink().visible
2098 );
2099 assert_eq!(sr.finished().calls.len(), 1, "call still parsed");
2100 }
2101
2102 #[test]
2103 fn show_tool_calls_false_suppresses_the_banner_but_keeps_visible_text() {
2104 let stanza = concat!(
2105 "answer before. ",
2106 "<|DSML|tool_calls>",
2107 "<|DSML|invoke name=\"bash\">",
2108 "<|DSML|parameter name=\"command\">ls -la</|DSML|parameter>",
2109 "</|DSML|invoke>",
2110 "</|DSML|tool_calls>",
2111 );
2112 let shown = run_chunked(stanza);
2114 assert!(
2115 shown.sink().visible.contains("🛠️"),
2116 "{:?}",
2117 shown.sink().visible
2118 );
2119
2120 let mut sr = StreamRenderer::new(Cap::default());
2123 sr.set_show_tool_calls(false);
2124 sr.push(stanza);
2125 sr.finish();
2126 let vis = &sr.sink().visible;
2127 assert!(vis.contains("answer before."), "prose kept: {vis:?}");
2128 assert!(!vis.contains("🛠️"), "banner suppressed: {vis:?}");
2129 assert!(!vis.contains("ls -la"), "params suppressed: {vis:?}");
2130 assert_eq!(sr.finished().calls.len(), 1, "call still parsed");
2131 }
2132
2133 #[test]
2134 fn show_thinking_false_suppresses_thinking_but_keeps_visible_text() {
2135 let text = "<think>secret reasoning</think>visible answer";
2136
2137 let shown = run_chunked(text);
2139 assert_eq!(shown.sink().think, "secret reasoning");
2140 assert_eq!(shown.sink().visible.trim(), "visible answer");
2141
2142 let mut sr = StreamRenderer::new(Cap::default());
2145 sr.set_show_thinking(false);
2146 sr.push(text);
2147 sr.finish();
2148 assert_eq!(sr.sink().think, "", "thinking suppressed");
2149 assert_eq!(sr.sink().visible.trim(), "visible answer", "prose kept");
2150 }
2151
2152 fn run_charwise(text: &str) -> StreamRenderer<Cap> {
2153 let mut sr = StreamRenderer::new(Cap::default());
2154 for ch in text.chars() {
2155 sr.push(ch.to_string());
2156 }
2157 sr.finish();
2158 sr
2159 }
2160
2161 const BASH_STANZA: &str = concat!(
2162 "<|DSML|tool_calls>",
2163 "<|DSML|invoke name=\"bash\">",
2164 "<|DSML|parameter name=\"command\">ls -la</|DSML|parameter|>",
2165 "</|DSML|invoke|>",
2166 "</|DSML|tool_calls|>",
2167 );
2168
2169 #[test]
2170 fn begin_in_think_routes_thinking_then_answer() {
2171 let mut sr = StreamRenderer::new(Cap::default());
2174 sr.begin_in_think();
2175 sr.push("weighing options</think>Final answer.");
2176 sr.finish();
2177 assert!(sr.sink().think.contains("weighing options"));
2178 assert!(sr.sink().visible.contains("Final answer."));
2179 assert!(!sr.sink().visible.contains("weighing options"));
2180 assert!(!sr.sink().visible.contains("think"));
2181 }
2182
2183 #[test]
2184 fn provider_explicit_think_tags_route_correctly() {
2185 let mut sr = StreamRenderer::new(Cap::default());
2190 sr.push("<think>reasoning here</think>");
2191 sr.push("visible answer");
2192 sr.finish();
2193 assert!(sr.sink().think.contains("reasoning here"));
2194 assert!(!sr.sink().think.contains("visible answer"));
2195 assert!(sr.sink().visible.contains("visible answer"));
2196
2197 let mut sr = StreamRenderer::new(Cap::default());
2200 sr.push("just an answer, no thinking");
2201 sr.finish();
2202 assert_eq!(sr.sink().visible, "just an answer, no thinking");
2203 assert_eq!(sr.sink().think, "");
2204 }
2205
2206 #[test]
2212 fn ssml_misspelling_is_accepted_as_an_alias() {
2213 let text = concat!(
2214 "Let me look at the documents module.\n",
2215 "<|SSML|tool_calls>\n",
2216 "<|SSML|invoke name=\"bash\">\n",
2217 "<|SSML|parameter name=\"command\" string=\"true\">cat documents.rs",
2218 "</|SSML|parameter>\n",
2219 "</|SSML|invoke>\n",
2220 "</|SSML|tool_calls>",
2221 );
2222 for sr in [run_chunked(text), run_charwise(text)] {
2223 let vis = &sr.sink().visible;
2224 assert!(vis.contains("🛠️ $ cat documents.rs"), "{vis:?}");
2225 assert!(!vis.contains("SSML"), "{vis:?}");
2226 let fin = sr.finished();
2227 assert_eq!(fin.calls.len(), 1);
2228 assert_eq!(fin.calls[0].name, "bash");
2229 assert_eq!(fin.calls[0].arg_value("command"), Some("cat documents.rs"));
2230 assert!(fin.error.is_none(), "{:?}", fin.error);
2231 }
2232 }
2233
2234 #[test]
2237 fn dsml_and_ssml_tags_mix_within_one_stanza() {
2238 let text = concat!(
2239 "<|DSML|tool_calls>",
2240 "<|SSML|invoke name=\"bash\">",
2241 "<|DSML|parameter name=\"command\">ls -la</|SSML|parameter>",
2242 "</|DSML|invoke>",
2243 "</|SSML|tool_calls>",
2244 );
2245 for sr in [run_chunked(text), run_charwise(text)] {
2246 let fin = sr.finished();
2247 assert_eq!(fin.calls.len(), 1);
2248 assert_eq!(fin.calls[0].arg_value("command"), Some("ls -la"));
2249 assert!(fin.error.is_none(), "{:?}", fin.error);
2250 }
2251 }
2252
2253 #[test]
2257 fn unrelated_marker_names_are_still_plain_text() {
2258 let text = "<|XSML|tool_calls><|XSML|invoke name=\"bash\">";
2259 for sr in [run_chunked(text), run_charwise(text)] {
2260 assert!(sr.finished().calls.is_empty());
2261 assert_eq!(sr.sink().visible, text);
2262 }
2263 }
2264
2265 #[test]
2266 fn prose_passes_through() {
2267 for sr in [run_chunked("Hello, world."), run_charwise("Hello, world.")] {
2268 assert_eq!(sr.sink().visible, "Hello, world.");
2269 assert_eq!(sr.sink().think, "");
2270 assert!(sr.finished().calls.is_empty());
2271 assert!(sr.finished().error.is_none());
2272 }
2273 }
2274
2275 #[test]
2276 fn bash_stanza_hides_dsml_and_shows_banner() {
2277 let text = format!("Let me look.\n{BASH_STANZA}");
2278 for sr in [run_chunked(&text), run_charwise(&text)] {
2279 let vis = &sr.sink().visible;
2280 assert!(vis.starts_with("Let me look.\n"), "{vis:?}");
2281 assert!(vis.contains("🛠️ $ ls -la"), "{vis:?}");
2282 assert!(!vis.contains("DSML"), "{vis:?}");
2283 let fin = sr.finished();
2284 assert_eq!(fin.calls.len(), 1);
2285 assert_eq!(fin.calls[0].name, "bash");
2286 assert_eq!(fin.calls[0].arg_value("command"), Some("ls -la"));
2287 assert!(fin.error.is_none());
2288 }
2289 }
2290
2291 #[test]
2292 fn read_banner_shows_path_and_range() {
2293 let stanza = concat!(
2294 "<|DSML|tool_calls>",
2295 "<|DSML|invoke name=\"read\">",
2296 "<|DSML|parameter name=\"path\" string=\"true\">src/main.rs</|DSML|parameter|>",
2297 "</|DSML|invoke|>",
2298 "</|DSML|tool_calls|>",
2299 );
2300 for sr in [run_chunked(stanza), run_charwise(stanza)] {
2301 let vis = &sr.sink().visible;
2302 assert!(vis.contains("🛠️ Reading src/main.rs 1:500...\n"), "{vis:?}");
2303 assert!(!vis.contains("DSML"), "{vis:?}");
2304 }
2305 }
2306
2307 #[test]
2308 fn read_banner_whole_file() {
2309 let stanza = concat!(
2310 "<|DSML|tool_calls>",
2311 "<|DSML|invoke name=\"read\">",
2312 "<|DSML|parameter name=\"path\" string=\"true\">a.c</|DSML|parameter|>",
2313 "<|DSML|parameter name=\"whole\">true</|DSML|parameter|>",
2314 "</|DSML|invoke|>",
2315 "</|DSML|tool_calls|>",
2316 );
2317 let sr = run_chunked(stanza);
2318 assert!(
2319 sr.sink()
2320 .visible
2321 .contains("🛠️ Reading a.c (whole file)...\n"),
2322 "{:?}",
2323 sr.sink().visible
2324 );
2325 }
2326
2327 #[test]
2328 fn edit_diff_uses_minus_plus_prefixes() {
2329 let stanza = concat!(
2330 "<|DSML|tool_calls>",
2331 "<|DSML|invoke name=\"edit\">",
2332 "<|DSML|parameter name=\"path\" string=\"true\">a.rs</|DSML|parameter|>",
2333 "<|DSML|parameter name=\"old\">let a = 1;</|DSML|parameter|>",
2334 "<|DSML|parameter name=\"new\">let a = 2;</|DSML|parameter|>",
2335 "</|DSML|invoke|>",
2336 "</|DSML|tool_calls|>",
2337 );
2338 for sr in [run_chunked(stanza), run_charwise(stanza)] {
2339 let vis = &sr.sink().visible;
2340 assert!(vis.contains("🛠️ edit path=a.rs"), "{vis:?}");
2341 assert!(vis.contains("- let a = 1;"), "{vis:?}");
2342 assert!(vis.contains("+ let a = 2;"), "{vis:?}");
2343 assert!(!vis.contains("DSML"), "{vis:?}");
2344 assert_eq!(sr.finished().calls[0].arg_value("new"), Some("let a = 2;"));
2345 }
2346 }
2347
2348 #[test]
2349 fn partial_marker_false_alarm_is_flushed() {
2350 let mut sr = StreamRenderer::new(Cap::default());
2351 sr.push("<|DSM");
2352 assert_eq!(sr.sink().visible, "");
2354 sr.push("ok");
2355 sr.finish();
2356 assert_eq!(sr.sink().visible, "<|DSMok");
2357 assert!(sr.finished().error.is_none());
2358 }
2359
2360 #[test]
2361 fn partial_marker_flushed_at_stream_end() {
2362 let mut sr = StreamRenderer::new(Cap::default());
2366 sr.push("done <|DSML|tool_c");
2367 sr.finish();
2368 assert!(
2369 sr.sink().visible.starts_with("done <|DSML|"),
2370 "{:?}",
2371 sr.sink().visible
2372 );
2373 assert!(
2374 sr.sink().visible.contains("[invalid tool call: "),
2375 "{:?}",
2376 sr.sink().visible
2377 );
2378 }
2379
2380 #[test]
2381 fn think_text_routes_to_think_sink() {
2382 let sr = run_chunked("<think>pondering</think>Answer.");
2383 assert_eq!(sr.sink().think, "pondering");
2384 assert!(
2385 sr.sink().visible.ends_with("Answer."),
2386 "{:?}",
2387 sr.sink().visible
2388 );
2389 assert!(!sr.sink().visible.contains("pondering"));
2390 }
2391
2392 #[test]
2393 fn think_tag_split_across_pushes() {
2394 let mut sr = StreamRenderer::new(Cap::default());
2395 sr.push("<th");
2396 sr.push("ink>hidden</th");
2397 sr.push("ink>shown");
2398 sr.finish();
2399 assert_eq!(sr.sink().think, "hidden");
2400 assert!(
2401 sr.sink().visible.ends_with("shown"),
2402 "{:?}",
2403 sr.sink().visible
2404 );
2405 }
2406
2407 #[test]
2408 fn dsml_inside_think_is_ignored_and_reported() {
2409 let text = format!("<think>{BASH_STANZA}</think>ok");
2410 for sr in [run_chunked(&text), run_charwise(&text)] {
2411 let fin = sr.finished();
2412 assert!(fin.dsml_in_think);
2413 assert!(fin.calls.is_empty(), "{:?}", fin.calls);
2418 assert!(
2419 sr.sink().visible.contains(
2420 "[tool call ignored: tool calling is not allowed inside <think></think>]"
2421 ),
2422 "{:?}",
2423 sr.sink().visible
2424 );
2425 assert!(!sr.sink().think.contains("DSML"), "{:?}", sr.sink().think);
2426 }
2427 }
2428
2429 fn run_allowing_in_think(text: &str) -> StreamRenderer<Cap> {
2430 let mut sr = StreamRenderer::new(Cap::default());
2431 sr.set_thinking_tool_calls(true);
2432 sr.push(text);
2433 sr.finish();
2434 sr
2435 }
2436
2437 fn run_allowing_in_think_charwise(text: &str) -> StreamRenderer<Cap> {
2438 let mut sr = StreamRenderer::new(Cap::default());
2439 sr.set_thinking_tool_calls(true);
2440 for ch in text.chars() {
2441 sr.push(ch.to_string());
2442 }
2443 sr.finish();
2444 sr
2445 }
2446
2447 #[test]
2452 fn an_unfinished_in_think_stanza_reports_placement_not_syntax() {
2453 let mut sr = StreamRenderer::new(Cap::default());
2454 sr.push("<think>let me look<|DSML|tool_calls><|DSML|invoke name=\"bash\">");
2456 sr.finish();
2457 let fin = sr.finished();
2458 assert!(fin.calls.is_empty(), "{:?}", fin.calls);
2459 assert!(fin.in_think_rejected, "rejected for placement");
2460 assert_eq!(fin.error, Some(IN_THINK_PROHIBITION));
2461 assert!(fin.ended_in_think);
2462 }
2463
2464 #[test]
2472 fn a_stanza_opened_in_think_but_closed_after_it_is_dispatched() {
2473 let mut sr = StreamRenderer::new(Cap::default());
2474 sr.push("<think>the correct format is:\n\n<|DSML|tool_calls>\n</think>\n\n");
2475 sr.push("<|DSML|invoke name=\"bash\">");
2476 sr.push("<|DSML|parameter name=\"command\">ls -la</|DSML|parameter|>");
2477 sr.push("</|DSML|invoke|></|DSML|tool_calls|>");
2478 sr.finish();
2479 let fin = sr.finished();
2480 assert_eq!(fin.calls.len(), 1, "{:?}", fin.calls);
2481 assert_eq!(fin.calls[0].arg_value("command"), Some("ls -la"));
2482 assert!(!fin.in_think_rejected, "the call closed outside thinking");
2483 assert!(fin.error.is_none(), "{:?}", fin.error);
2484 assert!(
2485 !sr.sink().visible.contains("[tool call ignored:"),
2486 "{:?}",
2487 sr.sink().visible
2488 );
2489 }
2490
2491 #[test]
2497 fn dsml_mentioned_while_thinking_without_a_stop_token_is_not_rejected() {
2498 for text in [
2499 "<think>each arg is <|DSML|parameter name=\"x\" string=\"true\">v</|DSML|parameter|> inside the invoke</think>the answer",
2502 "<think>and it ends with </|DSML|tool_calls|> of course</think>the answer",
2504 ] {
2505 let sr = run_chunked(text);
2506 let fin = sr.finished();
2507 assert!(fin.calls.is_empty(), "{:?}", fin.calls);
2508 assert!(!fin.in_think_rejected, "nothing completed; {text}");
2509 assert_eq!(fin.error, None, "{text}");
2510 assert!(fin.dsml_in_think, "{text}");
2513 assert!(sr.sink().visible.contains("the answer"), "{text}");
2514 }
2515 }
2516
2517 #[test]
2522 fn think_close_inside_a_parameter_value_is_payload_not_a_control_token() {
2523 let mut sr = StreamRenderer::new(Cap::default());
2524 sr.push("<think>writing it up</think>");
2525 sr.push("<|DSML|tool_calls><|DSML|invoke name=\"write\">");
2526 sr.push("<|DSML|parameter name=\"content\">close it with </think> when done");
2527 sr.push("</|DSML|parameter|></|DSML|invoke|></|DSML|tool_calls|>");
2528 sr.finish();
2529 let fin = sr.finished();
2530 assert_eq!(fin.calls.len(), 1, "{:?}", fin.calls);
2531 assert_eq!(
2532 fin.calls[0].arg_value("content"),
2533 Some("close it with </think> when done")
2534 );
2535 assert!(fin.error.is_none(), "{:?}", fin.error);
2536 }
2537
2538 #[test]
2542 fn a_stanza_wholly_inside_think_is_rejected_without_a_banner() {
2543 let mut sr = StreamRenderer::new(Cap::default());
2544 sr.push(format!("<think>thinking{BASH_STANZA}</think>done"));
2545 sr.finish();
2546 let fin = sr.finished();
2547 assert!(fin.calls.is_empty(), "{:?}", fin.calls);
2548 assert!(fin.in_think_rejected);
2549 assert_eq!(fin.error, Some(IN_THINK_PROHIBITION));
2550 assert!(
2551 !sr.sink().visible.contains("🛠️"),
2552 "no banner for a call that never happened: {:?}",
2553 sr.sink().visible
2554 );
2555 }
2556
2557 #[test]
2561 fn shorthand_invoke_renders_a_banner() {
2562 let mut sr = StreamRenderer::new(Cap::default());
2563 sr.push("<|DSML|tool_calls><|DSML|bash>");
2564 sr.push("<|DSML|parameter name=\"command\" string=\"true\">ls -la</|DSML|parameter|>");
2565 sr.push("</|DSML|invoke|></|DSML|tool_calls|>");
2566 sr.finish();
2567 let fin = sr.finished();
2568 assert_eq!(fin.calls.len(), 1, "{:?}", fin.calls);
2569 assert_eq!(fin.calls[0].name, "bash");
2570 assert!(
2571 sr.sink().visible.contains("🛠️ $ ls -la"),
2572 "{:?}",
2573 sr.sink().visible
2574 );
2575 assert!(
2576 !sr.sink().visible.contains("DSML"),
2577 "{:?}",
2578 sr.sink().visible
2579 );
2580 }
2581
2582 #[test]
2585 fn shorthand_parameter_renders_under_its_element_name() {
2586 let mut sr = StreamRenderer::new(Cap::default());
2587 sr.push("<|DSML|tool_calls><|DSML|invoke name=\"bash\">");
2588 sr.push("<|DSML|command string=\"true\">ls -la</|DSML|invoke>");
2589 sr.push("</|DSML|invoke|></|DSML|tool_calls|>");
2590 sr.finish();
2591 let fin = sr.finished();
2592 assert_eq!(fin.calls.len(), 1, "{:?}", fin.calls);
2593 assert_eq!(fin.calls[0].arg_value("command"), Some("ls -la"));
2594 assert!(
2595 sr.sink().visible.contains("🛠️ $ ls -la"),
2596 "{:?}",
2597 sr.sink().visible
2598 );
2599 }
2600
2601 #[test]
2604 fn a_completed_in_think_stanza_reports_placement() {
2605 let mut sr = StreamRenderer::new(Cap::default());
2606 sr.push(format!("<think>thinking{BASH_STANZA}"));
2607 sr.finish();
2608 let fin = sr.finished();
2609 assert!(fin.calls.is_empty(), "the call must not be dispatched");
2610 assert!(fin.in_think_rejected);
2611 assert_eq!(fin.error, Some(IN_THINK_PROHIBITION));
2612 }
2613
2614 #[test]
2618 fn allowing_in_think_calls_reports_no_placement_error() {
2619 let text = format!("<think>{BASH_STANZA}</think>ok");
2620 let sr = run_allowing_in_think(&text);
2621 let fin = sr.finished();
2622 assert_eq!(fin.calls.len(), 1);
2623 assert!(!fin.in_think_rejected, "nothing was rejected");
2624 assert!(fin.error.is_none(), "{:?}", fin.error);
2625 assert!(fin.dsml_in_think, "the marker was still seen in thinking");
2626 }
2627
2628 #[test]
2629 fn dsml_inside_think_is_executed_when_allowed() {
2630 let text = format!("<think>{BASH_STANZA}</think>ok");
2631 for sr in [
2632 run_allowing_in_think(&text),
2633 run_allowing_in_think_charwise(&text),
2634 ] {
2635 let fin = sr.finished();
2636 assert_eq!(fin.calls.len(), 1, "{:?}", fin.calls);
2637 assert_eq!(fin.calls[0].name, "bash");
2638 assert_eq!(fin.calls[0].arg_value("command"), Some("ls -la"));
2639 assert!(fin.error.is_none(), "{:?}", fin.error);
2640 assert!(
2641 !sr.sink().visible.contains("[tool call ignored:"),
2642 "{:?}",
2643 sr.sink().visible
2644 );
2645 assert!(
2648 sr.sink().visible.contains("🛠️ $ ls -la"),
2649 "{:?}",
2650 sr.sink().visible
2651 );
2652 assert!(
2653 !sr.sink().visible.contains("DSML"),
2654 "{:?}",
2655 sr.sink().visible
2656 );
2657 assert!(!sr.sink().think.contains("DSML"), "{:?}", sr.sink().think);
2658 }
2659 }
2660
2661 #[test]
2662 fn ended_in_think_reports_an_open_block() {
2663 let sr = run_allowing_in_think(&format!("<think>let me look{BASH_STANZA}"));
2665 assert!(sr.finished().ended_in_think);
2666
2667 let closed = run_allowing_in_think(&format!("<think>done</think>{BASH_STANZA}"));
2670 assert!(!closed.finished().ended_in_think);
2671 assert!(!run_chunked("plain answer").finished().ended_in_think);
2672
2673 let stanza_then_close =
2677 run_allowing_in_think(&format!("<think>a{BASH_STANZA}</think>tail"));
2678 assert!(!stanza_then_close.finished().ended_in_think);
2679 }
2680
2681 #[test]
2682 fn interrupted_stanza_reports_status() {
2683 let mut sr = StreamRenderer::new(Cap::default());
2684 sr.push("<|DSML|tool_calls><|DSML|invoke name=\"bash\">");
2685 sr.push("<|DSML|parameter name=\"command\">sleep 1");
2686 sr.finish();
2687 let vis = &sr.sink().visible;
2688 assert!(vis.contains("🛠️ $ sleep 1"), "{vis:?}");
2689 assert!(vis.contains("[tool call interrupted]\n"), "{vis:?}");
2690 assert!(sr.finished().calls.is_empty());
2691 }
2692
2693 #[test]
2694 fn incomplete_stanza_reports_incomplete_error() {
2695 let sr = run_chunked(concat!(
2696 "<|DSML|tool_calls>",
2697 "<|DSML|invoke name=\"bash\">",
2698 "<|DSML|parameter name=\"command\">ls",
2699 ));
2700 assert_eq!(sr.finished().error, Some("incomplete DSML tool call"));
2701 assert!(sr.finished().calls.is_empty());
2702 }
2703
2704 #[test]
2705 fn greedy_sampling_tracks_dsml_state() {
2706 let mut sr = StreamRenderer::new(Cap::default());
2707 sr.push("hello ");
2708 assert!(!sr.wants_greedy_sampling(), "prose");
2709 sr.push("<|DS");
2710 assert!(sr.wants_greedy_sampling(), "DSML-shaped held prefix");
2711 sr.push("ML|tool_calls><|DSML|invoke name=\"bash\">");
2712 assert!(sr.wants_greedy_sampling(), "structural markup");
2713 sr.push("<|DSML|parameter name=\"command\">ls -la");
2714 assert!(!sr.wants_greedy_sampling(), "free-form parameter value");
2715 sr.push("</|DSML|parameter");
2716 assert!(sr.wants_greedy_sampling(), "close tag streaming");
2717 sr.push("|></|DSML|invoke|></|DSML|tool_calls|>");
2718 assert!(!sr.wants_greedy_sampling(), "stanza done");
2719 }
2720
2721 #[test]
2722 fn edit_old_preflight_failure_is_reported_midstream() {
2723 let mut sr = StreamRenderer::new(Cap::default());
2724 sr.set_preflight(|call| {
2725 assert_eq!(call.name, "edit");
2726 assert_eq!(call.arg_value("path"), Some("src/a.rs"));
2727 Err("old text is not a unique match".to_string())
2728 });
2729 sr.push(concat!(
2730 "<|DSML|tool_calls>",
2731 "<|DSML|invoke name=\"edit\">",
2732 "<|DSML|parameter name=\"path\">src/a.rs</|DSML|parameter|>",
2733 "<|DSML|parameter name=\"old\">nope</|DSML|parameter|>",
2734 ));
2735 assert_eq!(
2737 sr.preflight_error(),
2738 Some(
2739 "edit old selector failed before new was generated: \
2740 old text is not a unique match"
2741 )
2742 );
2743 sr.finish();
2744 assert!(
2745 sr.sink()
2746 .errors
2747 .contains("[tool call stopped: edit old selector failed]"),
2748 "{:?}",
2749 sr.sink().errors
2750 );
2751 }
2752
2753 #[test]
2754 fn edit_old_preflight_pass_leaves_stream_clean() {
2755 let mut sr = StreamRenderer::new(Cap::default());
2756 sr.set_preflight(|_| Ok(()));
2757 sr.push(concat!(
2758 "<|DSML|tool_calls>",
2759 "<|DSML|invoke name=\"edit\">",
2760 "<|DSML|parameter name=\"path\">src/a.rs</|DSML|parameter|>",
2761 "<|DSML|parameter name=\"old\">a</|DSML|parameter|>",
2762 "<|DSML|parameter name=\"new\">b</|DSML|parameter|>",
2763 "</|DSML|invoke|>",
2764 "</|DSML|tool_calls|>",
2765 ));
2766 sr.finish();
2767 assert!(sr.preflight_error().is_none());
2768 assert_eq!(sr.finished().calls.len(), 1);
2769 assert!(sr.finished().error.is_none());
2770 }
2771
2772 #[test]
2773 fn malformed_stanza_suppresses_raw_and_reports_error() {
2774 let sr = run_chunked("<|DSML|tool_calls><b>");
2775 let vis = &sr.sink().visible;
2776 assert!(
2778 sr.sink()
2779 .errors
2780 .contains("[invalid tool call: unexpected DSML tag: <b>]"),
2781 "{:?}",
2782 sr.sink().errors
2783 );
2784 assert!(!vis.contains("tool_calls"), "{vis:?}");
2786 assert!(sr.finished().error.is_some());
2787 }
2788
2789 #[test]
2797 fn a_dsml_error_does_not_freeze_a_renderer_that_did_not_opt_in() {
2798 let mut sr = StreamRenderer::new(Cap::default());
2799 sr.push("junk \u{ff5c}DSML\u{ff5c} junk");
2800 sr.push("<think>reconsidering</think>Here is the corrected answer.");
2802 sr.finish();
2803 let vis = &sr.sink().visible;
2804 assert!(
2805 vis.contains("[invalid tool call: DSML markup outside a valid tool_calls block]"),
2806 "the error line must still render: {vis:?}"
2807 );
2808 assert!(
2809 vis.contains("Here is the corrected answer."),
2810 "output after the error must still render: {vis:?}"
2811 );
2812 assert_eq!(
2814 sr.finished().error,
2815 Some("DSML markup outside a valid tool_calls block")
2816 );
2817 }
2818
2819 #[test]
2822 fn an_opted_in_renderer_still_freezes_on_a_dsml_error() {
2823 let mut sr = StreamRenderer::new(Cap::default());
2824 sr.set_freeze_on_error(true);
2825 sr.push("junk \u{ff5c}DSML\u{ff5c} junk");
2826 sr.push("<think>reconsidering</think>Here is the corrected answer.");
2827 sr.finish();
2828 let vis = &sr.sink().visible;
2829 assert!(
2830 !vis.contains("Here is the corrected answer."),
2831 "frozen output must stay frozen: {vis:?}"
2832 );
2833 }
2834
2835 #[test]
2836 fn loose_dsml_marker_is_flagged() {
2837 let sr = run_chunked("junk |DSML| junk");
2838 assert!(
2839 sr.sink()
2840 .visible
2841 .contains("[invalid tool call: DSML markup outside a valid tool_calls block]"),
2842 "{:?}",
2843 sr.sink().visible
2844 );
2845 assert_eq!(
2846 sr.finished().error,
2847 Some("DSML markup outside a valid tool_calls block")
2848 );
2849 }
2850
2851 #[test]
2856 fn a_quoted_marker_in_think_does_not_disarm_the_loose_marker_validator() {
2857 let sr = run_chunked("<think>the |DSML| marker opens a call</think>junk |DSML| junk");
2858 assert_eq!(
2859 sr.finished().error,
2860 Some("DSML markup outside a valid tool_calls block"),
2861 "{:?}",
2862 sr.sink().visible
2863 );
2864 }
2865
2866 #[test]
2872 fn a_think_close_swallowed_by_an_unfinished_stanza_still_ends_thinking() {
2873 let sr = run_chunked(concat!(
2874 "<think>",
2875 "<|DSML|tool_calls>",
2876 "<|DSML|invoke name=\"write\">",
2877 "<|DSML|parameter name=\"content\">x</think>",
2878 ));
2879 let fin = sr.finished();
2880 assert!(!fin.ended_in_think, "in_think stayed stuck after </think>");
2881 assert!(
2882 !fin.in_think_rejected,
2883 "a call cut off after thinking closed is not an in-think call"
2884 );
2885 assert!(
2886 sr.sink().visible.contains("[tool call interrupted]"),
2887 "{:?}",
2888 sr.sink()
2889 );
2890 }
2891
2892 #[test]
2896 fn a_think_close_inside_a_valid_payload_stays_payload_text() {
2897 let sr = run_chunked(concat!(
2898 "<think>",
2899 "<|DSML|tool_calls>",
2900 "<|DSML|invoke name=\"bash\">",
2901 "<|DSML|parameter name=\"command\">echo </think></|DSML|parameter|>",
2902 "</|DSML|invoke|>",
2903 "</|DSML|tool_calls|>",
2904 ));
2905 let fin = sr.finished();
2906 assert!(fin.ended_in_think, "the thinking block never closed");
2907 assert_eq!(fin.calls.len(), 0, "an in-think stanza is not dispatched");
2908 }
2909
2910 #[test]
2911 fn implicit_invoke_opener_is_accepted() {
2912 let stanza = concat!(
2913 "<|DSML|invoke name=\"bash\">",
2914 "<|DSML|parameter name=\"command\">pwd</|DSML|parameter|>",
2915 "</|DSML|invoke|>",
2916 "</|DSML|tool_calls|>",
2917 );
2918 for sr in [run_chunked(stanza), run_charwise(stanza)] {
2919 assert!(
2920 sr.sink().visible.contains("🛠️ $ pwd"),
2921 "{:?}",
2922 sr.sink().visible
2923 );
2924 let fin = sr.finished();
2925 assert_eq!(fin.calls.len(), 1);
2926 assert_eq!(fin.calls[0].arg_value("command"), Some("pwd"));
2927 }
2928 }
2929
2930 #[test]
2935 fn opener_with_trailing_bar_is_accepted() {
2936 let stanza = concat!(
2937 "<|DSML|tool_calls|>",
2938 "<|DSML|invoke name=\"bash\">",
2939 "<|DSML|parameter name=\"command\" string=\"true\">pwd</|DSML|parameter|>",
2940 "</|DSML|invoke|>",
2941 "</|DSML|tool_calls|>",
2942 );
2943 for sr in [run_chunked(stanza), run_charwise(stanza)] {
2944 let fin = sr.finished();
2945 assert_eq!(fin.error, None);
2946 assert_eq!(fin.calls.len(), 1);
2947 assert_eq!(fin.calls[0].name, "bash");
2948 assert_eq!(fin.calls[0].arg_value("command"), Some("pwd"));
2949 }
2950 }
2951
2952 #[test]
2956 fn shorthand_parameter_element_dispatches() {
2957 let stanza = concat!(
2958 "<|DSML|tool_calls|>",
2959 "<|DSML|invoke name=\"bash\">",
2960 "<|DSML|command string=\"true\">ls</|DSML|invoke>",
2961 "</|DSML|invoke>",
2962 "</|DSML|tool_calls|>",
2963 );
2964 for sr in [run_chunked(stanza), run_charwise(stanza)] {
2965 let fin = sr.finished();
2966 assert_eq!(fin.error, None);
2967 assert_eq!(fin.calls.len(), 1);
2968 assert_eq!(fin.calls[0].name, "bash");
2969 assert_eq!(fin.calls[0].arg_value("command"), Some("ls"));
2970 }
2971 }
2972
2973 #[test]
2974 fn write_content_streams_without_label() {
2975 let stanza = concat!(
2976 "<|DSML|tool_calls>",
2977 "<|DSML|invoke name=\"write\">",
2978 "<|DSML|parameter name=\"path\" string=\"true\">x.txt</|DSML|parameter|>",
2979 "<|DSML|parameter name=\"content\">line one\nline two</|DSML|parameter|>",
2980 "</|DSML|invoke|>",
2981 "</|DSML|tool_calls|>",
2982 );
2983 for sr in [run_chunked(stanza), run_charwise(stanza)] {
2984 let vis = &sr.sink().visible;
2985 assert!(vis.contains("🛠️ write path=x.txt"), "{vis:?}");
2986 assert!(!vis.contains("line one"), "{vis:?}");
2988 assert!(
2989 sr.sink().think.contains("line one\nline two"),
2990 "{:?}",
2991 sr.sink().think
2992 );
2993 assert!(!vis.contains("content:"), "{vis:?}");
2994 assert!(!vis.contains("DSML"), "{vis:?}");
2995 }
2996 }
2997
2998 #[test]
2999 fn post_think_whitespace_gap_is_swallowed() {
3000 let sr = run_chunked("<think>x</think>\n\n Answer");
3001 assert!(
3002 sr.sink().visible.ends_with("Answer"),
3003 "{:?}",
3004 sr.sink().visible
3005 );
3006 assert!(!sr.sink().visible.contains("\n\n Answer"));
3007 }
3008
3009 #[test]
3010 fn charwise_and_chunked_agree() {
3011 let text = format!("hi <not dsml> there\n{BASH_STANZA}");
3012 let a = run_chunked(&text);
3013 let b = run_charwise(&text);
3014 assert_eq!(a.sink().visible, b.sink().visible);
3015 assert_eq!(a.finished().calls, b.finished().calls);
3016 }
3017
3018 #[test]
3019 fn tool_error_logging_honors_the_opt_out_env_var() {
3020 use std::ffi::OsStr;
3021 assert!(!super::logging_enabled_for(Some(OsStr::new("1"))));
3025 assert!(super::logging_enabled_for(None));
3026 assert!(super::logging_enabled_for(Some(OsStr::new("0"))));
3028 }
3029}