inline_tool_calls.rs · sigitcodereview/sigit

Code Commits Issues Pulls Cloud Agent CI/CD

fix/dsml-doubled-bar-markers

claude/github-issue-40-0au5ua claude/linked-issue-acp-tui-u4vvpa claude/minor-release-cut-vkugow claude/sigit-acp-local-chat-cx6380 claude/sigit-cloud-agent-expansion-reox0a claude/tool-permission-system claude/zen-feynman-0u78dk development feature/acp-cancel-request feature/acp-session-list-cursor feature/acp-v1-conformance feature/agent-tools-multiedit-glob-todos-remember feature/cargo-workspace feature/hooks feature/load-local-model-explicitly feature/sigit-agent-profile feature/sigit-code-cloud feature/xcode-mcp-server fix/acp-replay-loaded-session-history fix/acp-replay-reasoning-and-attachments fix/dsml-doubled-bar-markers main release/v1.5.10 release/v1.5.6 release/v1.6.1

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