From 83905690e60a32c0f296ca0ce5a5f53e64b4b03a Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Wed, 29 Jul 2026 18:47:14 +0200 Subject: [PATCH] Fix refactor recovery and metrics --- src/dsml.rs | 68 +++++++++++++++++++-------- src/runtime.rs | 122 ++++++++++++++++++++++--------------------------- 2 files changed, 103 insertions(+), 87 deletions(-) diff --git a/src/dsml.rs b/src/dsml.rs index 78169c8..9973c10 100644 --- a/src/dsml.rs +++ b/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) { - return Err("invalid or incomplete DSML 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())?; - let name = attribute(&raw[cursor..tag_end], "name") - .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?; - cursor = tag_end; - let mut arguments = Map::new(); - loop { - skip_whitespace(raw, &mut cursor); - if raw[cursor..].starts_with(syntax.invoke_end) { - cursor += syntax.invoke_end.len(); + let call = (|| { + if !raw[cursor..].starts_with(syntax.invoke_start) { + return Err(INCOMPLETE_TOOL_CALL.into()); + } + let tag_end = raw[cursor..] + .find('>') + .map(|offset| cursor + offset + 1) + .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?; + let name = attribute(&raw[cursor..tag_end], "name") + .ok_or_else(|| INCOMPLETE_TOOL_CALL.to_owned())?; + cursor = tag_end; + let mut arguments = Map::new(); + loop { + skip_whitespace(raw, &mut cursor); + 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)? - .ok_or_else(|| "invalid or incomplete DSML tool call".to_owned())?; - arguments.insert(name, value); + Err(error) => return Err(error), } - 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 = "truncated"; + + let (_, calls) = parse_tool_calls(raw).unwrap(); + + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "first"); + + let invalid = "invalid"; + assert!( + parse_tool_calls(invalid) + .unwrap_err() + .starts_with("invalid DSML tool arguments") + ); + } } diff --git a/src/runtime.rs b/src/runtime.rs index 5c6e5e8..04caab0 100644 --- a/src/runtime.rs +++ b/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 { - self.submit( - CommandRequest { - engine, - turn, - messages, - checkpoint: Some(checkpoint), - source, - operation: Operation::Generate, - idle_timeout, - }, - TrackingPolicy::QueuedAndRejected, - ) + self.submit(CommandRequest { + engine, + turn, + messages, + checkpoint: Some(checkpoint), + source, + operation: Operation::Generate, + idle_timeout, + }) } pub(crate) fn compact(&self, input: CompactionInput) -> Result { - self.submit( - CommandRequest { - engine: input.engine, - turn: input.turn, - messages: input.messages, - checkpoint: Some(CheckpointTarget::Local { - checkpoint: input.checkpoint, - bootstrap: None, - }), - source: WorkSource::LocalChat, - operation: Operation::Compact { - reason: input.reason, - rebuild_system_prompt: input.rebuild_system_prompt, - }, - idle_timeout: input.idle_timeout, + self.submit(CommandRequest { + engine: input.engine, + turn: input.turn, + messages: input.messages, + checkpoint: Some(CheckpointTarget::Local { + checkpoint: input.checkpoint, + bootstrap: None, + }), + source: WorkSource::LocalChat, + operation: Operation::Compact { + reason: input.reason, + rebuild_system_prompt: input.rebuild_system_prompt, }, - TrackingPolicy::None, - ) + idle_timeout: input.idle_timeout, + }) } pub(crate) fn measure_context( @@ -179,28 +160,22 @@ impl GenerationService { messages: Vec, idle_timeout: Duration, ) -> Result { - self.submit( - CommandRequest { - engine, - turn, - messages, - checkpoint: None, - source: WorkSource::LocalChat, - operation: Operation::Measure, - idle_timeout, - }, - TrackingPolicy::Queued, - ) + self.submit(CommandRequest { + engine, + turn, + messages, + checkpoint: None, + source: WorkSource::LocalChat, + operation: Operation::Measure, + idle_timeout, + }) } - fn submit( - &self, - request: CommandRequest, - tracking: TrackingPolicy, - ) -> Result { + fn submit(&self, request: CommandRequest) -> Result { 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, metrics: Arc) { 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()); + } }