fix: preserve pending usage across pause and compact

crates/agen/src/engine.rs

@@ -2275,6 +2275,9 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                         info!("Paused during response stream");
                     }
                     self.timeline.abort_current_block();
+                    // Preserve any UsageEvent already received for billing, but
+                    // do not treat this paused request as a normal completion.
+                    self.timeline.flush_usage();
                     return Err(EngineError::PauseRequested);
                 }
                 cancel = self.cancel_rx.recv() => {

crates/agen/tests/engine_state_test.rs

@@ -7,6 +7,7 @@ mod common;
 
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
+use std::time::Duration;
 
 use agen::Item;
 use agen::interceptor::{
@@ -18,16 +19,67 @@ use agen::interceptor::{
 };
 use agen::llm_client::{
     ClientError, LlmClient, Request, ResponseStream,
-    event::{Event, ResponseStatus, StatusEvent},
+    event::{Event, ResponseStatus, StatusEvent, UsageEvent},
 };
 use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
 use agen::{Engine, EngineError, EngineRunExit, History, RunInterruptionReason};
 use async_trait::async_trait;
 use common::MockLlmClient;
+use futures::StreamExt;
 
-// =============================================================================
-// Mutable State Tests
-// =============================================================================
+#[derive(Clone)]
+struct UsageThenPendingClient;
+
+#[async_trait]
+impl LlmClient for UsageThenPendingClient {
+    fn clone_boxed(&self) -> Box<dyn LlmClient> {
+        Box::new(self.clone())
+    }
+
+    async fn stream(&self, _request: Request) -> Result<ResponseStream, ClientError> {
+        let usage = Event::Usage(UsageEvent {
+            input_tokens: Some(90),
+            output_tokens: Some(2),
+            total_tokens: Some(92),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        });
+        Ok(Box::pin(
+            futures::stream::once(async move { Ok(usage) })
+                .chain(futures::stream::pending::<Result<Event, ClientError>>()),
+        ))
+    }
+}
+
+#[tokio::test]
+async fn pause_during_response_stream_flushes_observed_usage() {
+    let observed = Arc::new(Mutex::new(Vec::new()));
+    let observed_for_callback = Arc::clone(&observed);
+    let mut engine = Engine::new(UsageThenPendingClient);
+    engine.on_usage(move |usage| {
+        observed_for_callback.lock().unwrap().push(usage.clone());
+    });
+    let pause = engine.pause_sender();
+    let mut history = History::new();
+    let mut engine = engine.lock(&history);
+
+    tokio::spawn(async move {
+        tokio::time::sleep(Duration::from_millis(20)).await;
+        pause.send(()).await.unwrap();
+    });
+
+    let exit = tokio::time::timeout(
+        Duration::from_secs(2),
+        engine.run(&mut history, "pause after usage"),
+    )
+    .await
+    .expect("run should observe pause");
+    assert!(matches!(exit, EngineRunExit::Paused));
+    let observed = observed.lock().unwrap();
+    assert_eq!(observed.len(), 1);
+    assert_eq!(observed[0].input_tokens, Some(90));
+    assert_eq!(observed[0].output_tokens, Some(2));
+}
 
 /// Verify that system prompt can be set in Mutable state
 #[test]

crates/worker/src/compact/usage_tracker.rs

@@ -214,6 +214,14 @@ impl UsageTracker {
             .collect()
     }
 
+    /// Number of measurements that have not crossed the durable commit
+    /// boundary. Session rewrites must not activate a replacement Segment while
+    /// this is non-zero, because their history coordinates belong to the current
+    /// Segment.
+    pub(crate) fn pending_record_count(&self) -> usize {
+        self.state.lock().unwrap().pending_records.len()
+    }
+
     /// Remove the oldest measurement for one durable commit attempt.
     pub(crate) fn take_next(&self) -> Option<RecordedUsage> {
         self.state

crates/worker/src/worker.rs

@@ -3355,8 +3355,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
     /// Handle to the synchronous `MetricsTracker` buffer.
     ///
     /// Engine callbacks (e.g. the prune observer) clone this `Arc` and
-    /// `.push(metric)` into it; Worker drains it in `persist_turn` and
-    /// writes each metric via `session_metrics::record_metric`.
+    /// `.push(metric)` into it; Worker drains it at the next request accounting
+    /// boundary, with terminal persistence as fallback.
     pub(crate) fn metrics_tracker_handle(
         &self,
     ) -> Arc<crate::compact::metrics_tracker::MetricsTracker> {
@@ -5146,6 +5146,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
         mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
         trigger: CompactionTrigger,
     ) -> Result<SegmentId, WorkerError> {
+        let pending_records = self.usage_tracker.pending_record_count();
+        if pending_records > 0 {
+            return Err(WorkerError::PendingUsagePersistence { pending_records });
+        }
         self.release_pending_compaction_service().await?;
         let _rewrite_guard = self
             .prepare_session_rewrite(SessionRewriteKind::Compact)
@@ -7453,7 +7457,9 @@ fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory {
         | WorkerError::SegmentActivationIncomplete { .. } => {
             CompactFailureCategory::ActiveSegmentCommit
         }
-        WorkerError::Store(_) => CompactFailureCategory::Storage,
+        WorkerError::Store(_) | WorkerError::PendingUsagePersistence { .. } => {
+            CompactFailureCategory::Storage
+        }
         WorkerError::InvalidState(_) | WorkerError::Engine(_) => {
             CompactFailureCategory::InternalWorker
         }
@@ -7517,6 +7523,11 @@ pub enum WorkerError {
     #[error(transparent)]
     Provider(#[from] crate::model_client::ProviderError),
 
+    #[error(
+        "cannot replace the active Segment while {pending_records} usage record(s) await durable persistence"
+    )]
+    PendingUsagePersistence { pending_records: usize },
+
     #[error("active Segment changed before compaction commit")]
     CompactActiveSegmentChanged,
 
@@ -8574,6 +8585,90 @@ mod build_summary_prompt_tests {
         }
     }
 
+    #[tokio::test]
+    async fn failed_usage_append_fences_segment_rotation_until_retry_commits_once() {
+        let dir = tempfile::tempdir().unwrap();
+        let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
+        let mut worker = Worker::new(
+            minimal_manifest(),
+            Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
+            store.clone(),
+            WorkerWorkspaceContext::no_workspace(),
+            WorkerFilesystemAuthority::None,
+            Scope::empty(),
+        )
+        .await
+        .unwrap();
+        worker.ensure_segment_head().await.unwrap();
+        let session_id = worker.session_id();
+        let source_segment = worker.segment_id();
+
+        worker
+            .usage_tracker
+            .note_compaction_correlation_id("compact-id".into());
+        worker.usage_tracker.note_request(1);
+        worker.usage_tracker.record_usage(&agen::event::UsageEvent {
+            input_tokens: Some(10),
+            output_tokens: Some(1),
+            total_tokens: Some(11),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        });
+        worker.usage_tracker.request_completed();
+
+        let append_error = persist_pending_usage(
+            &worker.usage_tracker,
+            &worker.usage_history,
+            worker.compact_state.as_deref(),
+            |_, _, _| Err(StoreError::Io(std::io::Error::other("synthetic failure"))),
+        )
+        .expect_err("first usage append must fail");
+        assert!(append_error.to_string().contains("synthetic failure"));
+        assert_eq!(worker.usage_tracker.pending_record_count(), 1);
+
+        let error = worker
+            .compact(0)
+            .await
+            .expect_err("pending usage must fence Segment replacement");
+        assert!(matches!(
+            error,
+            WorkerError::PendingUsagePersistence { pending_records: 1 }
+        ));
+        assert_eq!(worker.segment_id(), source_segment);
+        assert_eq!(
+            store.list_segments(session_id).unwrap(),
+            vec![source_segment]
+        );
+
+        let metrics = persist_pending_usage(
+            &worker.usage_tracker,
+            &worker.usage_history,
+            worker.compact_state.as_deref(),
+            |entry, after_commit, _| {
+                store.append(session_id, source_segment, &entry)?;
+                after_commit();
+                Ok(())
+            },
+        )
+        .expect("retry should commit on the originating Segment");
+        assert_eq!(worker.usage_tracker.pending_record_count(), 0);
+        assert_eq!(worker.usage_history.lock().unwrap().len(), 1);
+        assert_eq!(metrics.len(), 1, "correlation must be emitted once");
+        assert_eq!(metrics[0].name, "compact.post_request");
+
+        let usage_entries = store
+            .read_all(session_id, source_segment)
+            .unwrap()
+            .into_iter()
+            .filter(|entry| matches!(entry, LogEntry::LlmUsage { .. }))
+            .count();
+        assert_eq!(usage_entries, 1, "retry must not duplicate accounting");
+        assert_eq!(
+            store.list_segments(session_id).unwrap(),
+            vec![source_segment]
+        );
+    }
+
     #[tokio::test]
     async fn fresh_run_clears_interrupted_budget_before_pre_run_compaction() {
         let dir = tempfile::tempdir().unwrap();

crates/worker/tests/controller_test.rs

@@ -3,7 +3,9 @@ use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 
 use agen::Engine;
-use agen::llm_client::event::{ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent};
+use agen::llm_client::event::{
+    ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent,
+};
 use agen::llm_client::types::Item;
 use agen::llm_client::{ClientError, LlmClient, Request};
 use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
@@ -2580,18 +2582,32 @@ async fn drain_until<F: FnMut(&Event) -> bool>(
 /// the interrupted turn's history consistency.
 #[tokio::test]
 async fn pause_then_resume_preserves_notifications_and_history_consistency() {
-    // Response 1: hang after opening a text block (no stop / completed),
-    // so the Engine is parked inside the stream read and `cancel_rx`
-    // races it cleanly on Method::Pause.
+    // Response 1: report billable partial usage, then hang after opening a text
+    // block (no stop / completed). Pause must flush that usage without treating
+    // the provider request as normally completed.
     let hang = MockResponse::Hang(vec![
         LlmEvent::text_block_start(0),
         LlmEvent::text_delta(0, "partial..."),
+        LlmEvent::Usage(UsageEvent {
+            input_tokens: Some(90),
+            output_tokens: Some(2),
+            total_tokens: Some(92),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        }),
     ]);
     // Response 2: a clean assistant reply delivered on Resume.
     let ok = MockResponse::Complete(vec![
         LlmEvent::text_block_start(0),
         LlmEvent::text_delta(0, "resumed output"),
         LlmEvent::text_block_stop(0, None),
+        LlmEvent::Usage(UsageEvent {
+            input_tokens: Some(100),
+            output_tokens: Some(3),
+            total_tokens: Some(103),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        }),
         LlmEvent::Status(StatusEvent {
             status: ResponseStatus::Completed,
         }),
@@ -2645,6 +2661,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() {
 
     tokio::time::sleep(std::time::Duration::from_millis(50)).await;
     assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Paused);
+    let (paused_entries, _) = handle.sink.subscribe_with_snapshot();
+    let paused_usage = paused_entries
+        .iter()
+        .filter_map(|entry| match entry {
+            LogEntry::LlmUsage {
+                input_total_tokens,
+                output_tokens,
+                ..
+            } => Some((*input_total_tokens, *output_tokens)),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(
+        paused_usage,
+        vec![(90, 2)],
+        "Pause must durably account for usage already observed on the interrupted request"
+    );
 
     handle
         .send(Method::Notify {
@@ -2729,6 +2762,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() {
         })
         .unwrap_or_default();
     assert_eq!(assistant_text, "resumed output");
+    let (final_entries, _) = handle.sink.subscribe_with_snapshot();
+    let final_usage = final_entries
+        .iter()
+        .filter_map(|entry| match entry {
+            LogEntry::LlmUsage {
+                input_total_tokens,
+                output_tokens,
+                ..
+            } => Some((*input_total_tokens, *output_tokens)),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(
+        final_usage,
+        vec![(90, 2), (100, 3)],
+        "the paused request must remain a distinct billing record rather than merging into resume"
+    );
     let has_tool_call = history.iter().any(|i| i.is_tool_call());
     assert!(!has_tool_call, "no orphan tool_call in history");
 }