Finish long-running agent parity

src/runtime.rs

@@ -46,6 +46,7 @@ impl CheckpointTarget {
 
 pub(crate) enum GenerationEvent {
     Loading,
+    Activity(&'static str),
     Chunk {
         reasoning: bool,
         content: String,
@@ -57,6 +58,23 @@ pub(crate) enum GenerationEvent {
     },
     Finished(Result<GenerationOutput, String>),
     Compacted(Result<CompactionOutput, String>),
+    Measured(Result<u32, String>),
+}
+
+enum Operation {
+    Generate,
+    Compact {
+        reason: String,
+        rebuild_system_prompt: String,
+    },
+    Measure,
+}
+
+#[derive(Clone, Copy)]
+enum ResponseKind {
+    Generation,
+    Compaction,
+    Measurement,
 }
 
 struct Command {
@@ -64,7 +82,7 @@ struct Command {
     turn: TurnSettings,
     messages: Vec<ChatTurn>,
     checkpoint: CheckpointTarget,
-    compact_reason: Option<String>,
+    operation: Operation,
     idle_timeout: Duration,
     cancel: Arc<AtomicBool>,
     events: Sender<GenerationEvent>,
@@ -100,7 +118,7 @@ impl GenerationService {
                 turn,
                 messages,
                 checkpoint,
-                compact_reason: None,
+                operation: Operation::Generate,
                 idle_timeout,
                 cancel: Arc::clone(&cancel),
                 events,
@@ -116,14 +134,46 @@ impl GenerationService {
         })
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn compact(
         &self,
         engine: EngineSettings,
         turn: TurnSettings,
         messages: Vec<ChatTurn>,
         reason: &str,
+        rebuild_system_prompt: String,
         checkpoint: PathBuf,
         idle_timeout: Duration,
+    ) -> Result<ActiveGeneration, String> {
+        let cancel = Arc::new(AtomicBool::new(false));
+        let (events, receiver) = mpsc::channel();
+        self.commands
+            .send(Command {
+                engine,
+                turn,
+                messages,
+                checkpoint: CheckpointTarget::Local(checkpoint),
+                operation: Operation::Compact {
+                    reason: reason.to_owned(),
+                    rebuild_system_prompt,
+                },
+                idle_timeout,
+                cancel: Arc::clone(&cancel),
+                events,
+            })
+            .map_err(|_| "The model runtime stopped unexpectedly.".to_owned())?;
+        Ok(ActiveGeneration {
+            events: receiver,
+            cancel,
+        })
+    }
+
+    pub(crate) fn measure_context(
+        &self,
+        engine: EngineSettings,
+        turn: TurnSettings,
+        messages: Vec<ChatTurn>,
+        idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
         let cancel = Arc::new(AtomicBool::new(false));
         let (events, receiver) = mpsc::channel();
@@ -133,8 +183,8 @@ impl GenerationService {
                 engine,
                 turn,
                 messages,
-                checkpoint: CheckpointTarget::Local(checkpoint),
-                compact_reason: Some(reason.to_owned()),
+                checkpoint: CheckpointTarget::OneShot(PathBuf::new()),
+                operation: Operation::Measure,
                 idle_timeout,
                 cancel: Arc::clone(&cancel),
                 events,
@@ -157,7 +207,11 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
                 let request_started = Instant::now();
                 let source = command.checkpoint.source();
                 let events = command.events.clone();
-                metrics.request_started(source);
+                let response = response_kind(&command.operation);
+                let tracked = !matches!(command.operation, Operation::Measure);
+                if tracked {
+                    metrics.request_started(source);
+                }
                 if let Err(error) = catch_runtime_panic(|| {
                     run_command(
                         command,
@@ -169,11 +223,13 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
                         &mut idle_timeout,
                     );
                 }) {
-                    metrics.request_failed(request_started.elapsed());
+                    if tracked {
+                        metrics.request_failed(request_started.elapsed());
+                    }
                     if loaded.take().is_some() {
                         metrics.unloaded();
                     }
-                    let _ = events.send(GenerationEvent::Finished(Err(error)));
+                    let _ = events.send(error_event(response, error));
                 }
             }
             Err(mpsc::RecvTimeoutError::Timeout) => {
@@ -197,12 +253,16 @@ fn run_command(
     last_used: &mut Instant,
     idle_timeout: &mut Duration,
 ) {
+    let response = response_kind(&command.operation);
+    let tracked = !matches!(command.operation, Operation::Measure);
     *idle_timeout = command.idle_timeout;
     if command.cancel.load(Ordering::Relaxed) {
-        metrics.request_failed(request_started.elapsed());
-        let _ = command.events.send(GenerationEvent::Finished(
-            Err("generation cancelled".into()),
-        ));
+        if tracked {
+            metrics.request_failed(request_started.elapsed());
+        }
+        let _ = command
+            .events
+            .send(error_event(response, "generation cancelled".into()));
         return;
     }
     if loaded
@@ -228,8 +288,10 @@ fn run_command(
                 Some((command.engine.clone(), generator))
             }
             Err(error) => {
-                metrics.request_failed(request_started.elapsed());
-                let _ = command.events.send(GenerationEvent::Finished(Err(error)));
+                if tracked {
+                    metrics.request_failed(request_started.elapsed());
+                }
+                let _ = command.events.send(error_event(response, error));
                 None
             }
         };
@@ -261,17 +323,25 @@ fn run_command(
                 tokens_per_second,
             });
         };
-        if let Some(reason) = &command.compact_reason {
+        if let Operation::Compact {
+            reason,
+            rebuild_system_prompt,
+        } = &command.operation
+        {
             let CheckpointTarget::Local(checkpoint) = &command.checkpoint else {
                 unreachable!("compaction checkpoints are local")
             };
             let result = generator.compact(
                 &command.messages,
                 &command.turn,
+                rebuild_system_prompt,
                 reason,
                 checkpoint,
                 &command.cancel,
                 &mut progress,
+                |activity| {
+                    let _ = command.events.send(GenerationEvent::Activity(activity));
+                },
             );
             match &result {
                 Ok(_) => {
@@ -283,6 +353,12 @@ fn run_command(
             *last_used = Instant::now();
             return;
         }
+        if matches!(command.operation, Operation::Measure) {
+            let result = generator.rendered_history_tokens(&command.messages, &command.turn);
+            let _ = command.events.send(GenerationEvent::Measured(result));
+            *last_used = Instant::now();
+            return;
+        }
         let result = match command.checkpoint {
             CheckpointTarget::Local(path) => generator.generate(
                 &path,
@@ -320,6 +396,22 @@ fn run_command(
     }
 }
 
+fn response_kind(operation: &Operation) -> ResponseKind {
+    match operation {
+        Operation::Generate => ResponseKind::Generation,
+        Operation::Compact { .. } => ResponseKind::Compaction,
+        Operation::Measure => ResponseKind::Measurement,
+    }
+}
+
+fn error_event(kind: ResponseKind, error: String) -> GenerationEvent {
+    match kind {
+        ResponseKind::Generation => GenerationEvent::Finished(Err(error)),
+        ResponseKind::Compaction => GenerationEvent::Compacted(Err(error)),
+        ResponseKind::Measurement => GenerationEvent::Measured(Err(error)),
+    }
+}
+
 fn catch_runtime_panic<T>(operation: impl FnOnce() -> T) -> Result<T, String> {
     // Generator owns GPU handles and is not unwind-safe; callers discard it on error.
     std::panic::catch_unwind(std::panic::AssertUnwindSafe(operation))
@@ -351,4 +443,16 @@ mod tests {
         );
         assert_eq!(catch_runtime_panic(|| 42), Ok(42));
     }
+
+    #[test]
+    fn operation_failures_use_the_matching_event() {
+        assert!(matches!(
+            error_event(ResponseKind::Compaction, "stopped".into()),
+            GenerationEvent::Compacted(Err(error)) if error == "stopped"
+        ));
+        assert!(matches!(
+            error_event(ResponseKind::Measurement, "stopped".into()),
+            GenerationEvent::Measured(Err(error)) if error == "stopped"
+        ));
+    }
 }