usageデータの永続化実装

crates/llm-worker/src/llm_client/scheme/anthropic/events.rs

@@ -265,12 +265,23 @@ impl AnthropicScheme {
     }
 
     fn convert_usage(&self, usage: &UsageData) -> UsageEvent {
-        let input = usage.input_tokens.unwrap_or(0);
+        // Anthropic の `input_tokens` は **キャッシュ外** の入力トークンのみで、
+        // プロンプト全長は input_tokens + cache_read + cache_creation。
+        // UsageEvent の `input_tokens` には「占有量（プロンプト全長）」を載せる
+        // 規約に合わせて、ここでキャッシュ分を足し込む。
+        // cache_read_input_tokens / cache_creation_input_tokens は内訳として
+        // 別フィールドに残るので、料金計算側で `input - cache_read - cache_creation`
+        // により非キャッシュ入力分は逆算可能。
+        let raw_input = usage.input_tokens.unwrap_or(0);
+        let cache_read = usage.cache_read_input_tokens.unwrap_or(0);
+        let cache_creation = usage.cache_creation_input_tokens.unwrap_or(0);
+        let input_total = raw_input + cache_read + cache_creation;
         let output = usage.output_tokens.unwrap_or(0);
+
         UsageEvent {
-            input_tokens: usage.input_tokens,
+            input_tokens: usage.input_tokens.map(|_| input_total),
             output_tokens: usage.output_tokens,
-            total_tokens: Some(input + output),
+            total_tokens: Some(input_total + output),
             cache_read_input_tokens: usage.cache_read_input_tokens,
             cache_creation_input_tokens: usage.cache_creation_input_tokens,
         }
@@ -289,12 +300,33 @@ mod tests {
         let event = scheme.parse_event("message_start", data).unwrap().unwrap();
         match event {
             Event::Usage(u) => {
+                // キャッシュなしなので input_total = raw_input = 10
                 assert_eq!(u.input_tokens, Some(10));
             }
             _ => panic!("Expected Usage event"),
         }
     }
 
+    #[test]
+    fn test_convert_usage_includes_cache_in_input_total() {
+        // Anthropic の input_tokens はキャッシュ外のみで、占有量は
+        // input + cache_read + cache_creation。
+        // UsageEvent.input_tokens は占有量に正規化される。
+        let scheme = AnthropicScheme::new();
+        let usage = UsageData {
+            input_tokens: Some(100),
+            output_tokens: Some(50),
+            cache_read_input_tokens: Some(800),
+            cache_creation_input_tokens: Some(200),
+        };
+        let event = scheme.convert_usage(&usage);
+        // 100 + 800 + 200 = 1100
+        assert_eq!(event.input_tokens, Some(1100));
+        assert_eq!(event.cache_read_input_tokens, Some(800));
+        assert_eq!(event.cache_creation_input_tokens, Some(200));
+        assert_eq!(event.total_tokens, Some(1150));
+    }
+
     #[test]
     fn test_parse_content_block_start_text() {
         let scheme = AnthropicScheme::new();

crates/llm-worker/src/timeline/timeline.rs

@@ -8,6 +8,33 @@ use std::marker::PhantomData;
 use super::event::*;
 use crate::handler::*;
 
+// =============================================================================
+// Helpers
+// =============================================================================
+
+/// 1リクエスト内で受信した複数 UsageEvent をマージする。
+/// 各フィールドについて新しい値が `Some` ならそれで上書き。
+/// プロバイダによっては input/cache 系を最初の event だけに載せ、
+/// output_tokens を後続 event で更新するため、最後の値だけを取るのではなく
+/// フィールド単位で latest-non-None を取る。
+fn merge_usage(acc: &mut UsageEvent, new: &UsageEvent) {
+    if new.input_tokens.is_some() {
+        acc.input_tokens = new.input_tokens;
+    }
+    if new.output_tokens.is_some() {
+        acc.output_tokens = new.output_tokens;
+    }
+    if new.total_tokens.is_some() {
+        acc.total_tokens = new.total_tokens;
+    }
+    if new.cache_read_input_tokens.is_some() {
+        acc.cache_read_input_tokens = new.cache_read_input_tokens;
+    }
+    if new.cache_creation_input_tokens.is_some() {
+        acc.cache_creation_input_tokens = new.cache_creation_input_tokens;
+    }
+}
+
 // =============================================================================
 // Type-erased Handler
 // =============================================================================
@@ -362,6 +389,12 @@ pub struct Timeline {
 
     // 現在アクティブなブロック
     current_block: Option<BlockType>,
+
+    // 1リクエスト内で受信した Usage event の集約バッファ。
+    // Anthropic は message_start と message_delta、Gemini は各チャンクと、
+    // 多くのプロバイダが複数 Usage を発行するため、リクエスト境界で
+    // 1度だけ発火するためにここでマージする。flush_usage() で発火する。
+    pending_usage: Option<UsageEvent>,
 }
 
 impl Default for Timeline {
@@ -381,6 +414,7 @@ impl Timeline {
             thinking_block_handlers: Vec::new(),
             tool_use_block_handlers: Vec::new(),
             current_block: None,
+            pending_usage: None,
         }
     }
 
@@ -491,9 +525,24 @@ impl Timeline {
         }
     }
 
+    /// Usage event を即時には dispatch せず、pending_usage にマージする。
+    /// 1リクエスト内で複数の Usage event が来ても、ハンドラには 1 度だけ
+    /// 最終値を渡したいため。flush_usage() で発火する。
     fn dispatch_usage(&mut self, event: &UsageEvent) {
-        for handler in &mut self.usage_handlers {
-            handler.dispatch(event);
+        match &mut self.pending_usage {
+            Some(acc) => merge_usage(acc, event),
+            None => self.pending_usage = Some(event.clone()),
+        }
+    }
+
+    /// pending_usage を usage_handlers に発火し、バッファをクリアする。
+    /// 1リクエスト分のストリーム終了時に1回だけ呼ぶ想定。
+    /// pending_usage が空ならば何もしない。
+    pub fn flush_usage(&mut self) {
+        if let Some(event) = self.pending_usage.take() {
+            for handler in &mut self.usage_handlers {
+                handler.dispatch(&event);
+            }
         }
     }
 
@@ -629,9 +678,63 @@ mod tests {
         timeline.on_usage(handler);
 
         timeline.dispatch(&Event::usage(100, 50));
+        // pending_usage に積まれているだけなのでまだ未発火
+        assert_eq!(calls.lock().unwrap().len(), 0);
 
+        // flush で 1 度だけ発火
+        timeline.flush_usage();
         let recorded = calls.lock().unwrap();
         assert_eq!(recorded.len(), 1);
         assert_eq!(recorded[0].input_tokens, Some(100));
     }
+
+    #[test]
+    fn test_usage_aggregation_and_flush() {
+        struct TestUsageHandler {
+            calls: Arc<Mutex<Vec<UsageEvent>>>,
+        }
+        impl Handler<UsageKind> for TestUsageHandler {
+            type Scope = ();
+            fn on_event(&mut self, _scope: &mut (), event: &UsageEvent) {
+                self.calls.lock().unwrap().push(event.clone());
+            }
+        }
+
+        let calls = Arc::new(Mutex::new(Vec::new()));
+        let mut timeline = Timeline::new();
+        timeline.on_usage(TestUsageHandler {
+            calls: calls.clone(),
+        });
+
+        // Anthropic 風: message_start で input + 暫定 output
+        timeline.dispatch(&Event::Usage(UsageEvent {
+            input_tokens: Some(409),
+            output_tokens: Some(1),
+            total_tokens: Some(410),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        }));
+        // message_delta で最終 output
+        timeline.dispatch(&Event::Usage(UsageEvent {
+            input_tokens: Some(409),
+            output_tokens: Some(71),
+            total_tokens: Some(480),
+            cache_read_input_tokens: Some(0),
+            cache_creation_input_tokens: Some(0),
+        }));
+
+        // 未 flush の段階では発火しない
+        assert_eq!(calls.lock().unwrap().len(), 0);
+
+        timeline.flush_usage();
+        let recorded = calls.lock().unwrap();
+        assert_eq!(recorded.len(), 1);
+        assert_eq!(recorded[0].input_tokens, Some(409));
+        assert_eq!(recorded[0].output_tokens, Some(71));
+
+        // flush 後にもう一度 flush しても何も起きない
+        drop(recorded);
+        timeline.flush_usage();
+        assert_eq!(calls.lock().unwrap().len(), 1);
+    }
 }

crates/llm-worker/src/worker.rs

@@ -803,7 +803,11 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
                                 }
                             }
                             let event = result
-                                .inspect_err(|_| self.last_run_interrupted = true)?;
+                                .inspect_err(|_| {
+                                    self.last_run_interrupted = true;
+                                    // 部分情報でも発火しておく（料金会計用）
+                                    self.timeline.flush_usage();
+                                })?;
                             self.timeline.dispatch(&event);
                         }
                         None => break,
@@ -814,11 +818,14 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
                         info!("Stream cancelled");
                     }
                     self.timeline.abort_current_block();
+                    self.timeline.flush_usage();
                     self.last_run_interrupted = true;
                     return Err(WorkerError::Cancelled);
                 }
             }
         }
+        // ストリーム完了時に集約済み Usage を 1 度だけ発火
+        self.timeline.flush_usage();
         debug!(event_count = event_count, "Stream completed");
         Ok(())
     }