feat: record compaction lifecycle metrics

This commit is contained in:
2026-09-16 06:01:04 +09:00
parent f40714efc2
commit c0a73c12ec
6 changed files with 900 additions and 43 deletions
+1
View File
@@ -1,6 +1,7 @@
pub(crate) mod metrics_tracker;
pub(crate) mod prune;
pub(crate) mod state;
pub(crate) mod telemetry;
pub(crate) mod token_counter;
pub(crate) mod usage_tracker;
pub(crate) mod worker;
+371
View File
@@ -0,0 +1,371 @@
use std::collections::BTreeMap;
use std::time::Duration;
use agen::token_counter::EstimateSource;
use session_metrics::Metric;
use session_store::{SegmentId, SessionId};
use super::usage_tracker::UsageSnapshot;
const MAX_SAFE_INTEGER: u64 = (1_u64 << 53) - 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CompactMode {
Manual,
Automatic,
}
impl CompactMode {
fn as_str(self) -> &'static str {
match self {
Self::Manual => "manual",
Self::Automatic => "automatic",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CompactThresholdPolicy {
Manual,
PreRun,
RequestThreshold,
}
impl CompactThresholdPolicy {
fn as_str(self) -> &'static str {
match self {
Self::Manual => "manual",
Self::PreRun => "pre_run",
Self::RequestThreshold => "request_threshold",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CompactFailureCategory {
Cancelled,
SummaryMissing,
SummaryTooLarge,
ResultContextTooLarge,
ActiveSegmentCommit,
Storage,
InternalWorker,
Preparation,
Other,
}
impl CompactFailureCategory {
fn as_str(self) -> &'static str {
match self {
Self::Cancelled => "cancelled",
Self::SummaryMissing => "summary_missing",
Self::SummaryTooLarge => "summary_too_large",
Self::ResultContextTooLarge => "result_context_too_large",
Self::ActiveSegmentCommit => "active_segment_commit",
Self::Storage => "storage",
Self::InternalWorker => "internal_worker",
Self::Preparation => "preparation",
Self::Other => "other",
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct CompactAttempt {
correlation_id: String,
session_id: SessionId,
source_segment_id: SegmentId,
mode: CompactMode,
threshold_policy: CompactThresholdPolicy,
pre_context_tokens: u64,
pre_context_source: EstimateSource,
retained_token_budget: u64,
}
impl CompactAttempt {
pub(crate) fn new(
correlation_id: String,
session_id: SessionId,
source_segment_id: SegmentId,
mode: CompactMode,
threshold_policy: CompactThresholdPolicy,
pre_context_tokens: u64,
pre_context_source: EstimateSource,
retained_token_budget: u64,
) -> Self {
debug_assert!(uuid::Uuid::parse_str(&correlation_id).is_ok());
Self {
correlation_id,
session_id,
source_segment_id,
mode,
threshold_policy,
pre_context_tokens,
pre_context_source,
retained_token_budget,
}
}
pub(crate) fn correlation_id(&self) -> &str {
&self.correlation_id
}
pub(crate) fn start_metric(&self) -> Metric {
self.metric("compact.start")
.with_value(safe_number(self.pre_context_tokens))
.with_dimension("occupancy_source", estimate_source(self.pre_context_source))
.with_dimension(
"retained_token_budget",
self.retained_token_budget.to_string(),
)
}
pub(crate) fn success_metrics(
&self,
result_segment_id: SegmentId,
elapsed: Duration,
stats: &CompactSuccessStats,
) -> Vec<Metric> {
let dimensions = self.base_dimensions();
let correlation_id = self.correlation_id.clone();
let mut metrics = vec![
metric_with_context(
"compact.finish",
elapsed.as_millis().min(u128::from(MAX_SAFE_INTEGER)) as u64,
&dimensions,
&correlation_id,
)
.with_dimension("outcome", "success")
.with_dimension("result_segment_id", result_segment_id.to_string())
.with_dimension("retained_items", stats.retained_items.to_string())
.with_dimension("summarized_items", stats.summarized_items.to_string()),
metric_with_context(
"compact.retained_tokens",
stats.retained_tokens,
&dimensions,
&correlation_id,
)
.with_dimension("source", estimate_source(stats.retained_tokens_source)),
metric_with_context(
"compact.overview_tokens",
stats.overview_tokens,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.summary_tokens",
stats.summary_tokens,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.auto_read_tokens",
stats.auto_read_tokens,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.result_context_tokens",
stats.result_context_tokens,
&dimensions,
&correlation_id,
)
.with_dimension("source", estimate_source(stats.result_context_source)),
metric_with_context(
"compact.compactor.input_tokens",
stats.usage.input_total_tokens,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.compactor.output_tokens",
stats.usage.output_tokens,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.compactor.cache_read_tokens",
stats.usage.cache_read_tokens,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.compactor.cache_write_tokens",
stats.usage.cache_write_tokens,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.compactor.requests",
stats.requests,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.compactor.turns",
stats.turns,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.compactor.tool_calls",
stats.tool_calls,
&dimensions,
&correlation_id,
),
metric_with_context(
"compact.duration_ms",
elapsed.as_millis().min(u128::from(MAX_SAFE_INTEGER)) as u64,
&dimensions,
&correlation_id,
),
];
for metric in &mut metrics {
metric
.dimensions
.insert("result_segment_id".into(), result_segment_id.to_string());
}
// Provider UsageEvent currently carries tokens but no price or cost. Keep
// the field explicit and valueless rather than fabricating a zero cost.
metrics.push(
self.metric("compact.compactor.cost_usd")
.with_dimension("status", "unavailable")
.with_dimension("reason", "provider_usage_unpriced")
.with_dimension("result_segment_id", result_segment_id.to_string()),
);
metrics
}
pub(crate) fn failure_metric(
&self,
observed_segment_id: SegmentId,
elapsed: Duration,
category: CompactFailureCategory,
) -> Metric {
let outcome = if category == CompactFailureCategory::Cancelled {
"cancelled"
} else {
"failure"
};
self.metric("compact.finish")
.with_value(elapsed.as_millis().min(u128::from(MAX_SAFE_INTEGER)) as f64)
.with_dimension("outcome", outcome)
.with_dimension("failure_category", category.as_str())
.with_dimension("observed_segment_id", observed_segment_id.to_string())
}
fn metric(&self, name: &'static str) -> Metric {
let mut metric = Metric::now(name).with_correlation_id(&self.correlation_id);
metric.dimensions = self.base_dimensions();
metric
}
fn base_dimensions(&self) -> BTreeMap<String, String> {
BTreeMap::from([
("session_id".into(), self.session_id.to_string()),
(
"source_segment_id".into(),
self.source_segment_id.to_string(),
),
("mode".into(), self.mode.as_str().into()),
(
"threshold_policy".into(),
self.threshold_policy.as_str().into(),
),
])
}
}
#[derive(Debug, Clone)]
pub(crate) struct CompactSuccessStats {
pub(crate) retained_items: u64,
pub(crate) summarized_items: u64,
pub(crate) retained_tokens: u64,
pub(crate) retained_tokens_source: EstimateSource,
pub(crate) overview_tokens: u64,
pub(crate) summary_tokens: u64,
pub(crate) auto_read_tokens: u64,
pub(crate) result_context_tokens: u64,
pub(crate) result_context_source: EstimateSource,
pub(crate) usage: UsageSnapshot,
pub(crate) requests: u64,
pub(crate) turns: u64,
pub(crate) tool_calls: u64,
}
fn metric_with_context(
name: &'static str,
value: u64,
dimensions: &BTreeMap<String, String>,
correlation_id: &str,
) -> Metric {
let mut metric = Metric::now(name)
.with_value(safe_number(value))
.with_correlation_id(correlation_id);
metric.dimensions = dimensions.clone();
metric
}
fn safe_number(value: u64) -> f64 {
value.min(MAX_SAFE_INTEGER) as f64
}
fn estimate_source(source: EstimateSource) -> &'static str {
match source {
EstimateSource::Measured => "measured",
EstimateSource::Interpolated => "interpolated",
EstimateSource::Extrapolated => "extrapolated",
EstimateSource::NoData => "no_data",
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn compact_metrics_use_fixed_bounded_labels_and_safe_numbers() {
let attempt = CompactAttempt::new(
uuid::Uuid::now_v7().to_string(),
uuid::Uuid::now_v7(),
uuid::Uuid::now_v7(),
CompactMode::Automatic,
CompactThresholdPolicy::RequestThreshold,
u64::MAX,
EstimateSource::Measured,
500,
);
let start = attempt.start_metric();
assert_eq!(start.name, "compact.start");
assert_eq!(start.value, Some(MAX_SAFE_INTEGER as f64));
assert_eq!(start.dimensions["mode"], "automatic");
assert_eq!(start.dimensions["threshold_policy"], "request_threshold");
assert_eq!(start.dimensions["occupancy_source"], "measured");
assert!(start.correlation_id.is_some());
assert!(start.dimensions.keys().all(|key| key.len() <= 32));
assert!(start.dimensions.values().all(|value| value.len() <= 64));
}
#[test]
fn failure_metrics_never_include_error_text() {
let attempt = CompactAttempt::new(
uuid::Uuid::now_v7().to_string(),
uuid::Uuid::now_v7(),
uuid::Uuid::now_v7(),
CompactMode::Manual,
CompactThresholdPolicy::Manual,
1,
EstimateSource::NoData,
1,
);
let metric = attempt.failure_metric(
uuid::Uuid::now_v7(),
Duration::from_millis(7),
CompactFailureCategory::InternalWorker,
);
let encoded = serde_json::to_string(&metric).unwrap();
assert!(encoded.contains("internal_worker"));
assert!(!encoded.contains("error"));
assert!(!encoded.contains("path"));
}
}
+76 -25
View File
@@ -19,14 +19,41 @@ use std::sync::Mutex;
use agen::UsageRecord;
use agen::timeline::event::UsageEvent;
/// One drained measurement: the underlying `UsageRecord` plus an optional
/// `correlation_id` stamped by the prune projection (or any other future
/// upstream observer) so that downstream metrics emitted alongside this
/// record can be joined to it after the fact.
/// The metric emitted after the next measured provider request.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PostRequestMetric {
Prune,
Compaction,
}
impl PostRequestMetric {
pub(crate) fn name(self) -> &'static str {
match self {
Self::Prune => "prune.post_request",
Self::Compaction => "compact.post_request",
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct PostRequestLink {
pub(crate) correlation_id: String,
pub(crate) metric: PostRequestMetric,
}
/// One drained measurement and its causal metric links.
#[derive(Debug, Clone)]
pub(crate) struct RecordedUsage {
pub(crate) record: UsageRecord,
pub(crate) correlation_id: Option<String>,
pub(crate) post_requests: Vec<PostRequestLink>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct UsageSnapshot {
pub(crate) input_total_tokens: u64,
pub(crate) cache_read_tokens: u64,
pub(crate) cache_write_tokens: u64,
pub(crate) output_tokens: u64,
}
/// Shared between the pre-request hook, the `on_usage` callback, and Worker.
@@ -34,11 +61,8 @@ 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 `correlation_id` set by an upstream observer (currently
/// the prune projection on `Fired`). Paired into the next
/// `RecordedUsage` and cleared. Skips that don't fire leave this
/// `None`, so the resulting record carries no correlation.
pending_correlation_id: Mutex<Option<String>>,
/// 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>>,
}
@@ -47,7 +71,7 @@ impl UsageTracker {
pub(crate) fn new() -> Self {
Self {
pending_history_len: Mutex::new(None),
pending_correlation_id: Mutex::new(None),
pending_correlations: Mutex::new(Vec::new()),
pending_records: Mutex::new(Vec::new()),
}
}
@@ -57,16 +81,23 @@ impl UsageTracker {
*self.pending_history_len.lock().unwrap() = Some(history_len);
}
/// Stash a `correlation_id` to be paired into the next `RecordedUsage`.
/// Currently invoked by the prune observer on `Fired` so that the
/// `prune.fire` metric and the `prune.post_request` metric (emitted
/// alongside the resulting `LlmUsage`) carry the same join key.
///
/// Overwrites any previous unconsumed value — by construction the
/// observer fires at most once per outgoing LLM request, immediately
/// before the pre-request hook captures `history_len`.
/// Pair a prune event with the next provider request.
pub(crate) fn note_correlation_id(&self, id: String) {
*self.pending_correlation_id.lock().unwrap() = Some(id);
self.note_post_request(id, PostRequestMetric::Prune);
}
/// Pair a completed compaction with the next normal provider request.
pub(crate) fn note_compaction_correlation_id(&self, id: String) {
self.note_post_request(id, PostRequestMetric::Compaction);
}
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 {
correlation_id: id,
metric,
});
}
/// Called from the `on_usage` callback with the aggregated final
@@ -79,7 +110,7 @@ impl UsageTracker {
Some(n) => n,
None => return,
};
let correlation_id = self.pending_correlation_id.lock().unwrap().take();
let post_requests = std::mem::take(&mut *self.pending_correlations.lock().unwrap());
// UsageEvent.input_tokens は scheme 層で「占有量(プロンプト全長)」に
// 正規化済みである前提(Anthropic は cache_read + cache_creation を
// 加算して emit する)。
@@ -95,7 +126,7 @@ impl UsageTracker {
cache_write_tokens: cache_write,
output_tokens: output,
},
correlation_id,
post_requests,
});
}
@@ -145,7 +176,7 @@ mod tests {
assert_eq!(records[0].record.cache_read_tokens, 800);
assert_eq!(records[0].record.cache_write_tokens, 100);
assert_eq!(records[0].record.output_tokens, 42);
assert!(records[0].correlation_id.is_none());
assert!(records[0].post_requests.is_empty());
}
#[test]
@@ -192,6 +223,24 @@ mod tests {
assert_eq!(records[1].record.cache_read_tokens, 50);
}
#[test]
fn prune_and_compaction_links_share_the_next_request() {
let tracker = UsageTracker::new();
tracker.note_compaction_correlation_id("compact-id".into());
tracker.note_correlation_id("prune-id".into());
tracker.note_request(5);
tracker.record_usage(&make_event(100, 10, 2, 20));
let records = tracker.drain();
assert_eq!(records[0].post_requests.len(), 2);
assert!(records[0].post_requests.iter().any(|link| {
link.correlation_id == "compact-id" && link.metric == PostRequestMetric::Compaction
}));
assert!(records[0].post_requests.iter().any(|link| {
link.correlation_id == "prune-id" && link.metric == PostRequestMetric::Prune
}));
}
#[test]
fn correlation_id_pairs_with_next_record_only() {
let tracker = UsageTracker::new();
@@ -205,7 +254,9 @@ mod tests {
let records = tracker.drain();
assert_eq!(records.len(), 2);
assert_eq!(records[0].correlation_id.as_deref(), Some("abc"));
assert!(records[1].correlation_id.is_none());
assert_eq!(records[0].post_requests.len(), 1);
assert_eq!(records[0].post_requests[0].correlation_id, "abc");
assert_eq!(records[0].post_requests[0].metric, PostRequestMetric::Prune);
assert!(records[1].post_requests.is_empty());
}
}
+156 -13
View File
@@ -4,7 +4,7 @@ use std::path::Path;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::time::{Duration, Instant};
use agen::llm_client::RequestConfig;
use agen::llm_client::client::LlmClient;
@@ -40,6 +40,10 @@ use manifest::{
};
use crate::compact::state::CompactState;
use crate::compact::telemetry::{
CompactAttempt, CompactFailureCategory, CompactMode, CompactSuccessStats,
CompactThresholdPolicy,
};
use crate::compact::usage_tracker::UsageTracker;
use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTaskRegistry};
use crate::feature::builtin::memory::WorkspaceMemoryBackendError;
@@ -3359,12 +3363,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
domain: session_metrics::DOMAIN.into(),
payload,
};
if let Err(err) = self.commit_entry(entry) {
warn!(name = %metric.name, error = %err, "failed to record session metric; dropping");
if self.commit_entry(entry).is_err() {
warn!(name = %metric.name, "failed to record session metric; dropping");
let bounded_name = metric.name.chars().take(64).collect::<String>();
self.alert(
AlertLevel::Warn,
AlertSource::Worker,
format!("failed to record metric `{}`: {}", metric.name, err),
format!("failed to record metric `{bounded_name}`"),
);
}
}
@@ -4890,7 +4895,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
for recorded in usage_records {
let crate::compact::usage_tracker::RecordedUsage {
record,
correlation_id,
post_requests,
} = recorded;
self.commit_entry(LogEntry::LlmUsage {
ts: segment_log::now_millis(),
@@ -4900,12 +4905,23 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
cache_write_tokens: record.cache_write_tokens,
output_tokens: record.output_tokens,
})?;
if let Some(id) = correlation_id {
let metric = session_metrics::Metric::now("prune.post_request")
.with_correlation_id(&id)
.with_value(record.cache_read_tokens as f64)
for link in post_requests {
let value = match link.metric {
crate::compact::usage_tracker::PostRequestMetric::Prune => {
record.cache_read_tokens
}
crate::compact::usage_tracker::PostRequestMetric::Compaction => {
record.input_total_tokens
}
};
let metric = session_metrics::Metric::now(link.metric.name())
.with_correlation_id(&link.correlation_id)
.with_value(value as f64)
.with_dimension("history_len", record.history_len.to_string())
.with_dimension("input_total_tokens", record.input_total_tokens.to_string())
.with_dimension("cache_read_tokens", record.cache_read_tokens.to_string())
.with_dimension("cache_write_tokens", record.cache_write_tokens.to_string())
.with_dimension("history_len", record.history_len.to_string());
.with_dimension("output_tokens", record.output_tokens.to_string());
self.try_record_metric(&metric);
}
self.usage_history
@@ -4986,6 +5002,31 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
error: None,
new_segment_id: None,
};
let started = Instant::now();
let source_location = self.segment_state.location();
let history_items = self.session.history().items_cloned();
let usage_history = self.usage_history();
let pre_context = agen::token_counter::total_tokens(&history_items, &usage_history);
let attempt = CompactAttempt::new(
lifecycle.compaction_id.clone(),
source_location.session_id,
source_location.segment_id,
match trigger {
CompactionTrigger::Manual => CompactMode::Manual,
CompactionTrigger::PreRun | CompactionTrigger::RequestThreshold => {
CompactMode::Automatic
}
},
match trigger {
CompactionTrigger::Manual => CompactThresholdPolicy::Manual,
CompactionTrigger::PreRun => CompactThresholdPolicy::PreRun,
CompactionTrigger::RequestThreshold => CompactThresholdPolicy::RequestThreshold,
},
pre_context.tokens,
pre_context.source,
retained_tokens,
);
self.try_record_metric(&attempt.start_metric());
let started_at_ms = lifecycle.started_at_ms;
self.set_compaction_progress(Some(InFlightCompaction {
phase: CompactionPhase::Preparing,
@@ -5006,12 +5047,24 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.await
};
match outcome {
Ok((new_segment_id, _summary)) => {
Ok((new_segment_id, _summary, stats)) => {
debug_assert_eq!(lifecycle.state, CompactionLifecycleState::Done);
for metric in attempt.success_metrics(new_segment_id, started.elapsed(), &stats) {
self.try_record_metric(&metric);
}
self.usage_tracker
.note_compaction_correlation_id(attempt.correlation_id().to_string());
self.release_compaction_service(&lifecycle).await;
Ok(new_segment_id)
}
Err(error) => {
let observed_segment_id = self.segment_state.location().segment_id;
let metric = attempt.failure_metric(
observed_segment_id,
started.elapsed(),
compact_failure_category(&error),
);
self.try_record_metric(&metric);
lifecycle.revision = lifecycle.revision.saturating_add(1);
lifecycle.state = if matches!(error, WorkerError::CompactCancelled) {
CompactionLifecycleState::Interrupted
@@ -5052,7 +5105,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
retained_tokens: u64,
lifecycle: &mut CompactionLifecycle,
trigger: CompactionTrigger,
) -> Result<(SegmentId, String), WorkerError> {
) -> Result<(SegmentId, String, CompactSuccessStats), WorkerError> {
use crate::compact::worker::{
CompactWorkerContext, CompactWorkerInterceptor, CompactionOutputFeature,
};
@@ -5414,6 +5467,57 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
} else {
Vec::new()
};
let auto_read_tokens = agen::token_counter::total_tokens(&auto_read_messages, &[]).tokens;
let retained_estimate = agen::token_counter::total_tokens(&retained_items, &[]);
let retained_item_count = u64::try_from(retained_items.len()).unwrap_or(u64::MAX);
let summarized_item_count = u64::try_from(items_to_summarise.len()).unwrap_or(u64::MAX);
let compactor_entries = handle.entries();
let compactor_turns = compactor_entries
.iter()
.filter_map(|entry| match entry {
LogEntry::TurnEnd { turn_count, .. } => u64::try_from(*turn_count).ok(),
_ => None,
})
.max()
.unwrap_or(0);
let compactor_tool_calls = u64::try_from(
compactor_entries
.iter()
.filter(|entry| {
matches!(
entry,
LogEntry::AnnotatedAssistantItem { entry, .. }
if matches!(&entry.item, session_store::LoggedItem::ToolCall { .. })
)
})
.count(),
)
.unwrap_or(u64::MAX);
let mut compactor_usage = crate::compact::usage_tracker::UsageSnapshot::default();
let mut compactor_requests = 0_u64;
for entry in &compactor_entries {
if let LogEntry::LlmUsage {
input_total_tokens,
cache_read_tokens,
cache_write_tokens,
output_tokens,
..
} = entry
{
compactor_requests = compactor_requests.saturating_add(1);
compactor_usage.input_total_tokens = compactor_usage
.input_total_tokens
.saturating_add(*input_total_tokens);
compactor_usage.cache_read_tokens = compactor_usage
.cache_read_tokens
.saturating_add(*cache_read_tokens);
compactor_usage.cache_write_tokens = compactor_usage
.cache_write_tokens
.saturating_add(*cache_write_tokens);
compactor_usage.output_tokens =
compactor_usage.output_tokens.saturating_add(*output_tokens);
}
}
// Reference list as a single system message; omitted when empty.
let reference_message = (!final_ctx.references.is_empty()).then(|| {
@@ -5651,7 +5755,25 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
lifecycle.state = CompactionLifecycleState::Done;
lifecycle.ended_at_ms = Some(segment_log::now_millis());
self.set_compaction_progress(None);
Ok((new_segment_id, summary_text))
Ok((
new_segment_id,
summary_text,
CompactSuccessStats {
retained_items: retained_item_count,
summarized_items: summarized_item_count,
retained_tokens: retained_estimate.tokens,
retained_tokens_source: retained_estimate.source,
overview_tokens: summary_input.overview_tokens,
summary_tokens,
auto_read_tokens,
result_context_tokens: result_estimate.tokens,
result_context_source: result_estimate.source,
usage: compactor_usage,
requests: compactor_requests,
turns: compactor_turns,
tool_calls: compactor_tool_calls,
},
))
}
/// Build the LlmClient for the compactor Engine.
@@ -7013,6 +7135,27 @@ fn restored_flow_runtime_state(
.transpose()
}
fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory {
match error {
WorkerError::CompactCancelled => CompactFailureCategory::Cancelled,
WorkerError::CompactSummaryMissing => CompactFailureCategory::SummaryMissing,
WorkerError::CompactSummaryTooLarge { .. } => CompactFailureCategory::SummaryTooLarge,
WorkerError::CompactResultContextTooLarge { .. } => {
CompactFailureCategory::ResultContextTooLarge
}
WorkerError::WorkerStore(_) => CompactFailureCategory::ActiveSegmentCommit,
WorkerError::Store(_) => CompactFailureCategory::Storage,
WorkerError::InvalidState(_) | WorkerError::Engine(_) => {
CompactFailureCategory::InternalWorker
}
WorkerError::Provider(_)
| WorkerError::PromptCatalog(_)
| WorkerError::FeatureLifecycle(_)
| WorkerError::FeatureInstall(_) => CompactFailureCategory::Preparation,
_ => CompactFailureCategory::Other,
}
}
#[derive(Debug, thiserror::Error)]
pub enum WorkerError {
#[error("invalid durable Worker state: {0}")]
+265 -5
View File
@@ -11,7 +11,7 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use agen::Engine;
use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
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 async_trait::async_trait;
@@ -234,6 +234,49 @@ fn write_summary_tool_use_events(call_id: &str, text: &str) -> Vec<LlmEvent> {
]
}
fn write_summary_tool_use_events_with_usage(
call_id: &str,
text: &str,
input_total: u64,
cache_read: u64,
cache_write: u64,
output: u64,
) -> Vec<LlmEvent> {
let mut events = write_summary_tool_use_events(call_id, text);
events.insert(
events.len() - 1,
LlmEvent::Usage(UsageEvent {
input_tokens: Some(input_total),
output_tokens: Some(output),
total_tokens: Some(input_total.saturating_add(output)),
cache_read_input_tokens: Some(cache_read),
cache_creation_input_tokens: Some(cache_write),
}),
);
events
}
fn text_events_with_full_usage(
text: &str,
input_total: u64,
cache_read: u64,
cache_write: u64,
output: u64,
) -> Vec<LlmEvent> {
let mut events = single_text_events(text);
events.insert(
events.len() - 1,
LlmEvent::Usage(UsageEvent {
input_tokens: Some(input_total),
output_tokens: Some(output),
total_tokens: Some(input_total.saturating_add(output)),
cache_read_input_tokens: Some(cache_read),
cache_creation_input_tokens: Some(cache_write),
}),
);
events
}
// A low compact_threshold guarantees `try_pre_run_compact` will fire
// the first time we check after a run.
const POST_RUN_MANIFEST_TOML: &str = r#"
@@ -257,6 +300,27 @@ target = "./"
permission = "write"
"#;
const MANUAL_ONLY_MANIFEST_TOML: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[compaction]
compact_threshold = 1000000000
compact_retained_tokens = 0
[[scope.allow]]
target = "./"
permission = "write"
"#;
// `compact_request_threshold` drives the WorkerInterceptor's mid-turn yield
// path. `compact_threshold` is left unset so the post-run check stays inert.
const MID_TURN_MANIFEST_TOML: &str = r#"
@@ -457,6 +521,28 @@ async fn failed_active_segment_commit_keeps_live_and_durable_history_on_old_segm
);
assert_eq!(worker.segment_id(), old_segment_id);
let failure_metrics =
session_metrics::read_segment_metrics(&segment_store, worker.session_id(), old_segment_id)
.unwrap();
let start = failure_metrics
.iter()
.find(|record| record.metric.name == "compact.start")
.unwrap();
let finish = failure_metrics
.iter()
.find(|record| record.metric.name == "compact.finish")
.unwrap();
assert_eq!(finish.metric.dimensions["outcome"], "failure");
assert_eq!(
finish.metric.dimensions["failure_category"],
"active_segment_commit"
);
assert_eq!(finish.metric.correlation_id, start.metric.correlation_id);
assert!(
!serde_json::to_string(&finish.metric)
.unwrap()
.contains("injected active Segment commit failure")
);
let metadata = metadata_store
.read_by_name("test-worker")
.unwrap()
@@ -605,16 +691,18 @@ permission = "write"
async fn compact_emits_session_start_carrying_summary_and_task_snapshot() {
let client = MockClient::new(vec![
single_text_events("hi"),
write_summary_tool_use_events("call-1", "summary"),
single_text_events("done"),
write_summary_tool_use_events_with_usage("call-1", "summary", 100, 10, 5, 20),
text_events_with_full_usage("done", 50, 3, 2, 10),
text_events_with_full_usage("after", 44, 4, 1, 6),
]);
let mut worker = make_worker(client).await;
let mut worker = make_worker_with_manifest(MANUAL_ONLY_MANIFEST_TOML, client).await;
let (tx, _rx_keep) = broadcast::channel::<Event>(64);
worker.attach_working_event_tx(tx);
worker.run_text("first").await.unwrap();
let session_id = worker.session_id();
let source_segment_id = worker.segment_id();
worker.compact(10_000).await.unwrap();
let compacted_segment_id = worker.segment_id();
let metadata = worker
@@ -642,6 +730,148 @@ async fn compact_emits_session_start_carrying_summary_and_task_snapshot() {
.any(|text| text.starts_with("[Session TaskStore snapshot]")),
"task snapshot system message missing from {system_texts:?}"
);
worker.run_text("after compaction").await.unwrap();
let metrics = session_metrics::read_session_metrics(worker.store(), session_id).unwrap();
let starts = metrics
.iter()
.filter(|record| record.metric.name == "compact.start")
.collect::<Vec<_>>();
assert_eq!(starts.len(), 1);
assert_eq!(starts[0].segment_id, source_segment_id);
assert_eq!(starts[0].metric.dimensions["mode"], "automatic");
assert_eq!(
starts[0].metric.dimensions["threshold_policy"],
"request_threshold"
);
let correlation_id = starts[0]
.metric
.correlation_id
.as_deref()
.expect("compact start must carry a correlation id");
let finish = metrics
.iter()
.find(|record| record.metric.name == "compact.finish")
.unwrap();
assert_eq!(finish.segment_id, compacted_segment_id);
assert_eq!(finish.metric.dimensions["outcome"], "success");
assert_eq!(
finish.metric.correlation_id.as_deref(),
Some(correlation_id)
);
assert_eq!(
finish.compacted_from.as_ref().unwrap().segment_id,
starts[0].segment_id
);
let value = |name: &str| {
metrics
.iter()
.find(|record| record.metric.name == name)
.and_then(|record| record.metric.value)
.unwrap() as u64
};
assert_eq!(value("compact.compactor.input_tokens"), 150);
assert_eq!(value("compact.compactor.cache_read_tokens"), 13);
assert_eq!(value("compact.compactor.cache_write_tokens"), 7);
assert_eq!(value("compact.compactor.output_tokens"), 30);
assert_eq!(value("compact.compactor.requests"), 2);
assert!(value("compact.compactor.tool_calls") >= 1);
assert!(value("compact.compactor.turns") >= 2);
assert!(value("compact.duration_ms") <= u64::MAX);
let cost = metrics
.iter()
.find(|record| record.metric.name == "compact.compactor.cost_usd")
.unwrap();
assert_eq!(cost.metric.value, None);
assert_eq!(cost.metric.dimensions["status"], "unavailable");
let post = metrics
.iter()
.find(|record| record.metric.name == "compact.post_request")
.unwrap();
assert_eq!(post.segment_id, compacted_segment_id);
assert_eq!(post.metric.correlation_id.as_deref(), Some(correlation_id));
assert_eq!(post.metric.dimensions["input_total_tokens"], "44");
assert_eq!(post.metric.dimensions["cache_read_tokens"], "4");
assert_eq!(post.metric.dimensions["cache_write_tokens"], "1");
assert_eq!(post.metric.dimensions["output_tokens"], "6");
}
#[tokio::test]
async fn manual_compact_metrics_identify_manual_mode() {
let client = MockClient::new(vec![
single_text_events("seed response"),
write_summary_tool_use_events("summary", "replacement summary"),
single_text_events("done"),
]);
let mut worker = make_worker_with_manifest(MANUAL_ONLY_MANIFEST_TOML, client).await;
worker.run_text("seed input").await.unwrap();
let session_id = worker.session_id();
worker.manual_compact().await.unwrap();
let metrics = session_metrics::read_session_metrics(worker.store(), session_id).unwrap();
let start = metrics
.iter()
.find(|record| record.metric.name == "compact.start")
.unwrap();
assert_eq!(start.metric.dimensions["mode"], "manual");
assert_eq!(start.metric.dimensions["threshold_policy"], "manual");
let finish = metrics
.iter()
.find(|record| record.metric.name == "compact.finish")
.unwrap();
assert_eq!(finish.metric.dimensions["outcome"], "success");
assert_eq!(finish.metric.correlation_id, start.metric.correlation_id);
}
#[tokio::test]
async fn compact_failure_and_cancellation_emit_bounded_categories() {
let client = MockClient::new(vec![
single_text_events("seed response"),
single_text_events("missing summary"),
single_text_events("still missing summary"),
]);
let mut worker = make_worker_with_manifest(MANUAL_ONLY_MANIFEST_TOML, client).await;
worker.run_text("seed input").await.unwrap();
let session_id = worker.session_id();
let source_segment_id = worker.segment_id();
let error = worker.manual_compact().await.unwrap_err();
assert!(matches!(error, worker::WorkerError::CompactSummaryMissing));
let metrics = session_metrics::read_session_metrics(worker.store(), session_id).unwrap();
let failure = metrics
.iter()
.find(|record| {
record.metric.name == "compact.finish"
&& record.metric.dimensions["outcome"] == "failure"
})
.unwrap();
assert_eq!(failure.segment_id, source_segment_id);
assert_eq!(
failure.metric.dimensions["failure_category"],
"summary_missing"
);
let encoded = serde_json::to_string(&failure.metric).unwrap();
assert!(!encoded.contains("missing summary"));
let (_cancel_tx, cancel_rx) = tokio::sync::watch::channel(true);
let error = worker
.manual_compact_with_cancel(cancel_rx)
.await
.unwrap_err();
assert!(matches!(error, worker::WorkerError::CompactCancelled));
let metrics = session_metrics::read_session_metrics(worker.store(), session_id).unwrap();
let cancelled = metrics
.iter()
.filter(|record| {
record.metric.name == "compact.finish"
&& record.metric.dimensions["outcome"] == "cancelled"
})
.last()
.unwrap();
assert_eq!(cancelled.segment_id, source_segment_id);
assert_eq!(cancelled.metric.dimensions["failure_category"], "cancelled");
}
#[tokio::test]
@@ -707,6 +937,20 @@ async fn pre_run_compact_publishes_runtime_progress_phases() {
entry,
LogEntry::Extension { domain, .. } if domain == "yoi.compaction"
)));
let metrics = session_metrics::read_session_metrics(worker.store(), session_before).unwrap();
let start = metrics
.iter()
.find(|record| record.metric.name == "compact.start")
.unwrap();
assert_eq!(start.segment_id, segment_before);
assert_eq!(start.metric.dimensions["mode"], "automatic");
assert_eq!(start.metric.dimensions["threshold_policy"], "pre_run");
let finish = metrics
.iter()
.find(|record| record.metric.name == "compact.finish")
.unwrap();
assert_eq!(finish.metric.dimensions["outcome"], "success");
assert_eq!(finish.metric.correlation_id, start.metric.correlation_id);
}
#[tokio::test]
@@ -723,7 +967,7 @@ async fn request_threshold_compact_publishes_runtime_progress() {
text_events_with_usage("a", 1000),
write_summary_tool_use_events("call-1", "summary"),
single_text_events("done"),
single_text_events("b"),
text_events_with_usage("b", 50),
]);
let mut worker = make_worker_with_manifest(MID_TURN_MANIFEST_TOML, client).await;
@@ -749,6 +993,22 @@ async fn request_threshold_compact_publishes_runtime_progress() {
.iter()
.any(|event| matches!(event, Event::CompactionProgress { compaction: None }))
);
let metrics =
session_metrics::read_session_metrics(worker.store(), worker.session_id()).unwrap();
let start = metrics
.iter()
.find(|record| record.metric.name == "compact.start")
.unwrap();
assert_eq!(
start.metric.dimensions["threshold_policy"],
"request_threshold"
);
let correlation_id = start.metric.correlation_id.as_deref().unwrap();
let post = metrics
.iter()
.find(|record| record.metric.name == "compact.post_request")
.unwrap();
assert_eq!(post.metric.correlation_id.as_deref(), Some(correlation_id));
}
#[tokio::test]