セッション関連の責務の分離

crates/llm-worker/src/lib.rs

@@ -48,12 +48,15 @@ pub mod llm_client;
 pub mod prune;
 pub mod state;
 pub mod timeline;
+pub mod token_counter;
 pub mod tool;
 pub mod tool_server;
+pub mod usage_record;
 
 pub use callback::{TextBlockScope, ToolUseBlockScope};
 pub use handler::ToolUseBlockStart;
 pub use interceptor::Interceptor;
 pub use message::{ContentPart, Item, Message, Role};
 pub use tool::{ToolCall, ToolOutputLimits, ToolResult};
+pub use usage_record::UsageRecord;
 pub use worker::{RunOutput, ToolRegistryError, Worker, WorkerConfig, WorkerError, WorkerResult};

crates/llm-worker/src/token_counter.rs

@@ -0,0 +1,222 @@
+//! Usage 履歴ベースのトークン会計（汎用部分）。
+//!
+//! `UsageRecord` の列（プロバイダ実測値）と現在の history から、
+//! 任意の history index 時点のプロンプト全長トークン数を pure に計算する。
+//!
+//! # 方針
+//!
+//! - ローカルトークナイザは持たない。実測値があればそれを採用し、
+//!   measurement 間はバイト数で按分、最新 measurement より先は最終 rate で外挿する
+//! - 推定の出どころは [`EstimateSource`] で呼び出し側に明示する。
+//!   課金判断には使えないが、compact / prune / memory phase 1 trigger 等の
+//!   閾値判定には十分な精度
+//! - `records` は `history_len` 昇順を仮定する（呼び出し側がそのように積む）
+
+use crate::{Item, UsageRecord};
+
+/// 推定の出どころ。
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum EstimateSource {
+    /// measurement の境界にちょうど一致（実測値そのもの）
+    Measured,
+    /// 連続する 2 つの measurement の間をバイト按分で計算
+    Interpolated,
+    /// 最後の measurement より新しい区間を最終 rate で外挿
+    Extrapolated,
+    /// measurement が 1 件も無く、バイト数のみのフォールバック
+    NoData,
+}
+
+/// トークン数の推定値。
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct TokenEstimate {
+    pub tokens: u64,
+    pub source: EstimateSource,
+}
+
+/// `items[..i]` までの累積バイト数（`prefix[i]`）を返す。長さは `items.len()+1`。
+pub fn prefix_bytes(items: &[Item]) -> Vec<u64> {
+    let mut prefix = Vec::with_capacity(items.len() + 1);
+    let mut acc: u64 = 0;
+    prefix.push(0);
+    for item in items {
+        acc = acc.saturating_add(item_bytes(item));
+        prefix.push(acc);
+    }
+    prefix
+}
+
+/// 1 Item の大きさ。JSON シリアライズ長を使う粗い近似。
+/// トークン数との絶対変換ではなく区間の按分にしか使わないので、
+/// プロバイダごとの overhead は比率でキャンセルされる。
+pub fn item_bytes(item: &Item) -> u64 {
+    serde_json::to_string(item)
+        .map(|s| s.len() as u64)
+        .unwrap_or(0)
+}
+
+/// `history[..index]` までのトークン数を推定する。
+///
+/// `prefix` は [`prefix_bytes`] で得た `history.len() + 1` 長の累積バイト列。
+/// 呼び出し側が 1 度だけ計算して使い回すことで、線形探索や複数回の推定が
+/// O(n) シリアライズで済む（内部で毎回再計算すると O(n²) になる）。
+pub fn tokens_at(
+    history: &[Item],
+    records: &[UsageRecord],
+    index: usize,
+    prefix: &[u64],
+) -> TokenEstimate {
+    debug_assert!(index <= history.len());
+    debug_assert_eq!(prefix.len(), history.len() + 1);
+
+    if index == 0 {
+        return TokenEstimate {
+            tokens: 0,
+            source: EstimateSource::Measured,
+        };
+    }
+
+    if records.is_empty() {
+        return TokenEstimate {
+            tokens: prefix[index] / 4,
+            source: EstimateSource::NoData,
+        };
+    }
+
+    // exact match（rev 走査で一番新しい record を採用）
+    if let Some(r) = records.iter().rev().find(|r| r.history_len == index) {
+        return TokenEstimate {
+            tokens: r.input_total_tokens,
+            source: EstimateSource::Measured,
+        };
+    }
+
+    let lower = records.iter().rev().find(|r| r.history_len < index);
+    let upper = records.iter().find(|r| r.history_len > index);
+    let cap = history.len();
+
+    match (lower, upper) {
+        (Some(lo), Some(up)) => {
+            let lo_bytes = prefix[lo.history_len.min(cap)];
+            let up_bytes = prefix[up.history_len.min(cap)];
+            let at_bytes = prefix[index];
+            let span_bytes = up_bytes.saturating_sub(lo_bytes);
+            let span_tokens = up.input_total_tokens.saturating_sub(lo.input_total_tokens);
+            if span_bytes == 0 || span_tokens == 0 {
+                return TokenEstimate {
+                    tokens: lo.input_total_tokens,
+                    source: EstimateSource::Interpolated,
+                };
+            }
+            let delta_bytes = at_bytes.saturating_sub(lo_bytes);
+            let delta_tokens =
+                (delta_bytes as u128 * span_tokens as u128 / span_bytes as u128) as u64;
+            TokenEstimate {
+                tokens: lo.input_total_tokens + delta_tokens,
+                source: EstimateSource::Interpolated,
+            }
+        }
+        (Some(lo), None) => {
+            let lo_bytes = prefix[lo.history_len.min(cap)];
+            let at_bytes = prefix[index];
+            if lo_bytes == 0 || lo.input_total_tokens == 0 {
+                return TokenEstimate {
+                    tokens: lo.input_total_tokens,
+                    source: EstimateSource::Extrapolated,
+                };
+            }
+            let delta_bytes = at_bytes.saturating_sub(lo_bytes);
+            let delta_tokens =
+                (delta_bytes as u128 * lo.input_total_tokens as u128 / lo_bytes as u128) as u64;
+            TokenEstimate {
+                tokens: lo.input_total_tokens + delta_tokens,
+                source: EstimateSource::Extrapolated,
+            }
+        }
+        (None, Some(up)) => {
+            let up_bytes = prefix[up.history_len.min(cap)];
+            let at_bytes = prefix[index];
+            if up_bytes == 0 {
+                return TokenEstimate {
+                    tokens: 0,
+                    source: EstimateSource::Interpolated,
+                };
+            }
+            let t = (at_bytes as u128 * up.input_total_tokens as u128 / up_bytes as u128) as u64;
+            TokenEstimate {
+                tokens: t,
+                source: EstimateSource::Interpolated,
+            }
+        }
+        (None, None) => unreachable!("records non-empty but neither lower nor upper matched"),
+    }
+}
+
+/// 現在の history 全体の推定トークン数。
+pub fn total_tokens(history: &[Item], records: &[UsageRecord]) -> TokenEstimate {
+    let prefix = prefix_bytes(history);
+    tokens_at(history, records, history.len(), &prefix)
+}
+
+/// 任意の history index 時点でのプロンプト全長推定。
+/// `history_len == 0` で 0 を返す。delta 計算 (extract trigger 等) で
+/// `total_tokens_at(now) - total_tokens_at(pointer)` の形で使う。
+pub fn total_tokens_at(
+    history: &[Item],
+    records: &[UsageRecord],
+    history_len: usize,
+) -> TokenEstimate {
+    let prefix = prefix_bytes(history);
+    tokens_at(history, records, history_len.min(history.len()), &prefix)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn msg(text: &str) -> Item {
+        Item::user_message(text)
+    }
+
+    fn record(history_len: usize, tokens: u64) -> UsageRecord {
+        UsageRecord {
+            history_len,
+            input_total_tokens: tokens,
+            cache_read_tokens: 0,
+            cache_write_tokens: 0,
+            output_tokens: 0,
+        }
+    }
+
+    #[test]
+    fn total_no_data_falls_back_to_byte_estimate() {
+        let history = vec![msg("hello world")];
+        let est = total_tokens(&history, &[]);
+        assert_eq!(est.source, EstimateSource::NoData);
+        assert!(est.tokens > 0);
+    }
+
+    #[test]
+    fn total_measured_when_last_record_matches_history_len() {
+        let history = vec![msg("a"), msg("b"), msg("c")];
+        let records = vec![record(3, 120)];
+        let est = total_tokens(&history, &records);
+        assert_eq!(est.source, EstimateSource::Measured);
+        assert_eq!(est.tokens, 120);
+    }
+
+    #[test]
+    fn total_extrapolated_when_history_grew_past_last_measurement() {
+        let history = vec![msg("a"), msg("b"), msg("c"), msg("d")];
+        let records = vec![record(3, 100)];
+        let est = total_tokens(&history, &records);
+        assert_eq!(est.source, EstimateSource::Extrapolated);
+        assert!(est.tokens > 100);
+    }
+
+    #[test]
+    fn total_zero_history_is_zero() {
+        let est = total_tokens(&[], &[]);
+        assert_eq!(est.tokens, 0);
+    }
+}

crates/llm-worker/src/usage_record.rs

@@ -0,0 +1,22 @@
+//! Per-LLM-request Usage measurement snapshot.
+//!
+//! 1 リクエストの送信時点での「ある history prefix 長で計測した占有量」を
+//! 1 件分にまとめたもの。`UsageEvent` (provider stream イベント) を
+//! 受けて呼び出し側 (typically Pod) が組み立て、永続化層
+//! (session-store) に流したり、token accounting (`token_counter`) で
+//! 履歴として参照したりする。
+
+/// LLM リクエスト送信時点での占有量スナップショット。
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct UsageRecord {
+    /// 送信時の history.len()
+    pub history_len: usize,
+    /// history[..history_len] の占有量（プロンプト全長、実測）
+    pub input_total_tokens: u64,
+    /// 上記のうちキャッシュから読み出された分
+    pub cache_read_tokens: u64,
+    /// 上記のうちこのリクエストでキャッシュに書かれた分
+    pub cache_write_tokens: u64,
+    /// このリクエストで生成された出力トークン数
+    pub output_tokens: u64,
+}

crates/llm-worker/src/worker.rs

@@ -66,7 +66,8 @@ pub struct WorkerConfig {
 }
 
 /// Worker execution result (status)
-#[derive(Debug)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
 pub enum WorkerResult {
     /// Completed (waiting for user input)
     Finished,