usageデータの永続化実装

crates/pod/src/lib.rs

@@ -10,6 +10,7 @@ mod compact_interceptor;
 mod compact_state;
 mod hook_interceptor;
 mod pod;
+mod usage_tracker;
 
 pub use controller::{PodController, PodHandle};
 pub use manifest::{PodManifest, ProviderConfig, ProviderKind, Scope};

crates/pod/src/pod.rs

@@ -20,6 +20,24 @@ use crate::hook::{
     PreToolCall,
 };
 use crate::hook_interceptor::HookInterceptor;
+use crate::usage_tracker::UsageTracker;
+use llm_worker::interceptor::PreRequestAction;
+use async_trait::async_trait;
+
+/// Pre-LLM-request hook that records `history.len()` at send time into a
+/// shared `UsageTracker`. The on_usage callback later pairs this with the
+/// aggregated UsageEvent to produce one `UsageRecord` per LLM call.
+struct UsageTrackingHook {
+    tracker: Arc<UsageTracker>,
+}
+
+#[async_trait]
+impl Hook<PreLlmRequest> for UsageTrackingHook {
+    async fn call(&self, context: &mut Vec<Item>) -> PreRequestAction {
+        self.tracker.note_request(context.len());
+        PreRequestAction::Continue
+    }
+}
 
 const SUMMARY_SYSTEM_PROMPT: &str = "\
 You are a context compaction assistant. \
@@ -53,6 +71,10 @@ pub struct Pod<C: LlmClient, St: Store> {
     manifest_dir: Option<PathBuf>,
     /// Shared compaction state (present when compact_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.
+    usage_tracker: Arc<UsageTracker>,
     /// Session-lifetime file-operation tracker from the builtin `tools`
     /// crate. Populated by the Controller when it registers the builtin
     /// tools so that Pod-owned operations (e.g. compaction) can consult
@@ -85,6 +107,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
             interceptor_installed: false,
             manifest_dir: None,
             compact_state: None,
+            usage_tracker: Arc::new(UsageTracker::new()),
             tracker: None,
         })
     }
@@ -118,6 +141,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
             interceptor_installed: false,
             manifest_dir: None,
             compact_state: None,
+            usage_tracker: Arc::new(UsageTracker::new()),
             tracker: None,
         })
     }
@@ -220,6 +244,14 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
     /// `on_usage` callback to track `input_tokens`.
     fn ensure_interceptor_installed(&mut self) {
         if !self.interceptor_installed {
+            // Pre-LLM-request hook: capture history.len() into the
+            // UsageTracker so the upcoming on_usage callback can pair
+            // it with the measured input_tokens.
+            self.hook_builder
+                .add_pre_llm_request(UsageTrackingHook {
+                    tracker: self.usage_tracker.clone(),
+                });
+
             let builder = std::mem::take(&mut self.hook_builder);
             let registry = Arc::new(builder.build());
             let hook_interceptor = HookInterceptor::new(registry);
@@ -230,6 +262,11 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
                 .as_ref()
                 .and_then(|c| c.compact_threshold);
 
+            // Usage tracking via on_usage callback. Independent of
+            // compact_threshold so that LlmUsage entries are persisted
+            // unconditionally.
+            let tracker_for_usage = self.usage_tracker.clone();
+
             if let Some(threshold) = compact_threshold {
                 let retained = self
                     .manifest
@@ -240,18 +277,23 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
 
                 let state = Arc::new(CompactState::new(threshold, retained));
 
-                // Track input_tokens via on_usage callback.
+                // Combined on_usage: feed both the legacy compact threshold
+                // tracker and the new UsageTracker.
                 let state_for_usage = state.clone();
                 self.worker_mut().on_usage(move |event| {
                     if let Some(tokens) = event.input_tokens {
                         state_for_usage.update_input_tokens(tokens);
                     }
+                    tracker_for_usage.record_usage(event);
                 });
 
                 let interceptor = CompactInterceptor::new(hook_interceptor, state.clone());
                 self.worker_mut().set_interceptor(interceptor);
                 self.compact_state = Some(state);
             } else {
+                self.worker_mut().on_usage(move |event| {
+                    tracker_for_usage.record_usage(event);
+                });
                 self.worker_mut().set_interceptor(hook_interceptor);
             }
 
@@ -439,6 +481,24 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
         )
         .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).
+        let usage_records = self.usage_tracker.drain();
+        for record in usage_records {
+            session_store::save_usage(
+                &self.store,
+                self.session_id,
+                &mut self.head_hash,
+                record.history_len,
+                record.input_total_tokens,
+                record.cache_read_tokens,
+                record.cache_write_tokens,
+                record.output_tokens,
+            )
+            .await?;
+        }
+
         let interrupted = self.worker.as_ref().unwrap().last_run_interrupted();
         let outcome = match result {
             Ok(WorkerResult::Finished) => Outcome::Finished,
@@ -597,6 +657,7 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
             interceptor_installed: false,
             manifest_dir,
             compact_state: None,
+            usage_tracker: Arc::new(UsageTracker::new()),
             tracker: None,
         })
     }

crates/pod/src/usage_tracker.rs

@@ -0,0 +1,136 @@
+//! Tracks per-LLM-request Usage measurements within a Pod run.
+//!
+//! Bridge between two sync touchpoints in the Worker 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. Pod drains them
+//! in `persist_turn` and writes them as `LogEntry::LlmUsage` entries.
+//!
+//! Multiple LLM calls per Pod 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_worker::timeline::event::UsageEvent;
+use session_store::UsageRecord;
+
+/// 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>>,
+    /// Records accumulated during the current run; drained by Pod.
+    pending_records: Mutex<Vec<UsageRecord>>,
+}
+
+impl UsageTracker {
+    pub(crate) fn new() -> Self {
+        Self {
+            pending_history_len: 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);
+    }
+
+    /// 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.
+    pub(crate) fn record_usage(&self, event: &UsageEvent) {
+        let history_len = match self.pending_history_len.lock().unwrap().take() {
+            Some(n) => n,
+            None => return,
+        };
+        // 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(UsageRecord {
+            history_len,
+            input_total_tokens: input_total,
+            cache_read_tokens: cache_read,
+            cache_write_tokens: cache_write,
+            output_tokens: output,
+        });
+    }
+
+    /// Drain accumulated records. Called by Pod after a run completes,
+    /// before persisting the turn.
+    pub(crate) fn drain(&self) -> Vec<UsageRecord> {
+        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].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);
+    }
+
+    #[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].history_len, 5);
+        assert_eq!(records[1].history_len, 10);
+        assert_eq!(records[1].cache_read_tokens, 50);
+    }
+}