From 2dda4b90b4d6c7f7029342d51c2705983299855a Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 20 Sep 2026 19:17:33 +0900 Subject: [PATCH] fix: preserve pending usage across pause and compact --- crates/agen/src/engine.rs | 3 + crates/agen/tests/engine_state_test.rs | 60 +++++++++++- crates/worker/src/compact/usage_tracker.rs | 8 ++ crates/worker/src/worker.rs | 101 ++++++++++++++++++++- crates/worker/tests/controller_test.rs | 58 +++++++++++- 5 files changed, 219 insertions(+), 11 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index da3cead9d..2235614bc 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -2275,6 +2275,9 @@ impl Engine { 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() => { diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index 5060015d5..051f94394 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/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 { + Box::new(self.clone()) + } + + async fn stream(&self, _request: Request) -> Result { + 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::>()), + )) + } +} + +#[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] diff --git a/crates/worker/src/compact/usage_tracker.rs b/crates/worker/src/compact/usage_tracker.rs index 24323e9e9..b0d068f5c 100644 --- a/crates/worker/src/compact/usage_tracker.rs +++ b/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 { self.state diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index cde3d8b3a..298067a13 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -3355,8 +3355,8 @@ impl Worker { /// 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 { @@ -5146,6 +5146,10 @@ impl Worker { mut cancel: Option>, trigger: CompactionTrigger, ) -> Result { + 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(); diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 2193945fa..849291347 100644 --- a/crates/worker/tests/controller_test.rs +++ b/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 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::>(); + 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::>(); + 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"); }