fix: persist post-compact usage at request boundary

This commit is contained in:
2026-09-20 18:57:11 +09:00
parent 0c80f84e9c
commit cd648a9abd
5 changed files with 731 additions and 96 deletions
+4 -5
View File
@@ -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<Metric> {
std::mem::take(&mut *self.pending.lock().unwrap())
}
+332 -40
View File
@@ -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<PostRequestLink>,
/// 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<PendingRequest>,
pending_correlations: Vec<PostRequestLink>,
pending_records: VecDeque<PendingRecordedUsage>,
}
/// 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<Option<usize>>,
/// Optional causal link consumed by the next measured request.
pending_correlations: Mutex<Vec<PostRequestLink>>,
/// Records accumulated during the current run; drained by Worker.
pending_records: Mutex<Vec<RecordedUsage>>,
state: Mutex<UsageTrackerState>,
}
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<UsageRecord> {
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<RecordedUsage> {
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<RecordedUsage> {
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<RecordedUsage> {
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<Vec<UsageRecord>>,
compact_state: Option<&CompactState>,
mut commit: impl FnMut(LogEntry, &mut dyn FnMut(), &[Metric]) -> Result<(), StoreError>,
) -> Result<Vec<Metric>, 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::<Vec<_>>();
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(_)
));
}
}
+150 -3
View File
@@ -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<Arc<UsageTracker>>,
/// 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<Arc<MetricsTracker>>,
/// 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<MetricsTracker>) -> Self {
self.metrics_tracker = Some(metrics_tracker);
self
}
pub(crate) fn with_prompt_workspace_id(mut self, workspace_id: Option<String>) -> Self {
self.prompt_workspace_id = workspace_id;
self
@@ -272,6 +283,56 @@ impl WorkerInterceptor {
}
}
}
fn try_persist_metric(writer: &Arc<dyn SystemItemCommitter>, 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<SessionHistoryMetadata> for WorkerInterceptor {
&self,
context: AssistantTurnEndContext<'_, SessionHistoryMetadata>,
) -> InterceptorResult<TurnEndAction> {
// `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<AtomicUsize>,
}
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<AtomicBool>,
}
@@ -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());
+85 -45
View File
@@ -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<session_store::FsStore> {
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<St: Store + Clone> {
@@ -2154,16 +2200,16 @@ pub struct Worker<C: LlmClient, St: Store> {
/// Shared compaction state (present when threshold is configured).
compact_state: Option<Arc<CompactState>>,
/// 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<UsageTracker>,
/// 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<crate::compact::metrics_tracker::MetricsTracker>,
/// 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<Mutex>` so that callbacks injected into the
@@ -3546,6 +3592,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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!(
+160 -3
View File
@@ -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<dyn Fn(usize) + Send + Sync>;
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
observer: Arc<Mutex<Option<CallObserver>>>,
}
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<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + 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<ToolOutput, ToolError> {
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<dyn Tool>,
)
})
}
fn grow_tool_use_events(call_id: &str, input_tokens: u64) -> Vec<LlmEvent> {
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<AtomicUsize>,
@@ -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::<Vec<_>>();
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::<Vec<_>>();
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::<Vec<_>>();
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![