token-counter実装
This commit is contained in:
parent
f607a52fbb
commit
2edc2dc245
4
TODO.md
4
TODO.md
|
|
@ -2,8 +2,8 @@
|
|||
- [ ] ツール設計
|
||||
- [ ] Bash ツール (Permission 層と統合) → [tickets/bash-tool.md](tickets/bash-tool.md)
|
||||
- [ ] Scope の再設計 (pwd + writable、必須化) → [tickets/scope-redesign.md](tickets/scope-redesign.md)
|
||||
- [ ] Usage 履歴の永続化 (session-store に LlmUsage entry 追加) → [tickets/usage-history.md](tickets/usage-history.md)
|
||||
- [ ] トークン会計 (Usage 履歴ベースの retained/savings 計算) → [tickets/token-counter.md](tickets/token-counter.md)
|
||||
- [ ] Prune をコンテキスト射影に変更 → [tickets/prune-projection.md](tickets/prune-projection.md)
|
||||
- [ ] Prune の savings 推定を正確にする → [tickets/prune-savings-estimation.md](tickets/prune-savings-estimation.md)
|
||||
- [ ] Compact の改善（要約品質 + 挙動詳細） → [tickets/compact-improvements.md](tickets/compact-improvements.md)
|
||||
- [ ] Protocol の設計 → [tickets/protocol-design.md](tickets/protocol-design.md)
|
||||
- [ ] パーミッション: パターンベースのツール実行制御 → [tickets/permission-extension-point.md](tickets/permission-extension-point.md)
|
||||
|
|
|
|||
|
|
@ -95,8 +95,18 @@ fn item_bytes(item: &Item) -> u64 {
|
|||
}
|
||||
|
||||
/// `history[..index]` までのトークン数を推定する。
|
||||
fn tokens_at(history: &[Item], records: &[UsageRecord], index: usize) -> TokenEstimate {
|
||||
///
|
||||
/// `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 {
|
||||
|
|
@ -106,7 +116,6 @@ fn tokens_at(history: &[Item], records: &[UsageRecord], index: usize) -> TokenEs
|
|||
}
|
||||
|
||||
if records.is_empty() {
|
||||
let prefix = prefix_bytes(history);
|
||||
return TokenEstimate {
|
||||
tokens: prefix[index] / 4,
|
||||
source: EstimateSource::NoData,
|
||||
|
|
@ -123,7 +132,6 @@ fn tokens_at(history: &[Item], records: &[UsageRecord], index: usize) -> TokenEs
|
|||
|
||||
let lower = records.iter().rev().find(|r| r.history_len < index);
|
||||
let upper = records.iter().find(|r| r.history_len > index);
|
||||
let prefix = prefix_bytes(history);
|
||||
let cap = history.len();
|
||||
|
||||
match (lower, upper) {
|
||||
|
|
@ -186,7 +194,8 @@ fn tokens_at(history: &[Item], records: &[UsageRecord], index: usize) -> TokenEs
|
|||
}
|
||||
|
||||
fn total_tokens_impl(history: &[Item], records: &[UsageRecord]) -> TokenEstimate {
|
||||
tokens_at(history, records, history.len())
|
||||
let prefix = prefix_bytes(history);
|
||||
tokens_at(history, records, history.len(), &prefix)
|
||||
}
|
||||
|
||||
fn split_for_retained_impl(
|
||||
|
|
@ -194,7 +203,8 @@ fn split_for_retained_impl(
|
|||
records: &[UsageRecord],
|
||||
retained: u64,
|
||||
) -> SplitPoint {
|
||||
let current = total_tokens_impl(history, records);
|
||||
let prefix = prefix_bytes(history);
|
||||
let current = tokens_at(history, records, history.len(), &prefix);
|
||||
if current.tokens <= retained {
|
||||
return SplitPoint {
|
||||
index: 0,
|
||||
|
|
@ -204,11 +214,12 @@ fn split_for_retained_impl(
|
|||
let target = current.tokens - retained;
|
||||
|
||||
// `tokens_at` が target 以上になる最小の idx を線形探索。
|
||||
// history.len() は高々数百〜数千なので十分速い。将来ボトルネックになれば
|
||||
// prefix を使い回すので 1 回の split 呼び出しあたり O(n) で済む
|
||||
// （内部で毎回再計算すると O(n²) になる）。将来ボトルネックになれば
|
||||
// record 境界で二分探索に置き換える。
|
||||
let mut chosen_source = current.source;
|
||||
for idx in 1..=history.len() {
|
||||
let est = tokens_at(history, records, idx);
|
||||
let est = tokens_at(history, records, idx, &prefix);
|
||||
if est.tokens >= target {
|
||||
chosen_source = est.source;
|
||||
return SplitPoint {
|
||||
|
|
@ -234,8 +245,9 @@ pub(crate) fn savings_for_drop_impl(
|
|||
source: EstimateSource::Measured,
|
||||
};
|
||||
}
|
||||
let s = tokens_at(history, records, range.start);
|
||||
let e = tokens_at(history, records, range.end);
|
||||
let prefix = prefix_bytes(history);
|
||||
let s = tokens_at(history, records, range.start, &prefix);
|
||||
let e = tokens_at(history, records, range.end, &prefix);
|
||||
TokenEstimate {
|
||||
tokens: e.tokens.saturating_sub(s.tokens),
|
||||
source: s.source.worst(e.source),
|
||||
|
|
|
|||
43
tickets/prune-projection.md
Normal file
43
tickets/prune-projection.md
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
# Prune をコンテキスト射影に変更
|
||||
|
||||
## 背景
|
||||
|
||||
現状の `apply_prune` は Worker が保持する history の `Item::ToolResult { content }`
|
||||
を直接 `None` に書き換えている。PreLlmRequest hook の `context: &mut Vec<Item>` は
|
||||
Worker の history そのものなので、prune すると元の content が永久に失われる。
|
||||
|
||||
問題:
|
||||
- session-store に persist される history が prune 済みになり、restore しても content が戻らない
|
||||
- compact worker が要約を作る際に、本来参照できるはずの content が消えている
|
||||
- 「どこまで prune したか」という状態が暗黙的（content が None かどうか）
|
||||
|
||||
## 方針
|
||||
|
||||
永続化された記録（Worker の history / session-store のログ）は変えず、
|
||||
LLM に送るコンテキストを組み立てる段階で content を省く。
|
||||
|
||||
Prune は「どの ToolResult の content を省略するか」を決める射影ロジックであり、
|
||||
history の変換ではない。
|
||||
|
||||
### 設計の方向
|
||||
|
||||
1. prune の判定結果を「省略対象の index 集合」として保持する
|
||||
（現在の `prunable_indices` がそのまま使える）
|
||||
2. PreLlmRequest hook でコンテキストを構築する際に、省略対象の
|
||||
ToolResult は content を含めずに組み立てる
|
||||
3. Worker の history は一切変更しない
|
||||
|
||||
### 影響範囲
|
||||
|
||||
- `crates/llm-worker/src/prune.rs`: `apply_prune` を廃止または射影版に置き換え
|
||||
- `crates/pod/src/prune_hook.rs`: history を mutate せず、射影したコンテキストを返す
|
||||
- PreLlmRequest hook の戻り値: 現在 `PreRequestAction::Continue` で history を
|
||||
そのまま使う設計。射影したコンテキストを渡す方法の検討が必要
|
||||
|
||||
## 依存
|
||||
|
||||
- なし
|
||||
|
||||
## ブロックする後続
|
||||
|
||||
- [compact-improvements.md](compact-improvements.md) — compact worker が content を参照する前提
|
||||
40
tickets/prune-savings-estimation.md
Normal file
40
tickets/prune-savings-estimation.md
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
# Prune の savings 推定を正確にする
|
||||
|
||||
## 背景
|
||||
|
||||
現在の PruneHook は `savings_for_drop_impl(context, &snapshot, first..last)` で
|
||||
候補範囲を「丸ごと drop した場合」の savings を計算し、`min_savings` と比較している。
|
||||
|
||||
しかし prune が実際に行うのは ToolResult の content を省略するだけで、
|
||||
item 自体（summary、メタデータ）は残る。そのため `savings_for_drop` は
|
||||
実際の節約量を過大評価しており、本来 prune 不要な場面でも発動しうる。
|
||||
|
||||
savings の推定は prune 側の責務であり、token_counter の汎用 API に
|
||||
prune 固有の挙動を押し込むべきではない。
|
||||
|
||||
## 方針
|
||||
|
||||
PruneHook が「content 部分だけの savings」を計算する。
|
||||
|
||||
### 計算方法
|
||||
|
||||
候補の各 ToolResult について:
|
||||
- content ありの item のトークン推定
|
||||
- content を None にした場合（summary のみ）のトークン推定
|
||||
- 差分が prune による実際の savings
|
||||
|
||||
バイト数の差分を measurement 由来の rate で換算するか、
|
||||
`tokens_at` を使った前後比較にするかは実装時判断。
|
||||
|
||||
### 影響範囲
|
||||
|
||||
- `crates/pod/src/prune_hook.rs`: savings 計算ロジックの置き換え
|
||||
- `crates/pod/src/token_counter.rs`: 必要に応じて content-level の推定ヘルパーを追加
|
||||
|
||||
## 依存
|
||||
|
||||
- [prune-projection.md](prune-projection.md) — prune が射影ベースになった後の方が設計しやすい
|
||||
|
||||
## ブロックする後続
|
||||
|
||||
- なし（チューニングの精度改善）
|
||||
|
|
@ -1,155 +0,0 @@
|
|||
# トークン会計 (Usage 履歴ベース)
|
||||
|
||||
## 背景
|
||||
|
||||
Compact / Prune の挙動改善に「**履歴上の任意位置のトークン数**」と
|
||||
「**ある変更でどれだけトークンが浮くか**」を答えられる仕組みが要る。
|
||||
|
||||
ローカル近似（`len/4` や BPE テーブル）はモデル/言語/ツール overhead の
|
||||
誤差が大きく、Anthropic に至ってはオフラインで正確に数える手段が無い。
|
||||
|
||||
一方、LLM レスポンスの `Usage` は **送信した history prefix に対して
|
||||
プロバイダが実測したトークン数** であり、これが手元にある最も正確な情報源。
|
||||
|
||||
リクエスト毎にスナップショットを蓄積すれば、ターンより細かい粒度で
|
||||
履歴の任意 suffix のトークン数を実測ベースで逆算できる。
|
||||
|
||||
## 動機
|
||||
|
||||
正確なトークン数（推定でも実測由来）が要る箇所:
|
||||
|
||||
- **Compact 閾値判定** — 現状 `CompactState::last_input_tokens` (`AtomicU64`) が
|
||||
on_usage callback で更新されているが、これは usage_history と情報源が二重化
|
||||
している。本チケットで `Session::total_tokens()` を生やせば、`compact_interceptor.rs` /
|
||||
`controller.rs` から閾値判定がこの API 経由になり、`last_input_tokens` 経路を
|
||||
撤去できる（撤去自体は compact-improvements 側で実施）
|
||||
- **Compact の retained_tokens 切り出し** — 末尾から N トークン残す cut 位置を決める
|
||||
- **Prune の `min_savings` 判定** — 「この content を捨てたら何トークン浮くか」を見積もる
|
||||
- **Compact worker の auto-read budget 判定** — `mark_read_required` の累計
|
||||
- **UI 向けトークン表示**（将来）
|
||||
|
||||
ターン境界での切り出しでは粒度が粗すぎる。長く自走するエージェントは
|
||||
1ターン内で多数のリクエストを回し、ターン長が大きくバラつくため。
|
||||
|
||||
## 前提
|
||||
|
||||
- [usage-history.md](usage-history.md) — session-store に `LogEntry::LlmUsage`
|
||||
を追加し、Worker からリクエスト送信時の prefix と実測値を組で記録する基盤。
|
||||
本チケットはその履歴を消費する側。
|
||||
|
||||
## 方針
|
||||
|
||||
ローカルなトークナイザは持たない。プロバイダ実測値の履歴 (`UsageRecord` 列)
|
||||
と現在の history items から、ピュアな計算で答えを返す。
|
||||
|
||||
API は **session 概念を持つ型のメソッド** として生やす。両方を所有する
|
||||
オーナーが呼ぶ形になるので引数はゼロ。具体的な配置先は実装時に確定する
|
||||
（候補: `RestoredState` / Worker 内の history 所有者 / 薄い view 型）。
|
||||
|
||||
```rust
|
||||
impl Session {
|
||||
/// 現在の history 全体の推定トークン数。
|
||||
/// 最後の measurement + その後追加された未測定分の按分。
|
||||
pub fn total_tokens(&self) -> TokenEstimate;
|
||||
|
||||
/// 末尾から `retained_tokens` 以上を残すための分割位置 (history index)。
|
||||
/// `items[..cut.index]` が捨てる/要約される側、`items[cut.index..]` が残る側。
|
||||
pub fn split_for_retained(&self, retained_tokens: u64) -> SplitPoint;
|
||||
|
||||
/// 指定範囲の items を drop した場合の推定節約トークン数。
|
||||
pub fn savings_for_drop(&self, range: Range<usize>) -> TokenEstimate;
|
||||
}
|
||||
|
||||
pub struct TokenEstimate {
|
||||
pub tokens: u64,
|
||||
pub source: EstimateSource,
|
||||
}
|
||||
|
||||
pub struct SplitPoint {
|
||||
pub index: usize,
|
||||
pub source: EstimateSource,
|
||||
}
|
||||
|
||||
/// 推定の出どころ。呼び出し側が「概算である」ことを認識して扱えるよう明示する。
|
||||
pub enum EstimateSource {
|
||||
/// measurement の境界とちょうど一致（実測そのもの）
|
||||
Measured,
|
||||
/// 連続する 2 measurement の間をバイト按分で計算
|
||||
Interpolated,
|
||||
/// 最後の measurement より新しい区間 → 最後の rate で外挿
|
||||
Extrapolated,
|
||||
/// measurement ゼロ件 → バイト数のみのフォールバック
|
||||
NoData,
|
||||
}
|
||||
```
|
||||
|
||||
呼び出し側:
|
||||
```rust
|
||||
let cut = session.split_for_retained(8_000);
|
||||
let saved = session.savings_for_drop(0..cut.index);
|
||||
if saved.tokens >= min_savings {
|
||||
// prune を実行
|
||||
}
|
||||
```
|
||||
|
||||
## 設計ポイント
|
||||
|
||||
- **状態を持たない**: 計算は所有 history と所有 measurements を見るだけの pure 関数。
|
||||
trait もインスタンスも tokenizer も要らない
|
||||
- **概算であることを返り値で明示**: `EstimateSource` で呼び出し側が
|
||||
measurement 直上 / 按分 / 外挿 / 履歴無しを区別できる。課金判断には使えない
|
||||
- **provider 非依存**: tiktoken やプロバイダ別実装は一切不要。
|
||||
実測値に provider/model/言語/ツール overhead が全て込み
|
||||
- **キャッシュヒットの扱い**: `cache_read_input_tokens` は「コンテキスト占有量」
|
||||
には含めるが「実コスト」とは区別する。Compact/Prune の判定は占有量基準
|
||||
（= raw input_tokens、cache_read 込みの total）
|
||||
|
||||
## 計算アルゴリズム
|
||||
|
||||
`measurements: &[UsageRecord]` （`(history_len_at_send, input_tokens)` の昇順列）と
|
||||
`history: &[Item]` から:
|
||||
|
||||
### `total_tokens`
|
||||
1. measurement が無い → `NoData`、history のバイト数で粗い概算
|
||||
2. 最新 measurement の `tokens` を起点に、`history.len() > last.history_len` の
|
||||
差分があれば最終 rate (tokens / total_bytes) で按分して足す → `Extrapolated`
|
||||
3. ぴったり history が一致 → `Measured`
|
||||
|
||||
### `split_for_retained(retained)`
|
||||
1. measurements を末尾から走査
|
||||
2. `current_total - measurement.tokens >= retained` を満たす最小の measurement を見つける
|
||||
3. 一致する measurement あり → `index = measurement.history_len`, `Measured`
|
||||
4. 2つの measurement の間に境界がある → 区間内のアイテムをバイト按分して `Interpolated`
|
||||
5. 最後の measurement より新しい区間で境界 → 最終 rate で外挿、`Extrapolated`
|
||||
6. measurements 不足 → `NoData` でバイト按分フォールバック
|
||||
|
||||
### `savings_for_drop(range)`
|
||||
1. range と measurement の境界の包含関係を見て、区間を完全に含む measurement の
|
||||
差分を取る → `Measured` or `Interpolated`
|
||||
2. range が最後の measurement より後ろを含む → 外挿で `Extrapolated`
|
||||
3. measurements 無し → バイト数のみ、`NoData`
|
||||
|
||||
## 実装対象
|
||||
|
||||
- session-store または最終的な配置先クレート:
|
||||
- `Session`（または既存の history+measurements 所有者）に
|
||||
`total_tokens` / `split_for_retained` / `savings_for_drop` を実装
|
||||
- `TokenEstimate` / `SplitPoint` / `EstimateSource` 型
|
||||
- 単体テスト: measurement 0/1/N 件、ぴったり境界、按分、外挿、prune 後の整合性
|
||||
- `crates/llm-worker/src/prune.rs`:
|
||||
- `estimate_tokens` を削除し、`min_savings` 判定を `Session::savings_for_drop`
|
||||
呼び出しに置き換え（呼び出し側で渡す）
|
||||
- prune の API シグネチャ調整は最小限に
|
||||
|
||||
## レビュー状態
|
||||
|
||||
Reviewed — [token-counter.review.md](token-counter.review.md)
|
||||
|
||||
## 依存
|
||||
|
||||
- [usage-history.md](usage-history.md) — Usage を session-store に積む基盤
|
||||
|
||||
## ブロックする後続
|
||||
|
||||
- [compact-improvements.md](compact-improvements.md) — retained_tokens 化、
|
||||
auto-read budget、prune の min_savings 精度向上が依存
|
||||
|
|
@ -1,60 +0,0 @@
|
|||
# token-counter レビュー
|
||||
|
||||
## 要件の充足
|
||||
|
||||
チケットが定義した 3 API・型・アルゴリズムは全て実装されている:
|
||||
|
||||
- `Pod::total_tokens()` → `TokenEstimate`
|
||||
- `Pod::split_for_retained(retained)` → `SplitPoint`
|
||||
- `Pod::savings_for_drop(range)` → `TokenEstimate`
|
||||
- `EstimateSource`: `Measured / Interpolated / Extrapolated / NoData`
|
||||
|
||||
設計方針（状態を持たない pure 関数、provider 非依存、ローカルトークナイザ不要）も
|
||||
満たされている。`_impl` 関数群は `(&[Item], &[UsageRecord])` だけを受け取り、
|
||||
Pod メソッドは history と usage_history を渡すだけの薄いラッパー。
|
||||
|
||||
## アーキテクチャ
|
||||
|
||||
| レイヤー | 変更内容 |
|
||||
|---------|---------|
|
||||
| pod::token_counter | pure な計算関数 + Pod のメソッドとして公開 |
|
||||
| pod::Pod | `usage_history: Arc<Mutex<Vec<UsageRecord>>>` を追加。restore で復元、persist_turn で追記、compact で clear |
|
||||
| pod::PruneHook | min_savings 判定を `savings_for_drop_impl` に委譲。usage_history の shared handle を保持 |
|
||||
| llm-worker::prune | `prune()` → `prunable_indices()` + `apply_prune()` に分解。min_savings 判定とトークン会計への依存を除去 |
|
||||
|
||||
prune の責務分離が適切。llm-worker 側は pure な候補抽出と適用のみ、
|
||||
トークン会計への依存は pod 層に閉じている。
|
||||
|
||||
## 指摘と対処
|
||||
|
||||
### 1. split_for_retained_impl の O(n²) シリアライズ（非ブロッカー、未対処）
|
||||
|
||||
`tokens_at` を `1..=history.len()` で毎回呼び、内部で `prefix_bytes`（history 全体の
|
||||
JSON シリアライズ）を都度計算。長大セッションでは item 数に対して二乗になる。
|
||||
`prefix_bytes` をループ外で 1 回だけ計算して渡す形にリファクタリングすべきだが、
|
||||
現時点の history サイズでは実害なし。パフォーマンスが問題になった段階で対処。
|
||||
|
||||
### 2. PruneHook の savings 過大評価（認識済み、未対処）
|
||||
|
||||
prune は content を None にするだけで item を消さないため、`savings_for_drop`
|
||||
（範囲全体の drop を仮定）は実際の節約量より大きい値を返す。閾値判定としては
|
||||
prune を発動しやすい方向＝安全側。ログの `estimated_savings_tokens` が過大になる
|
||||
点はチューニング時に注意。
|
||||
|
||||
### 3. compact 後の usage_history.clear()（後続チケットで対処）
|
||||
|
||||
compact 直後は measurement が空になり `total_tokens()` が `NoData` を返す。
|
||||
compact-improvements で `last_input_tokens` を撤去して閾値判定を usage 経由に
|
||||
一本化する際、この NoData 期間の扱いを設計する必要がある。
|
||||
|
||||
## テスト
|
||||
|
||||
token_counter: 13 件（NoData / Measured / Extrapolated / Interpolated 各ケース、
|
||||
split の境界、savings の measurement 差分、空 range、out-of-range）。
|
||||
|
||||
prune (llm-worker): `prunable_indices` + `apply_prune` に分解後のテスト 5 件。
|
||||
候補抽出、適用、冪等性、既 prune 済み除外、境界。
|
||||
|
||||
## 判定
|
||||
|
||||
承認。
|
||||
Loading…
Reference in New Issue
Block a user