0438851 · sigitcodereview/sigit · siGit Code & Deploy

@@ -124,7 +124,7 @@ fn parse_k3_tools_block(inner: &str, tools: &[ToolSpec]) -> Option<Vec<Recovered

124 }

125

126 let (call_inner, after_call) = split_once(after_header, K3_CALL_CLOSE)?;

127 - let arguments = parse_k3_arguments(call_inner)?;

127 + let arguments = parse_k3_arguments(call_inner, tools, name)?;

128 calls.push(Recovered {

129 name: name.to_string(),

130 arguments: serde_json::Value::Object(arguments).to_string(),

@@ -135,7 +135,11 @@ fn parse_k3_tools_block(inner: &str, tools: &[ToolSpec]) -> Option<Vec<Recovered

135 Some(calls)

136 }

137

138 -fn parse_k3_arguments(mut rest: &str) -> Option<serde_json::Map<String, serde_json::Value>> {

138 +fn parse_k3_arguments(

139 + mut rest: &str,

140 + tools: &[ToolSpec],

141 + tool_name: &str,

142 +) -> Option<serde_json::Map<String, serde_json::Value>> {

143 let mut args = serde_json::Map::new();

144

145 while !rest.trim().is_empty() {

@@ -144,10 +148,18 @@ fn parse_k3_arguments(mut rest: &str) -> Option<serde_json::Map<String, serde_js

148 let (attr_text, after_header) = split_once(header, K3_SEP)?;

149 let attrs = parse_k3_attributes(attr_text)?;

150 let key = attrs.get("key")?.trim();

147 - let value_type = attrs.get("type").map(String::as_str).unwrap_or("string");

151 if key.is_empty() {

152 return None;

153 }

154 + // K3 normally emits `type` for every argument. If an endpoint omits

155 + // it, retain the legacy recovery path's schema-based typing instead of

156 + // turning e.g. an integer `task_id` into a JSON string.

157 + let schema_type = declared_type(tools, tool_name, key);

158 + let value_type = attrs

159 + .get("type")

160 + .map(String::as_str)

161 + .or(schema_type.as_deref())

162 + .unwrap_or("string");

163

164 let (raw_value, after_argument) = split_once(after_header, K3_ARGUMENT_CLOSE)?;

165 let value = if value_type == "string" {

@@ -190,6 +202,9 @@ fn parse_k3_attributes(mut text: &str) -> Option<HashMap<String, String>> {

202 }

203

204 fn decode_k3_attribute(value: &str) -> String {

205 + // K3's encoder escapes only `&` and `"` in attributes. These replacements

206 + // deliberately mirror that format rather than implementing general HTML

207 + // entity decoding (which could alter literal argument names).

208 value.replace(""", "\"").replace("&", "&")

209 }

210

@@ -323,23 +338,20 @@ impl<'a> StreamScanner<'a> {

338 None => break,

339 },

340 K3_RESPONSE_OPEN => {

326 - let Some(event) =

341 + let Some(mut unwrapped) =

342 self.drain_wrapped_text(K3_RESPONSE_OPEN, K3_RESPONSE_CLOSE, true)

343 else {

344 break;

345 };

331 - if let Some(event) = event {

332 - events.push(event);

333 - }

346 + events.append(&mut unwrapped);

347 }

348 K3_THINK_OPEN => {

336 - let Some(event) = self.drain_wrapped_text(K3_THINK_OPEN, K3_THINK_CLOSE, false)

349 + let Some(mut unwrapped) =

350 + self.drain_wrapped_text(K3_THINK_OPEN, K3_THINK_CLOSE, false)

351 else {

352 break;

353 };

340 - if let Some(event) = event {

341 - events.push(event);

342 - }

354 + events.append(&mut unwrapped);

355 }

356 _ => unreachable!("unknown inline marker"),

357 }

@@ -381,20 +393,45 @@ impl<'a> StreamScanner<'a> {

393 open: &str,

394 close: &str,

395 emit_inner: bool,

384 - ) -> Option<Option<ScanEvent>> {

385 - let close_rel = self.pending[open.len()..].find(close)?;

396 + ) -> Option<Vec<ScanEvent>> {

397 + let after_open = &self.pending[open.len()..];

398 + let close_rel = match after_open.find(close) {

399 + Some(close_rel) => close_rel,

400 + None => {

401 + // A malformed response/think wrapper must not swallow a later

402 + // valid tool block forever. Preserve the malformed prefix as

403 + // text, leave the next marker pending, and resume scanning.

404 + let (next_rel, _) = find_next_marker(after_open)?;

405 + return Some(vec![ScanEvent::Text(

406 + self.pending

407 + .drain(..open.len() + next_rel)

408 + .collect::<String>(),

409 + )]);

410 + }

411 + };

412 let close_idx = open.len() + close_rel;

413 let block: String = self.pending.drain(..close_idx + close.len()).collect();

414

415 if emit_inner {

416 let inner = &block[open.len()..block.len() - close.len()];

391 - Some(Some(ScanEvent::Text(inner.to_string())))

417 + Some(scan_complete_text(inner, self.tools))

418 } else {

393 - Some(None)

419 + Some(Vec::new())

420 }

421 }

422 }

423

424 +/// Scan text from a wrapper whose closing marker has already arrived. This

425 +/// preserves event order when a response body itself contains an inline block.

426 +fn scan_complete_text(text: &str, tools: &[ToolSpec]) -> Vec<ScanEvent> {

427 + let mut scanner = StreamScanner::new(tools);

428 + let mut events = scanner.push(text);

429 + if let Some(rest) = scanner.take_pending() {

430 + events.push(ScanEvent::Text(rest));

431 + }

432 + events

433 +}

434 +

435 fn find_next_marker<'a>(text: &str) -> Option<(usize, &'a str)> {

436 START_MARKERS

437 .iter()

@@ -539,6 +576,18 @@ mod tests {

576 assert!(args["task_id"].is_number());

577 }

578

579 + #[test]

580 + fn kimi_k3_missing_argument_type_uses_the_tool_schema() {

581 + let tools = vec![command_output_spec()];

582 + let (_, calls) = extract(

583 + "<|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|>",

584 + &tools,

585 + );

586 + let args: serde_json::Value = serde_json::from_str(&calls[0].arguments).unwrap();

587 + assert_eq!(args["task_id"], 2);

588 + assert!(args["task_id"].is_number());

589 + }

590 +

591 #[test]

592 fn kimi_k3_multiple_calls_are_recovered_in_order() {

593 let tools = vec![command_output_spec()];

@@ -563,11 +612,24 @@ mod tests {

612 assert!(calls.is_empty());

613 }

614

615 + #[test]

616 + fn kimi_k3_response_content_is_scanned_for_nested_tool_calls() {

617 + let tools = vec![command_output_spec()];

618 + let (text, calls) = extract(

619 + "<|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|>",

620 + &tools,

621 + );

622 + assert_eq!(text, "before after");

623 + assert_eq!(calls.len(), 1);

624 + assert_eq!(calls[0].name, "command_output");

625 + }

626 +

627 #[test]

628 fn kimi_k3_attribute_entities_are_decoded() {

629 let attrs = parse_k3_attributes(r#"tool="run"command" key="a&b""#).unwrap();

630 assert_eq!(attrs["tool"], "run\"command");

631 assert_eq!(attrs["key"], "a&b");

632 + assert_eq!(decode_k3_attribute("line
break"), "line
break");

633 }

634

635 #[test]

@@ -656,6 +718,45 @@ mod tests {

718 assert_eq!(calls[0].name, "run_command");

719 }

720

721 + #[test]

722 + fn scanner_recovers_tools_after_an_unclosed_k3_wrapper() {

723 + let tools = vec![command_output_spec()];

724 + let mut scanner = StreamScanner::new(&tools);

725 + let mut text = String::new();

726 + let mut calls = Vec::new();

727 +

728 + for chunk in [

729 + "before <|open|>response<|sep|>broken ",

730 + "<|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",

731 + ] {

732 + for event in scanner.push(chunk) {

733 + match event {

734 + ScanEvent::Text(chunk) => text.push_str(&chunk),

735 + ScanEvent::ToolCall(call) => calls.push(call),

736 + }

737 + }

738 + }

739 + if let Some(rest) = scanner.take_pending() {

740 + text.push_str(&rest);

741 + }

742 +

743 + assert_eq!(text, "before <|open|>response<|sep|>broken after");

744 + assert_eq!(calls.len(), 1);

745 + assert_eq!(calls[0].name, "command_output");

746 + }

747 +

748 + #[test]

749 + fn scanner_recovers_tools_after_an_unclosed_k3_think_wrapper() {

750 + let tools = vec![command_output_spec()];

751 + let (text, calls) = extract(

752 + "<|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|>",

753 + &tools,

754 + );

755 + assert_eq!(text, "<|open|>think<|sep|>broken ");

756 + assert_eq!(calls.len(), 1);

757 + assert_eq!(calls[0].name, "command_output");

758 + }

759 +

760 #[test]

761 fn scanner_forwards_ordinary_text_without_holding_it_back() {

762 let mut scanner = StreamScanner::new(&[]);

Harden Kimi K3 inline recovery