inline_tool_calls.rs · sigitcode-old/sigit

Fork 1

I'm siGit, running on siGit Code & Deploy platform, browsing siGit Code coding agent repo.

Code Commits Issues 15 Pulls 1 Cloud Agent CICD

feature/account-commands-system-status

claude/github-issue-40-0au5ua claude/sigit-acp-local-chat-cx6380 claude/zen-feynman-0u78dk development feature/account-commands-system-status feature/skill-slash-commands feature/tui-slash-autocomplete feature/xcode-mcp-server fix/zed-commands-after-session-response main

rs 1,943 lines 76.3 KB

Raw

1	//! Recovery for tool calls a model emits as literal text instead of the
2	//! endpoint's structured `tool_calls` field.
3	//!
4	//! Several open-weight families siGit Code can drive are fine-tuned on chat
5	//! templates that render a tool call inside assistant content. A serving stack
6	//! is supposed to parse that back into the OpenAI-shaped `tool_calls` field
7	//! before it reaches us. When it doesn't, the raw tag arrives as ordinary
8	//! content.
9	//!
10	//! Nothing then executes it. Worse, `consume_stream` forwards content to the
11	//! UI as it arrives, so the tag is rendered verbatim in the editor and the
12	//! turn ends with no tool calls, which reads to the agent loop as "the model
13	//! chose to answer in prose". Whatever it was actually trying to do is
14	//! dropped.
15	//!
16	//! This lives here rather than only in the gateway because siGit Code talks to
17	//! more than one kind of endpoint: siGit Code Cloud, an arbitrary
18	//! OpenAI-compatible base URL from `providers.toml` or `OPENAI_BASE_URL`, and
19	//! on-device models. A fix in any single upstream leaves the others exposed.
20	//!
21	//! Only well-formed shapes are recovered. A block that doesn't match is never
22	//! guessed at: reissuing a `run_command` is cheap, but guessing wrong at a
23	//! half-parsed `edit_file` would write the wrong change to a file. It isn't
24	//! shown either. A tool-call block that fails to parse, or that is still open
25	//! when the reply ends, is reported as [`ScanEvent::Malformed`] and kept out of
26	//! both the UI and the conversation history. Left in the history, the model
27	//! reads back a call it "made" with no result attached, and later turns start
28	//! inventing `<function_results>` blocks of their own. The backend tells the
29	//! model the call didn't run and asks for it again instead.
30	//!
31	//! A marker the model merely mentions in prose, say while explaining this very
32	//! scanner, is not a call. So a block that fails to parse is only treated as
33	//! malformed when what follows the marker looks like the start of a call (see
34	//! [`looks_like_xml_call`] and [`looks_like_k3_call`]); otherwise it stays
35	//! text, as it always did.
36	//!
37	//! One spelling difference is folded away before matching: DSML markers are
38	//! built from fullwidth bars, and some stacks hand them back with each bar
39	//! doubled (`<｜｜DSML｜｜invoke …>`). See [`fold_dsml_bars`]. DSML envelopes
40	//! come as `tool_calls`, `function_calls` or bare `calls`, and a closer may put
41	//! its slash before the bar (`</｜DSML｜invoke>`); see [`fold_dsml_closers`].
42	//!
43	//! Some models wrap the call in `<tool_call>` or `<minimax:tool_call>` and
44	//! write `<invoke name="…">` with parameters as `<parameter name="k">v</parameter>`
45	//! or `<k>v</k>`. The tokenizer's `]<]minimax[>[` special token also leaks
46	//! into content, between tags and in prose; it is stripped before matching.
47
48	use std::collections::HashMap;
49
50	use crate::backend::ToolSpec;
51
52	const XML_OPEN_TAG: &str = "<tool_call>";
53	const XML_CLOSE_TAG: &str = "</tool_call>";
54
55	const NAMESPACED_OPEN_TAG: &str = "<minimax:tool_call>";
56	const NAMESPACED_CLOSE_TAG: &str = "</minimax:tool_call>";
57	/// A tokenizer special token that leaks into content. Never meant to
58	/// be visible, so it is removed wherever it appears.
59	const LEAKED_SENTINEL: &str = "]<]minimax[>[";
60	const INVOKE_OPEN: &str = "<invoke";
61	const INVOKE_CLOSE: &str = "</invoke>";
62	const PARAMETER_OPEN: &str = "<parameter";
63	const PARAMETER_CLOSE: &str = "</parameter>";
64
65	const K3_TOOLS_OPEN: &str = "<\|open\|>tools<\|sep\|>";
66	const K3_TOOLS_CLOSE: &str = "<\|close\|>tools<\|sep\|>";
67	const K3_CALL_OPEN: &str = "<\|open\|>call";
68	const K3_CALL_CLOSE: &str = "<\|close\|>call<\|sep\|>";
69	const K3_ARGUMENT_OPEN: &str = "<\|open\|>argument";
70	const K3_ARGUMENT_CLOSE: &str = "<\|close\|>argument<\|sep\|>";
71	const K3_RESPONSE_OPEN: &str = "<\|open\|>response<\|sep\|>";
72	const K3_RESPONSE_CLOSE: &str = "<\|close\|>response<\|sep\|>";
73	const K3_THINK_OPEN: &str = "<\|open\|>think<\|sep\|>";
74	const K3_THINK_CLOSE: &str = "<\|close\|>think<\|sep\|>";
75	const K3_SEP: &str = "<\|sep\|>";
76
77	/// `｜` (U+FF5C), the bar every DSML marker above is built from.
78	const FULLWIDTH_BAR: char = '｜';
79	/// Two of them — the spelling [`fold_dsml_bars`] rewrites into one.
80	const DOUBLED_BAR: &str = "｜｜";
81	/// The name those markers wrap, and the only place [`fold_dsml_bars`] folds.
82	const DSML_TAG_NAME: &str = "DSML";
83
84	const DSML_TOOL_CALLS_OPEN: &str = "<｜DSML｜tool_calls>";
85	const DSML_TOOL_CALLS_CLOSE: &str = "<｜DSML｜/tool_calls>";
86	// The published DSML template names the envelope `function_calls`; one stack
87	// was seen writing bare `calls`.
88	const DSML_FUNCTION_CALLS_OPEN: &str = "<｜DSML｜function_calls>";
89	const DSML_FUNCTION_CALLS_CLOSE: &str = "<｜DSML｜/function_calls>";
90	const DSML_CALLS_OPEN: &str = "<｜DSML｜calls>";
91	const DSML_CALLS_CLOSE: &str = "<｜DSML｜/calls>";
92	/// Every DSML envelope spelling as `(open, close)`.
93	const DSML_ENVELOPES: &[(&str, &str)] = &[
94	(DSML_TOOL_CALLS_OPEN, DSML_TOOL_CALLS_CLOSE),
95	(DSML_FUNCTION_CALLS_OPEN, DSML_FUNCTION_CALLS_CLOSE),
96	(DSML_CALLS_OPEN, DSML_CALLS_CLOSE),
97	];
98	/// A closer with the slash before the bar, and the canonical spelling the
99	/// scanner matches instead.
100	const DSML_SLASH_FIRST_CLOSER: &str = "</｜DSML｜";
101	const DSML_CANONICAL_CLOSER: &str = "<｜DSML｜/";
102	const DSML_INVOKE_OPEN: &str = "<｜DSML｜invoke";
103	const DSML_INVOKE_CLOSE: &str = "<｜DSML｜/invoke>";
104	const DSML_PARAMETER_OPEN: &str = "<｜DSML｜parameter";
105	const DSML_PARAMETER_CLOSE: &str = "<｜DSML｜/parameter>";
106
107	const START_MARKERS: &[&str] = &[
108	XML_OPEN_TAG,
109	NAMESPACED_OPEN_TAG,
110	K3_TOOLS_OPEN,
111	K3_RESPONSE_OPEN,
112	K3_THINK_OPEN,
113	DSML_TOOL_CALLS_OPEN,
114	DSML_FUNCTION_CALLS_OPEN,
115	DSML_CALLS_OPEN,
116	];
117
118	/// One tool call recovered from inline text.
119	#[derive(Debug, Clone, PartialEq, Eq)]
120	pub struct Recovered {
121	pub name: String,
122	/// Arguments as a JSON-encoded string, matching `ToolCall::arguments`.
123	pub arguments: String,
124	}
125
126	/// What [`extract`] found in a complete text blob.
127	#[derive(Debug, Default, PartialEq, Eq)]
128	pub struct Extracted {
129	/// The text with every tool-call block removed, recovered or not.
130	pub text: String,
131	/// Calls recovered from well-formed blocks, in order.
132	pub calls: Vec<Recovered>,
133	/// Tool-call blocks that were dropped because they couldn't be parsed.
134	pub malformed: usize,
135	}
136
137	/// Scan a complete text blob for inline tool-call blocks. Recovered blocks
138	/// become calls; blocks that don't parse cleanly, including an unterminated
139	/// trailing one, are dropped and counted in [`Extracted::malformed`].
140	pub fn extract(text: &str, tools: &[ToolSpec]) -> Extracted {
141	let mut scanner = StreamScanner::new(tools);
142	let mut extracted = Extracted {
143	text: String::with_capacity(text.len()),
144	..Extracted::default()
145	};
146
147	for event in scanner.push(text).into_iter().chain(scanner.finish()) {
148	match event {
149	ScanEvent::Text(text) => extracted.text.push_str(&text),
150	ScanEvent::ToolCall(call) => extracted.calls.push(call),
151	ScanEvent::Malformed(_) => extracted.malformed += 1,
152	}
153	}
154
155	extracted
156	}
157
158	/// Parse one legacy XML block's inner text: an offered tool name, then zero or
159	/// more `<arg_key>K</arg_key><arg_value>V</arg_value>` pairs. The malformed
160	/// opening marker observed from GLM is accepted only for the first argument;
161	/// later arguments stay strict so text that merely resembles a call cannot be
162	/// recovered as one.
163	fn parse_xml_block(inner: &str, tools: &[ToolSpec]) -> Option<Recovered> {
164	let key_idx = inner.find("<arg_key>");
165	let malformed_key_idx = inner.find(XML_OPEN_TAG);
166	let check_status_name = key_idx.and_then(\|idx\| parse_check_status_preamble(&inner[..idx]));
167	let (name, mut rest, malformed_first_key) =
168	if let (Some(idx), Some(name)) = (key_idx, check_status_name) {
169	(name, &inner[idx..], false)
170	} else {
171	match (key_idx, malformed_key_idx) {
172	(Some(key_idx), Some(malformed_key_idx)) if malformed_key_idx < key_idx => (
173	inner[..malformed_key_idx].trim(),
174	&inner[malformed_key_idx..],
175	true,
176	),
177	(Some(idx), _) => (inner[..idx].trim(), &inner[idx..], false),
178	(None, Some(idx)) => (inner[..idx].trim(), &inner[idx..], true),
179	(None, None) => (inner.trim(), "", false),
180	}
181	};
182	if name.is_empty() \|\| name.contains(['<', '>']) \|\| !offered_tool(tools, name) {
183	return None;
184	}
185
186	let mut args = serde_json::Map::new();
187	let mut first_key = true;
188	while !rest.is_empty() {
189	// The GLM alias is permitted exactly once, at the first key. A later
190	// `<tool_call>` must reject the complete block rather than letting a
191	// partial call escape recovery.
192	if !first_key && rest.starts_with(XML_OPEN_TAG) {
193	return None;
194	}
195	let key_open = if first_key && malformed_first_key {
196	XML_OPEN_TAG
197	} else {
198	"<arg_key>"
199	};
200	let (key, after_key) = split_once(rest.strip_prefix(key_open)?, "</arg_key>")?;
201	let (value, after_value) =
202	split_once(after_key.strip_prefix("<arg_value>")?, "</arg_value>")?;
203	let key = key.trim();
204	if key.is_empty() {
205	return None;
206	}
207	args.insert(
208	key.to_string(),
209	coerce(value, declared_type(tools, name, key).as_deref()),
210	);
211	rest = after_value;
212	first_key = false;
213	}
214
215	Some(Recovered {
216	name: name.to_string(),
217	arguments: serde_json::Value::Object(args).to_string(),
218	})
219	}
220
221	/// Some GLM-family responses prefix the real arguments with generation-control
222	/// metadata shaped like `tool_name CheckStatus=value</arg_value>`. It is not a
223	/// tool argument: accept that exact bounded preamble and discard it, while
224	/// leaving every other malformed prefix untouched.
225	fn parse_check_status_preamble(prefix: &str) -> Option<&str> {
226	let (name, status) = prefix.trim().split_once(" CheckStatus=")?;
227	let status = status.strip_suffix("</arg_value>")?;
228	if name.is_empty()
229	\|\| status.is_empty()
230	\|\| !status
231	.chars()
232	.all(\|character\| character.is_ascii_alphanumeric() \|\| matches!(character, '_' \| '-'))
233	{
234	return None;
235	}
236	Some(name)
237	}
238
239	fn parse_k3_tools_block(inner: &str, tools: &[ToolSpec]) -> Option<Vec<Recovered>> {
240	let mut calls = Vec::new();
241	let mut rest = inner;
242
243	while !rest.trim().is_empty() {
244	rest = rest.trim_start();
245	let header = rest.strip_prefix(K3_CALL_OPEN)?;
246	let (attr_text, after_header) = split_once(header, K3_SEP)?;
247	let attrs = parse_k3_attributes(attr_text)?;
248	let name = attrs.get("tool")?;
249	if name.is_empty() \|\| !offered_tool(tools, name) {
250	return None;
251	}
252
253	let (call_inner, after_call) = split_once(after_header, K3_CALL_CLOSE)?;
254	let arguments = parse_k3_arguments(call_inner, tools, name)?;
255	calls.push(Recovered {
256	name: name.to_string(),
257	arguments: serde_json::Value::Object(arguments).to_string(),
258	});
259	rest = after_call;
260	}
261
262	Some(calls)
263	}
264
265	fn parse_k3_arguments(
266	mut rest: &str,
267	tools: &[ToolSpec],
268	tool_name: &str,
269	) -> Option<serde_json::Map<String, serde_json::Value>> {
270	let mut args = serde_json::Map::new();
271
272	while !rest.trim().is_empty() {
273	rest = rest.trim_start();
274	let header = rest.strip_prefix(K3_ARGUMENT_OPEN)?;
275	let (attr_text, after_header) = split_once(header, K3_SEP)?;
276	let attrs = parse_k3_attributes(attr_text)?;
277	let key = attrs.get("key")?.trim();
278	if key.is_empty() {
279	return None;
280	}
281	// K3 normally emits `type` for every argument. If an endpoint omits
282	// it, retain the legacy recovery path's schema-based typing instead of
283	// turning e.g. an integer `task_id` into a JSON string.
284	let schema_type = declared_type(tools, tool_name, key);
285	let value_type = attrs
286	.get("type")
287	.map(String::as_str)
288	.or(schema_type.as_deref())
289	.unwrap_or("string");
290
291	let (raw_value, after_argument) = split_once(after_header, K3_ARGUMENT_CLOSE)?;
292	let value = if value_type == "string" {
293	serde_json::Value::String(raw_value.to_string())
294	} else {
295	serde_json::from_str(raw_value).ok()?
296	};
297	args.insert(key.to_string(), value);
298	rest = after_argument;
299	}
300
301	Some(args)
302	}
303
304	fn parse_k3_attributes(mut text: &str) -> Option<HashMap<String, String>> {
305	let mut attrs = HashMap::new();
306	text = text.trim();
307
308	while !text.is_empty() {
309	let eq_idx = text.find('=')?;
310	let key = text[..eq_idx].trim();
311	if key.is_empty()
312	\|\| !key
313	.chars()
314	.all(\|ch\| ch.is_ascii_alphanumeric() \|\| ch == '_' \|\| ch == '-')
315	{
316	return None;
317	}
318
319	let mut value_part = text[eq_idx + 1..].trim_start();
320	value_part = value_part.strip_prefix('"')?;
321	let end_idx = value_part.find('"')?;
322	let raw_value = &value_part[..end_idx];
323	attrs.insert(key.to_string(), decode_k3_attribute(raw_value));
324
325	text = value_part[end_idx + 1..].trim_start();
326	}
327
328	Some(attrs)
329	}
330
331	fn decode_k3_attribute(value: &str) -> String {
332	// K3's encoder escapes only `&` and `"` in attributes. These replacements
333	// deliberately mirror that format rather than implementing general HTML
334	// entity decoding (which could alter literal argument names).
335	value.replace(""", "\"").replace("&", "&")
336	}
337
338	/// Parse a DSML `<｜DSML｜tool_calls>` envelope's inner text: zero or
339	/// more `<｜DSML｜invoke name="$TOOL">...</｜DSML｜invoke>` blocks, each
340	/// holding `<｜DSML｜parameter name="$KEY" string="true\|false">...` entries.
341	fn parse_dsml_block(inner: &str, tools: &[ToolSpec]) -> Option<Vec<Recovered>> {
342	let mut calls = Vec::new();
343	let mut rest = inner;
344
345	while !rest.trim().is_empty() {
346	rest = rest.trim_start();
347	let header = rest.strip_prefix(DSML_INVOKE_OPEN)?;
348	let (attr_text, after_header) = split_once(header, ">")?;
349	let attrs = parse_k3_attributes(attr_text)?;
350	let name = attrs.get("name")?;
351	if name.is_empty() \|\| !offered_tool(tools, name) {
352	return None;
353	}
354
355	let (call_inner, after_call) = split_once(after_header, DSML_INVOKE_CLOSE)?;
356	let arguments = parse_dsml_parameters(call_inner, tools, name)?;
357	calls.push(Recovered {
358	name: name.to_string(),
359	arguments: serde_json::Value::Object(arguments).to_string(),
360	});
361	rest = after_call;
362	}
363
364	Some(calls)
365	}
366
367	fn parse_dsml_parameters(
368	mut rest: &str,
369	tools: &[ToolSpec],
370	tool_name: &str,
371	) -> Option<serde_json::Map<String, serde_json::Value>> {
372	let mut args = serde_json::Map::new();
373
374	while !rest.trim().is_empty() {
375	rest = rest.trim_start();
376	let header = rest.strip_prefix(DSML_PARAMETER_OPEN)?;
377	let (attr_text, after_header) = split_once(header, ">")?;
378	let attrs = parse_k3_attributes(attr_text)?;
379	let key = attrs.get("name")?.trim();
380	if key.is_empty() {
381	return None;
382	}
383	// DSML marks each parameter's encoding with `string="true\|false"`. If
384	// an endpoint omits it, fall back to the schema-based typing the
385	// other recovery paths use instead of mistyping the value.
386	let is_string = match attrs.get("string").map(String::as_str) {
387	Some("true") => true,
388	Some("false") => false,
389	_ => {
390	declared_type(tools, tool_name, key)
391	.as_deref()
392	.unwrap_or("string")
393	== "string"
394	}
395	};
396
397	let (raw_value, after_param) = split_once(after_header, DSML_PARAMETER_CLOSE)?;
398	let value = if is_string {
399	serde_json::Value::String(raw_value.to_string())
400	} else {
401	serde_json::from_str(raw_value).ok()?
402	};
403	args.insert(key.to_string(), value);
404	rest = after_param;
405	}
406
407	Some(args)
408	}
409
410	/// Parse an `<invoke>`-shaped `<tool_call>` body of one or more
411	/// `<invoke name="$TOOL">…</invoke>` blocks. A parameter is either
412	/// `<parameter name="k">v</parameter>` or `<k>v</k>`.
413	fn parse_invoke_block(inner: &str, tools: &[ToolSpec]) -> Option<Vec<Recovered>> {
414	let mut calls = Vec::new();
415	let mut rest = inner;
416
417	while !rest.trim().is_empty() {
418	rest = rest.trim_start();
419	let header = rest.strip_prefix(INVOKE_OPEN)?;
420	let (attr_text, after_header) = split_once(header, ">")?;
421	let attrs = parse_k3_attributes(attr_text)?;
422	let name = attrs.get("name")?;
423	if name.is_empty() \|\| !offered_tool(tools, name) {
424	return None;
425	}
426
427	let (call_inner, after_call) = split_once(after_header, INVOKE_CLOSE)?;
428	let arguments = parse_invoke_parameters(call_inner, tools, name)?;
429	calls.push(Recovered {
430	name: name.to_string(),
431	arguments: serde_json::Value::Object(arguments).to_string(),
432	});
433	rest = after_call;
434	}
435
436	if calls.is_empty() { None } else { Some(calls) }
437	}
438
439	fn parse_invoke_parameters(
440	mut rest: &str,
441	tools: &[ToolSpec],
442	tool_name: &str,
443	) -> Option<serde_json::Map<String, serde_json::Value>> {
444	let mut args = serde_json::Map::new();
445
446	while !rest.trim().is_empty() {
447	rest = rest.trim_start();
448	let (key, raw_value, after) = if let Some(header) = rest.strip_prefix(PARAMETER_OPEN) {
449	let (attr_text, after_header) = split_once(header, ">")?;
450	let attrs = parse_k3_attributes(attr_text)?;
451	let key = attrs.get("name")?.trim().to_string();
452	let (raw_value, after) = split_once(after_header, PARAMETER_CLOSE)?;
453	(key, raw_value, after)
454	} else {
455	let header = rest.strip_prefix('<')?;
456	let (key, after_header) = split_once(header, ">")?;
457	if !key
458	.chars()
459	.all(\|ch\| ch.is_ascii_alphanumeric() \|\| ch == '_' \|\| ch == '-')
460	{
461	return None;
462	}
463	let (raw_value, after) = split_once(after_header, &format!("</{key}>"))?;
464	(key.to_string(), raw_value, after)
465	};
466	if key.is_empty() {
467	return None;
468	}
469
470	let declared = declared_type(tools, tool_name, &key);
471	let value = match declared.as_deref() {
472	Some("array" \| "object") => serde_json::from_str(raw_value).ok()?,
473	other => coerce(raw_value, other),
474	};
475	args.insert(key, value);
476	rest = after;
477	}
478
479	Some(args)
480	}
481
482	fn split_once<'a>(s: &'a str, delim: &str) -> Option<(&'a str, &'a str)> {
483	let idx = s.find(delim)?;
484	Some((&s[..idx], &s[idx + delim.len()..]))
485	}
486
487	fn offered_tool(tools: &[ToolSpec], tool_name: &str) -> bool {
488	tools.iter().any(\|spec\| spec.name == tool_name)
489	}
490
491	/// The JSON Schema `type` declared for `tool_name`'s `key` parameter, e.g.
492	/// `"integer"` for `command_output`'s `task_id`. [`ToolSpec`] carries the
493	/// schema as an unparsed string, so this parses it per lookup; the argument
494	/// counts involved are tiny and this only runs on the recovery path.
495	fn declared_type(tools: &[ToolSpec], tool_name: &str, key: &str) -> Option<String> {
496	let spec = tools.iter().find(\|spec\| spec.name == tool_name)?;
497	let schema: serde_json::Value = serde_json::from_str(&spec.parameters_schema).ok()?;
498	Some(
499	schema
500	.get("properties")?
501	.get(key)?
502	.get("type")?
503	.as_str()?
504	.to_string(),
505	)
506	}
507
508	/// Convert a raw legacy text value to JSON per its declared schema type.
509	/// Falls back to a plain string for `"string"`, an unknown type, or a value
510	/// that doesn't parse as what it claims to be.
511	fn coerce(raw: &str, declared: Option<&str>) -> serde_json::Value {
512	match declared {
513	Some("integer") => raw
514	.trim()
515	.parse::<i64>()
516	.map(serde_json::Value::from)
517	.unwrap_or_else(\|_\| serde_json::Value::String(raw.to_string())),
518	Some("number") => raw
519	.trim()
520	.parse::<f64>()
521	.ok()
522	.and_then(serde_json::Number::from_f64)
523	.map(serde_json::Value::Number)
524	.unwrap_or_else(\|\| serde_json::Value::String(raw.to_string())),
525	Some("boolean") => match raw.trim() {
526	"true" => serde_json::Value::Bool(true),
527	"false" => serde_json::Value::Bool(false),
528	_ => serde_json::Value::String(raw.to_string()),
529	},
530	_ => serde_json::Value::String(raw.to_string()),
531	}
532	}
533
534	/// One event from incrementally scanning a stream of content deltas.
535	#[derive(Debug, Clone, PartialEq, Eq)]
536	pub enum ScanEvent {
537	/// Plain text, safe to forward to the UI.
538	Text(String),
539	/// A tool call recovered from a completed block.
540	ToolCall(Recovered),
541	/// A tool-call block that couldn't be parsed, or never closed. Carries the
542	/// raw block for logging; it must not be shown or executed.
543	Malformed(String),
544	}
545
546	/// Incremental scanner for a live delta stream.
547	///
548	/// Holds back only the text that could still turn into a known inline protocol
549	/// marker, so ordinary answers still reach the UI token by token. Buffering
550	/// starts once an opening marker appears, and lasts only until the matching
551	/// close arrives.
552	pub struct StreamScanner<'a> {
553	tools: &'a [ToolSpec],
554	pending: String,
555	}
556
557	impl<'a> StreamScanner<'a> {
558	pub fn new(tools: &'a [ToolSpec]) -> Self {
559	Self {
560	tools,
561	pending: String::new(),
562	}
563	}
564
565	/// Feed the next content delta. Returns events to emit now, in order; may
566	/// be empty if the chunk was absorbed into a tag that hasn't closed yet.
567	pub fn push(&mut self, chunk: &str) -> Vec<ScanEvent> {
568	self.pending.push_str(chunk);
569	// Like the bar fold, this runs on the whole buffer so a sentinel split
570	// across deltas is seen joined; `drain` holds back a half sentinel.
571	if self.pending.contains(LEAKED_SENTINEL) {
572	self.pending = self.pending.replace(LEAKED_SENTINEL, "");
573	}
574	// The fold runs on the whole buffer, not the chunk: a doubled bar split
575	// across two deltas is only visible once they are joined, and a marker
576	// is never emitted before then because a partial prefix is held back.
577	// Both the single- and doubled-bar spellings are partial prefixes
578	// (START_MARKERS is single-bar), so a doubled-bar prefix is flushed as
579	// text on this push; the fold then rewrites the buffer and the next
580	// push re-runs `drain` against canonical markers.
581	if let Some(folded) = fold_dsml_bars(&self.pending) {
582	self.pending = folded;
583	}
584	if let Some(folded) = fold_dsml_closers(&self.pending) {
585	self.pending = folded;
586	}
587	self.drain()
588	}
589
590	/// Call once the stream has ended to flush whatever is still held back.
591	/// A tool-call block that opened but never closed is [`ScanEvent::Malformed`]:
592	/// it can't be executed safely, and it isn't prose either. Anything else,
593	/// such as a partial marker prefix or an unclosed Kimi response wrapper, is
594	/// handed back as text.
595	pub fn finish(&mut self) -> Option<ScanEvent> {
596	if self.pending.is_empty() {
597	return None;
598	}
599	let rest = std::mem::take(&mut self.pending);
600	Some(
601	if rest
602	.strip_prefix(K3_TOOLS_OPEN)
603	.is_some_and(looks_like_k3_call)
604	\|\| [XML_OPEN_TAG, NAMESPACED_OPEN_TAG]
605	.iter()
606	.any(\|open\| rest.strip_prefix(open).is_some_and(looks_like_xml_call))
607	\|\| DSML_ENVELOPES
608	.iter()
609	.any(\|(open, _)\| rest.strip_prefix(open).is_some_and(looks_like_dsml_call))
610	{
611	ScanEvent::Malformed(rest)
612	} else {
613	ScanEvent::Text(rest)
614	},
615	)
616	}
617
618	fn drain(&mut self) -> Vec<ScanEvent> {
619	let mut events = Vec::new();
620	loop {
621	let Some((open_idx, marker)) = find_next_marker(&self.pending) else {
622	let keep = max_partial_prefix_len(&self.pending, START_MARKERS)
623	.max(partial_prefix_len(&self.pending, LEAKED_SENTINEL));
624	let flush_len = self.pending.len() - keep;
625	if flush_len > 0 {
626	events.push(ScanEvent::Text(
627	self.pending.drain(..flush_len).collect::<String>(),
628	));
629	}
630	break;
631	};
632
633	if open_idx > 0 {
634	events.push(ScanEvent::Text(
635	self.pending.drain(..open_idx).collect::<String>(),
636	));
637	}
638
639	match marker {
640	XML_OPEN_TAG \| NAMESPACED_OPEN_TAG => {
641	let close = if marker == XML_OPEN_TAG {
642	XML_CLOSE_TAG
643	} else {
644	NAMESPACED_CLOSE_TAG
645	};
646	match self.drain_xml(marker, close) {
647	Some(mut recovered) => events.append(&mut recovered),
648	None => break,
649	}
650	}
651	K3_TOOLS_OPEN => match self.drain_k3_tools() {
652	Some(mut recovered) => events.append(&mut recovered),
653	None => break,
654	},
655	DSML_TOOL_CALLS_OPEN \| DSML_FUNCTION_CALLS_OPEN \| DSML_CALLS_OPEN => {
656	let close = DSML_ENVELOPES
657	.iter()
658	.find_map(\|(open, close)\| (open == marker).then_some(close))
659	.expect("every DSML open marker has a close");
660	match self.drain_dsml(marker, close) {
661	Some(mut recovered) => events.append(&mut recovered),
662	None => break,
663	}
664	}
665	K3_RESPONSE_OPEN => {
666	let Some(mut unwrapped) =
667	self.drain_wrapped_text(K3_RESPONSE_OPEN, K3_RESPONSE_CLOSE, true)
668	else {
669	break;
670	};
671	events.append(&mut unwrapped);
672	}
673	K3_THINK_OPEN => {
674	let Some(mut unwrapped) =
675	self.drain_wrapped_text(K3_THINK_OPEN, K3_THINK_CLOSE, false)
676	else {
677	break;
678	};
679	events.append(&mut unwrapped);
680	}
681	_ => unreachable!("unknown inline marker"),
682	}
683	}
684	events
685	}
686
687	fn drain_xml(&mut self, open: &str, close: &str) -> Option<Vec<ScanEvent>> {
688	let close_rel = self.pending[open.len()..].find(close)?;
689	let close_idx = open.len() + close_rel;
690	let block: String = self.pending.drain(..close_idx + close.len()).collect();
691	let inner = &block[open.len()..block.len() - close.len()];
692	if inner.trim_start().starts_with(INVOKE_OPEN) {
693	return Some(match parse_invoke_block(inner, self.tools) {
694	Some(calls) => calls.into_iter().map(ScanEvent::ToolCall).collect(),
695	None => vec![ScanEvent::Malformed(block)],
696	});
697	}
698	Some(vec![match parse_xml_block(inner, self.tools) {
699	Some(call) => ScanEvent::ToolCall(call),
700	None if looks_like_xml_call(inner) => ScanEvent::Malformed(block),
701	None => ScanEvent::Text(block),
702	}])
703	}
704
705	fn drain_k3_tools(&mut self) -> Option<Vec<ScanEvent>> {
706	let close_rel = self.pending[K3_TOOLS_OPEN.len()..].find(K3_TOOLS_CLOSE)?;
707	let close_idx = K3_TOOLS_OPEN.len() + close_rel;
708	let block: String = self
709	.pending
710	.drain(..close_idx + K3_TOOLS_CLOSE.len())
711	.collect();
712	let inner = &block[K3_TOOLS_OPEN.len()..block.len() - K3_TOOLS_CLOSE.len()];
713
714	Some(match parse_k3_tools_block(inner, self.tools) {
715	Some(calls) => calls.into_iter().map(ScanEvent::ToolCall).collect(),
716	None if looks_like_k3_call(inner) => vec![ScanEvent::Malformed(block)],
717	None => vec![ScanEvent::Text(block)],
718	})
719	}
720
721	fn drain_dsml(&mut self, open: &str, close: &str) -> Option<Vec<ScanEvent>> {
722	let close_rel = self.pending[open.len()..].find(close)?;
723	let close_idx = open.len() + close_rel;
724	let block: String = self.pending.drain(..close_idx + close.len()).collect();
725	let inner = &block[open.len()..block.len() - close.len()];
726
727	Some(match parse_dsml_block(inner, self.tools) {
728	Some(calls) => calls.into_iter().map(ScanEvent::ToolCall).collect(),
729	None if looks_like_dsml_call(inner) => vec![ScanEvent::Malformed(block)],
730	None => vec![ScanEvent::Text(block)],
731	})
732	}
733
734	fn drain_wrapped_text(
735	&mut self,
736	open: &str,
737	close: &str,
738	emit_inner: bool,
739	) -> Option<Vec<ScanEvent>> {
740	let after_open = &self.pending[open.len()..];
741	let close_rel = match after_open.find(close) {
742	Some(close_rel) => close_rel,
743	None => {
744	// A malformed response/think wrapper must not swallow a later
745	// valid tool block forever. Preserve the malformed prefix as
746	// text, leave the next marker pending, and resume scanning.
747	let (next_rel, _) = find_next_marker(after_open)?;
748	return Some(vec![ScanEvent::Text(
749	self.pending
750	.drain(..open.len() + next_rel)
751	.collect::<String>(),
752	)]);
753	}
754	};
755	let close_idx = open.len() + close_rel;
756	let block: String = self.pending.drain(..close_idx + close.len()).collect();
757
758	if emit_inner {
759	let inner = &block[open.len()..block.len() - close.len()];
760	Some(scan_complete_text(inner, self.tools))
761	} else {
762	Some(Vec::new())
763	}
764	}
765	}
766
767	/// Scan text from a wrapper whose closing marker has already arrived. This
768	/// preserves event order when a response body itself contains an inline block.
769	/// Whether the text after a `<tool_call>` tag reads as the start of a call:
770	/// a tool name straight after the tag or on the next line, a JSON object or
771	/// another tag after any whitespace, or nothing at all (a reply cut off right
772	/// after the tag). Prose that mentions the tag, like "`<tool_call>` wraps..."
773	/// or "<tool_call> and </tool_call> tags", fails every one of those.
774	fn looks_like_xml_call(after_tag: &str) -> bool {
775	let after_newlines = after_tag.trim_start_matches(['\n', '\r']);
776	if after_newlines
777	.chars()
778	.next()
779	.is_some_and(\|c\| c.is_ascii_alphanumeric() \|\| c == '_')
780	{
781	return true;
782	}
783	let rest = after_newlines.trim_start();
784	rest.is_empty() \|\| rest.starts_with(['{', '<'])
785	}
786
787	/// Whether the text after a Kimi `<\|open\|>tools<\|sep\|>` marker reads as the
788	/// start of a call: another `<\|open\|>call` marker after any whitespace, or
789	/// nothing at all. A quoted marker, like the constant in this file, fails.
790	fn looks_like_k3_call(after_marker: &str) -> bool {
791	let rest = after_marker.trim_start();
792	rest.is_empty() \|\| rest.starts_with(K3_CALL_OPEN)
793	}
794
795	/// Whether the text after a DSML `<｜DSML｜tool_calls>` marker reads as
796	/// the start of a call: a `<｜DSML｜invoke` marker after any whitespace, or
797	/// nothing at all. A quoted marker, like the constant in this file, fails.
798	fn looks_like_dsml_call(after_marker: &str) -> bool {
799	let rest = after_marker.trim_start();
800	rest.is_empty() \|\| rest.starts_with(DSML_INVOKE_OPEN)
801	}
802
803	/// Every DSML marker above carries exactly one fullwidth bar on each side of
804	/// `DSML`. Some serving stacks emit each bar twice (`<｜｜DSML｜｜invoke …>`),
805	/// which no marker matches, so the block falls through as text and renders raw
806	/// in the editor.
807	///
808	/// Folding the run on both sides of `DSML` down to a single bar, rather than
809	/// listing a doubled-bar constant next to each one, keeps the marker set at six
810	/// and covers a run of any length. A bar run elsewhere is never touched, so a
811	/// fullwidth bar used as prose or as a table rule passes through unchanged. The
812	/// fold is not conditional on the block parsing, so a doubled-bar marker that
813	/// turns out not to be a call comes back with single bars — markup either way.
814	///
815	/// Returns `None` when there is nothing to fold — in particular for any buffer
816	/// without a doubled bar, which is the common case and costs no allocation.
817	fn fold_dsml_bars(text: &str) -> Option<String> {
818	if !text.contains(DSML_TAG_NAME) {
819	return None;
820	}
821	// A DSML tag with only single bars is canonical; only the doubled-bar
822	// spelling needs rewriting, so the buffer-walking allocation can be
823	// skipped when the only defect isn't present.
824	if !text.contains(DOUBLED_BAR) {
825	return None;
826	}
827	let mut folded = String::with_capacity(text.len());
828	let mut changed = false;
829	let mut rest = text;
830	while let Some(idx) = rest.find(DSML_TAG_NAME) {
831	let (before, from_name) = rest.split_at(idx);
832	// The run of bars closing the tag being opened, i.e. `<｜｜` in
833	// `<｜｜DSML｜｜invoke`. Counted in bars, not bytes: a lone fullwidth bar
834	// is three bytes, so a byte count would read as a run of its own.
835	let opening_run = bar_run_len(before, false);
836	folded.push_str(&before[..before.len() - opening_run * FULLWIDTH_BAR.len_utf8()]);
837	if opening_run > 0 {
838	changed \|= opening_run > 1;
839	folded.push(FULLWIDTH_BAR);
840	}
841	folded.push_str(DSML_TAG_NAME);
842	let after = &from_name[DSML_TAG_NAME.len()..];
843	// The run opening the tag being closed, i.e. `｜｜>` in
844	// `<｜DSML｜｜/invoke>`.
845	let closing_run = bar_run_len(after, true);
846	if closing_run > 0 {
847	changed \|= closing_run > 1;
848	folded.push(FULLWIDTH_BAR);
849	}
850	rest = &after[closing_run * FULLWIDTH_BAR.len_utf8()..];
851	}
852	folded.push_str(rest);
853	changed.then_some(folded)
854	}
855
856	/// Number of fullwidth bars in the run at the end of `text`, or at its start
857	/// when `leading` is set. The caller passes the slice anchored to `DSML`, so
858	/// the run being counted always sits at the boundary, never in the middle.
859	fn bar_run_len(text: &str, leading: bool) -> usize {
860	let bytes = if leading {
861	text.len() - text.trim_start_matches(FULLWIDTH_BAR).len()
862	} else {
863	text.len() - text.trim_end_matches(FULLWIDTH_BAR).len()
864	};
865	bytes / FULLWIDTH_BAR.len_utf8()
866	}
867
868	/// Some stacks write a DSML closer with the slash before the bar
869	/// (`</｜DSML｜invoke>`), as the published template does, instead of after it.
870	/// Rewritten to the canonical `<｜DSML｜/invoke>` once bars are folded. Returns
871	/// `None` when nothing needs rewriting.
872	fn fold_dsml_closers(text: &str) -> Option<String> {
873	text.contains(DSML_SLASH_FIRST_CLOSER)
874	.then(\|\| text.replace(DSML_SLASH_FIRST_CLOSER, DSML_CANONICAL_CLOSER))
875	}
876
877	fn scan_complete_text(text: &str, tools: &[ToolSpec]) -> Vec<ScanEvent> {
878	let mut scanner = StreamScanner::new(tools);
879	let mut events = scanner.push(text);
880	events.extend(scanner.finish());
881	events
882	}
883
884	fn find_next_marker<'a>(text: &str) -> Option<(usize, &'a str)> {
885	START_MARKERS
886	.iter()
887	.filter_map(\|marker\| text.find(marker).map(\|idx\| (idx, *marker)))
888	.min_by_key(\|(idx, _)\| *idx)
889	}
890
891	/// Length of the longest suffix of `buf` that is a prefix of any known opening
892	/// marker. Lets a split opening tag be caught without delaying text that can't
893	/// possibly be part of one.
894	fn max_partial_prefix_len(buf: &str, needles: &[&str]) -> usize {
895	needles
896	.iter()
897	.map(\|needle\| partial_prefix_len(buf, needle))
898	.max()
899	.unwrap_or(0)
900	}
901
902	fn partial_prefix_len(buf: &str, needle: &str) -> usize {
903	let max = buf.len().min(needle.len() - 1);
904	for len in (1..=max).rev() {
905	let start = buf.len() - len;
906	if buf.is_char_boundary(start) && needle.starts_with(&buf[start..]) {
907	return len;
908	}
909	}
910	0
911	}
912
913	#[cfg(test)]
914	mod tests {
915	use super::*;
916
917	fn spec(name: &str, schema: serde_json::Value) -> ToolSpec {
918	ToolSpec {
919	name: name.to_string(),
920	description: String::new(),
921	parameters_schema: schema.to_string(),
922	}
923	}
924
925	fn command_output_spec() -> ToolSpec {
926	spec(
927	"command_output",
928	serde_json::json!({
929	"type": "object",
930	"properties": { "task_id": { "type": "integer" } }
931	}),
932	)
933	}
934
935	fn run_command_spec() -> ToolSpec {
936	spec(
937	"run_command",
938	serde_json::json!({
939	"type": "object",
940	"properties": {
941	"command": { "type": "string" },
942	"cwd": { "type": "string" },
943	"run_in_background": { "type": "boolean" }
944	}
945	}),
946	)
947	}
948
949	#[test]
950	fn an_integer_argument_is_recovered_as_a_json_number() {
951	let tools = vec![command_output_spec()];
952	let Extracted { text, calls, .. } = extract(
953	"Checking.<tool_call>command_output<arg_key>task_id</arg_key><arg_value>2</arg_value></tool_call>",
954	&tools,
955	);
956	assert_eq!(text, "Checking.");
957	assert_eq!(calls.len(), 1);
958	assert_eq!(calls[0].name, "command_output");
959	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
960	assert_eq!(args["task_id"], 2);
961	assert!(args["task_id"].is_number());
962	}
963
964	#[test]
965	fn recovers_multiple_string_arguments_in_order() {
966	let tools = vec![run_command_spec()];
967	let Extracted { text, calls, .. } = extract(
968	"<tool_call>run_command<arg_key>command</arg_key><arg_value>cargo test --locked 2>&1 \| tail -20</arg_value><arg_key>cwd</arg_key><arg_value>/Users/x/sigit</arg_value><arg_key>run_in_background</arg_key><arg_value>true</arg_value></tool_call>",
969	&tools,
970	);
971	assert_eq!(text, "");
972	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
973	assert_eq!(args["command"], "cargo test --locked 2>&1 \| tail -20");
974	assert_eq!(args["cwd"], "/Users/x/sigit");
975	assert_eq!(args["run_in_background"], true);
976	}
977
978	#[test]
979	fn recovers_glm_call_with_a_malformed_first_argument_marker() {
980	let tools = vec![run_command_spec()];
981	let Extracted { text, calls, .. } = extract(
982	"<tool_call>run_command<tool_call>command</arg_key><arg_value>pwd</arg_value><arg_key>cwd</arg_key><arg_value>/tmp</arg_value></tool_call>",
983	&tools,
984	);
985
986	assert_eq!(text, "");
987	assert_eq!(calls.len(), 1);
988	assert_eq!(calls[0].name, "run_command");
989	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
990	assert_eq!(args["command"], "pwd");
991	assert_eq!(args["cwd"], "/tmp");
992	}
993
994	#[test]
995	fn recovers_glm_call_with_check_status_before_real_arguments() {
996	let tools = vec![command_output_spec()];
997	let Extracted { text, calls, .. } = extract(
998	"polling <tool_call>command_output CheckStatus=true_or_poll_again_with_different_params</arg_value><arg_key>task_id</arg_key><arg_value>1</arg_value></tool_call>",
999	&tools,
1000	);
1001
1002	assert_eq!(text, "polling ");
1003	assert_eq!(calls.len(), 1);
1004	assert_eq!(calls[0].name, "command_output");
1005	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1006	assert_eq!(args["task_id"], 1);
1007	assert!(args["task_id"].is_number());
1008	assert!(args.get("CheckStatus").is_none());
1009	}
1010
1011	#[test]
1012	fn malformed_check_status_preamble_is_dropped() {
1013	let text = "<tool_call>command_output CheckStatus=keep going</arg_value><arg_key>task_id</arg_key><arg_value>1</arg_value></tool_call>";
1014	let Extracted {
1015	text: out,
1016	calls,
1017	malformed,
1018	} = extract(text, &[command_output_spec()]);
1019	assert_eq!(out, "");
1020	assert_eq!(malformed, 1);
1021	assert!(calls.is_empty());
1022	}
1023
1024	#[test]
1025	fn check_status_after_a_real_argument_is_dropped() {
1026	let text = "<tool_call>command_output<arg_key>task_id</arg_key><arg_value>1</arg_value> CheckStatus=true_or_poll_again_with_different_params</arg_value></tool_call>";
1027	let Extracted {
1028	text: out,
1029	calls,
1030	malformed,
1031	} = extract(text, &[command_output_spec()]);
1032	assert_eq!(out, "");
1033	assert_eq!(malformed, 1);
1034	assert!(calls.is_empty());
1035	}
1036
1037	#[test]
1038	fn check_status_for_an_unknown_tool_is_dropped() {
1039	let text = "<tool_call>not_offered CheckStatus=true_or_poll_again_with_different_params</arg_value><arg_key>task_id</arg_key><arg_value>1</arg_value></tool_call>";
1040	let Extracted {
1041	text: out,
1042	calls,
1043	malformed,
1044	} = extract(text, &[command_output_spec()]);
1045	assert_eq!(out, "");
1046	assert_eq!(malformed, 1);
1047	assert!(calls.is_empty());
1048	}
1049
1050	#[test]
1051	fn unknown_glm_tool_is_dropped() {
1052	let text = "<tool_call>not_offered<tool_call>command</arg_key><arg_value>pwd</arg_value></tool_call>";
1053	let Extracted {
1054	text: out,
1055	calls,
1056	malformed,
1057	} = extract(text, &[run_command_spec()]);
1058	assert_eq!(out, "");
1059	assert_eq!(malformed, 1);
1060	assert!(calls.is_empty());
1061	}
1062
1063	#[test]
1064	fn malformed_glm_marker_is_accepted_only_for_the_first_argument() {
1065	let text = "<tool_call>run_command<arg_key>command</arg_key><arg_value>pwd</arg_value><tool_call>cwd</arg_key><arg_value>/tmp</arg_value></tool_call>";
1066	let Extracted {
1067	text: out,
1068	calls,
1069	malformed,
1070	} = extract(text, &[run_command_spec()]);
1071	assert_eq!(out, "");
1072	assert_eq!(malformed, 1);
1073	assert!(calls.is_empty());
1074	}
1075
1076	#[test]
1077	fn a_malformed_block_is_dropped_rather_than_guessed_at() {
1078	let text = "<tool_call>edit_file<arg_value>new_text</arg_key><arg_value>def x; end</arg_value></tool_call>";
1079	let Extracted {
1080	text: out,
1081	calls,
1082	malformed,
1083	} = extract(text, &[run_command_spec()]);
1084	assert_eq!(out, "");
1085	assert_eq!(malformed, 1);
1086	assert!(calls.is_empty());
1087	}
1088
1089	#[test]
1090	fn an_unterminated_tag_is_dropped() {
1091	let text = "working on it <tool_call>run_command<arg_key>command</arg_key>";
1092	let Extracted {
1093	text: out,
1094	calls,
1095	malformed,
1096	} = extract(text, &[run_command_spec()]);
1097	assert_eq!(out, "working on it ");
1098	assert_eq!(malformed, 1);
1099	assert!(calls.is_empty());
1100	}
1101
1102	#[test]
1103	fn ordinary_text_passes_through_untouched() {
1104	let Extracted {
1105	text: out, calls, ..
1106	} = extract("if x < y then a<b, nothing to recover", &[]);
1107	assert_eq!(out, "if x < y then a<b, nothing to recover");
1108	assert!(calls.is_empty());
1109	}
1110
1111	#[test]
1112	fn kimi_k3_run_command_block_is_recovered_and_removed() {
1113	let tools = vec![run_command_spec()];
1114	let Extracted { text, calls, .. } = extract(
1115	"Working <\|open\|>tools<\|sep\|><\|open\|>call tool=\"run_command\" index=\"1\"<\|sep\|><\|open\|>argument key=\"command\" type=\"string\"<\|sep\|>sleep 180; echo waited<\|close\|>argument<\|sep\|><\|open\|>argument key=\"cwd\" type=\"string\"<\|sep\|>/Users/setoelkahfi/Repositories/onde-ed<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>",
1116	&tools,
1117	);
1118	assert_eq!(text, "Working ");
1119	assert_eq!(calls.len(), 1);
1120	assert_eq!(calls[0].name, "run_command");
1121	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1122	assert_eq!(args["command"], "sleep 180; echo waited");
1123	assert_eq!(args["cwd"], "/Users/setoelkahfi/Repositories/onde-ed");
1124	}
1125
1126	#[test]
1127	fn kimi_k3_non_string_arguments_are_json_decoded() {
1128	let tools = vec![command_output_spec()];
1129	let Extracted { calls, .. } = extract(
1130	"<\|open\|>tools<\|sep\|><\|open\|>call tool=\"command_output\" index=\"1\"<\|sep\|><\|open\|>argument key=\"task_id\" type=\"integer\"<\|sep\|>2<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>",
1131	&tools,
1132	);
1133	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1134	assert_eq!(args["task_id"], 2);
1135	assert!(args["task_id"].is_number());
1136	}
1137
1138	#[test]
1139	fn kimi_k3_missing_argument_type_uses_the_tool_schema() {
1140	let tools = vec![command_output_spec()];
1141	let Extracted { calls, .. } = extract(
1142	"<\|open\|>tools<\|sep\|><\|open\|>call tool=\"command_output\" index=\"1\"<\|sep\|><\|open\|>argument key=\"task_id\"<\|sep\|>2<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>",
1143	&tools,
1144	);
1145	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1146	assert_eq!(args["task_id"], 2);
1147	assert!(args["task_id"].is_number());
1148	}
1149
1150	#[test]
1151	fn kimi_k3_multiple_calls_are_recovered_in_order() {
1152	let tools = vec![command_output_spec()];
1153	let Extracted { calls, .. } = extract(
1154	"<\|open\|>tools<\|sep\|><\|open\|>call tool=\"command_output\" index=\"1\"<\|sep\|><\|open\|>argument key=\"task_id\" type=\"integer\"<\|sep\|>1<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|open\|>call tool=\"command_output\" index=\"2\"<\|sep\|><\|open\|>argument key=\"task_id\" type=\"integer\"<\|sep\|>2<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>",
1155	&tools,
1156	);
1157	let first: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1158	let second: serde_json::Value = serde_json::from_str(&calls[1].arguments).unwrap();
1159	assert_eq!(calls.len(), 2);
1160	assert_eq!(first["task_id"], 1);
1161	assert_eq!(second["task_id"], 2);
1162	}
1163
1164	#[test]
1165	fn kimi_k3_response_is_unwrapped_and_think_is_hidden() {
1166	let Extracted { text, calls, .. } = extract(
1167	"a<\|open\|>think<\|sep\|>private<\|close\|>think<\|sep\|>b<\|open\|>response<\|sep\|>visible<\|close\|>response<\|sep\|>c",
1168	&[],
1169	);
1170	assert_eq!(text, "abvisiblec");
1171	assert!(calls.is_empty());
1172	}
1173
1174	#[test]
1175	fn kimi_k3_response_content_is_scanned_for_nested_tool_calls() {
1176	let tools = vec![command_output_spec()];
1177	let Extracted { text, calls, .. } = extract(
1178	"<\|open\|>response<\|sep\|>before <\|open\|>think<\|sep\|>private<\|close\|>think<\|sep\|><\|open\|>tools<\|sep\|><\|open\|>call tool=\"command_output\" index=\"1\"<\|sep\|><\|open\|>argument key=\"task_id\" type=\"integer\"<\|sep\|>2<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|> after<\|close\|>response<\|sep\|>",
1179	&tools,
1180	);
1181	assert_eq!(text, "before after");
1182	assert_eq!(calls.len(), 1);
1183	assert_eq!(calls[0].name, "command_output");
1184	}
1185
1186	#[test]
1187	fn kimi_k3_attribute_entities_are_decoded() {
1188	let attrs = parse_k3_attributes(r#"tool="run"command" key="a&b""#).unwrap();
1189	assert_eq!(attrs["tool"], "run\"command");
1190	assert_eq!(attrs["key"], "a&b");
1191	assert_eq!(decode_k3_attribute("line break"), "line break");
1192	}
1193
1194	#[test]
1195	fn unterminated_kimi_k3_tools_block_is_dropped() {
1196	let text = "<\|open\|>tools<\|sep\|><\|open\|>call tool=\"run_command\" index=\"1\"<\|sep\|><\|open\|>argument key=\"command\" type=\"string\"<\|sep\|>echo hi<\|close\|>argument<\|sep\|>";
1197	let Extracted {
1198	text: out,
1199	calls,
1200	malformed,
1201	} = extract(text, &[run_command_spec()]);
1202	assert_eq!(out, "");
1203	assert_eq!(malformed, 1);
1204	assert!(calls.is_empty());
1205	}
1206
1207	#[test]
1208	fn unknown_kimi_k3_tool_block_is_dropped() {
1209	let text = "<\|open\|>tools<\|sep\|><\|open\|>call tool=\"delete_everything\" index=\"1\"<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>";
1210	let Extracted {
1211	text: out,
1212	calls,
1213	malformed,
1214	} = extract(text, &[run_command_spec()]);
1215	assert_eq!(out, "");
1216	assert_eq!(malformed, 1);
1217	assert!(calls.is_empty());
1218	}
1219
1220	#[test]
1221	fn invalid_kimi_k3_non_string_argument_is_dropped() {
1222	let text = "<\|open\|>tools<\|sep\|><\|open\|>call tool=\"command_output\" index=\"1\"<\|sep\|><\|open\|>argument key=\"task_id\" type=\"integer\"<\|sep\|>not-json<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>";
1223	let Extracted {
1224	text: out,
1225	calls,
1226	malformed,
1227	} = extract(text, &[command_output_spec()]);
1228	assert_eq!(out, "");
1229	assert_eq!(malformed, 1);
1230	assert!(calls.is_empty());
1231	}
1232
1233	#[test]
1234	fn scanner_recovers_a_legacy_tag_split_across_chunk_boundaries() {
1235	let tools = vec![command_output_spec()];
1236	let mut scanner = StreamScanner::new(&tools);
1237	let mut text = String::new();
1238	let mut calls = Vec::new();
1239
1240	for chunk in [
1241	"polling ",
1242	"<tool_c",
1243	"all>command_output<arg_",
1244	"key>task_id</arg_key><arg_value>2</arg_v",
1245	"alue></tool_call>",
1246	" done",
1247	] {
1248	for event in scanner.push(chunk) {
1249	match event {
1250	ScanEvent::Text(t) => text.push_str(&t),
1251	ScanEvent::ToolCall(c) => calls.push(c),
1252	ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1253	}
1254	}
1255	}
1256	if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1257	text.push_str(&rest);
1258	}
1259
1260	assert_eq!(text, "polling done", "no tag fragment may reach the UI");
1261	assert_eq!(calls.len(), 1);
1262	assert_eq!(calls[0].name, "command_output");
1263	}
1264
1265	#[test]
1266	fn scanner_recovers_glm_call_with_malformed_marker_split_across_chunks() {
1267	let tools = vec![run_command_spec()];
1268	let mut scanner = StreamScanner::new(&tools);
1269	let mut text = String::new();
1270	let mut calls = Vec::new();
1271
1272	for chunk in [
1273	"before ",
1274	"<tool_call>run_command<tool_",
1275	"call>command</arg_key><arg_value>pwd</arg_value></tool_call>",
1276	" after",
1277	] {
1278	for event in scanner.push(chunk) {
1279	match event {
1280	ScanEvent::Text(value) => text.push_str(&value),
1281	ScanEvent::ToolCall(call) => calls.push(call),
1282	ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1283	}
1284	}
1285	}
1286	if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1287	text.push_str(&rest);
1288	}
1289
1290	assert_eq!(text, "before after");
1291	assert_eq!(calls.len(), 1);
1292	assert_eq!(calls[0].name, "run_command");
1293	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1294	assert_eq!(args["command"], "pwd");
1295	}
1296
1297	#[test]
1298	fn scanner_recovers_check_status_call_split_across_chunks() {
1299	let tools = vec![command_output_spec()];
1300	let mut scanner = StreamScanner::new(&tools);
1301	let mut text = String::new();
1302	let mut calls = Vec::new();
1303
1304	for chunk in [
1305	"before <tool_call>command_output CheckStatus=true_or_poll_",
1306	"again_with_different_params</arg_",
1307	"value><arg_key>task_id</arg_key><arg_value>1</arg_value></tool_",
1308	"call> after",
1309	] {
1310	for event in scanner.push(chunk) {
1311	match event {
1312	ScanEvent::Text(value) => text.push_str(&value),
1313	ScanEvent::ToolCall(call) => calls.push(call),
1314	ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1315	}
1316	}
1317	}
1318	if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1319	text.push_str(&rest);
1320	}
1321
1322	assert_eq!(text, "before after");
1323	assert_eq!(calls.len(), 1);
1324	assert_eq!(calls[0].name, "command_output");
1325	}
1326
1327	#[test]
1328	fn scanner_recovers_kimi_k3_tools_split_across_chunk_boundaries() {
1329	let tools = vec![run_command_spec()];
1330	let mut scanner = StreamScanner::new(&tools);
1331	let mut text = String::new();
1332	let mut calls = Vec::new();
1333
1334	for chunk in [
1335	"pre ",
1336	"<\|open\|>too",
1337	"ls<\|sep\|><\|open\|>call tool=\"run_command\" index=\"1\"<\|sep\|>",
1338	"<\|open\|>argument key=\"command\" type=\"string\"<\|sep\|>echo hi",
1339	"<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>",
1340	" post",
1341	] {
1342	for event in scanner.push(chunk) {
1343	match event {
1344	ScanEvent::Text(t) => text.push_str(&t),
1345	ScanEvent::ToolCall(c) => calls.push(c),
1346	ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1347	}
1348	}
1349	}
1350	if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1351	text.push_str(&rest);
1352	}
1353
1354	assert_eq!(text, "pre post");
1355	assert_eq!(calls.len(), 1);
1356	assert_eq!(calls[0].name, "run_command");
1357	}
1358
1359	#[test]
1360	fn scanner_recovers_tools_after_an_unclosed_k3_wrapper() {
1361	let tools = vec![command_output_spec()];
1362	let mut scanner = StreamScanner::new(&tools);
1363	let mut text = String::new();
1364	let mut calls = Vec::new();
1365
1366	for chunk in [
1367	"before <\|open\|>response<\|sep\|>broken ",
1368	"<\|open\|>tools<\|sep\|><\|open\|>call tool=\"command_output\" index=\"1\"<\|sep\|><\|open\|>argument key=\"task_id\" type=\"integer\"<\|sep\|>2<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|> after",
1369	] {
1370	for event in scanner.push(chunk) {
1371	match event {
1372	ScanEvent::Text(chunk) => text.push_str(&chunk),
1373	ScanEvent::ToolCall(call) => calls.push(call),
1374	ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1375	}
1376	}
1377	}
1378	if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1379	text.push_str(&rest);
1380	}
1381
1382	assert_eq!(text, "before <\|open\|>response<\|sep\|>broken after");
1383	assert_eq!(calls.len(), 1);
1384	assert_eq!(calls[0].name, "command_output");
1385	}
1386
1387	#[test]
1388	fn scanner_recovers_tools_after_an_unclosed_k3_think_wrapper() {
1389	let tools = vec![command_output_spec()];
1390	let Extracted { text, calls, .. } = extract(
1391	"<\|open\|>think<\|sep\|>broken <\|open\|>tools<\|sep\|><\|open\|>call tool=\"command_output\" index=\"1\"<\|sep\|><\|open\|>argument key=\"task_id\" type=\"integer\"<\|sep\|>2<\|close\|>argument<\|sep\|><\|close\|>call<\|sep\|><\|close\|>tools<\|sep\|>",
1392	&tools,
1393	);
1394	assert_eq!(text, "<\|open\|>think<\|sep\|>broken ");
1395	assert_eq!(calls.len(), 1);
1396	assert_eq!(calls[0].name, "command_output");
1397	}
1398
1399	#[test]
1400	fn scanner_forwards_ordinary_text_without_holding_it_back() {
1401	let mut scanner = StreamScanner::new(&[]);
1402	let events = scanner.push("a normal streamed answer");
1403	assert_eq!(
1404	events,
1405	vec![ScanEvent::Text("a normal streamed answer".to_string())]
1406	);
1407	assert!(scanner.finish().is_none());
1408	}
1409
1410	#[test]
1411	fn scanner_releases_a_near_miss_prefix_once_it_cannot_match() {
1412	let mut scanner = StreamScanner::new(&[]);
1413	assert_eq!(
1414	scanner.push("see <tool"),
1415	vec![ScanEvent::Text("see ".to_string())]
1416	);
1417	assert_eq!(
1418	scanner.push("box for details"),
1419	vec![ScanEvent::Text("<toolbox for details".to_string())]
1420	);
1421	assert!(scanner.finish().is_none());
1422	}
1423
1424	#[test]
1425	fn scanner_reports_an_unclosed_tag_at_stream_end_as_malformed() {
1426	let tools = vec![run_command_spec()];
1427	let mut scanner = StreamScanner::new(&tools);
1428	let events = scanner.push("partial <tool_call>run_command<arg_key>cmd");
1429	assert_eq!(events, vec![ScanEvent::Text("partial ".to_string())]);
1430	assert_eq!(
1431	scanner.finish(),
1432	Some(ScanEvent::Malformed(
1433	"<tool_call>run_command<arg_key>cmd".to_string()
1434	))
1435	);
1436	}
1437
1438	#[test]
1439	fn scanner_hands_back_a_partial_marker_prefix_at_stream_end_as_text() {
1440	let mut scanner = StreamScanner::new(&[]);
1441	assert_eq!(
1442	scanner.push("see <tool_"),
1443	vec![ScanEvent::Text("see ".to_string())]
1444	);
1445	assert_eq!(
1446	scanner.finish(),
1447	Some(ScanEvent::Text("<tool_".to_string()))
1448	);
1449
1450	let mut scanner = StreamScanner::new(&[]);
1451	assert_eq!(
1452	scanner.push("see <\|open\|>to"),
1453	vec![ScanEvent::Text("see ".to_string())]
1454	);
1455	assert_eq!(
1456	scanner.finish(),
1457	Some(ScanEvent::Text("<\|open\|>to".to_string()))
1458	);
1459	}
1460
1461	#[test]
1462	fn a_tag_mentioned_in_prose_stays_text() {
1463	let tools = vec![run_command_spec()];
1464	for text in [
1465	"Wrap it in `<tool_call>` and `</tool_call>` tags.",
1466	"The scanner holds <tool_call> and </tool_call> blocks back.",
1467	"The model opens with `<tool_call>` and never closes it.",
1468	] {
1469	assert_eq!(
1470	extract(text, &tools),
1471	Extracted {
1472	text: text.to_string(),
1473	..Extracted::default()
1474	},
1475	"{text}"
1476	);
1477	}
1478	}
1479
1480	#[test]
1481	fn a_kimi_marker_mentioned_in_prose_stays_text() {
1482	let tools = vec![run_command_spec()];
1483	for text in [
1484	"The marker is `<\|open\|>tools<\|sep\|>` and it never closes here.",
1485	"Quote `<\|open\|>tools<\|sep\|>` then `<\|close\|>tools<\|sep\|>` in docs.",
1486	] {
1487	assert_eq!(
1488	extract(text, &tools),
1489	Extracted {
1490	text: text.to_string(),
1491	..Extracted::default()
1492	},
1493	"{text}"
1494	);
1495	}
1496	}
1497
1498	#[test]
1499	fn a_json_shaped_block_is_dropped() {
1500	let text = "<tool_call>\n{\"name\": \"run_command\", \"arguments\": {}}\n</tool_call>";
1501	let Extracted {
1502	text: out,
1503	calls,
1504	malformed,
1505	} = extract(text, &[run_command_spec()]);
1506	assert_eq!(out, "");
1507	assert_eq!(malformed, 1);
1508	assert!(calls.is_empty());
1509	}
1510
1511	/// Issue #105: a GLM-family reply that opened a call, then wrote an
1512	/// invented invoke/result transcript instead of closing it. None of that
1513	/// may reach the editor, and it must not count as a recovered call.
1514	#[test]
1515	fn a_hallucinated_invoke_transcript_is_dropped() {
1516	let tools = vec![command_output_spec()];
1517	let reply = "Checking.<tool_call>command_output\n\n<invokeID>reply_A</invokeID>\n\
1518	<parameter_name>task_id</parameter_name>\n<parameter>6</parameter>\n\
1519	</invoke>\n\n<function_results>PyPI and NuGet are queued.\n\
1520	<tool_call>command_output<arg_key>task_id</arg_key><arg_value>5</task_id>\n\
1521	</invoke>\n</function_results>";
1522	let Extracted {
1523	text,
1524	calls,
1525	malformed,
1526	} = extract(reply, &tools);
1527	assert_eq!(text, "Checking.");
1528	assert!(calls.is_empty());
1529	assert_eq!(malformed, 1);
1530	}
1531
1532	/// Issue #97: the well-formed blocks from that thread are recovered as
1533	/// real calls, multi-line argument values included.
1534	#[test]
1535	fn well_formed_blocks_from_issue_97_are_recovered() {
1536	let tools = vec![
1537	run_command_spec(),
1538	spec(
1539	"edit_file",
1540	serde_json::json!({
1541	"type": "object",
1542	"properties": {
1543	"path": { "type": "string" },
1544	"old_text": { "type": "string" },
1545	"new_text": { "type": "string" }
1546	}
1547	}),
1548	),
1549	];
1550	let reply = "Now wire `body_path` into the step:<tool_call>edit_file<arg_key>new_text</arg_key><arg_value> - name: Release\n with:\n body_path: release_notes.md</arg_value><arg_key>old_text</arg_key><arg_value> - name: Release\n with:</arg_value><arg_key>path</arg_key><arg_value>/repo/.github/workflows/release-github.yml</arg_value></tool_call>";
1551	let Extracted {
1552	text,
1553	calls,
1554	malformed,
1555	} = extract(reply, &tools);
1556	assert_eq!(text, "Now wire `body_path` into the step:");
1557	assert_eq!(malformed, 0);
1558	assert_eq!(calls.len(), 1);
1559	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1560	assert_eq!(
1561	args["new_text"],
1562	" - name: Release\n with:\n body_path: release_notes.md"
1563	);
1564	assert_eq!(args["path"], "/repo/.github/workflows/release-github.yml");
1565
1566	let Extracted { calls, .. } = extract(
1567	"<tool_call>run_command<arg_key>command</arg_key><arg_value>sleep 90</arg_value><arg_key>cwd</arg_key><arg_value>/repo</arg_value></tool_call>",
1568	&tools,
1569	);
1570	assert_eq!(calls.len(), 1);
1571	assert_eq!(calls[0].name, "run_command");
1572	}
1573
1574	#[test]
1575	fn dsml_multiple_invokes_are_recovered_in_order() {
1576	let tools = vec![command_output_spec()];
1577	let Extracted { text, calls, .. } = extract(
1578	"Working <｜DSML｜tool_calls><｜DSML｜invoke name=\"command_output\"><｜DSML｜parameter name=\"task_id\" string=\"false\">1<｜DSML｜/parameter><｜DSML｜/invoke><｜DSML｜invoke name=\"command_output\"><｜DSML｜parameter name=\"task_id\" string=\"false\">2<｜DSML｜/parameter><｜DSML｜/invoke><｜DSML｜/tool_calls>",
1579	&tools,
1580	);
1581	assert_eq!(text, "Working ");
1582	assert_eq!(calls.len(), 2);
1583	let first: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1584	let second: serde_json::Value = serde_json::from_str(&calls[1].arguments).unwrap();
1585	assert_eq!(first["task_id"], 1);
1586	assert_eq!(second["task_id"], 2);
1587	}
1588
1589	#[test]
1590	fn dsml_string_true_is_kept_verbatim_and_string_false_is_json_typed() {
1591	let tools = vec![run_command_spec()];
1592	let Extracted { calls, .. } = extract(
1593	"<｜DSML｜tool_calls><｜DSML｜invoke name=\"run_command\"><｜DSML｜parameter name=\"command\" string=\"true\">echo {\"a\":1}<｜DSML｜/parameter><｜DSML｜parameter name=\"run_in_background\" string=\"false\">true<｜DSML｜/parameter><｜DSML｜/invoke><｜DSML｜/tool_calls>",
1594	&tools,
1595	);
1596	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1597	assert_eq!(args["command"], "echo {\"a\":1}");
1598	assert_eq!(args["run_in_background"], true);
1599	}
1600
1601	#[test]
1602	fn dsml_missing_string_attribute_uses_the_tool_schema() {
1603	let tools = vec![command_output_spec()];
1604	let Extracted { calls, .. } = extract(
1605	"<｜DSML｜tool_calls><｜DSML｜invoke name=\"command_output\"><｜DSML｜parameter name=\"task_id\">2<｜DSML｜/parameter><｜DSML｜/invoke><｜DSML｜/tool_calls>",
1606	&tools,
1607	);
1608	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1609	assert_eq!(args["task_id"], 2);
1610	assert!(args["task_id"].is_number());
1611	}
1612
1613	#[test]
1614	fn unterminated_dsml_tool_calls_block_is_dropped() {
1615	let text = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"run_command\"><｜DSML｜parameter name=\"command\" string=\"true\">echo hi<｜DSML｜/parameter>";
1616	let Extracted {
1617	text: out,
1618	calls,
1619	malformed,
1620	} = extract(text, &[run_command_spec()]);
1621	assert_eq!(out, "");
1622	assert_eq!(malformed, 1);
1623	assert!(calls.is_empty());
1624	}
1625
1626	#[test]
1627	fn unknown_dsml_tool_is_dropped() {
1628	let text = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"delete_everything\"><｜DSML｜/invoke><｜DSML｜/tool_calls>";
1629	let Extracted {
1630	text: out,
1631	calls,
1632	malformed,
1633	} = extract(text, &[run_command_spec()]);
1634	assert_eq!(out, "");
1635	assert_eq!(malformed, 1);
1636	assert!(calls.is_empty());
1637	}
1638
1639	#[test]
1640	fn invalid_dsml_json_typed_parameter_is_dropped() {
1641	let text = "<｜DSML｜tool_calls><｜DSML｜invoke name=\"command_output\"><｜DSML｜parameter name=\"task_id\" string=\"false\">not-json<｜DSML｜/parameter><｜DSML｜/invoke><｜DSML｜/tool_calls>";
1642	let Extracted {
1643	text: out,
1644	calls,
1645	malformed,
1646	} = extract(text, &[command_output_spec()]);
1647	assert_eq!(out, "");
1648	assert_eq!(malformed, 1);
1649	assert!(calls.is_empty());
1650	}
1651
1652	#[test]
1653	fn a_dsml_marker_mentioned_in_prose_stays_text() {
1654	let tools = vec![run_command_spec()];
1655	for text in [
1656	"The marker is `<｜DSML｜tool_calls>` and it never closes here.",
1657	"Quote `<｜DSML｜tool_calls>` then `<｜DSML｜/tool_calls>` in docs.",
1658	] {
1659	assert_eq!(
1660	extract(text, &tools),
1661	Extracted {
1662	text: text.to_string(),
1663	..Extracted::default()
1664	},
1665	"{text}"
1666	);
1667	}
1668	}
1669
1670	#[test]
1671	fn scanner_recovers_dsml_tool_calls_split_across_chunk_boundaries_including_inside_the_marker()
1672	{
1673	let tools = vec![run_command_spec()];
1674	let mut scanner = StreamScanner::new(&tools);
1675	let mut text = String::new();
1676	let mut calls = Vec::new();
1677
1678	for chunk in [
1679	"pre ",
1680	"<｜DSML",
1681	"｜tool_calls><｜DSML｜invoke name=\"run_command\"><｜DSML｜paramet",
1682	"er name=\"command\" string=\"true\">echo hi<｜DSML｜/parameter><｜DSML｜/invoke><｜DSML｜/tool_calls>",
1683	" post",
1684	] {
1685	for event in scanner.push(chunk) {
1686	match event {
1687	ScanEvent::Text(t) => text.push_str(&t),
1688	ScanEvent::ToolCall(c) => calls.push(c),
1689	ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1690	}
1691	}
1692	}
1693	if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1694	text.push_str(&rest);
1695	}
1696
1697	assert_eq!(text, "pre post");
1698	assert_eq!(calls.len(), 1);
1699	assert_eq!(calls[0].name, "run_command");
1700	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1701	assert_eq!(args["command"], "echo hi");
1702	}
1703
1704	#[test]
1705	fn scanner_reports_an_unclosed_dsml_tag_at_stream_end_as_malformed() {
1706	let tools = vec![run_command_spec()];
1707	let mut scanner = StreamScanner::new(&tools);
1708	let events =
1709	scanner.push("partial <｜DSML｜tool_calls><｜DSML｜invoke name=\"run_command\">");
1710	assert_eq!(events, vec![ScanEvent::Text("partial ".to_string())]);
1711	assert_eq!(
1712	scanner.finish(),
1713	Some(ScanEvent::Malformed(
1714	"<｜DSML｜tool_calls><｜DSML｜invoke name=\"run_command\">".to_string()
1715	))
1716	);
1717	}
1718
1719	#[test]
1720	fn fold_dsml_bars_collapses_runs_glued_to_dsml_only() {
1721	assert_eq!(
1722	fold_dsml_bars("<｜｜DSML｜｜invoke name=\"run_command\">"),
1723	Some("<｜DSML｜invoke name=\"run_command\">".to_string())
1724	);
1725	assert_eq!(
1726	fold_dsml_bars("<｜｜｜｜DSML｜/tool_calls>"),
1727	Some("<｜DSML｜/tool_calls>".to_string())
1728	);
1729	// Only the bar run adjacent to DSML folds; a doubled bar further along
1730	// in the tag name or body is not part of the marker grammar.
1731	assert_eq!(
1732	fold_dsml_bars("<｜DSML｜｜param｜｜eter"),
1733	Some("<｜DSML｜param｜｜eter".to_string())
1734	);
1735	// A bar run that isn't a marker, and single-bar markup, are left alone.
1736	assert_eq!(fold_dsml_bars("a ｜｜ b <｜DSML｜invoke>"), None);
1737	assert_eq!(fold_dsml_bars("<｜DSML｜invoke>"), None);
1738	assert_eq!(fold_dsml_bars("plain text"), None);
1739	}
1740
1741	#[test]
1742	fn a_doubled_bar_dsml_call_is_recovered() {
1743	let tools = vec![run_command_spec()];
1744	let Extracted {
1745	text,
1746	calls,
1747	malformed,
1748	} = extract(
1749	"before <｜｜DSML｜｜tool_calls><｜｜DSML｜｜invoke name=\"run_command\">\
1750	<｜｜DSML｜｜parameter name=\"command\" string=\"true\">echo hi\
1751	<｜｜DSML｜｜/parameter><｜｜DSML｜｜/invoke><｜｜DSML｜｜/tool_calls> after",
1752	&tools,
1753	);
1754
1755	assert_eq!(text, "before after");
1756	assert_eq!(malformed, 0);
1757	assert_eq!(calls.len(), 1);
1758	assert_eq!(calls[0].name, "run_command");
1759	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1760	assert_eq!(args["command"], "echo hi");
1761	}
1762
1763	#[test]
1764	fn a_doubled_bar_split_across_pushes_is_recovered() {
1765	let tools = vec![run_command_spec()];
1766	let mut scanner = StreamScanner::new(&tools);
1767	let mut text = String::new();
1768	let mut calls = Vec::new();
1769
1770	for chunk in [
1771	"pre <｜",
1772	"｜DSML｜｜tool_calls><｜｜DSML｜｜invoke name=\"run_com",
1773	"mand\"><｜｜DSML｜｜parameter name=\"command\" string=\"true\">echo hi\
1774	<｜｜DSML｜｜/parameter><｜｜DSML｜｜/invoke><｜｜DSML｜｜/tool_calls>",
1775	" post",
1776	] {
1777	for event in scanner.push(chunk) {
1778	match event {
1779	ScanEvent::Text(t) => text.push_str(&t),
1780	ScanEvent::ToolCall(call) => calls.push(call),
1781	ScanEvent::Malformed(block) => panic!("unexpected malformed block: {block}"),
1782	}
1783	}
1784	}
1785	if let Some(ScanEvent::Text(rest)) = scanner.finish() {
1786	text.push_str(&rest);
1787	}
1788
1789	assert_eq!(text, "pre post");
1790	assert_eq!(calls.len(), 1);
1791	assert_eq!(calls[0].name, "run_command");
1792	let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
1793	assert_eq!(args["command"], "echo hi");
1794	}
1795
1796	#[test]
1797	fn doubled_bar_markup_outside_a_call_returns_folded_as_text() {
1798	let tools = vec![run_command_spec()];
1799	let Extracted {
1800	text,
1801	calls,
1802	malformed,
1803	} = extract("the tag <｜｜DSML｜｜done> is not a call", &tools);
1804
1805	assert_eq!(text, "the tag <｜DSML｜done> is not a call");
1806	assert!(calls.is_empty());
1807	assert_eq!(malformed, 0);
1808	}
1809
1810	const ISSUE_246: &str = "Now let me run the lint/typecheck:]<]minimax[>[<tool_call>\n\
1811	]<]minimax[>[<invoke name=\"run_command\">]<]minimax[>[<command>cd app && npm run lint \| tail -60]<]minimax[>[</command>\
1812	]<]minimax[>[<timeout>120000]<]minimax[>[</timeout>]<]minimax[>[</invoke>\n]<]minimax[>[";
1813
1814	fn args_of(call: &Recovered) -> serde_json::Value {
1815	serde_json::from_str(&call.arguments).unwrap()
1816	}
1817
1818	fn feed(tools: &[ToolSpec], chunks: &[&str]) -> (String, Vec<Recovered>, usize) {
1819	let mut scanner = StreamScanner::new(tools);
1820	let (mut text, mut calls, mut malformed) = (String::new(), Vec::new(), 0);
1821	let mut handle = \|event\| match event {
1822	ScanEvent::Text(t) => text.push_str(&t),
1823	ScanEvent::ToolCall(c) => calls.push(c),
1824	ScanEvent::Malformed(_) => malformed += 1,
1825	};
1826	for chunk in chunks {
1827	scanner.push(chunk).into_iter().for_each(&mut handle);
1828	}
1829	scanner.finish().into_iter().for_each(&mut handle);
1830	(text, calls, malformed)
1831	}
1832
1833	#[test]
1834	fn an_unclosed_invoke_call_with_sentinels_is_dropped() {
1835	let tools = vec![run_command_spec()];
1836	let extracted = extract(ISSUE_246, &tools);
1837	assert_eq!(extracted.text, "Now let me run the lint/typecheck:");
1838	assert_eq!(extracted.malformed, 1);
1839	assert!(extracted.calls.is_empty());
1840	assert!(!extracted.text.contains("minimax"));
1841	}
1842
1843	#[test]
1844	fn a_closed_invoke_call_with_sentinels_is_recovered() {
1845	let tools = vec![run_command_spec()];
1846	let text = format!("{ISSUE_246}</tool_call>");
1847	let extracted = extract(&text, &tools);
1848	assert_eq!(extracted.malformed, 0);
1849	assert_eq!(extracted.calls.len(), 1);
1850	assert_eq!(extracted.calls[0].name, "run_command");
1851	assert_eq!(
1852	args_of(&extracted.calls[0])["command"],
1853	"cd app && npm run lint \| tail -60"
1854	);
1855	}
1856
1857	#[test]
1858	fn a_namespaced_envelope_is_recovered_with_typed_parameters() {
1859	let tools = vec![run_command_spec(), command_output_spec()];
1860	let extracted = extract(
1861	"<minimax:tool_call><invoke name=\"run_command\"><parameter name=\"command\">echo hi</parameter></invoke>\
1862	<invoke name=\"command_output\"><parameter name=\"task_id\">3</parameter></invoke></minimax:tool_call>",
1863	&tools,
1864	);
1865	assert_eq!(extracted.malformed, 0);
1866	assert_eq!(extracted.calls.len(), 2);
1867	assert_eq!(args_of(&extracted.calls[0])["command"], "echo hi");
1868	assert_eq!(args_of(&extracted.calls[1])["task_id"], 3);
1869	}
1870
1871	#[test]
1872	fn a_sentinel_split_across_pushes_is_removed() {
1873	let tools = vec![run_command_spec()];
1874	let (text, calls, malformed) = feed(&tools, &["a]<]mini", "max[>[b"]);
1875	assert_eq!(text, "ab");
1876	assert!(calls.is_empty());
1877	assert_eq!(malformed, 0);
1878	}
1879
1880	#[test]
1881	fn an_unknown_tool_in_an_invoke_block_is_malformed() {
1882	let tools = vec![run_command_spec()];
1883	let extracted = extract(
1884	"<tool_call><invoke name=\"delete_everything\"><x>1</x></invoke></tool_call>",
1885	&tools,
1886	);
1887	assert_eq!(extracted.malformed, 1);
1888	assert!(extracted.calls.is_empty());
1889	assert_eq!(extracted.text, "");
1890	}
1891
1892	const ISSUE_236_BODY: &str = "<｜｜DSML｜｜invoke name=\"run_command\">\
1893	<｜｜DSML｜｜parameter name=\"command\" string=\"true\">echo hi</｜｜DSML｜｜parameter>\
1894	</｜｜DSML｜｜invoke>";
1895
1896	#[test]
1897	fn doubled_bar_dsml_with_slash_first_closers_is_recovered_for_every_envelope() {
1898	let tools = vec![run_command_spec()];
1899	for name in ["function_calls", "calls", "tool_calls"] {
1900	let text =
1901	format!("before <｜｜DSML｜｜{name}>{ISSUE_236_BODY}</｜｜DSML｜｜{name}> after");
1902	let extracted = extract(&text, &tools);
1903	assert_eq!(extracted.text, "before after", "{name}");
1904	assert_eq!(extracted.malformed, 0, "{name}");
1905	assert_eq!(extracted.calls.len(), 1, "{name}");
1906	assert_eq!(args_of(&extracted.calls[0])["command"], "echo hi");
1907	}
1908	}
1909
1910	#[test]
1911	fn a_dsml_function_calls_open_marker_split_mid_name_is_recovered() {
1912	let tools = vec![run_command_spec()];
1913	let (text, calls, malformed) = feed(
1914	&tools,
1915	&[
1916	"pre <｜DSML｜func",
1917	"tion_calls><｜DSML｜invoke name=\"run_command\"><｜DSML｜parameter name=\"command\" string=\"true\">echo hi</｜DSML｜parameter></｜DSML｜invoke></｜DSML｜function_calls>",
1918	" post",
1919	],
1920	);
1921	assert_eq!(text, "pre post");
1922	assert_eq!(malformed, 0);
1923	assert_eq!(calls.len(), 1);
1924	}
1925
1926	#[test]
1927	fn a_dsml_function_calls_marker_mentioned_in_prose_stays_text() {
1928	let tools = vec![run_command_spec()];
1929	for text in [
1930	"The marker is `<｜DSML｜function_calls>` and it never closes here.",
1931	"Quote `<｜DSML｜function_calls>` then `<｜DSML｜/function_calls>` in docs.",
1932	] {
1933	assert_eq!(
1934	extract(text, &tools),
1935	Extracted {
1936	text: text.to_string(),
1937	..Extracted::default()
1938	},
1939	"{text}"
1940	);
1941	}
1942	}
1943	}