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 left
22 //! in the text untouched, deliberately. Reissuing a `run_command` is cheap,
23 //! but guessing wrong at a half-parsed `edit_file` would write the wrong
24 //! change to a file, so a malformed block stays visible rather than being
25 //! silently misinterpreted.
26
27 use std::collections::HashMap;
28
29 use crate::backend::ToolSpec;
30
31 const XML_OPEN_TAG: &str = "<tool_call>";
32 const XML_CLOSE_TAG: &str = "</tool_call>";
33
34 const K3_TOOLS_OPEN: &str = "<|open|>tools<|sep|>";
35 const K3_TOOLS_CLOSE: &str = "<|close|>tools<|sep|>";
36 const K3_CALL_OPEN: &str = "<|open|>call";
37 const K3_CALL_CLOSE: &str = "<|close|>call<|sep|>";
38 const K3_ARGUMENT_OPEN: &str = "<|open|>argument";
39 const K3_ARGUMENT_CLOSE: &str = "<|close|>argument<|sep|>";
40 const K3_RESPONSE_OPEN: &str = "<|open|>response<|sep|>";
41 const K3_RESPONSE_CLOSE: &str = "<|close|>response<|sep|>";
42 const K3_THINK_OPEN: &str = "<|open|>think<|sep|>";
43 const K3_THINK_CLOSE: &str = "<|close|>think<|sep|>";
44 const K3_SEP: &str = "<|sep|>";
45
46 const START_MARKERS: &[&str] = &[XML_OPEN_TAG, K3_TOOLS_OPEN, K3_RESPONSE_OPEN, K3_THINK_OPEN];
47
48 /// One tool call recovered from inline text.
49 #[derive(Debug, Clone, PartialEq, Eq)]
50 pub struct Recovered {
51 pub name: String,
52 /// Arguments as a JSON-encoded string, matching `ToolCall::arguments`.
53 pub arguments: String,
54 }
55
56 /// Scan a complete text blob for inline tool-call blocks. Returns the text
57 /// with recovered blocks removed, plus the calls recovered from them. Blocks
58 /// that don't parse cleanly, and unterminated trailing tags, are left in the
59 /// returned text untouched.
60 pub fn extract(text: &str, tools: &[ToolSpec]) -> (String, Vec<Recovered>) {
61 let mut scanner = StreamScanner::new(tools);
62 let mut out = String::with_capacity(text.len());
63 let mut calls = Vec::new();
64
65 for event in scanner.push(text) {
66 match event {
67 ScanEvent::Text(text) => out.push_str(&text),
68 ScanEvent::ToolCall(call) => calls.push(call),
69 }
70 }
71 if let Some(rest) = scanner.take_pending() {
72 out.push_str(&rest);
73 }
74
75 (out, calls)
76 }
77
78 /// Parse one legacy XML block's inner text: an offered tool name, then zero or
79 /// more `<arg_key>K</arg_key><arg_value>V</arg_value>` pairs. The malformed
80 /// opening marker observed from GLM is accepted only for the first argument;
81 /// later arguments stay strict so text that merely resembles a call cannot be
82 /// recovered as one.
83 fn parse_xml_block(inner: &str, tools: &[ToolSpec]) -> Option<Recovered> {
84 let key_idx = inner.find("<arg_key>");
85 let malformed_key_idx = inner.find(XML_OPEN_TAG);
86 let check_status_name = key_idx.and_then(|idx| parse_check_status_preamble(&inner[..idx]));
87 let (name, mut rest, malformed_first_key) =
88 if let (Some(idx), Some(name)) = (key_idx, check_status_name) {
89 (name, &inner[idx..], false)
90 } else {
91 match (key_idx, malformed_key_idx) {
92 (Some(key_idx), Some(malformed_key_idx)) if malformed_key_idx < key_idx => (
93 inner[..malformed_key_idx].trim(),
94 &inner[malformed_key_idx..],
95 true,
96 ),
97 (Some(idx), _) => (inner[..idx].trim(), &inner[idx..], false),
98 (None, Some(idx)) => (inner[..idx].trim(), &inner[idx..], true),
99 (None, None) => (inner.trim(), "", false),
100 }
101 };
102 if name.is_empty() || name.contains(['<', '>']) || !offered_tool(tools, name) {
103 return None;
104 }
105
106 let mut args = serde_json::Map::new();
107 let mut first_key = true;
108 while !rest.is_empty() {
109 // The GLM alias is permitted exactly once, at the first key. A later
110 // `<tool_call>` must reject the complete block rather than letting a
111 // partial call escape recovery.
112 if !first_key && rest.starts_with(XML_OPEN_TAG) {
113 return None;
114 }
115 let key_open = if first_key && malformed_first_key {
116 XML_OPEN_TAG
117 } else {
118 "<arg_key>"
119 };
120 let (key, after_key) = split_once(rest.strip_prefix(key_open)?, "</arg_key>")?;
121 let (value, after_value) =
122 split_once(after_key.strip_prefix("<arg_value>")?, "</arg_value>")?;
123 let key = key.trim();
124 if key.is_empty() {
125 return None;
126 }
127 args.insert(
128 key.to_string(),
129 coerce(value, declared_type(tools, name, key).as_deref()),
130 );
131 rest = after_value;
132 first_key = false;
133 }
134
135 Some(Recovered {
136 name: name.to_string(),
137 arguments: serde_json::Value::Object(args).to_string(),
138 })
139 }
140
141 /// Some GLM-family responses prefix the real arguments with generation-control
142 /// metadata shaped like `tool_name CheckStatus=value</arg_value>`. It is not a
143 /// tool argument: accept that exact bounded preamble and discard it, while
144 /// leaving every other malformed prefix untouched.
145 fn parse_check_status_preamble(prefix: &str) -> Option<&str> {
146 let (name, status) = prefix.trim().split_once(" CheckStatus=")?;
147 let status = status.strip_suffix("</arg_value>")?;
148 if name.is_empty()
149 || status.is_empty()
150 || !status
151 .chars()
152 .all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '-'))
153 {
154 return None;
155 }
156 Some(name)
157 }
158
159 fn parse_k3_tools_block(inner: &str, tools: &[ToolSpec]) -> Option<Vec<Recovered>> {
160 let mut calls = Vec::new();
161 let mut rest = inner;
162
163 while !rest.trim().is_empty() {
164 rest = rest.trim_start();
165 let header = rest.strip_prefix(K3_CALL_OPEN)?;
166 let (attr_text, after_header) = split_once(header, K3_SEP)?;
167 let attrs = parse_k3_attributes(attr_text)?;
168 let name = attrs.get("tool")?;
169 if name.is_empty() || !offered_tool(tools, name) {
170 return None;
171 }
172
173 let (call_inner, after_call) = split_once(after_header, K3_CALL_CLOSE)?;
174 let arguments = parse_k3_arguments(call_inner, tools, name)?;
175 calls.push(Recovered {
176 name: name.to_string(),
177 arguments: serde_json::Value::Object(arguments).to_string(),
178 });
179 rest = after_call;
180 }
181
182 Some(calls)
183 }
184
185 fn parse_k3_arguments(
186 mut rest: &str,
187 tools: &[ToolSpec],
188 tool_name: &str,
189 ) -> Option<serde_json::Map<String, serde_json::Value>> {
190 let mut args = serde_json::Map::new();
191
192 while !rest.trim().is_empty() {
193 rest = rest.trim_start();
194 let header = rest.strip_prefix(K3_ARGUMENT_OPEN)?;
195 let (attr_text, after_header) = split_once(header, K3_SEP)?;
196 let attrs = parse_k3_attributes(attr_text)?;
197 let key = attrs.get("key")?.trim();
198 if key.is_empty() {
199 return None;
200 }
201 // K3 normally emits `type` for every argument. If an endpoint omits
202 // it, retain the legacy recovery path's schema-based typing instead of
203 // turning e.g. an integer `task_id` into a JSON string.
204 let schema_type = declared_type(tools, tool_name, key);
205 let value_type = attrs
206 .get("type")
207 .map(String::as_str)
208 .or(schema_type.as_deref())
209 .unwrap_or("string");
210
211 let (raw_value, after_argument) = split_once(after_header, K3_ARGUMENT_CLOSE)?;
212 let value = if value_type == "string" {
213 serde_json::Value::String(raw_value.to_string())
214 } else {
215 serde_json::from_str(raw_value).ok()?
216 };
217 args.insert(key.to_string(), value);
218 rest = after_argument;
219 }
220
221 Some(args)
222 }
223
224 fn parse_k3_attributes(mut text: &str) -> Option<HashMap<String, String>> {
225 let mut attrs = HashMap::new();
226 text = text.trim();
227
228 while !text.is_empty() {
229 let eq_idx = text.find('=')?;
230 let key = text[..eq_idx].trim();
231 if key.is_empty()
232 || !key
233 .chars()
234 .all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '-')
235 {
236 return None;
237 }
238
239 let mut value_part = text[eq_idx + 1..].trim_start();
240 value_part = value_part.strip_prefix('"')?;
241 let end_idx = value_part.find('"')?;
242 let raw_value = &value_part[..end_idx];
243 attrs.insert(key.to_string(), decode_k3_attribute(raw_value));
244
245 text = value_part[end_idx + 1..].trim_start();
246 }
247
248 Some(attrs)
249 }
250
251 fn decode_k3_attribute(value: &str) -> String {
252 // K3's encoder escapes only `&` and `"` in attributes. These replacements
253 // deliberately mirror that format rather than implementing general HTML
254 // entity decoding (which could alter literal argument names).
255 value.replace("&quot;", "\"").replace("&amp;", "&")
256 }
257
258 fn split_once<'a>(s: &'a str, delim: &str) -> Option<(&'a str, &'a str)> {
259 let idx = s.find(delim)?;
260 Some((&s[..idx], &s[idx + delim.len()..]))
261 }
262
263 fn offered_tool(tools: &[ToolSpec], tool_name: &str) -> bool {
264 tools.iter().any(|spec| spec.name == tool_name)
265 }
266
267 /// The JSON Schema `type` declared for `tool_name`'s `key` parameter, e.g.
268 /// `"integer"` for `command_output`'s `task_id`. [`ToolSpec`] carries the
269 /// schema as an unparsed string, so this parses it per lookup; the argument
270 /// counts involved are tiny and this only runs on the recovery path.
271 fn declared_type(tools: &[ToolSpec], tool_name: &str, key: &str) -> Option<String> {
272 let spec = tools.iter().find(|spec| spec.name == tool_name)?;
273 let schema: serde_json::Value = serde_json::from_str(&spec.parameters_schema).ok()?;
274 Some(
275 schema
276 .get("properties")?
277 .get(key)?
278 .get("type")?
279 .as_str()?
280 .to_string(),
281 )
282 }
283
284 /// Convert a raw legacy text value to JSON per its declared schema type.
285 /// Falls back to a plain string for `"string"`, an unknown type, or a value
286 /// that doesn't parse as what it claims to be.
287 fn coerce(raw: &str, declared: Option<&str>) -> serde_json::Value {
288 match declared {
289 Some("integer") => raw
290 .trim()
291 .parse::<i64>()
292 .map(serde_json::Value::from)
293 .unwrap_or_else(|_| serde_json::Value::String(raw.to_string())),
294 Some("number") => raw
295 .trim()
296 .parse::<f64>()
297 .ok()
298 .and_then(serde_json::Number::from_f64)
299 .map(serde_json::Value::Number)
300 .unwrap_or_else(|| serde_json::Value::String(raw.to_string())),
301 Some("boolean") => match raw.trim() {
302 "true" => serde_json::Value::Bool(true),
303 "false" => serde_json::Value::Bool(false),
304 _ => serde_json::Value::String(raw.to_string()),
305 },
306 _ => serde_json::Value::String(raw.to_string()),
307 }
308 }
309
310 /// One event from incrementally scanning a stream of content deltas.
311 #[derive(Debug, Clone, PartialEq, Eq)]
312 pub enum ScanEvent {
313 /// Plain text, safe to forward to the UI.
314 Text(String),
315 /// A tool call recovered from a completed block.
316 ToolCall(Recovered),
317 }
318
319 /// Incremental scanner for a live delta stream.
320 ///
321 /// Holds back only the text that could still turn into a known inline protocol
322 /// marker, so ordinary answers still reach the UI token by token. Buffering
323 /// starts once an opening marker appears, and lasts only until the matching
324 /// close arrives.
325 pub struct StreamScanner<'a> {
326 tools: &'a [ToolSpec],
327 pending: String,
328 }
329
330 impl<'a> StreamScanner<'a> {
331 pub fn new(tools: &'a [ToolSpec]) -> Self {
332 Self {
333 tools,
334 pending: String::new(),
335 }
336 }
337
338 /// Feed the next content delta. Returns events to emit now, in order; may
339 /// be empty if the chunk was absorbed into a tag that hasn't closed yet.
340 pub fn push(&mut self, chunk: &str) -> Vec<ScanEvent> {
341 self.pending.push_str(chunk);
342 self.drain()
343 }
344
345 /// Take whatever is still held back, e.g. a `<tool_call` or Kimi marker
346 /// prefix that never closed. Text we never saw the end of can't be
347 /// interpreted safely, so the honest thing is to hand it back as-is.
348 pub fn take_pending(&mut self) -> Option<String> {
349 if self.pending.is_empty() {
350 None
351 } else {
352 Some(std::mem::take(&mut self.pending))
353 }
354 }
355
356 fn drain(&mut self) -> Vec<ScanEvent> {
357 let mut events = Vec::new();
358 loop {
359 let Some((open_idx, marker)) = find_next_marker(&self.pending) else {
360 let keep = max_partial_prefix_len(&self.pending, START_MARKERS);
361 let flush_len = self.pending.len() - keep;
362 if flush_len > 0 {
363 events.push(ScanEvent::Text(
364 self.pending.drain(..flush_len).collect::<String>(),
365 ));
366 }
367 break;
368 };
369
370 if open_idx > 0 {
371 events.push(ScanEvent::Text(
372 self.pending.drain(..open_idx).collect::<String>(),
373 ));
374 }
375
376 match marker {
377 XML_OPEN_TAG => {
378 let Some(event) = self.drain_xml() else {
379 break;
380 };
381 events.push(event);
382 }
383 K3_TOOLS_OPEN => match self.drain_k3_tools() {
384 Some(mut recovered) => events.append(&mut recovered),
385 None => break,
386 },
387 K3_RESPONSE_OPEN => {
388 let Some(mut unwrapped) =
389 self.drain_wrapped_text(K3_RESPONSE_OPEN, K3_RESPONSE_CLOSE, true)
390 else {
391 break;
392 };
393 events.append(&mut unwrapped);
394 }
395 K3_THINK_OPEN => {
396 let Some(mut unwrapped) =
397 self.drain_wrapped_text(K3_THINK_OPEN, K3_THINK_CLOSE, false)
398 else {
399 break;
400 };
401 events.append(&mut unwrapped);
402 }
403 _ => unreachable!("unknown inline marker"),
404 }
405 }
406 events
407 }
408
409 fn drain_xml(&mut self) -> Option<ScanEvent> {
410 let close_rel = self.pending[XML_OPEN_TAG.len()..].find(XML_CLOSE_TAG)?;
411 let close_idx = XML_OPEN_TAG.len() + close_rel;
412 let block: String = self
413 .pending
414 .drain(..close_idx + XML_CLOSE_TAG.len())
415 .collect();
416 let inner = &block[XML_OPEN_TAG.len()..block.len() - XML_CLOSE_TAG.len()];
417 Some(match parse_xml_block(inner, self.tools) {
418 Some(call) => ScanEvent::ToolCall(call),
419 None => ScanEvent::Text(block),
420 })
421 }
422
423 fn drain_k3_tools(&mut self) -> Option<Vec<ScanEvent>> {
424 let close_rel = self.pending[K3_TOOLS_OPEN.len()..].find(K3_TOOLS_CLOSE)?;
425 let close_idx = K3_TOOLS_OPEN.len() + close_rel;
426 let block: String = self
427 .pending
428 .drain(..close_idx + K3_TOOLS_CLOSE.len())
429 .collect();
430 let inner = &block[K3_TOOLS_OPEN.len()..block.len() - K3_TOOLS_CLOSE.len()];
431
432 Some(match parse_k3_tools_block(inner, self.tools) {
433 Some(calls) => calls.into_iter().map(ScanEvent::ToolCall).collect(),
434 None => vec![ScanEvent::Text(block)],
435 })
436 }
437
438 fn drain_wrapped_text(
439 &mut self,
440 open: &str,
441 close: &str,
442 emit_inner: bool,
443 ) -> Option<Vec<ScanEvent>> {
444 let after_open = &self.pending[open.len()..];
445 let close_rel = match after_open.find(close) {
446 Some(close_rel) => close_rel,
447 None => {
448 // A malformed response/think wrapper must not swallow a later
449 // valid tool block forever. Preserve the malformed prefix as
450 // text, leave the next marker pending, and resume scanning.
451 let (next_rel, _) = find_next_marker(after_open)?;
452 return Some(vec![ScanEvent::Text(
453 self.pending
454 .drain(..open.len() + next_rel)
455 .collect::<String>(),
456 )]);
457 }
458 };
459 let close_idx = open.len() + close_rel;
460 let block: String = self.pending.drain(..close_idx + close.len()).collect();
461
462 if emit_inner {
463 let inner = &block[open.len()..block.len() - close.len()];
464 Some(scan_complete_text(inner, self.tools))
465 } else {
466 Some(Vec::new())
467 }
468 }
469 }
470
471 /// Scan text from a wrapper whose closing marker has already arrived. This
472 /// preserves event order when a response body itself contains an inline block.
473 fn scan_complete_text(text: &str, tools: &[ToolSpec]) -> Vec<ScanEvent> {
474 let mut scanner = StreamScanner::new(tools);
475 let mut events = scanner.push(text);
476 if let Some(rest) = scanner.take_pending() {
477 events.push(ScanEvent::Text(rest));
478 }
479 events
480 }
481
482 fn find_next_marker<'a>(text: &str) -> Option<(usize, &'a str)> {
483 START_MARKERS
484 .iter()
485 .filter_map(|marker| text.find(marker).map(|idx| (idx, *marker)))
486 .min_by_key(|(idx, _)| *idx)
487 }
488
489 /// Length of the longest suffix of `buf` that is a prefix of any known opening
490 /// marker. Lets a split opening tag be caught without delaying text that can't
491 /// possibly be part of one.
492 fn max_partial_prefix_len(buf: &str, needles: &[&str]) -> usize {
493 needles
494 .iter()
495 .map(|needle| partial_prefix_len(buf, needle))
496 .max()
497 .unwrap_or(0)
498 }
499
500 fn partial_prefix_len(buf: &str, needle: &str) -> usize {
501 let max = buf.len().min(needle.len() - 1);
502 for len in (1..=max).rev() {
503 let start = buf.len() - len;
504 if buf.is_char_boundary(start) && needle.starts_with(&buf[start..]) {
505 return len;
506 }
507 }
508 0
509 }
510
511 #[cfg(test)]
512 mod tests {
513 use super::*;
514
515 fn spec(name: &str, schema: serde_json::Value) -> ToolSpec {
516 ToolSpec {
517 name: name.to_string(),
518 description: String::new(),
519 parameters_schema: schema.to_string(),
520 }
521 }
522
523 fn command_output_spec() -> ToolSpec {
524 spec(
525 "command_output",
526 serde_json::json!({
527 "type": "object",
528 "properties": { "task_id": { "type": "integer" } }
529 }),
530 )
531 }
532
533 fn run_command_spec() -> ToolSpec {
534 spec(
535 "run_command",
536 serde_json::json!({
537 "type": "object",
538 "properties": {
539 "command": { "type": "string" },
540 "cwd": { "type": "string" },
541 "run_in_background": { "type": "boolean" }
542 }
543 }),
544 )
545 }
546
547 #[test]
548 fn an_integer_argument_is_recovered_as_a_json_number() {
549 let tools = vec![command_output_spec()];
550 let (text, calls) = extract(
551 "Checking.<tool_call>command_output<arg_key>task_id</arg_key><arg_value>2</arg_value></tool_call>",
552 &tools,
553 );
554 assert_eq!(text, "Checking.");
555 assert_eq!(calls.len(), 1);
556 assert_eq!(calls[0].name, "command_output");
557 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
558 assert_eq!(args["task_id"], 2);
559 assert!(args["task_id"].is_number());
560 }
561
562 #[test]
563 fn recovers_multiple_string_arguments_in_order() {
564 let tools = vec![run_command_spec()];
565 let (text, calls) = extract(
566 "<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>",
567 &tools,
568 );
569 assert_eq!(text, "");
570 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
571 assert_eq!(args["command"], "cargo test --locked 2>&1 | tail -20");
572 assert_eq!(args["cwd"], "/Users/x/sigit");
573 assert_eq!(args["run_in_background"], true);
574 }
575
576 #[test]
577 fn recovers_glm_call_with_a_malformed_first_argument_marker() {
578 let tools = vec![run_command_spec()];
579 let (text, calls) = extract(
580 "<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>",
581 &tools,
582 );
583
584 assert_eq!(text, "");
585 assert_eq!(calls.len(), 1);
586 assert_eq!(calls[0].name, "run_command");
587 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
588 assert_eq!(args["command"], "pwd");
589 assert_eq!(args["cwd"], "/tmp");
590 }
591
592 #[test]
593 fn recovers_glm_call_with_check_status_before_real_arguments() {
594 let tools = vec![command_output_spec()];
595 let (text, calls) = extract(
596 "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>",
597 &tools,
598 );
599
600 assert_eq!(text, "polling ");
601 assert_eq!(calls.len(), 1);
602 assert_eq!(calls[0].name, "command_output");
603 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
604 assert_eq!(args["task_id"], 1);
605 assert!(args["task_id"].is_number());
606 assert!(args.get("CheckStatus").is_none());
607 }
608
609 #[test]
610 fn malformed_check_status_preamble_is_left_as_text() {
611 let text = "<tool_call>command_output CheckStatus=keep going</arg_value><arg_key>task_id</arg_key><arg_value>1</arg_value></tool_call>";
612 let (out, calls) = extract(text, &[command_output_spec()]);
613 assert_eq!(out, text);
614 assert!(calls.is_empty());
615 }
616
617 #[test]
618 fn check_status_after_a_real_argument_is_left_as_text() {
619 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>";
620 let (out, calls) = extract(text, &[command_output_spec()]);
621 assert_eq!(out, text);
622 assert!(calls.is_empty());
623 }
624
625 #[test]
626 fn check_status_for_an_unknown_tool_is_left_as_text() {
627 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>";
628 let (out, calls) = extract(text, &[command_output_spec()]);
629 assert_eq!(out, text);
630 assert!(calls.is_empty());
631 }
632
633 #[test]
634 fn unknown_glm_tool_is_left_as_text() {
635 let text = "<tool_call>not_offered<tool_call>command</arg_key><arg_value>pwd</arg_value></tool_call>";
636 let (out, calls) = extract(text, &[run_command_spec()]);
637 assert_eq!(out, text);
638 assert!(calls.is_empty());
639 }
640
641 #[test]
642 fn malformed_glm_marker_is_accepted_only_for_the_first_argument() {
643 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>";
644 let (out, calls) = extract(text, &[run_command_spec()]);
645 assert_eq!(out, text);
646 assert!(calls.is_empty());
647 }
648
649 #[test]
650 fn a_malformed_block_is_left_alone_rather_than_guessed_at() {
651 let text = "<tool_call>edit_file<arg_value>new_text</arg_key><arg_value>def x; end</arg_value></tool_call>";
652 let (out, calls) = extract(text, &[run_command_spec()]);
653 assert_eq!(out, text);
654 assert!(calls.is_empty());
655 }
656
657 #[test]
658 fn an_unterminated_tag_is_left_alone() {
659 let text = "working on it <tool_call>run_command<arg_key>command</arg_key>";
660 let (out, calls) = extract(text, &[run_command_spec()]);
661 assert_eq!(out, text);
662 assert!(calls.is_empty());
663 }
664
665 #[test]
666 fn ordinary_text_passes_through_untouched() {
667 let (out, calls) = extract("if x < y then a<b, nothing to recover", &[]);
668 assert_eq!(out, "if x < y then a<b, nothing to recover");
669 assert!(calls.is_empty());
670 }
671
672 #[test]
673 fn kimi_k3_run_command_block_is_recovered_and_removed() {
674 let tools = vec![run_command_spec()];
675 let (text, calls) = extract(
676 "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|>",
677 &tools,
678 );
679 assert_eq!(text, "Working ");
680 assert_eq!(calls.len(), 1);
681 assert_eq!(calls[0].name, "run_command");
682 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
683 assert_eq!(args["command"], "sleep 180; echo waited");
684 assert_eq!(args["cwd"], "/Users/setoelkahfi/Repositories/onde-ed");
685 }
686
687 #[test]
688 fn kimi_k3_non_string_arguments_are_json_decoded() {
689 let tools = vec![command_output_spec()];
690 let (_, calls) = extract(
691 "<|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|>",
692 &tools,
693 );
694 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
695 assert_eq!(args["task_id"], 2);
696 assert!(args["task_id"].is_number());
697 }
698
699 #[test]
700 fn kimi_k3_missing_argument_type_uses_the_tool_schema() {
701 let tools = vec![command_output_spec()];
702 let (_, calls) = extract(
703 "<|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|>",
704 &tools,
705 );
706 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
707 assert_eq!(args["task_id"], 2);
708 assert!(args["task_id"].is_number());
709 }
710
711 #[test]
712 fn kimi_k3_multiple_calls_are_recovered_in_order() {
713 let tools = vec![command_output_spec()];
714 let (_, calls) = extract(
715 "<|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|>",
716 &tools,
717 );
718 let first: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
719 let second: serde_json::Value = serde_json::from_str(&calls[1].arguments).unwrap();
720 assert_eq!(calls.len(), 2);
721 assert_eq!(first["task_id"], 1);
722 assert_eq!(second["task_id"], 2);
723 }
724
725 #[test]
726 fn kimi_k3_response_is_unwrapped_and_think_is_hidden() {
727 let (text, calls) = extract(
728 "a<|open|>think<|sep|>private<|close|>think<|sep|>b<|open|>response<|sep|>visible<|close|>response<|sep|>c",
729 &[],
730 );
731 assert_eq!(text, "abvisiblec");
732 assert!(calls.is_empty());
733 }
734
735 #[test]
736 fn kimi_k3_response_content_is_scanned_for_nested_tool_calls() {
737 let tools = vec![command_output_spec()];
738 let (text, calls) = extract(
739 "<|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|>",
740 &tools,
741 );
742 assert_eq!(text, "before after");
743 assert_eq!(calls.len(), 1);
744 assert_eq!(calls[0].name, "command_output");
745 }
746
747 #[test]
748 fn kimi_k3_attribute_entities_are_decoded() {
749 let attrs = parse_k3_attributes(r#"tool="run&quot;command" key="a&amp;b""#).unwrap();
750 assert_eq!(attrs["tool"], "run\"command");
751 assert_eq!(attrs["key"], "a&b");
752 assert_eq!(decode_k3_attribute("line&#10;break"), "line&#10;break");
753 }
754
755 #[test]
756 fn malformed_kimi_k3_tools_block_is_left_alone() {
757 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|>";
758 let (out, calls) = extract(text, &[run_command_spec()]);
759 assert_eq!(out, text);
760 assert!(calls.is_empty());
761 }
762
763 #[test]
764 fn unknown_kimi_k3_tool_block_is_left_alone() {
765 let text = "<|open|>tools<|sep|><|open|>call tool=\"delete_everything\" index=\"1\"<|sep|><|close|>call<|sep|><|close|>tools<|sep|>";
766 let (out, calls) = extract(text, &[run_command_spec()]);
767 assert_eq!(out, text);
768 assert!(calls.is_empty());
769 }
770
771 #[test]
772 fn invalid_kimi_k3_non_string_argument_is_left_alone() {
773 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|>";
774 let (out, calls) = extract(text, &[command_output_spec()]);
775 assert_eq!(out, text);
776 assert!(calls.is_empty());
777 }
778
779 #[test]
780 fn scanner_recovers_a_legacy_tag_split_across_chunk_boundaries() {
781 let tools = vec![command_output_spec()];
782 let mut scanner = StreamScanner::new(&tools);
783 let mut text = String::new();
784 let mut calls = Vec::new();
785
786 for chunk in [
787 "polling ",
788 "<tool_c",
789 "all>command_output<arg_",
790 "key>task_id</arg_key><arg_value>2</arg_v",
791 "alue></tool_call>",
792 " done",
793 ] {
794 for event in scanner.push(chunk) {
795 match event {
796 ScanEvent::Text(t) => text.push_str(&t),
797 ScanEvent::ToolCall(c) => calls.push(c),
798 }
799 }
800 }
801 if let Some(rest) = scanner.take_pending() {
802 text.push_str(&rest);
803 }
804
805 assert_eq!(text, "polling done", "no tag fragment may reach the UI");
806 assert_eq!(calls.len(), 1);
807 assert_eq!(calls[0].name, "command_output");
808 }
809
810 #[test]
811 fn scanner_recovers_glm_call_with_malformed_marker_split_across_chunks() {
812 let tools = vec![run_command_spec()];
813 let mut scanner = StreamScanner::new(&tools);
814 let mut text = String::new();
815 let mut calls = Vec::new();
816
817 for chunk in [
818 "before ",
819 "<tool_call>run_command<tool_",
820 "call>command</arg_key><arg_value>pwd</arg_value></tool_call>",
821 " after",
822 ] {
823 for event in scanner.push(chunk) {
824 match event {
825 ScanEvent::Text(value) => text.push_str(&value),
826 ScanEvent::ToolCall(call) => calls.push(call),
827 }
828 }
829 }
830 if let Some(rest) = scanner.take_pending() {
831 text.push_str(&rest);
832 }
833
834 assert_eq!(text, "before after");
835 assert_eq!(calls.len(), 1);
836 assert_eq!(calls[0].name, "run_command");
837 let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();
838 assert_eq!(args["command"], "pwd");
839 }
840
841 #[test]
842 fn scanner_recovers_check_status_call_split_across_chunks() {
843 let tools = vec![command_output_spec()];
844 let mut scanner = StreamScanner::new(&tools);
845 let mut text = String::new();
846 let mut calls = Vec::new();
847
848 for chunk in [
849 "before <tool_call>command_output CheckStatus=true_or_poll_",
850 "again_with_different_params</arg_",
851 "value><arg_key>task_id</arg_key><arg_value>1</arg_value></tool_",
852 "call> after",
853 ] {
854 for event in scanner.push(chunk) {
855 match event {
856 ScanEvent::Text(value) => text.push_str(&value),
857 ScanEvent::ToolCall(call) => calls.push(call),
858 }
859 }
860 }
861 if let Some(rest) = scanner.take_pending() {
862 text.push_str(&rest);
863 }
864
865 assert_eq!(text, "before after");
866 assert_eq!(calls.len(), 1);
867 assert_eq!(calls[0].name, "command_output");
868 }
869
870 #[test]
871 fn scanner_recovers_kimi_k3_tools_split_across_chunk_boundaries() {
872 let tools = vec![run_command_spec()];
873 let mut scanner = StreamScanner::new(&tools);
874 let mut text = String::new();
875 let mut calls = Vec::new();
876
877 for chunk in [
878 "pre ",
879 "<|open|>too",
880 "ls<|sep|><|open|>call tool=\"run_command\" index=\"1\"<|sep|>",
881 "<|open|>argument key=\"command\" type=\"string\"<|sep|>echo hi",
882 "<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|>",
883 " post",
884 ] {
885 for event in scanner.push(chunk) {
886 match event {
887 ScanEvent::Text(t) => text.push_str(&t),
888 ScanEvent::ToolCall(c) => calls.push(c),
889 }
890 }
891 }
892 if let Some(rest) = scanner.take_pending() {
893 text.push_str(&rest);
894 }
895
896 assert_eq!(text, "pre post");
897 assert_eq!(calls.len(), 1);
898 assert_eq!(calls[0].name, "run_command");
899 }
900
901 #[test]
902 fn scanner_recovers_tools_after_an_unclosed_k3_wrapper() {
903 let tools = vec![command_output_spec()];
904 let mut scanner = StreamScanner::new(&tools);
905 let mut text = String::new();
906 let mut calls = Vec::new();
907
908 for chunk in [
909 "before <|open|>response<|sep|>broken ",
910 "<|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",
911 ] {
912 for event in scanner.push(chunk) {
913 match event {
914 ScanEvent::Text(chunk) => text.push_str(&chunk),
915 ScanEvent::ToolCall(call) => calls.push(call),
916 }
917 }
918 }
919 if let Some(rest) = scanner.take_pending() {
920 text.push_str(&rest);
921 }
922
923 assert_eq!(text, "before <|open|>response<|sep|>broken after");
924 assert_eq!(calls.len(), 1);
925 assert_eq!(calls[0].name, "command_output");
926 }
927
928 #[test]
929 fn scanner_recovers_tools_after_an_unclosed_k3_think_wrapper() {
930 let tools = vec![command_output_spec()];
931 let (text, calls) = extract(
932 "<|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|>",
933 &tools,
934 );
935 assert_eq!(text, "<|open|>think<|sep|>broken ");
936 assert_eq!(calls.len(), 1);
937 assert_eq!(calls[0].name, "command_output");
938 }
939
940 #[test]
941 fn scanner_forwards_ordinary_text_without_holding_it_back() {
942 let mut scanner = StreamScanner::new(&[]);
943 let events = scanner.push("a normal streamed answer");
944 assert_eq!(
945 events,
946 vec![ScanEvent::Text("a normal streamed answer".to_string())]
947 );
948 assert!(scanner.take_pending().is_none());
949 }
950
951 #[test]
952 fn scanner_releases_a_near_miss_prefix_once_it_cannot_match() {
953 let mut scanner = StreamScanner::new(&[]);
954 assert_eq!(
955 scanner.push("see <tool"),
956 vec![ScanEvent::Text("see ".to_string())]
957 );
958 assert_eq!(
959 scanner.push("box for details"),
960 vec![ScanEvent::Text("<toolbox for details".to_string())]
961 );
962 assert!(scanner.take_pending().is_none());
963 }
964
965 #[test]
966 fn scanner_hands_back_an_unclosed_tag_at_stream_end() {
967 let tools = vec![run_command_spec()];
968 let mut scanner = StreamScanner::new(&tools);
969 let events = scanner.push("partial <tool_call>run_command<arg_key>cmd");
970 assert_eq!(events, vec![ScanEvent::Text("partial ".to_string())]);
971 assert_eq!(
972 scanner.take_pending().as_deref(),
973 Some("<tool_call>run_command<arg_key>cmd")
974 );
975 }
976 }