feat: session-metrics実装
This commit is contained in:
@@ -0,0 +1,50 @@
|
||||
//! Sync buffer for `session_metrics::Metric` values queued from inside
|
||||
//! Worker callbacks (which run synchronously and cannot themselves
|
||||
//! perform `async` store writes).
|
||||
//!
|
||||
//! Pod drains this buffer in `persist_turn` and writes each metric via
|
||||
//! `session_metrics::record_metric`, alongside the regular `LlmUsage`
|
||||
//! entries.
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
use session_metrics::Metric;
|
||||
|
||||
pub(crate) struct MetricsTracker {
|
||||
pending: Mutex<Vec<Metric>>,
|
||||
}
|
||||
|
||||
impl MetricsTracker {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
pending: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Queue a metric for the next `persist_turn` flush.
|
||||
pub(crate) fn push(&self, metric: Metric) {
|
||||
self.pending.lock().unwrap().push(metric);
|
||||
}
|
||||
|
||||
/// Drain all queued metrics. Called by Pod after a run completes.
|
||||
pub(crate) fn drain(&self) -> Vec<Metric> {
|
||||
std::mem::take(&mut *self.pending.lock().unwrap())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn push_then_drain_returns_in_order_and_clears() {
|
||||
let t = MetricsTracker::new();
|
||||
t.push(Metric::now("a"));
|
||||
t.push(Metric::now("b"));
|
||||
let drained = t.drain();
|
||||
assert_eq!(drained.len(), 2);
|
||||
assert_eq!(drained[0].name, "a");
|
||||
assert_eq!(drained[1].name, "b");
|
||||
assert!(t.drain().is_empty());
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub(crate) mod metrics_tracker;
|
||||
pub(crate) mod prune;
|
||||
pub(crate) mod state;
|
||||
pub(crate) mod token_counter;
|
||||
|
||||
@@ -5,10 +5,16 @@
|
||||
//! 直後)。Worker は usage 履歴を知らないので、`min_savings` 判定に使う savings
|
||||
//! の見積もりはコールバックで外部から注入する。このモジュールはそのコールバック
|
||||
//! を組み立てて Worker に差し込むための `impl Pod` を提供する。
|
||||
//!
|
||||
//! 同じ経路で `PruneObserver` も install し、評価のたびに `prune.fire` /
|
||||
//! `prune.skip` metric を `MetricsTracker` に積む。`Fired` 時は uuid を
|
||||
//! `UsageTracker` にも stash しておき、後続の `LlmUsage` と組で
|
||||
//! `prune.post_request` を吐けるようにする。
|
||||
|
||||
use llm_worker::Item;
|
||||
use llm_worker::llm_client::client::LlmClient;
|
||||
use llm_worker::prune::{PruneConfig, SavingsEstimator};
|
||||
use llm_worker::prune::{PruneConfig, PruneDecision, PruneObserver, SavingsEstimator};
|
||||
use session_metrics::Metric;
|
||||
use session_store::Store;
|
||||
|
||||
use crate::Pod;
|
||||
@@ -24,6 +30,12 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
/// Measurement-less estimates (before the first LLM call, or immediately
|
||||
/// after a compact) return `0` from the estimator, which naturally
|
||||
/// prevents the prune projection from firing until usage data exists.
|
||||
///
|
||||
/// Also installs a [`PruneObserver`] that pushes `prune.fire` /
|
||||
/// `prune.skip` metrics into the shared [`MetricsTracker`]. On `Fired`
|
||||
/// the observer additionally stashes a fresh correlation_id in
|
||||
/// [`UsageTracker`] so the next `LlmUsage` can be paired with a
|
||||
/// `prune.post_request` metric carrying the same id.
|
||||
pub fn attach_prune(&mut self, config: PruneConfig) {
|
||||
let usage = self.usage_history_handle();
|
||||
let estimator: SavingsEstimator = Box::new(move |history: &[Item], indices| {
|
||||
@@ -34,9 +46,43 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
_ => est.tokens,
|
||||
}
|
||||
});
|
||||
|
||||
let metrics = self.metrics_tracker_handle();
|
||||
let usage_tracker = self.usage_tracker_handle();
|
||||
let observer: PruneObserver = Box::new(move |eval| {
|
||||
match &eval.decision {
|
||||
PruneDecision::Fired { .. } => {
|
||||
let correlation_id = uuid::Uuid::now_v7().to_string();
|
||||
let mut metric = Metric::now("prune.fire")
|
||||
.with_value(eval.estimated_savings as f64)
|
||||
.with_correlation_id(&correlation_id)
|
||||
.with_dimension("candidate_count", eval.candidate_count.to_string());
|
||||
if let Some(border) = eval.border_turn {
|
||||
metric = metric.with_dimension("border_turn", border.to_string());
|
||||
}
|
||||
metrics.push(metric);
|
||||
usage_tracker.note_correlation_id(correlation_id);
|
||||
}
|
||||
PruneDecision::SkippedNoCandidates => {
|
||||
metrics.push(
|
||||
Metric::now("prune.skip").with_dimension("reason", "no_candidates"),
|
||||
);
|
||||
}
|
||||
PruneDecision::SkippedBelowMinSavings => {
|
||||
metrics.push(
|
||||
Metric::now("prune.skip")
|
||||
.with_dimension("reason", "below_min_savings")
|
||||
.with_dimension("candidate_count", eval.candidate_count.to_string())
|
||||
.with_value(eval.estimated_savings as f64),
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let worker = self.worker_mut();
|
||||
worker.set_prune_config(Some(config));
|
||||
worker.set_savings_estimator(Some(estimator));
|
||||
worker.set_prune_observer(Some(observer));
|
||||
}
|
||||
|
||||
/// If the manifest has a `[compaction]` section, build a `PruneConfig`
|
||||
|
||||
@@ -19,19 +19,35 @@ use std::sync::Mutex;
|
||||
use llm_worker::UsageRecord;
|
||||
use llm_worker::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.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct RecordedUsage {
|
||||
pub(crate) record: UsageRecord,
|
||||
pub(crate) correlation_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Shared between the pre-request hook, the `on_usage` callback, and Pod.
|
||||
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>>,
|
||||
/// Records accumulated during the current run; drained by Pod.
|
||||
pending_records: Mutex<Vec<UsageRecord>>,
|
||||
pending_records: Mutex<Vec<RecordedUsage>>,
|
||||
}
|
||||
|
||||
impl UsageTracker {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
pending_history_len: Mutex::new(None),
|
||||
pending_correlation_id: Mutex::new(None),
|
||||
pending_records: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
@@ -41,16 +57,29 @@ 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`.
|
||||
pub(crate) fn note_correlation_id(&self, id: String) {
|
||||
*self.pending_correlation_id.lock().unwrap() = Some(id);
|
||||
}
|
||||
|
||||
/// Called from the `on_usage` callback with the aggregated final
|
||||
/// UsageEvent. If a `history_len` was previously stashed via
|
||||
/// `note_request`, builds a `UsageRecord` and pushes it onto the buffer.
|
||||
/// If not (e.g. test code that fires Usage outside a request), drops
|
||||
/// the event.
|
||||
/// `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.
|
||||
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 correlation_id = self.pending_correlation_id.lock().unwrap().take();
|
||||
// UsageEvent.input_tokens は scheme 層で「占有量(プロンプト全長)」に
|
||||
// 正規化済みである前提(Anthropic は cache_read + cache_creation を
|
||||
// 加算して emit する)。
|
||||
@@ -58,18 +87,21 @@ 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(UsageRecord {
|
||||
history_len,
|
||||
input_total_tokens: input_total,
|
||||
cache_read_tokens: cache_read,
|
||||
cache_write_tokens: cache_write,
|
||||
output_tokens: output,
|
||||
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,
|
||||
},
|
||||
correlation_id,
|
||||
});
|
||||
}
|
||||
|
||||
/// Drain accumulated records. Called by Pod after a run completes,
|
||||
/// before persisting the turn.
|
||||
pub(crate) fn drain(&self) -> Vec<UsageRecord> {
|
||||
pub(crate) fn drain(&self) -> Vec<RecordedUsage> {
|
||||
std::mem::take(&mut *self.pending_records.lock().unwrap())
|
||||
}
|
||||
}
|
||||
@@ -96,11 +128,12 @@ mod tests {
|
||||
|
||||
let records = tracker.drain();
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].history_len, 5);
|
||||
assert_eq!(records[0].input_total_tokens, 1000);
|
||||
assert_eq!(records[0].cache_read_tokens, 800);
|
||||
assert_eq!(records[0].cache_write_tokens, 100);
|
||||
assert_eq!(records[0].output_tokens, 42);
|
||||
assert_eq!(records[0].record.history_len, 5);
|
||||
assert_eq!(records[0].record.input_total_tokens, 1000);
|
||||
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());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -129,8 +162,25 @@ mod tests {
|
||||
|
||||
let records = tracker.drain();
|
||||
assert_eq!(records.len(), 2);
|
||||
assert_eq!(records[0].history_len, 5);
|
||||
assert_eq!(records[1].history_len, 10);
|
||||
assert_eq!(records[1].cache_read_tokens, 50);
|
||||
assert_eq!(records[0].record.history_len, 5);
|
||||
assert_eq!(records[1].record.history_len, 10);
|
||||
assert_eq!(records[1].record.cache_read_tokens, 50);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn correlation_id_pairs_with_next_record_only() {
|
||||
let tracker = UsageTracker::new();
|
||||
// Stash an ID, then run a request → the ID should land on this record.
|
||||
tracker.note_correlation_id("abc".into());
|
||||
tracker.note_request(5);
|
||||
tracker.record_usage(&make_event(100, 0, 0, 20));
|
||||
// Next request without a fresh stash → no correlation_id.
|
||||
tracker.note_request(10);
|
||||
tracker.record_usage(&make_event(200, 50, 0, 30));
|
||||
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
||||
+68
-2
@@ -77,6 +77,11 @@ pub struct Pod<C: LlmClient, St: Store> {
|
||||
/// Captures `(history_len, UsageEvent)` pairs during a run; drained
|
||||
/// in `persist_turn` and persisted as `LogEntry::LlmUsage` entries.
|
||||
usage_tracker: Arc<UsageTracker>,
|
||||
/// Sync-side buffer for `Metric` values queued from inside Worker
|
||||
/// callbacks (currently the prune observer). Drained in `persist_turn`
|
||||
/// and written via `session_metrics::record_metric` alongside
|
||||
/// `LogEntry::LlmUsage`. Always present after construction.
|
||||
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.
|
||||
/// Read by token-accounting APIs (`Pod::total_tokens`, etc.).
|
||||
@@ -203,6 +208,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
interceptor_installed: false,
|
||||
compact_state: None,
|
||||
usage_tracker: Arc::new(UsageTracker::new()),
|
||||
metrics_tracker: Arc::new(crate::compact::metrics_tracker::MetricsTracker::new()),
|
||||
usage_history: Arc::new(Mutex::new(Vec::<UsageRecord>::new())),
|
||||
tracker: None,
|
||||
system_prompt_template: None,
|
||||
@@ -391,6 +397,26 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
self.usage_history.clone()
|
||||
}
|
||||
|
||||
/// Handle to the per-LLM-request `UsageTracker`.
|
||||
///
|
||||
/// Sibling modules (e.g. the prune observer) clone this `Arc` to stash
|
||||
/// per-request side state (e.g. a `correlation_id`) that pairs with
|
||||
/// the next `LlmUsage`.
|
||||
pub(crate) fn usage_tracker_handle(&self) -> Arc<UsageTracker> {
|
||||
self.usage_tracker.clone()
|
||||
}
|
||||
|
||||
/// Handle to the synchronous `MetricsTracker` buffer.
|
||||
///
|
||||
/// Worker callbacks (e.g. the prune observer) clone this `Arc` and
|
||||
/// `.push(metric)` into it; Pod drains it in `persist_turn` and
|
||||
/// writes each metric via `session_metrics::record_metric`.
|
||||
pub(crate) fn metrics_tracker_handle(
|
||||
&self,
|
||||
) -> Arc<crate::compact::metrics_tracker::MetricsTracker> {
|
||||
self.metrics_tracker.clone()
|
||||
}
|
||||
|
||||
/// Attach the session-scoped file-operation tracker from the builtin
|
||||
/// `tools` crate. Called by the Controller immediately after it
|
||||
/// registers the builtin tools on the Worker. Overwrites any
|
||||
@@ -1068,13 +1094,36 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Flush any sync-buffered metrics from this run first
|
||||
// (currently `prune.fire` / `prune.skip` from the prune observer).
|
||||
// Ordered before LlmUsage so that a `prune.fire` and the
|
||||
// `prune.post_request` derived from the matching usage record
|
||||
// appear in the log close together.
|
||||
let pending_metrics = self.metrics_tracker.drain();
|
||||
for metric in pending_metrics {
|
||||
session_metrics::record_metric(
|
||||
&self.store,
|
||||
self.session_id,
|
||||
&mut self.head_hash,
|
||||
&metric,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
// 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 Pod::run). Each is also appended to
|
||||
// the in-memory `usage_history` so token-accounting APIs see it
|
||||
// before the next run.
|
||||
// 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 record in usage_records {
|
||||
for recorded in usage_records {
|
||||
let crate::compact::usage_tracker::RecordedUsage {
|
||||
record,
|
||||
correlation_id,
|
||||
} = recorded;
|
||||
session_store::save_usage(
|
||||
&self.store,
|
||||
self.session_id,
|
||||
@@ -1086,6 +1135,20 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
record.output_tokens,
|
||||
)
|
||||
.await?;
|
||||
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)
|
||||
.with_dimension("cache_write_tokens", record.cache_write_tokens.to_string())
|
||||
.with_dimension("history_len", record.history_len.to_string());
|
||||
session_metrics::record_metric(
|
||||
&self.store,
|
||||
self.session_id,
|
||||
&mut self.head_hash,
|
||||
&metric,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
self.usage_history
|
||||
.lock()
|
||||
.expect("usage_history poisoned")
|
||||
@@ -1895,6 +1958,7 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
|
||||
interceptor_installed: false,
|
||||
compact_state: None,
|
||||
usage_tracker: Arc::new(UsageTracker::new()),
|
||||
metrics_tracker: Arc::new(crate::compact::metrics_tracker::MetricsTracker::new()),
|
||||
usage_history: Arc::new(Mutex::new(Vec::new())),
|
||||
tracker: None,
|
||||
system_prompt_template: common.system_prompt_template,
|
||||
@@ -1956,6 +2020,7 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
|
||||
interceptor_installed: false,
|
||||
compact_state: None,
|
||||
usage_tracker: Arc::new(UsageTracker::new()),
|
||||
metrics_tracker: Arc::new(crate::compact::metrics_tracker::MetricsTracker::new()),
|
||||
usage_history: Arc::new(Mutex::new(Vec::new())),
|
||||
tracker: None,
|
||||
system_prompt_template: common.system_prompt_template,
|
||||
@@ -2067,6 +2132,7 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
|
||||
interceptor_installed: false,
|
||||
compact_state: None,
|
||||
usage_tracker: Arc::new(UsageTracker::new()),
|
||||
metrics_tracker: Arc::new(crate::compact::metrics_tracker::MetricsTracker::new()),
|
||||
usage_history: Arc::new(Mutex::new(state.usage_history)),
|
||||
tracker: None,
|
||||
// Restore replays the saved system_prompt verbatim — no
|
||||
|
||||
Reference in New Issue
Block a user