Harden inline tool call recovery
Co-Authored-By: siGit Code <noreply@sigit.si>
Seto Elkahfi committed
Sep 13, 2026 at 13:43 UTC
a5b47bec64fedb7b7acc981cef17b0af8fef2b53
1 file changed
+28
-4
src/inline_tool_calls.rs
+28
-4
@@ -77,9 +77,11 @@ pub fn extract(text: &str, tools: &[ToolSpec]) -> (String, Vec<Recovered>) {
77
(out, calls)
78
}
79
80
-/// Parse one block's inner text: a name, then zero or more
80
+/// Parse one block's inner text: an offered tool name, then zero or more
81
/// `<arg_key>K</arg_key><arg_value>V</arg_value>` pairs. Returns `None` the
82
-/// moment anything departs from that shape.
82
+/// moment anything departs from that shape. The malformed opening marker is
83
+/// accepted only for the first argument; later arguments stay strict so text
84
+/// that merely resembles a call cannot be recovered as one.
85
fn parse_block(inner: &str, tools: &[ToolSpec]) -> Option<Recovered> {
86
let key_idx = inner.find("<arg_key>");
87
let malformed_key_idx = inner.find(OPEN_TAG);
@@ -93,7 +95,7 @@ fn parse_block(inner: &str, tools: &[ToolSpec]) -> Option<Recovered> {
95
(None, Some(idx)) => (inner[..idx].trim(), &inner[idx..], true),
96
(None, None) => (inner.trim(), "", false),
97
};
96
- if name.is_empty() || name.contains(['<', '>']) || !tools.iter().any(|tool| tool.name == name) {
98
+ if name.is_empty() || name.contains(['<', '>']) || !offered_tool(tools, name) {
99
return None;
100
}
101
@@ -126,6 +128,10 @@ fn parse_block(inner: &str, tools: &[ToolSpec]) -> Option<Recovered> {
128
})
129
}
130
131
+fn offered_tool(tools: &[ToolSpec], name: &str) -> bool {
132
+ tools.iter().any(|tool| tool.name == name)
133
+}
134
+
135
fn split_once<'a>(s: &'a str, delim: &str) -> Option<(&'a str, &'a str)> {
136
let idx = s.find(delim)?;
137
Some((&s[..idx], &s[idx + delim.len()..]))
@@ -370,7 +376,25 @@ mod tests {
376
#[test]
377
fn an_unknown_malformed_call_is_left_alone() {
378
let text = "<tool_call>not_offered<tool_call>command</arg_key><arg_value>pwd</arg_value></tool_call>";
373
- let (out, calls) = extract(text, &[]);
379
+ let (out, calls) = extract(text, &[command_output_spec()]);
380
+ assert_eq!(out, text);
381
+ assert!(calls.is_empty());
382
+ }
383
+
384
+ #[test]
385
+ fn malformed_marker_is_accepted_only_for_the_first_argument() {
386
+ let tools = vec![spec(
387
+ "run_command",
388
+ serde_json::json!({
389
+ "type": "object",
390
+ "properties": {
391
+ "command": { "type": "string" },
392
+ "cwd": { "type": "string" }
393
+ }
394
+ }),
395
+ )];
396
+ 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>";
397
+ let (out, calls) = extract(text, &tools);
398
assert_eq!(out, text);
399
assert!(calls.is_empty());
400
}