212 lines
8.3 KiB
Rust
212 lines
8.3 KiB
Rust
//! Tracks per-LLM-request Usage measurements within a Worker run.
|
|
//!
|
|
//! Bridge between two sync touchpoints in the Engine lifecycle:
|
|
//!
|
|
//! - **`pre_llm_request` hook** (async, but synchronously accessed via the
|
|
//! tracker): captures `history.len()` at the moment a request goes out.
|
|
//! - **`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.
|
|
//!
|
|
//! 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::sync::Mutex;
|
|
|
|
use llm_engine::UsageRecord;
|
|
use llm_engine::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 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 `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 Worker.
|
|
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()),
|
|
}
|
|
}
|
|
|
|
/// Called from a `pre_llm_request` hook with the current history length.
|
|
pub(crate) fn note_request(&self, history_len: usize) {
|
|
*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 `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 する）。
|
|
let input_total = event.input_tokens.unwrap_or(0);
|
|
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,
|
|
},
|
|
correlation_id,
|
|
});
|
|
}
|
|
|
|
/// Return a clone of the accumulated `UsageRecord`s 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
|
|
.lock()
|
|
.unwrap()
|
|
.iter()
|
|
.map(|r| r.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())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn make_event(input: u64, cache_read: u64, cache_write: u64, output: u64) -> UsageEvent {
|
|
UsageEvent {
|
|
input_tokens: Some(input),
|
|
output_tokens: Some(output),
|
|
total_tokens: Some(input + output),
|
|
cache_read_input_tokens: Some(cache_read),
|
|
cache_creation_input_tokens: Some(cache_write),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn pairs_history_len_with_usage_event() {
|
|
let tracker = UsageTracker::new();
|
|
tracker.note_request(5);
|
|
tracker.record_usage(&make_event(1000, 800, 100, 42));
|
|
|
|
let records = tracker.drain();
|
|
assert_eq!(records.len(), 1);
|
|
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]
|
|
fn records_clones_without_clearing() {
|
|
let tracker = UsageTracker::new();
|
|
tracker.note_request(1);
|
|
tracker.record_usage(&make_event(10, 0, 0, 5));
|
|
|
|
let records = tracker.records();
|
|
assert_eq!(records.len(), 1);
|
|
assert_eq!(records[0].history_len, 1);
|
|
assert_eq!(records[0].input_total_tokens, 10);
|
|
assert_eq!(tracker.records().len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn drain_clears_buffer() {
|
|
let tracker = UsageTracker::new();
|
|
tracker.note_request(1);
|
|
tracker.record_usage(&make_event(10, 0, 0, 5));
|
|
assert_eq!(tracker.drain().len(), 1);
|
|
assert_eq!(tracker.drain().len(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn usage_without_pending_history_len_is_dropped() {
|
|
let tracker = UsageTracker::new();
|
|
tracker.record_usage(&make_event(10, 0, 0, 5));
|
|
assert_eq!(tracker.drain().len(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn multiple_requests_in_one_run() {
|
|
let tracker = UsageTracker::new();
|
|
tracker.note_request(5);
|
|
tracker.record_usage(&make_event(100, 0, 0, 20));
|
|
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].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());
|
|
}
|
|
}
|