src/dsml.rs

@@ -45,10 +45,11 @@ pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value
     else {
         return Ok((text.to_owned(), Vec::new()));
     };
-    let Some(relative_end) = text[start..].find(syntax.tool_end) else {
+    let relative_end = text[start..].find(syntax.tool_end);
-        return Err("invalid or incomplete DSML tool call".into());
+    let end = relative_end
-    };
+        .map(|relative_end| start + relative_end + syntax.tool_end.len())
-    let end = start + relative_end + syntax.tool_end.len();
+        .unwrap_or(text.len());
+    let outer_call_complete = relative_end.is_some();
     let content = text[..start].trim_end().to_owned();
     let raw = &text[start..end];
     let mut cursor = syntax.tool_start.len();
@@ -58,6 +59,11 @@ pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value
         if raw[cursor..].starts_with(syntax.tool_end) {
             break;
         }
+        // DS4 recovers calls whose invokes are complete even if generation
+        // stopped before emitting the outer tool-calls closing tag.
+        if !outer_call_complete && cursor == raw.len() && !calls.is_empty() {
+            break;
+        }
         if !raw[cursor..].starts_with(syntax.invoke_start) {
             return Err("invalid or incomplete DSML tool call".into());
         }
@@ -207,4 +213,16 @@ mod tests {
             assert_eq!(calls[0].1["path"], "src/main.rs");
         }
     }
+
+    #[test]
+    fn recovers_complete_invokes_without_the_outer_closing_tag() {
+        let raw = "before<｜DSML｜tool_calls><｜DSML｜invoke name=\"read\"><｜DSML｜parameter name=\"path\" string=\"true\">src/main.rs</｜DSML｜parameter></｜DSML｜invoke>";
+
+        let (content, calls) = parse_tool_calls(raw).unwrap();
+
+        assert_eq!(content, "before");
+        assert_eq!(calls.len(), 1);
+        assert_eq!(calls[0].0, "read");
+        assert_eq!(calls[0].1["path"], "src/main.rs");
+    }
 }