Fix refactor recovery and metrics

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,31 +66,40 @@ pub(crate) fn parse_tool_calls(text: &str) -> Result<(String, Vec<(String, Value
         if !outer_call_complete && cursor == raw.len() && !calls.is_empty() {
             break;
         }
-        if !raw[cursor..].starts_with(syntax.invoke_start) {
+        let call = (|| {
-            return Err("invalid or incomplete DSML tool call".into());
+            if !raw[cursor..].starts_with(syntax.invoke_start) {
-        }
+                return Err(INCOMPLETE_TOOL_CALL.into());
-        let tag_end = raw[cursor..]
+            }
-            .find('>')
+            let tag_end = raw[cursor..]
-            .map(|offset| cursor + offset + 1)
+                .find('>')
-            .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+                .map(|offset| cursor + offset + 1)
-        let name = attribute(&raw[cursor..tag_end], "name")
+                .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?;
-            .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
+            let name = attribute(&raw[cursor..tag_end], "name")
-        cursor = tag_end;
+                .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?;
-        let mut arguments = Map::new();
+            cursor = tag_end;
-        loop {
+            let mut arguments = Map::new();
-            skip_whitespace(raw, &mut cursor);
+            loop {
-            if raw[cursor..].starts_with(syntax.invoke_end) {
+                skip_whitespace(raw, &mut cursor);
-                cursor += syntax.invoke_end.len();
+                if raw[cursor..].starts_with(syntax.invoke_end) {
+                    cursor += syntax.invoke_end.len();
+                    break;
+                }
+                let (name, value) = parse_parameter(raw, &mut cursor, syntax)?
+                    .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?;
+                arguments.insert(name, value);
+            }
+            Ok((name, Value::Object(arguments)))
+        })();
+        match call {
+            Ok(call) => calls.push(call),
+            Err(error) if !calls.is_empty() && error == INCOMPLETE_TOOL_CALL => {
                 break;
             }
-            let (name, value) = parse_parameter(raw, &mut cursor, syntax)?
+            Err(error) => return Err(error),
-                .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?;
-            arguments.insert(name, value);
         }
-        calls.push((name, Value::Object(arguments)));
     }
     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,39 +124,33 @@ impl GenerationService {
         source: WorkSource,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
-        self.submit(
+        self.submit(CommandRequest {
-            CommandRequest {
+            engine,
-                engine,
+            turn,
-                turn,
+            messages,
-                messages,
+            checkpoint: Some(checkpoint),
-                checkpoint: Some(checkpoint),
+            source,
-                source,
+            operation: Operation::Generate,
-                operation: Operation::Generate,
+            idle_timeout,
-                idle_timeout,
+        })
-            },
-            TrackingPolicy::QueuedAndRejected,
-        )
     }
 
     pub(crate) fn compact(&self, input: CompactionInput) -> Result<ActiveGeneration, String> {
-        self.submit(
+        self.submit(CommandRequest {
-            CommandRequest {
+            engine: input.engine,
-                engine: input.engine,
+            turn: input.turn,
-                turn: input.turn,
+            messages: input.messages,
-                messages: input.messages,
+            checkpoint: Some(CheckpointTarget::Local {
-                checkpoint: Some(CheckpointTarget::Local {
+                checkpoint: input.checkpoint,
-                    checkpoint: input.checkpoint,
+                bootstrap: None,
-                    bootstrap: None,
+            }),
-                }),
+            source: WorkSource::LocalChat,
-                source: WorkSource::LocalChat,
+            operation: Operation::Compact {
-                operation: Operation::Compact {
+                reason: input.reason,
-                    reason: input.reason,
+                rebuild_system_prompt: input.rebuild_system_prompt,
-                    rebuild_system_prompt: input.rebuild_system_prompt,
-                },
-                idle_timeout: input.idle_timeout,
             },
-            TrackingPolicy::None,
+            idle_timeout: input.idle_timeout,
-        )
+        })
     }
 
     pub(crate) fn measure_context(
@@ -179,28 +160,22 @@ impl GenerationService {
         messages: Vec<ChatTurn>,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
-        self.submit(
+        self.submit(CommandRequest {
-            CommandRequest {
+            engine,
-                engine,
+            turn,
-                turn,
+            messages,
-                messages,
+            checkpoint: None,
-                checkpoint: None,
+            source: WorkSource::LocalChat,
-                source: WorkSource::LocalChat,
+            operation: Operation::Measure,
-                operation: Operation::Measure,
+            idle_timeout,
-                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());
+    }
 }