Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

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

Browse Source
This commit is contained in:
Hare
2026-04-28 15:43:34 +09:00
parent e49fb3f1a0
commit 6fe19b84ce
16 changed files with 449 additions and 436 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pod/src/compact/token_counter.rs
+14 -201
View File
@@ -1,47 +1,27 @@
//! Usage 履歴ベースのトークン会計。
//! Compact / prune 専用のトークン会計補助。
//!
//! `UsageRecord` の列（プロバイダ実測値）と現在の history から、
//! 「末尾 N トークン残すための split 位置」「prune 射影で節約される
//! トークン数」などを pure に計算する。
//! 汎用部分（`prefix_bytes`, `tokens_at`, `total_tokens`, `total_tokens_at`）は
//! [`llm_worker::token_counter`] にあり、`UsageRecord` の列と現在の history から
//! pure に推定する。本モジュールは compact / prune 固有のロジック
//! （`split_for_retained`, `savings_for_prune`）と、Pod 上の公開 API に
//! 限定する。
//!
//! # 方針
//!
//! - ローカルトークナイザは持たない。実測値があればそれを採用し、
//! measurement 間はバイト数で按分、最新 measurement より先は最終 rate で外挿する
//! - 推定の出どころは [`EstimateSource`] で呼び出し側に明示する。
//! 課金判断には使えないが、compact/prune の閾値判定には十分な精度
//! - `records` は `history_len` 昇順を仮定する（`collect_state` と
//! `UsageTracker` がそのように積む）
//!
//! 公開 API は本ファイル内の `impl Pod` で [`Pod`](crate::Pod) のメソッドとして
//! 生やしている。pure な補助関数はこのモジュール内に private に閉じる。
//! 課金判断には使えないが、compact / prune の閾値判定には十分な精度
use llm_worker::Item;
use llm_worker::llm_client::client::LlmClient;
use session_store::{Store, UsageRecord};
use llm_worker::token_counter::{item_bytes, prefix_bytes, tokens_at};
use llm_worker::{Item, UsageRecord};
use session_store::Store;
pub use llm_worker::token_counter::{EstimateSource, TokenEstimate};
use crate::Pod;
/// 推定の出どころ。
#[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,
}
/// history を分割する位置。
///
/// `items[..index]` が捨てる／要約される側、`items[index..]` が残る側。
@@ -51,141 +31,6 @@ pub struct SplitPoint {
pub source: EstimateSource,
}
/// `items[..i]` までの累積バイト数（`prefix[i]`）を返す。長さは `items.len()+1`。
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 は比率でキャンセルされる。
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²) になる）。
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"),
}
}
pub(crate) fn total_tokens_impl(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(crate) fn total_tokens_at_impl(
history: &[Item],
records: &[UsageRecord],
history_len: usize,
) -> TokenEstimate {
let prefix = prefix_bytes(history);
tokens_at(history, records, history_len.min(history.len()), &prefix)
}
fn split_for_retained_impl(history: &[Item], records: &[UsageRecord], retained: u64) -> SplitPoint {
let prefix = prefix_bytes(history);
let current = tokens_at(history, records, history.len(), &prefix);
@@ -300,7 +145,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// 最後の measurement と、その後に追加された未測定分のバイト按分／外挿。
pub fn total_tokens(&self) -> TokenEstimate {
let usage = self.usage_history();
total_tokens_impl(self.history(), &usage)
llm_worker::token_counter::total_tokens(self.history(), &usage)
}
/// 任意の history index 時点でのプロンプト全長推定。
@@ -311,7 +156,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// pointer 以降に増えたプロンプト長を測るのに使う。
pub fn total_tokens_at(&self, history_len: usize) -> TokenEstimate {
let usage = self.usage_history();
total_tokens_at_impl(self.history(), &usage, history_len)
llm_worker::token_counter::total_tokens_at(self.history(), &usage, history_len)
}
/// 末尾から `retained` トークン以上を残すための分割位置。
@@ -341,38 +186,6 @@ mod tests {
}
}
#[test]
fn total_no_data_falls_back_to_byte_estimate() {
let history = vec![msg("hello world")];
let est = total_tokens_impl(&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_impl(&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_impl(&history, &records);
assert_eq!(est.source, EstimateSource::Extrapolated);
assert!(est.tokens > 100);
}
#[test]
fn total_zero_history_is_zero() {
let est = total_tokens_impl(&[], &[]);
assert_eq!(est.tokens, 0);
}
#[test]
fn split_returns_zero_when_current_below_retained() {
let history = vec![msg("a"), msg("b")];
crates/pod/src/compact/usage_tracker.rs
+1 -1
View File
@@ -16,8 +16,8 @@
use std::sync::Mutex;
use llm_worker::UsageRecord;
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 {
crates/pod/src/ipc/interceptor.rs
+3 -3
View File
@@ -16,12 +16,12 @@ use llm_worker::interceptor::{
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo,
ToolResultInfo, TurnEndAction,
};
use llm_worker::UsageRecord;
use llm_worker::tool::ToolOutput;
use session_store::UsageRecord;
use tracing::info;
use crate::compact::state::CompactState;
use crate::compact::token_counter::total_tokens_impl;
use llm_worker::token_counter::total_tokens;
use crate::hook::{
AbortInfo, HookRegistry, PreRequestInfo, PromptSubmitInfo, ToolCallSummary, ToolResultSummary,
TurnEndInfo,
@@ -82,7 +82,7 @@ impl PodInterceptor {
fn estimated_tokens(&self, context: &[Item]) -> Option<u64> {
let handle = self.usage_history.as_ref()?;
let records = handle.lock().expect("usage_history poisoned").clone();
Some(total_tokens_impl(context, &records).tokens)
Some(total_tokens(context, &records).tokens)
}
}
crates/pod/src/pod.rs
+24 -21
View File
@@ -6,10 +6,8 @@ use llm_worker::Item;
use llm_worker::llm_client::RequestConfig;
use llm_worker::llm_client::client::LlmClient;
use llm_worker::state::Mutable;
use llm_worker::{ToolOutputLimits, Worker, WorkerError, WorkerResult};
use session_store::{
EntryHash, Outcome, SessionId, SessionStartState, Store, StoreError, UsageRecord,
};
use llm_worker::{ToolOutputLimits, UsageRecord, Worker, WorkerError, WorkerResult};
use session_store::{EntryHash, SessionId, SessionStartState, Store, StoreError};
use tracing::{info, warn};
use manifest::{PodManifest, PodManifestConfig, ResolveError, Scope, ScopeError, WorkerManifest};
@@ -963,23 +961,28 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
}
let interrupted = self.worker.as_ref().unwrap().last_run_interrupted();
let outcome = match result {
Ok(WorkerResult::Finished) => Outcome::Finished,
Ok(WorkerResult::Paused) => Outcome::Paused,
Ok(WorkerResult::LimitReached) => Outcome::LimitReached,
Ok(WorkerResult::Yielded) => Outcome::Yielded,
Err(e) => Outcome::Error {
message: e.to_string(),
},
};
session_store::save_outcome(
&self.store,
self.session_id,
&mut self.head_hash,
outcome,
interrupted,
)
.await?;
match result {
Ok(r) => {
session_store::save_run_completed(
&self.store,
self.session_id,
&mut self.head_hash,
r.clone(),
interrupted,
)
.await?;
}
Err(e) => {
session_store::save_run_errored(
&self.store,
self.session_id,
&mut self.head_hash,
e.to_string(),
interrupted,
)
.await?;
}
}
Ok(())
}