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

pruneのトークン計算置き換え・Podに接続

Browse Source
This commit is contained in:
Hare
2026-04-14 02:35:35 +09:00
parent 5a995cf099
commit 2e004161e4
9 changed files with 310 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pod/src/compact_interceptor.rs
+1 -1
View File
@@ -39,7 +39,7 @@ impl Interceptor for CompactInterceptor {
}
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
// Step 1: Delegate to inner (PruneHook and other hooks run first).
// Step 1: Delegate to inner hooks first.
let inner_action = self.inner.pre_llm_request(context).await;
if !matches!(inner_action, PreRequestAction::Continue) {
return inner_action;
crates/pod/src/lib.rs
+1 -3
View File
@@ -4,12 +4,11 @@ pub mod runtime_dir;
pub mod shared_state;
pub mod socket_server;
pub mod prune_hook;
mod compact_interceptor;
mod compact_state;
mod hook_interceptor;
mod pod;
mod prune;
mod token_counter;
mod usage_tracker;
@@ -18,7 +17,6 @@ pub use token_counter::{EstimateSource, SplitPoint, TokenEstimate};
pub use controller::{PodController, PodHandle};
pub use manifest::{PodManifest, ProviderConfig, ProviderKind, Scope};
pub use hook::{Hook, HookEventKind, HookRegistryBuilder};
pub use prune_hook::PruneHook;
pub use pod::{Pod, PodError, PodRunResult, apply_worker_manifest};
pub use protocol::{ErrorCode, Event, Method, TurnResult};
pub use provider::{ProviderError, build_client};
crates/pod/src/pod.rs
+24 -11
View File
@@ -79,9 +79,9 @@ pub struct Pod<C: LlmClient, St: Store> {
/// Restored from session log on `restore`, appended on each persist.
/// Read by token-accounting APIs (`Pod::total_tokens`, etc.).
///
/// Wrapped in `Arc<Mutex>` so that hooks living on the Worker
/// (e.g. `PruneHook`) can share the same view via
/// [`Pod::usage_history_handle`].
/// Wrapped in `Arc<Mutex>` so that callbacks injected into the
/// Worker (e.g. the savings estimator used by the prune projection)
/// can share the same view via [`Pod::usage_history_handle`].
usage_history: Arc<Mutex<Vec<UsageRecord>>>,
/// Session-lifetime file-operation tracker from the builtin `tools`
/// crate. Populated by the Controller when it registers the builtin
@@ -104,7 +104,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
history: worker.history(),
};
let (session_id, head_hash) = session_store::create_session(&store, state).await?;
Ok(Self {
let mut pod = Self {
manifest,
worker: Some(worker),
store,
@@ -118,7 +118,9 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
usage_tracker: Arc::new(UsageTracker::new()),
usage_history: Arc::new(Mutex::new(Vec::<UsageRecord>::new())),
tracker: None,
})
};
pod.apply_prune_from_manifest();
Ok(pod)
}
/// Restore a Pod from a persisted session.
@@ -139,7 +141,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
worker.set_turn_count(state.turn_count);
worker.set_last_run_interrupted(state.last_run_interrupted);
Ok(Self {
let mut pod = Self {
manifest,
worker: Some(worker),
store,
@@ -153,7 +155,9 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
usage_tracker: Arc::new(UsageTracker::new()),
usage_history: Arc::new(Mutex::new(state.usage_history)),
tracker: None,
})
};
pod.apply_prune_from_manifest();
Ok(pod)
}
/// The session ID used for persistence.
@@ -206,9 +210,16 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// Shared handle to the cumulative Usage history.
///
/// Hooks (e.g. `PruneHook`) take a clone of this `Arc` so they can
/// read the latest measurements at request time. The handle outlives
/// Callbacks that need live access to the latest measurements (e.g.
/// the savings estimator that `attach_prune` installs on the Worker)
/// clone this `Arc` and read it at request time. The handle outlives
/// any individual run.
///
/// **Locking contract:** the inner `Mutex` is held only for a short
/// clone (`lock().unwrap().clone()`) and released immediately.
/// Callers must not hold the guard across `.await` points, I/O, or
/// long computations — the guard is implicitly assumed to be
/// non-contended at every Pod lifecycle event.
pub fn usage_history_handle(&self) -> Arc<Mutex<Vec<UsageRecord>>> {
self.usage_history.clone()
}
@@ -686,7 +697,7 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
history: worker.history(),
};
let (session_id, head_hash) = session_store::create_session(&store, state).await?;
Ok(Self {
let mut pod = Self {
manifest,
worker: Some(worker),
store,
@@ -700,7 +711,9 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
usage_tracker: Arc::new(UsageTracker::new()),
usage_history: Arc::new(Mutex::new(Vec::new())),
tracker: None,
})
};
pod.apply_prune_from_manifest();
Ok(pod)
}
}
crates/pod/src/prune.rs
+56
View File
@@ -0,0 +1,56 @@
//! Prune integration — wires the Worker's prune projection to the Pod's
//! usage-history-backed token accounting.
//!
//! Worker 自身がコンテキスト射影を行う（`worker.rs` の `request_context` 構築
//! 直後）。Worker は usage 履歴を知らないので、`min_savings` 判定に使う savings
//! の見積もりはコールバックで外部から注入する。このモジュールはそのコールバック
//! を組み立てて Worker に差し込むための `impl Pod` を提供する。
use llm_worker::Item;
use llm_worker::llm_client::client::LlmClient;
use llm_worker::prune::{PruneConfig, SavingsEstimator};
use session_store::Store;
use crate::Pod;
use crate::token_counter::{EstimateSource, savings_for_drop_impl};
impl<C: LlmClient, St: Store> Pod<C, St> {
/// Enable prune projection on the underlying Worker.
///
/// Registers the config and a savings-estimator closure on the Worker.
/// The estimator captures a shared handle to [`Pod::usage_history_handle`]
/// so that every LLM request sees the latest measurements.
///
/// Measurement-less ranges (before the first LLM call, or immediately
/// after a compact) return `0` from the estimator, which naturally
/// prevents the prune projection from firing until usage data exists.
pub fn attach_prune(&mut self, config: PruneConfig) {
let usage = self.usage_history_handle();
let estimator: SavingsEstimator = Box::new(move |history: &[Item], range| {
let snapshot = usage.lock().expect("usage_history poisoned").clone();
let est = savings_for_drop_impl(history, &snapshot, range);
match est.source {
EstimateSource::NoData => 0,
_ => est.tokens,
}
});
let worker = self.worker_mut();
worker.set_prune_config(Some(config));
worker.set_savings_estimator(Some(estimator));
}
/// If the manifest has a `[compaction]` section, build a `PruneConfig`
/// from its `prune_*` fields and call [`attach_prune`](Self::attach_prune).
/// Otherwise no-op. Called from all Pod constructors so prune is
/// active whenever the manifest asks for it.
pub(crate) fn apply_prune_from_manifest(&mut self) {
let Some(compaction) = self.manifest().compaction.as_ref() else {
return;
};
let config = PruneConfig {
protected_turns: compaction.prune_protected_turns,
min_savings: compaction.prune_min_savings,
};
self.attach_prune(config);
}
}
crates/pod/src/prune_hook.rs
-95
View File
@@ -1,95 +0,0 @@
//! PruneHook — projects the LLM request context before each call.
//!
//! Prune は **コンテキスト射影** として実装する。`PreLlmRequest` hook に
//! 渡される `context: &mut Vec<Item>` は Worker が毎 turn 冒頭で history を
//! clone した一時配列 (`worker.rs:701`)。ここで ToolResult.content を省いても
//! Worker の永続履歴には影響しない。`prunable_indices` で候補を抽出し、
//! `min_savings` を満たせば content を `None` に射影する。
//!
//! `min_savings` の判定は usage 履歴ベースのトークン会計
//! ([`crate::token_counter::savings_for_drop_impl`]) で行う。
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use llm_worker::Item;
use llm_worker::interceptor::PreRequestAction;
use llm_worker::prune::{PruneConfig, prunable_indices};
use session_store::UsageRecord;
use tracing::debug;
use crate::hook::{Hook, PreLlmRequest};
use crate::token_counter::{EstimateSource, savings_for_drop_impl};
/// Hook that conditionally prunes old tool-result content before each
/// LLM request, reclaiming context-window tokens.
///
/// `usage_history` は [`crate::Pod::usage_history_handle`] から共有された
/// `Arc<Mutex<_>>`。リクエスト直前に snapshot を取って savings を見積もる。
pub struct PruneHook {
config: PruneConfig,
usage_history: Arc<Mutex<Vec<UsageRecord>>>,
}
impl PruneHook {
pub fn new(config: PruneConfig, usage_history: Arc<Mutex<Vec<UsageRecord>>>) -> Self {
Self {
config,
usage_history,
}
}
}
#[async_trait]
impl Hook<PreLlmRequest> for PruneHook {
async fn call(&self, context: &mut Vec<Item>) -> PreRequestAction {
let candidates = prunable_indices(context, self.config.protected_turns);
if candidates.is_empty() {
return PreRequestAction::Continue;
}
// 候補範囲のトークン節約量を usage 履歴ベースで見積もる。
// content だけ削除する場合の上限値（範囲全体を消した場合の savings）として
// 近似する。実際の content drop は items 数を変えないので、本来の savings
// はこの値以下。閾値判定は上振れ方向＝「やや prune を発動しやすい」側で安全。
let first = *candidates.first().unwrap();
let last = *candidates.last().unwrap() + 1;
let snapshot = self
.usage_history
.lock()
.expect("usage_history poisoned")
.clone();
let savings = savings_for_drop_impl(context, &snapshot, first..last);
// measurement が無い場合 (NoData) は判定材料がないので prune を見送る。
// 最初の LLM call が走るまでは usage_history が空なのでこのパスを通る。
if matches!(savings.source, EstimateSource::NoData) {
return PreRequestAction::Continue;
}
if savings.tokens < self.config.min_savings {
return PreRequestAction::Continue;
}
// 射影: context (= history の clone) 上の対象 ToolResult だけ content を
// drop する。Worker の永続履歴は別インスタンスなので影響を受けない。
let mut projected = 0usize;
for &i in &candidates {
if let Item::ToolResult { content, .. } = &mut context[i] {
if content.is_some() {
*content = None;
projected += 1;
}
}
}
if projected > 0 {
debug!(
pruned = projected,
estimated_savings_tokens = savings.tokens,
source = ?savings.source,
"Projected old tool-result content out of request context"
);
}
PreRequestAction::Continue
}
}