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/metrics_tracker.rs b/crates/worker/src/compact/metrics_tracker.rs index c0e81033c..dcea2131a 100644 --- a/crates/worker/src/compact/metrics_tracker.rs +++ b/crates/worker/src/compact/metrics_tracker.rs @@ -2,9 +2,8 @@ //! Engine callbacks (which run synchronously and cannot themselves //! perform `async` store writes). //! -//! Worker drains this buffer in `persist_turn` and writes each metric via -//! `session_metrics::record_metric`, alongside the regular `LlmUsage` -//! entries. +//! Worker drains this buffer at the next successful request accounting +//! boundary, with terminal `persist_turn` as a fallback. use std::sync::Mutex; @@ -21,12 +20,12 @@ impl MetricsTracker { } } - /// Queue a metric for the next `persist_turn` flush. + /// Queue a metric for the next request or terminal accounting flush. pub(crate) fn push(&self, metric: Metric) { self.pending.lock().unwrap().push(metric); } - /// Drain all queued metrics. Called by Worker after a run completes. + /// Drain all queued metrics at an accounting boundary. pub(crate) fn drain(&self) -> Vec { std::mem::take(&mut *self.pending.lock().unwrap()) } diff --git a/crates/worker/src/compact/usage_tracker.rs b/crates/worker/src/compact/usage_tracker.rs index 3310ca338..b0d068f5c 100644 --- a/crates/worker/src/compact/usage_tracker.rs +++ b/crates/worker/src/compact/usage_tracker.rs @@ -7,17 +7,24 @@ //! - **`on_usage` callback** (sync closure): receives the aggregated final //! `UsageEvent` for that request after the stream completes. //! -//! Pairing the two yields one `UsageRecord` per LLM call. Worker drains them -//! in `persist_turn` and writes them as `LogEntry::LlmUsage` entries. +//! Pairing the two yields one `UsageRecord` per LLM call. Worker commits +//! successful requests from the assistant-turn boundary and uses terminal +//! `persist_turn` only for records that have not already crossed that boundary. //! //! Multiple LLM calls per Worker run (tool loop) are supported: each call //! produces its own `(history_len, UsageEvent)` pair, and the records are //! buffered in chronological order. +use std::collections::VecDeque; use std::sync::Mutex; use agen::UsageRecord; use agen::timeline::event::UsageEvent; +use session_metrics::Metric; +use session_store::{LogEntry, StoreError, segment_log}; + +use crate::compact::state::CompactState; +use crate::compact::telemetry::correlated_post_request_metric; /// The metric emitted after the next measured provider request. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -41,11 +48,15 @@ pub(crate) struct PostRequestLink { pub(crate) metric: PostRequestMetric, } -/// One drained measurement and its causal metric links. +/// One measured request and its causal metric links. #[derive(Debug, Clone)] pub(crate) struct RecordedUsage { pub(crate) record: UsageRecord, pub(crate) post_requests: Vec, + /// True only when the provider stream completed normally and reported an + /// input occupancy. Partial/error/cancelled streams remain billable usage + /// records, but cannot satisfy a post-compaction rearm. + pub(crate) completed_with_occupancy: bool, } #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -56,29 +67,52 @@ pub(crate) struct UsageSnapshot { pub(crate) output_tokens: u64, } +#[derive(Debug)] +struct PendingRequest { + id: u64, + history_len: usize, + usage_recorded: bool, +} + +#[derive(Debug)] +struct PendingRecordedUsage { + request_id: u64, + has_input_occupancy: bool, + usage: RecordedUsage, +} + +#[derive(Debug, Default)] +struct UsageTrackerState { + next_request_id: u64, + pending_request: Option, + pending_correlations: Vec, + pending_records: VecDeque, +} + /// Shared between the pre-request hook, the `on_usage` callback, and Worker. pub(crate) struct UsageTracker { - /// `history.len()` captured at the most recent `pre_llm_request`. - /// Cleared when paired with an incoming `on_usage` event. - pending_history_len: Mutex>, - /// Optional causal link consumed by the next measured request. - pending_correlations: Mutex>, - /// Records accumulated during the current run; drained by Worker. - pending_records: Mutex>, + state: Mutex, } impl UsageTracker { pub(crate) fn new() -> Self { Self { - pending_history_len: Mutex::new(None), - pending_correlations: Mutex::new(Vec::new()), - pending_records: Mutex::new(Vec::new()), + state: Mutex::new(UsageTrackerState::default()), } } /// Called from a `pre_llm_request` hook with the current history length. + /// Replacing an older pending request classifies that request as incomplete; + /// any partial usage already recorded for it remains available for billing. pub(crate) fn note_request(&self, history_len: usize) { - *self.pending_history_len.lock().unwrap() = Some(history_len); + let mut state = self.state.lock().unwrap(); + let id = state.next_request_id; + state.next_request_id = state.next_request_id.wrapping_add(1); + state.pending_request = Some(PendingRequest { + id, + history_len, + usage_recorded: false, + }); } /// Pair a prune event with the next provider request. @@ -92,25 +126,31 @@ impl UsageTracker { } fn note_post_request(&self, id: String, metric: PostRequestMetric) { - let mut pending = self.pending_correlations.lock().unwrap(); - pending.retain(|link| link.metric != metric); - pending.push(PostRequestLink { + let mut state = self.state.lock().unwrap(); + state + .pending_correlations + .retain(|link| link.metric != metric); + state.pending_correlations.push(PostRequestLink { correlation_id: id, metric, }); } /// Called from the `on_usage` callback with the aggregated final - /// UsageEvent. If a `history_len` was previously stashed via - /// `note_request`, builds a `RecordedUsage` and pushes it onto the - /// buffer. If not (e.g. test code that fires Usage outside a request), - /// drops the event. + /// UsageEvent. The record remains associated with the current request until + /// [`Self::request_completed`] proves that the provider stream completed + /// normally. Error/cancel paths still retain the measurement for billing. pub(crate) fn record_usage(&self, event: &UsageEvent) { - let history_len = match self.pending_history_len.lock().unwrap().take() { - Some(n) => n, - None => return, + let mut state = self.state.lock().unwrap(); + let Some(pending) = state.pending_request.as_mut() else { + return; }; - let post_requests = std::mem::take(&mut *self.pending_correlations.lock().unwrap()); + if pending.usage_recorded { + return; + } + pending.usage_recorded = true; + let request_id = pending.id; + let history_len = pending.history_len; // UsageEvent.input_tokens は scheme 層で「占有量(プロンプト全長)」に // 正規化済みである前提(Anthropic は cache_read + cache_creation を // 加算して emit する)。 @@ -118,35 +158,166 @@ impl UsageTracker { let cache_read = event.cache_read_input_tokens.unwrap_or(0); let cache_write = event.cache_creation_input_tokens.unwrap_or(0); let output = event.output_tokens.unwrap_or(0); - self.pending_records.lock().unwrap().push(RecordedUsage { - record: UsageRecord { - history_len, - input_total_tokens: input_total, - cache_read_tokens: cache_read, - cache_write_tokens: cache_write, - output_tokens: output, + state.pending_records.push_back(PendingRecordedUsage { + request_id, + has_input_occupancy: event.input_tokens.is_some(), + usage: RecordedUsage { + record: UsageRecord { + history_len, + input_total_tokens: input_total, + cache_read_tokens: cache_read, + cache_write_tokens: cache_write, + output_tokens: output, + }, + post_requests: Vec::new(), + completed_with_occupancy: false, }, - post_requests, }); } - /// Return a clone of the accumulated `UsageRecord`s without clearing them. + /// Mark the current request as normally completed. A post-request causal + /// link is consumed only when that exact request also reported a valid + /// input occupancy; callbacks from partial/error/cancelled streams cannot + /// rearm the compaction guard. + pub(crate) fn request_completed(&self) { + let mut state = self.state.lock().unwrap(); + let Some(pending) = state.pending_request.take() else { + return; + }; + let Some(recorded_index) = state + .pending_records + .iter() + .rposition(|recorded| recorded.request_id == pending.id) + else { + return; + }; + let has_occupancy = state.pending_records[recorded_index].has_input_occupancy; + if !has_occupancy { + return; + } + let post_requests = std::mem::take(&mut state.pending_correlations); + let recorded = &mut state.pending_records[recorded_index].usage; + recorded.completed_with_occupancy = true; + recorded.post_requests = post_requests; + } + + /// Return a clone of accumulated measurements without clearing them. /// Used by request-time circuit breakers that need the same occupancy /// projection as Worker persistence while the run is still active. pub(crate) fn records(&self) -> Vec { - self.pending_records + self.state .lock() .unwrap() + .pending_records .iter() - .map(|r| r.record.clone()) + .map(|r| r.usage.record.clone()) .collect() } - /// Drain accumulated records. Called by Worker after a run completes, - /// before persisting the turn. - pub(crate) fn drain(&self) -> Vec { - std::mem::take(&mut *self.pending_records.lock().unwrap()) + /// 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 + .lock() + .unwrap() + .pending_records + .pop_front() + .map(|recorded| recorded.usage) + } + + /// Restore a measurement after its durable commit failed. + pub(crate) fn restore_front(&self, usage: RecordedUsage) { + self.state + .lock() + .unwrap() + .pending_records + .push_front(PendingRecordedUsage { + // Persist ordering is the only remaining identity requirement. + request_id: u64::MAX, + has_input_occupancy: usage.completed_with_occupancy, + usage, + }); + } + + /// Drain accumulated records at a terminal Worker boundary. + pub(crate) fn drain(&self) -> Vec { + let mut state = self.state.lock().unwrap(); + state.pending_request = None; + state + .pending_records + .drain(..) + .map(|recorded| recorded.usage) + .collect() + } +} + +/// Commit all currently measured usage records in request order. +/// +/// A record is removed from the tracker only after its `LlmUsage` append +/// succeeds. This is the single accounting boundary shared by the in-run +/// interceptor and terminal `persist_turn`, preventing both loss on append +/// failure and duplicate terminal writes. Correlated metrics are handed to the +/// same commit boundary so they stay on the usage record's Segment; the return +/// value is only useful as test/diagnostic evidence and must not be re-emitted. +pub(crate) fn persist_pending_usage( + tracker: &UsageTracker, + usage_history: &Mutex>, + compact_state: Option<&CompactState>, + mut commit: impl FnMut(LogEntry, &mut dyn FnMut(), &[Metric]) -> Result<(), StoreError>, +) -> Result, StoreError> { + let mut metrics = Vec::new(); + while let Some(recorded) = tracker.take_next() { + let RecordedUsage { + record, + post_requests, + completed_with_occupancy, + } = recorded; + let post_metrics = post_requests + .iter() + .map(|link| correlated_post_request_metric(link.metric, &link.correlation_id, &record)) + .collect::>(); + let usage_entry = LogEntry::LlmUsage { + ts: segment_log::now_millis(), + history_len: record.history_len, + input_total_tokens: record.input_total_tokens, + cache_read_tokens: record.cache_read_tokens, + cache_write_tokens: record.cache_write_tokens, + output_tokens: record.output_tokens, + }; + let mut publish_committed_usage = || { + usage_history + .lock() + .expect("usage_history poisoned") + .push(record.clone()); + + if completed_with_occupancy + && post_requests + .iter() + .any(|link| link.metric == PostRequestMetric::Compaction) + { + if let Some(state) = compact_state { + state.post_compact_request_committed(); + } + } + }; + if let Err(error) = commit(usage_entry, &mut publish_committed_usage, &post_metrics) { + tracker.restore_front(RecordedUsage { + record, + post_requests, + completed_with_occupancy, + }); + return Err(error); + } + metrics.extend(post_metrics); + } + Ok(metrics) } #[cfg(test)] @@ -230,6 +401,7 @@ mod tests { tracker.note_correlation_id("prune-id".into()); tracker.note_request(5); tracker.record_usage(&make_event(100, 10, 2, 20)); + tracker.request_completed(); let records = tracker.drain(); assert_eq!(records[0].post_requests.len(), 2); @@ -248,9 +420,11 @@ mod tests { tracker.note_correlation_id("abc".into()); tracker.note_request(5); tracker.record_usage(&make_event(100, 0, 0, 20)); + tracker.request_completed(); // Next request without a fresh stash → no correlation_id. tracker.note_request(10); tracker.record_usage(&make_event(200, 50, 0, 30)); + tracker.request_completed(); let records = tracker.drain(); assert_eq!(records.len(), 2); @@ -259,4 +433,130 @@ mod tests { assert_eq!(records[0].post_requests[0].metric, PostRequestMetric::Prune); assert!(records[1].post_requests.is_empty()); } + + #[test] + fn partial_usage_does_not_consume_post_compact_link() { + let tracker = UsageTracker::new(); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(5); + tracker.record_usage(&make_event(90, 0, 0, 2)); + + // Starting another request proves the prior stream did not complete. + tracker.note_request(6); + tracker.record_usage(&make_event(100, 0, 0, 3)); + tracker.request_completed(); + + let records = tracker.drain(); + assert_eq!(records.len(), 2); + assert!(!records[0].completed_with_occupancy); + assert!(records[0].post_requests.is_empty()); + assert!(records[1].completed_with_occupancy); + assert_eq!(records[1].post_requests.len(), 1); + assert_eq!(records[1].post_requests[0].correlation_id, "compact-id"); + } + + #[test] + fn completed_request_without_input_occupancy_does_not_consume_link() { + let tracker = UsageTracker::new(); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(5); + let mut missing_input = make_event(0, 0, 0, 2); + missing_input.input_tokens = None; + tracker.record_usage(&missing_input); + tracker.request_completed(); + + tracker.note_request(6); + tracker.record_usage(&make_event(100, 0, 0, 3)); + tracker.request_completed(); + + let records = tracker.drain(); + assert!(!records[0].completed_with_occupancy); + assert!(records[0].post_requests.is_empty()); + assert!(records[1].completed_with_occupancy); + assert_eq!(records[1].post_requests[0].correlation_id, "compact-id"); + } + + #[test] + fn completed_request_without_usage_does_not_complete_previous_partial_record() { + let tracker = UsageTracker::new(); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(5); + tracker.record_usage(&make_event(90, 0, 0, 2)); + + tracker.note_request(6); + tracker.request_completed(); + + let records = tracker.drain(); + assert_eq!(records.len(), 1); + assert!(!records[0].completed_with_occupancy); + assert!(records[0].post_requests.is_empty()); + } + + #[test] + fn failed_commit_can_restore_record_without_reordering() { + let tracker = UsageTracker::new(); + tracker.note_request(1); + tracker.record_usage(&make_event(10, 0, 0, 1)); + tracker.request_completed(); + tracker.note_request(2); + tracker.record_usage(&make_event(20, 0, 0, 2)); + tracker.request_completed(); + + let first = tracker.take_next().expect("first record"); + tracker.restore_front(first); + let records = tracker.drain(); + assert_eq!(records.len(), 2); + assert_eq!(records[0].record.input_total_tokens, 10); + assert_eq!(records[1].record.input_total_tokens, 20); + } + + #[test] + fn compaction_rearms_only_after_usage_append_succeeds() { + use crate::compact::state::{AutomaticCompactDecision, CompactionOutcome}; + + let tracker = UsageTracker::new(); + let state = CompactState::new(None, Some(10), 0); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Succeeded)); + + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(1); + tracker.record_usage(&make_event(5, 0, 0, 1)); + tracker.request_completed(); + let usage_history = Mutex::new(Vec::new()); + + let error = persist_pending_usage(&tracker, &usage_history, Some(&state), |_, _, _| { + Err(StoreError::Io(std::io::Error::other("synthetic failure"))) + }) + .expect_err("usage append must fail"); + assert!(error.to_string().contains("synthetic failure")); + assert!(usage_history.lock().unwrap().is_empty()); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Block(_) + )); + + let mut entries = Vec::new(); + let metrics = persist_pending_usage( + &tracker, + &usage_history, + Some(&state), + |entry, after_commit, _post_metrics| { + entries.push(entry); + after_commit(); + Ok(()) + }, + ) + .expect("retry should commit"); + assert_eq!(entries.len(), 1); + assert_eq!(usage_history.lock().unwrap().len(), 1); + assert_eq!(metrics.len(), 1); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + } } diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 773c2298a..802d0c7d7 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -23,11 +23,12 @@ use agen::interceptor::{ use agen::tool::ToolOutput; use arc_swap::ArcSwap; use async_trait::async_trait; -use tracing::info; +use tracing::{info, warn}; +use crate::compact::metrics_tracker::MetricsTracker; use crate::compact::state::{AutomaticCompactDecision, CompactState}; -use crate::compact::usage_tracker::UsageTracker; -use session_store::SystemItem; +use crate::compact::usage_tracker::{UsageTracker, persist_pending_usage}; +use session_store::{SystemItem, segment_log}; use crate::hook::{ HookEventKind, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, @@ -81,6 +82,10 @@ pub(crate) struct WorkerInterceptor { /// persisted into `usage_history`. Subsequent tool-loop LLM calls must /// see these records during pre-request safety accounting. usage_tracker: Option>, + /// Metrics produced synchronously while projecting the request (notably + /// prune.fire/skip). Flushed before the matching request usage so metric + /// correlation order remains stable at the in-run commit boundary. + metrics_tracker: Option>, /// Pending-notification buffer drained into `worker.history` /// via [`Self::pending_history_appends`] just before the next LLM /// request. The Engine `extend`s these into its persistent history @@ -152,6 +157,7 @@ impl WorkerInterceptor { compact_state, usage_history, usage_tracker: None, + metrics_tracker: None, pending_notifies, pending_attachments, prompts, @@ -168,6 +174,11 @@ impl WorkerInterceptor { self } + pub(crate) fn with_metrics_tracker(mut self, metrics_tracker: Arc) -> Self { + self.metrics_tracker = Some(metrics_tracker); + self + } + pub(crate) fn with_prompt_workspace_id(mut self, workspace_id: Option) -> Self { self.prompt_workspace_id = workspace_id; self @@ -272,6 +283,56 @@ impl WorkerInterceptor { } } } + fn try_persist_metric(writer: &Arc, metric: &session_metrics::Metric) { + let payload = serde_json::to_value(metric).expect("Metric is Serialize"); + if writer + .commit_log_entry(session_store::LogEntry::Extension { + ts: segment_log::now_millis(), + domain: session_metrics::DOMAIN.into(), + payload, + }) + .is_err() + { + warn!(name = %metric.name, "failed to record session metric; dropping"); + } + } + + fn persist_completed_request_usage(&self) -> InterceptorResult<()> { + let Some(tracker) = self.usage_tracker.as_ref() else { + return Ok(()); + }; + tracker.request_completed(); + + let (Some(writer), Some(usage_history)) = + (self.log_writer.as_ref(), self.usage_history.as_ref()) + else { + // Low-level Worker tests may omit the type-erased writer. Their + // terminal persist_turn remains the durable fallback. + return Ok(()); + }; + if let Some(metrics_tracker) = self.metrics_tracker.as_ref() { + for metric in metrics_tracker.drain() { + Self::try_persist_metric(writer, &metric); + } + } + + persist_pending_usage( + tracker, + usage_history, + self.compact_state.as_deref(), + |entry, after_commit, post_metrics| { + writer.commit_usage_boundary(entry, after_commit, post_metrics) + }, + ) + .map_err(|error| { + InterceptorError::new( + InterceptorErrorCategory::Dependency, + format!("usage persistence failed: {error}"), + ) + })?; + Ok(()) + } + fn attach_prompt_provenance(&self, items: &mut [SystemItem]) { let prompts = self.prompts.load(); let projection = prompts.projection(); @@ -596,6 +657,11 @@ impl Interceptor for WorkerInterceptor { &self, context: AssistantTurnEndContext<'_, SessionHistoryMetadata>, ) -> InterceptorResult { + // `on_assistant_turn_end` is reached only after a normal provider + // stream completed and its usage callback was flushed. Commit that + // exact request before tools can extend the same logical run. + self.persist_completed_request_usage()?; + let history = context.history; let final_text_preview = history .iter() @@ -752,6 +818,25 @@ mod tests { } } + struct FailingUsageCommitter { + attempts: Arc, + } + + impl SystemItemCommitter for FailingUsageCommitter { + fn commit_log_entry( + &self, + entry: session_store::LogEntry, + ) -> Result<(), session_store::StoreError> { + if matches!(entry, session_store::LogEntry::LlmUsage { .. }) { + self.attempts.fetch_add(1, Ordering::SeqCst); + return Err(session_store::StoreError::Io(std::io::Error::other( + "synthetic usage failure", + ))); + } + Ok(()) + } + } + struct AppendingPreRequestHook { saw_handle: Arc, } @@ -955,6 +1040,68 @@ mod tests { ); } + #[tokio::test] + async fn completed_request_usage_failure_stops_before_rearm() { + use crate::compact::state::{AutomaticCompactDecision, CompactionOutcome}; + + let registry = Arc::new(HookRegistryBuilder::new().build()); + let state = Arc::new(CompactState::new(None, Some(10), 0)); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Succeeded)); + + let usage_history = Arc::new(Mutex::new(Vec::new())); + let tracker = Arc::new(UsageTracker::new()); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(1); + tracker.record_usage(&agen::event::UsageEvent { + input_tokens: Some(5), + output_tokens: Some(1), + total_tokens: Some(6), + cache_read_input_tokens: Some(0), + cache_creation_input_tokens: Some(0), + }); + let attempts = Arc::new(AtomicUsize::new(0)); + let interceptor = WorkerInterceptor::new( + registry, + Some(Arc::clone(&state)), + Some(Arc::clone(&usage_history)), + NotifyBuffer::new(), + Arc::new(Mutex::new(Vec::new())), + test_prompts(), + Some(Arc::new(FailingUsageCommitter { + attempts: Arc::clone(&attempts), + })), + ) + .with_usage_tracker(Arc::clone(&tracker)); + + let error = interceptor + .on_assistant_turn_end(AssistantTurnEndContext { + invocation: Default::default(), + assistant_entries: &[], + history: &[], + tool_calls: &[], + }) + .await + .expect_err("usage storage failure must stop the request boundary"); + + assert_eq!(error.category(), InterceptorErrorCategory::Dependency); + assert!(error.diagnostic().contains("usage persistence failed")); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + assert_eq!( + tracker.records().len(), + 1, + "failed record must remain pending" + ); + assert!(usage_history.lock().unwrap().is_empty()); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Block(_) + )); + } + #[tokio::test] async fn pre_llm_request_counts_in_flight_usage_records() { let registry = Arc::new(HookRegistryBuilder::new().build()); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index e1a90d770..9463dcede 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -45,9 +45,9 @@ use crate::compact::state::{ }; use crate::compact::telemetry::{ CompactAttempt, CompactFailureCategory, CompactMode, CompactSuccessStats, - CompactThresholdPolicy, correlated_post_request_metric, new_compact_metric_correlation_id, + CompactThresholdPolicy, new_compact_metric_correlation_id, }; -use crate::compact::usage_tracker::UsageTracker; +use crate::compact::usage_tracker::{UsageTracker, persist_pending_usage}; use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTaskRegistry}; use crate::feature::builtin::memory::WorkspaceMemoryBackendError; use crate::feature::builtin::{TaskFeature, WorkerObservationProvider}; @@ -1947,6 +1947,28 @@ impl PendingSubmissionHandle { pub trait SystemItemCommitter: Send + Sync { fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError>; + /// Commit one authoritative usage record while holding the same append + /// barrier across its in-memory publication, guard rearm, and correlated + /// post-request metrics. Metric appends remain best-effort. + fn commit_usage_boundary( + &self, + entry: LogEntry, + after_commit: &mut dyn FnMut(), + post_metrics: &[session_metrics::Metric], + ) -> Result<(), StoreError> { + self.commit_log_entry(entry)?; + after_commit(); + for metric in post_metrics { + let payload = serde_json::to_value(metric).expect("Metric is Serialize"); + let _ = self.commit_log_entry(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: session_metrics::DOMAIN.into(), + payload, + }); + } + Ok(()) + } + fn commit_system_item_with_extensions( &self, item: SystemItem, @@ -1978,6 +2000,30 @@ where fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError> { self.append_entry(entry) } + + fn commit_usage_boundary( + &self, + entry: LogEntry, + after_commit: &mut dyn FnMut(), + post_metrics: &[session_metrics::Metric], + ) -> Result<(), StoreError> { + let _append_guard = self + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + self.append_entry_locked(entry)?; + after_commit(); + for metric in post_metrics { + let payload = serde_json::to_value(metric).expect("Metric is Serialize"); + let _ = self.append_entry_locked(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: session_metrics::DOMAIN.into(), + payload, + }); + } + Ok(()) + } } struct SessionFlowRuntimeStateCommitter { @@ -2155,16 +2201,16 @@ pub struct Worker { /// Shared compaction state (present when threshold is configured). compact_state: Option>, /// Per-LLM-request Usage tracker. Always present after construction. - /// Captures `(history_len, UsageEvent)` pairs during a run; drained - /// in `persist_turn` and persisted as `LogEntry::LlmUsage` entries. + /// Captures `(history_len, UsageEvent)` pairs until their authoritative + /// request-boundary or terminal `LlmUsage` append succeeds. usage_tracker: Arc, /// Sync-side buffer for `Metric` values queued from inside Engine - /// callbacks (currently the prune observer). Drained in `persist_turn` - /// and written via `session_metrics::record_metric` alongside - /// `LogEntry::LlmUsage`. Always present after construction. + /// callbacks (currently the prune observer). Drained at the same request + /// boundary as usage, with terminal `persist_turn` as the fallback. metrics_tracker: Arc, /// Cumulative Usage measurement timeline, one entry per LLM call. - /// Restored from session log on `restore`, appended on each persist. + /// Restored from session log on `restore`, appended after each durable + /// request-boundary usage commit. /// Read by token-accounting APIs (`Worker::total_tokens`, etc.). /// /// Wrapped in `Arc` so that callbacks injected into the @@ -3330,8 +3376,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 { @@ -3567,6 +3613,7 @@ impl Worker { self.pending_committed_history.clone(), ) .with_usage_tracker(self.usage_tracker.clone()) + .with_metrics_tracker(self.metrics_tracker.clone()) .with_prompt_workspace_id( self.workspace_context .workspace_id() @@ -4452,6 +4499,15 @@ impl Worker { self.handle_worker_result(result, history_before).await } + fn ensure_no_pending_usage_for_segment_replacement(&self) -> Result<(), WorkerError> { + let pending_records = self.usage_tracker.pending_record_count(); + if pending_records == 0 { + Ok(()) + } else { + Err(WorkerError::PendingUsagePersistence { pending_records }) + } + } + /// Ensure the session exists and the writer's tally still matches /// the on-disk entry count. /// @@ -4493,6 +4549,7 @@ impl Worker { if store_count == entries_written { return Ok(()); } + self.ensure_no_pending_usage_for_segment_replacement()?; // Auto-fork within the same Session: mint a fresh Segment and // switch to it. The source segment is left immutable (no terminal // marker is written back); the fork relationship is recorded @@ -5029,45 +5086,38 @@ impl Worker { self.try_record_metric(&metric); } - // Persist any LLM Usage measurements collected during this run. - // One LogEntry::LlmUsage per LLM call (the tool loop may have run - // many calls within a single Worker::run). Each is also appended to - // the in-memory `usage_history` so token-accounting APIs see it - // before the next run. Records carrying a `correlation_id` (set - // by an upstream observer such as the prune projection) also get - // a paired `prune.post_request` metric so cache_read/write can be - // joined back to the originating event. - let usage_records = self.usage_tracker.drain(); - for recorded in usage_records { - let crate::compact::usage_tracker::RecordedUsage { - record, - post_requests, - } = recorded; - let committed_post_compact_request = post_requests.iter().any(|link| { - link.metric == crate::compact::usage_tracker::PostRequestMetric::Compaction - }); - self.commit_entry(LogEntry::LlmUsage { - ts: segment_log::now_millis(), - history_len: record.history_len, - input_total_tokens: record.input_total_tokens, - cache_read_tokens: record.cache_read_tokens, - cache_write_tokens: record.cache_write_tokens, - output_tokens: record.output_tokens, - })?; - if committed_post_compact_request { - if let Some(state) = &self.compact_state { - state.post_compact_request_committed(); - } - } - for link in post_requests { - let metric = - correlated_post_request_metric(link.metric, &link.correlation_id, &record); + // Persist any request measurements that did not already cross the + // in-run request boundary (for example partial usage from an errored or + // cancelled stream, or low-level Workers without an attached writer). + // `persist_pending_usage` removes each record only after its authoritative + // LlmUsage append succeeds, so the in-run and terminal paths cannot + // double-write or silently lose a failed append. + let usage_tracker = self.usage_tracker.clone(); + let usage_history = self.usage_history.clone(); + let compact_state = self.compact_state.clone(); + if let Some(writer) = self.log_writer.clone() { + let _ = persist_pending_usage( + &usage_tracker, + &usage_history, + compact_state.as_deref(), + |entry, after_commit, post_metrics| { + writer.commit_usage_boundary(entry, after_commit, post_metrics) + }, + )?; + } else { + let post_request_metrics = persist_pending_usage( + &usage_tracker, + &usage_history, + compact_state.as_deref(), + |entry, after_commit, _post_metrics| { + self.commit_entry(entry)?; + after_commit(); + Ok(()) + }, + )?; + for metric in post_request_metrics { self.try_record_metric(&metric); } - self.usage_history - .lock() - .expect("usage_history poisoned") - .push(record); } let interrupted = matches!( @@ -5127,6 +5177,7 @@ impl Worker { mut cancel: Option>, trigger: CompactionTrigger, ) -> Result { + self.ensure_no_pending_usage_for_segment_replacement()?; self.release_pending_compaction_service().await?; let _rewrite_guard = self .prepare_session_rewrite(SessionRewriteKind::Compact) @@ -7438,7 +7489,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 } @@ -7502,6 +7555,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, @@ -8567,6 +8625,117 @@ 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); + + // Simulate another writer extending the active Segment. The normal run + // prelude must not auto-fork while the failed record still belongs to + // this Segment. + store + .append( + session_id, + source_segment, + &LogEntry::TurnEnd { + ts: segment_log::now_millis(), + turn_count: 99, + }, + ) + .unwrap(); + let run_error = worker + .run_text("must stop before provider ownership") + .await + .expect_err("head drift must not fork past pending usage"); + assert!(matches!( + run_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 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/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index c6c37367c..7e086533f 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -7,15 +7,16 @@ //! (driven by `compact_request_threshold` → `PreRequestAction::Yield`) use std::pin::Pin; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use agen::Engine; use agen::llm_client::event::{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}; use async_trait::async_trait; -use futures::Stream; +use futures::{Stream, StreamExt}; use protocol::{Event, Method, RunResult}; use session_store::{ CombinedStore, FsStore, FsWorkerStore, LogEntry, Store, WorkerMetadata, WorkerMetadataStore, @@ -125,10 +126,13 @@ fn annotated(item: Item) -> session_store::LoggedHistoryEntry { } } +type CallObserver = Arc; + #[derive(Clone)] struct MockClient { responses: Arc>>, call_count: Arc, + observer: Arc>>, } impl MockClient { @@ -136,8 +140,13 @@ impl MockClient { Self { responses: Arc::new(responses), call_count: Arc::new(AtomicUsize::new(0)), + observer: Arc::new(Mutex::new(None)), } } + + fn set_observer(&self, observer: impl Fn(usize) + Send + Sync + 'static) { + *self.observer.lock().unwrap() = Some(Arc::new(observer)); + } } #[async_trait] @@ -152,6 +161,9 @@ impl LlmClient for MockClient { ) -> Result> + Send>>, ClientError> { let count = self.call_count.fetch_add(1, Ordering::SeqCst); + if let Some(observer) = self.observer.lock().unwrap().clone() { + observer(count); + } if count >= self.responses.len() { return Err(ClientError::Config("mock client exhausted".into())); } @@ -161,6 +173,90 @@ impl LlmClient for MockClient { } } +#[derive(Clone)] +struct PauseAfterUsageClient { + responses: Arc>>, + call_count: Arc, + hanging_call: usize, +} + +impl PauseAfterUsageClient { + fn new(responses: Vec>, hanging_call: usize) -> Self { + Self { + responses: Arc::new(responses), + call_count: Arc::new(AtomicUsize::new(0)), + hanging_call, + } + } +} + +#[async_trait] +impl LlmClient for PauseAfterUsageClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + Send>>, ClientError> + { + let count = self.call_count.fetch_add(1, Ordering::SeqCst); + let events = self + .responses + .get(count) + .cloned() + .ok_or_else(|| ClientError::Config("pause mock client exhausted".into()))?; + let finite = futures::stream::iter(events.into_iter().map(Ok)); + if count == self.hanging_call { + Ok(Box::pin(finite.chain(futures::stream::pending()))) + } else { + Ok(Box::pin(finite)) + } + } +} + +#[derive(Clone)] +struct GrowTool; + +#[async_trait] +impl Tool for GrowTool { + async fn execute( + &self, + _input: &str, + _ctx: agen::tool::ToolExecutionContext, + ) -> Result { + Ok(ToolOutput { + summary: "grew context".into(), + content: Some("x".repeat(400_000)), + attachments: Vec::new(), + }) + } +} + +fn grow_tool_definition() -> ToolDefinition { + Arc::new(|| { + ( + ToolMeta::new("grow") + .description("test tool that grows request context") + .input_schema(serde_json::json!({"type": "object"})), + Arc::new(GrowTool) as Arc, + ) + }) +} + +fn grow_tool_use_events(call_id: &str, input_tokens: u64) -> Vec { + vec![ + LlmEvent::tool_use_start(0, call_id, "grow"), + LlmEvent::tool_input_delta(0, "{}"), + LlmEvent::tool_use_stop(0), + LlmEvent::usage(input_tokens, 1), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] +} + #[derive(Clone)] struct BlockingCompactClient { calls: Arc, @@ -365,12 +461,13 @@ where let scope = worker::Scope::writable(&pwd).unwrap(); std::mem::forget(pwd_tmp); - let worker = + let mut engine = Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); + engine.register_tool(grow_tool_definition()); let observed_store = store.clone(); let mut worker = Worker::new( manifest, - worker, + engine, store, worker::WorkerWorkspaceContext::local_filesystem(None), worker::WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()), @@ -379,6 +476,7 @@ where .await .unwrap(); worker.enable_worker_metadata_write_through().unwrap(); + worker.attach_log_writer(Arc::new(worker.log_writer_handle())); (worker, observed_store) } @@ -1057,6 +1155,289 @@ async fn request_threshold_compact_publishes_runtime_progress() { assert_eq!(post.metric.correlation_id.as_deref(), Some(correlation_id)); } +#[tokio::test] +async fn completed_post_compact_usage_rearms_within_one_tool_loop() { + let client = MockClient::new(vec![ + text_events_with_usage("seed", 100_000), + write_summary_tool_use_events("compact-1", "first summary"), + single_text_events("done"), + grow_tool_use_events("grow-1", 1_000), + write_summary_tool_use_events("compact-2", "second summary"), + single_text_events("done"), + text_events_with_usage("finished", 800), + ]); + let call_count = Arc::clone(&client.call_count); + let observed_client = client.clone(); + let manifest = MID_TURN_MANIFEST_TOML.replace( + "compact_request_threshold = 100", + "compact_threshold = 80000\ncompact_request_threshold = 50000", + ); + let (mut worker, store) = make_worker_with_manifest_and_store(&manifest, client).await; + let session_id = worker.session_id(); + let durable_at_next_request = Arc::new(AtomicBool::new(false)); + let observed_barrier = Arc::clone(&durable_at_next_request); + let observed_store = store.clone(); + observed_client.set_observer(move |call_index| { + if call_index != 4 { + return; + } + let usage_is_durable = observed_store + .list_segments(session_id) + .unwrap() + .into_iter() + .flat_map(|segment_id| observed_store.read_all(session_id, segment_id).unwrap()) + .any(|entry| { + matches!( + entry, + LogEntry::LlmUsage { + input_total_tokens: 1_000, + .. + } + ) + }); + assert!( + usage_is_durable, + "the first post-compact usage must be durable before the next provider request" + ); + observed_barrier.store(true, Ordering::SeqCst); + }); + + worker.run_text("seed input").await.unwrap(); + worker + .run_text("continue through two compactions") + .await + .expect("a successful post-compact request must rearm before the tool-loop threshold"); + + assert_eq!( + call_count.load(Ordering::SeqCst), + 7, + "both compactor cycles and both normal post-compact requests must run" + ); + assert!(durable_at_next_request.load(Ordering::SeqCst)); + + let metrics = session_metrics::read_session_metrics(&store, session_id).unwrap(); + let compact_starts = metrics + .iter() + .filter(|record| record.metric.name == "compact.start") + .collect::>(); + assert_eq!(compact_starts.len(), 2); + assert!(compact_starts.iter().any(|record| { + record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run") + })); + assert!(compact_starts.iter().any(|record| { + record.metric.dimensions.get("trigger").map(String::as_str) == Some("request_threshold") + })); + let compact_posts = metrics + .iter() + .filter(|record| record.metric.name == "compact.post_request") + .collect::>(); + assert_eq!(compact_posts.len(), 2); + assert_ne!( + compact_posts[0].metric.correlation_id, + compact_posts[1].metric.correlation_id + ); + + let usage = store + .list_segments(session_id) + .unwrap() + .into_iter() + .flat_map(|segment_id| store.read_all(session_id, segment_id).unwrap()) + .filter_map(|entry| match entry { + LogEntry::LlmUsage { + input_total_tokens, .. + } => Some(input_total_tokens), + _ => None, + }) + .collect::>(); + assert_eq!(usage.iter().filter(|tokens| **tokens == 1_000).count(), 1); + assert_eq!(usage.iter().filter(|tokens| **tokens == 800).count(), 1); + + let restored = session_store::restore(&store, session_id, worker.segment_id()).unwrap(); + assert_eq!(restored.usage_history.len(), 1); + assert_eq!(restored.usage_history[0].input_total_tokens, 800); +} + +#[tokio::test] +async fn post_compact_pause_bills_partial_without_rearming_or_consuming_correlation() { + let client = PauseAfterUsageClient::new( + vec![ + text_events_with_usage("seed", 100_000), + write_summary_tool_use_events("manual-compact", "pause test summary"), + single_text_events("done"), + vec![ + LlmEvent::text_block_start(0), + LlmEvent::text_delta(0, "partial"), + LlmEvent::usage(90, 2), + ], + grow_tool_use_events("grow-after-resume", 80), + write_summary_tool_use_events("automatic-compact", "rearmed summary"), + single_text_events("done"), + text_events_with_usage("finished", 50), + ], + 3, + ); + let call_count = Arc::clone(&client.call_count); + let manifest = MID_TURN_MANIFEST_TOML.replace( + "compact_request_threshold = 100", + "compact_threshold = 80000\ncompact_request_threshold = 50000", + ); + let (worker, store) = make_worker_with_manifest_and_store(&manifest, client).await; + let session_id = worker.session_id(); + + let runtime_tmp = tempfile::tempdir().unwrap(); + let bash_output_dir = runtime_tmp.path().join("bash-output"); + let (handle, shutdown_receiver) = + WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir) + .await + .unwrap(); + let mut rx = handle.subscribe(); + handle + .send(Method::submit_text( + protocol::new_submission_request_id(), + "seed history", + )) + .await + .unwrap(); + loop { + if matches!( + tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("timeout waiting for seed") + .expect("event"), + Event::RunEnd { + result: RunResult::Finished + } + ) { + break; + } + } + handle + .send(Method::submit_text( + protocol::new_submission_request_id(), + "pause the first post-compact request", + )) + .await + .unwrap(); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while call_count.load(Ordering::SeqCst) < 4 { + tokio::task::yield_now().await; + } + }) + .await + .expect("post-compact provider request should start"); + handle + .send(Method::Pause { + command: protocol::WorkerCommandEnvelope::new(1), + }) + .await + .unwrap(); + loop { + let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("timeout waiting for pause") + .expect("event"); + if matches!( + event, + Event::RunEnd { + result: RunResult::Paused + } + ) { + break; + } + } + + let paused_usage = store + .list_segments(session_id) + .unwrap() + .into_iter() + .flat_map(|segment_id| store.read_all(session_id, segment_id).unwrap()) + .filter_map(|entry| match entry { + LogEntry::LlmUsage { + input_total_tokens, .. + } => Some(input_total_tokens), + _ => None, + }) + .collect::>(); + assert_eq!( + paused_usage.iter().filter(|tokens| **tokens == 90).count(), + 1 + ); + let paused_metrics = session_metrics::read_session_metrics(&store, session_id).unwrap(); + assert!( + paused_metrics + .iter() + .all(|record| record.metric.name != "compact.post_request"), + "partial Pause usage must not consume compact correlation or rearm" + ); + + handle + .send(Method::Resume { + command: protocol::WorkerCommandEnvelope::new(2), + }) + .await + .unwrap(); + loop { + if matches!( + tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("timeout waiting for resume completion") + .expect("event"), + Event::RunEnd { + result: RunResult::Finished + } + ) { + break; + } + } + + assert_eq!( + call_count.load(Ordering::SeqCst), + 8, + "the completed resumed request must rearm and permit request-threshold compaction" + ); + let metrics = session_metrics::read_session_metrics(&store, session_id).unwrap(); + let compact_starts = metrics + .iter() + .filter(|record| record.metric.name == "compact.start") + .collect::>(); + assert_eq!(compact_starts.len(), 2); + let pre_run_correlation = compact_starts + .iter() + .find(|record| { + record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run") + }) + .and_then(|record| record.metric.correlation_id.as_deref()) + .expect("pre-run automatic compact correlation"); + let pre_run_posts = metrics + .iter() + .filter(|record| { + record.metric.name == "compact.post_request" + && record.metric.correlation_id.as_deref() == Some(pre_run_correlation) + }) + .collect::>(); + assert_eq!(pre_run_posts.len(), 1); + assert_eq!( + pre_run_posts[0].metric.dimensions["input_total_tokens"], + "80" + ); + assert!(metrics.iter().all(|record| { + record.metric.name != "compact.post_request" + || record.metric.dimensions["input_total_tokens"] != "90" + })); + + handle + .send(Method::Shutdown { + command: protocol::WorkerCommandEnvelope::new(3), + }) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), shutdown_receiver) + .await + .expect("controller shutdown timeout") + .expect("shutdown confirmation"); +} + #[tokio::test] async fn compacted_context_above_request_threshold_fails_before_provider_request() { let client = MockClient::new(vec![ diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index d92c539ae..0f996122b 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"); }