From cd648a9abd1b6fc1996faceffb8b9e6f770d8bc9 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 20 Sep 2026 18:57:11 +0900 Subject: [PATCH] fix: persist post-compact usage at request boundary --- crates/worker/src/compact/metrics_tracker.rs | 9 +- crates/worker/src/compact/usage_tracker.rs | 372 +++++++++++++++++-- crates/worker/src/ipc/interceptor.rs | 153 +++++++- crates/worker/src/worker.rs | 130 ++++--- crates/worker/tests/compact_events_test.rs | 163 +++++++- 5 files changed, 731 insertions(+), 96 deletions(-) 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..24323e9e9 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,158 @@ 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()) + /// 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 +393,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 +412,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 +425,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 255933ab2..cde3d8b3a 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}; @@ -1946,6 +1946,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, @@ -1977,6 +1999,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 { @@ -2154,16 +2200,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 @@ -3546,6 +3592,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() @@ -5008,45 +5055,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!( diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index c6c37367c..51ec551cb 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -7,13 +7,14 @@ //! (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 protocol::{Event, Method, RunResult}; @@ -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,47 @@ impl LlmClient for MockClient { } } +#[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 +418,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 +433,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 +1112,108 @@ 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 compacted_context_above_request_threshold_fails_before_provider_request() { let client = MockClient::new(vec![