src/dsml.rs

@@ -1,5 +1,7 @@
 use serde_json::{Map, Value};
 
+const INCOMPLETE_TOOL_CALL: &str = "invalid or incomplete DSML tool call";
+
 #[derive(Clone, Copy)]
 pub(crate) struct Syntax {
     pub(crate) tool_start: &'static str,
@@ -64,15 +66,16 @@ pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value
         if !outer_call_complete && cursor == raw.len() && !calls.is_empty() {
             break;
         }
+        let call = (|| {
             if !raw[cursor..].starts_with(syntax.invoke_start) {
-            return Err("invalid or incomplete DSML tool call".into());
+                return Err(INCOMPLETE_TOOL_CALL.into());
             }
             let tag_end = raw[cursor..]
                 .find('>')
                 .map(|offset| cursor + offset + 1)
-            .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+                .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?;
             let name = attribute(&raw[cursor..tag_end], "name")
-            .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+                .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?;
             cursor = tag_end;
             let mut arguments = Map::new();
             loop {
@@ -82,13 +85,21 @@ pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value
                     break;
                 }
                 let (name, value) = parse_parameter(raw, &mut cursor, syntax)?
-                .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+                    .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?;
                 arguments.insert(name, value);
             }
-        calls.push((name, Value::Object(arguments)));
+            Ok((name, Value::Object(arguments)))
+        })();
+        match call {
+            Ok(call) => calls.push(call),
+            Err(error) if !calls.is_empty() && error == INCOMPLETE_TOOL_CALL => {
+                break;
+            }
+            Err(error) => return Err(error),
+        }
     }
     if calls.is_empty() {
-        return Err("invalid or incomplete DSML tool call".into());
+        return Err(INCOMPLETE_TOOL_CALL.into());
     }
     Ok((content, calls))
 }
@@ -225,4 +236,21 @@ mod tests {
         assert_eq!(calls[0].0, "read");
         assert_eq!(calls[0].1["path"], "src/main.rs");
     }
+
+    #[test]
+    fn recovers_complete_invokes_before_an_incomplete_one() {
+        let raw = "<tool_calls><invoke name=\"first\"></invoke><invoke name=\"second\"><parameter name=\"value\">truncated";
+
+        let (_, calls) = parse_tool_calls(raw).unwrap();
+
+        assert_eq!(calls.len(), 1);
+        assert_eq!(calls[0].0, "first");
+
+        let invalid = "<tool_calls><invoke name=\"first\"></invoke><invoke name=\"second\"><parameter name=\"value\" string=\"false\">invalid</parameter></invoke></tool_calls>";
+        assert!(
+            parse_tool_calls(invalid)
+                .unwrap_err()
+                .starts_with("invalid DSML tool arguments")
+        );
+    }
 }

src/runtime.rs

@@ -62,24 +62,11 @@ enum Operation {
     Measure,
 }
 
-#[derive(Clone, Copy)]
-enum TrackingPolicy {
-    None,
-    Queued,
-    QueuedAndRejected,
-}
-
-impl TrackingPolicy {
-    fn records_queue(self) -> bool {
-        matches!(self, Self::Queued | Self::QueuedAndRejected)
-    }
-
-    fn records_rejection(self) -> bool {
-        matches!(self, Self::QueuedAndRejected)
-    }
-}
-
 impl Operation {
+    fn tracks_metrics(&self) -> bool {
+        !matches!(self, Self::Measure)
+    }
+
     fn error_handler(&self) -> fn(String) -> GenerationEvent {
         match self {
             Self::Generate => |error| GenerationEvent::Finished(Err(error)),
@@ -137,8 +124,7 @@ impl GenerationService {
         source: WorkSource,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
-        self.submit(
+        self.submit(CommandRequest {
-            CommandRequest {
             engine,
             turn,
             messages,
@@ -146,14 +132,11 @@ impl GenerationService {
             source,
             operation: Operation::Generate,
             idle_timeout,
-            },
+        })
-            TrackingPolicy::QueuedAndRejected,
-        )
     }
 
     pub(crate) fn compact(&self, input: CompactionInput) -> Result<ActiveGeneration, String> {
-        self.submit(
+        self.submit(CommandRequest {
-            CommandRequest {
             engine: input.engine,
             turn: input.turn,
             messages: input.messages,
@@ -167,9 +150,7 @@ impl GenerationService {
                 rebuild_system_prompt: input.rebuild_system_prompt,
             },
             idle_timeout: input.idle_timeout,
-            },
+        })
-            TrackingPolicy::None,
-        )
     }
 
     pub(crate) fn measure_context(
@@ -179,8 +160,7 @@ impl GenerationService {
         messages: Vec<ChatTurn>,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
-        self.submit(
+        self.submit(CommandRequest {
-            CommandRequest {
             engine,
             turn,
             messages,
@@ -188,19 +168,14 @@ impl GenerationService {
             source: WorkSource::LocalChat,
             operation: Operation::Measure,
             idle_timeout,
-            },
+        })
-            TrackingPolicy::Queued,
-        )
     }
 
-    fn submit(
+    fn submit(&self, request: CommandRequest) -> Result<ActiveGeneration, String> {
-        &self,
-        request: CommandRequest,
-        tracking: TrackingPolicy,
-    ) -> Result<ActiveGeneration, String> {
         let cancel = Arc::new(AtomicBool::new(false));
         let (events, receiver) = mpsc::channel();
-        if tracking.records_queue() {
+        let tracked = request.operation.tracks_metrics();
+        if tracked {
             self.metrics.request_queued(request.source);
         }
         let command = Command {
@@ -215,7 +190,7 @@ impl GenerationService {
             events,
         };
         if self.commands.send(command).is_err() {
-            if tracking.records_rejection() {
+            if tracked {
                 self.metrics.request_rejected();
             }
             return Err("The model runtime stopped unexpectedly.".to_owned());
@@ -250,7 +225,7 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
                 let source = command.source;
                 let events = command.events.clone();
                 let error_event = command.operation.error_handler();
-                let tracked = !matches!(command.operation, Operation::Measure);
+                let tracked = command.operation.tracks_metrics();
                 if tracked {
                     metrics.request_started(source);
                 }
@@ -285,7 +260,7 @@ fn run_command(
     source: WorkSource,
 ) {
     let error_event = command.operation.error_handler();
-    let tracked = !matches!(command.operation, Operation::Measure);
+    let tracked = command.operation.tracks_metrics();
     state.idle_timeout = command.idle_timeout;
     if command.cancel.load(Ordering::Relaxed) {
         if tracked {
@@ -483,4 +458,17 @@ mod tests {
             GenerationEvent::Measured(Err(error)) if error == "stopped"
         ));
     }
+
+    #[test]
+    fn operation_tracking_is_consistent() {
+        assert!(Operation::Generate.tracks_metrics());
+        assert!(
+            Operation::Compact {
+                reason: String::new(),
+                rebuild_system_prompt: String::new(),
+            }
+            .tracks_metrics()
+        );
+        assert!(!Operation::Measure.tracks_metrics());
+    }
 }