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

feat: protect prune tail by token budget

Browse Source
This commit is contained in:
Hare
2026-05-23 05:00:06 +09:00
parent 4072d35f81
commit 9ee7f04805
12 changed files with 423 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pod/src/compact/prune.rs
+42 -16
View File
@@ -13,19 +13,24 @@
use llm_worker::Item;
use llm_worker::llm_client::client::LlmClient;
use llm_worker::prune::{PruneConfig, PruneDecision, PruneObserver, SavingsEstimator};
use llm_worker::prune::{
PruneConfig, PruneDecision, PruneObserver, SavingsEstimator, TokenEstimator,
};
use session_metrics::Metric;
use session_store::Store;
use crate::Pod;
use crate::compact::token_counter::{EstimateSource, savings_for_prune_impl};
use crate::compact::token_counter::{
EstimateSource, savings_for_prune_impl, token_estimates_for_prune_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.
/// Registers the config and token/savings-estimator closures on the Worker.
/// The estimators combine persisted [`Pod::usage_history_handle`] records
/// with in-flight `UsageTracker` records so multi-request tool loops can
/// prune before the surrounding Pod run finishes.
///
/// Measurement-less estimates (before the first LLM call, or immediately
/// after a compact) return `0` from the estimator, which naturally
@@ -37,9 +42,25 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// [`UsageTracker`] so the next `LlmUsage` can be paired with a
/// `prune.post_request` metric carrying the same id.
pub fn attach_prune(&mut self, config: PruneConfig) {
let usage = self.usage_history_handle();
let usage_history_for_tokens = self.usage_history_handle();
let usage_tracker_for_tokens = self.usage_tracker_handle();
let token_estimator: TokenEstimator = Box::new(move |history: &[Item]| {
let mut snapshot = usage_history_for_tokens
.lock()
.expect("usage_history poisoned")
.clone();
snapshot.extend(usage_tracker_for_tokens.records());
token_estimates_for_prune_impl(history, &snapshot)
});
let usage_history_for_savings = self.usage_history_handle();
let usage_tracker_for_savings = self.usage_tracker_handle();
let estimator: SavingsEstimator = Box::new(move |history: &[Item], indices| {
let snapshot = usage.lock().expect("usage_history poisoned").clone();
let mut snapshot = usage_history_for_savings
.lock()
.expect("usage_history poisoned")
.clone();
snapshot.extend(usage_tracker_for_savings.records());
let est = savings_for_prune_impl(history, &snapshot, indices);
match est.source {
EstimateSource::NoData => 0,
@@ -56,8 +77,9 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
.with_value(eval.estimated_savings as f64)
.with_correlation_id(&correlation_id)
.with_dimension("candidate_count", eval.candidate_count.to_string());
if let Some(border) = eval.border_turn {
metric = metric.with_dimension("border_turn", border.to_string());
if let Some(protected_start) = eval.protected_start_index {
metric =
metric.with_dimension("protected_start_index", protected_start.to_string());
}
metrics.push(metric);
usage_tracker.note_correlation_id(correlation_id);
@@ -66,17 +88,21 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
metrics.push(Metric::now("prune.skip").with_dimension("reason", "no_candidates"));
}
PruneDecision::SkippedBelowMinSavings => {
metrics.push(
Metric::now("prune.skip")
.with_dimension("reason", "below_min_savings")
.with_dimension("candidate_count", eval.candidate_count.to_string())
.with_value(eval.estimated_savings as f64),
);
let mut metric = Metric::now("prune.skip")
.with_dimension("reason", "below_min_savings")
.with_dimension("candidate_count", eval.candidate_count.to_string())
.with_value(eval.estimated_savings as f64);
if let Some(protected_start) = eval.protected_start_index {
metric =
metric.with_dimension("protected_start_index", protected_start.to_string());
}
metrics.push(metric);
}
});
let worker = self.worker_mut();
worker.set_prune_config(Some(config));
worker.set_token_estimator(Some(token_estimator));
worker.set_savings_estimator(Some(estimator));
worker.set_prune_observer(Some(observer));
}
@@ -90,7 +116,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
return;
};
let config = PruneConfig {
protected_turns: compaction.prune_protected_turns,
protected_tokens: compaction.prune_protected_tokens,
min_savings: compaction.prune_min_savings,
};
self.attach_prune(config);
crates/pod/src/compact/token_counter.rs
+35
View File
@@ -132,6 +132,21 @@ fn tool_result_content_bytes(item: &Item) -> u64 {
item_bytes(item).saturating_sub(item_bytes(&cleared))
}
/// Prefix-boundary token estimates used by Prune to find its protected suffix.
///
/// Returns `history.len() + 1` entries where entry `i` estimates
/// `history[..i]`. This shares the same [`tokens_at`] accounting as compact's
/// retained-tail split and prune's savings estimate.
pub(crate) fn token_estimates_for_prune_impl(
history: &[Item],
records: &[UsageRecord],
) -> Vec<TokenEstimate> {
let prefix = prefix_bytes(history);
(0..=history.len())
.map(|idx| tokens_at(history, records, idx, &prefix))
.collect()
}
/// Prune 射影（`ToolResult.content = None`）で節約されるトークン数の推定。
///
/// `indices` は [`llm_worker::prune::prunable_indices`] が返す候補列を
@@ -278,6 +293,26 @@ mod tests {
}
}
#[test]
fn token_estimates_for_prune_returns_every_prefix_boundary() {
let history = vec![msg("a"), msg("b"), msg("c")];
let estimates = token_estimates_for_prune_impl(&history, &[record(3, 300)]);
assert_eq!(estimates.len(), history.len() + 1);
assert_eq!(estimates[0].tokens, 0);
assert_eq!(estimates[3].tokens, 300);
assert_eq!(estimates[3].source, EstimateSource::Measured);
}
#[test]
fn token_estimates_for_prune_propagates_no_data() {
let history = vec![msg("a"), msg("b")];
let estimates = token_estimates_for_prune_impl(&history, &[]);
assert_eq!(estimates.len(), history.len() + 1);
assert_eq!(estimates[0].source, EstimateSource::Measured);
assert_eq!(estimates[1].source, EstimateSource::NoData);
assert_eq!(estimates[2].source, EstimateSource::NoData);
}
#[test]
fn savings_for_prune_skips_non_toolresult_indices() {
let history = vec![msg("a"), msg("b"), msg("c")];
crates/pod/tests/session_metrics_test.rs
+41 -11
View File
@@ -4,10 +4,10 @@
//! returns a long `ToolOutput.content`, then inspects the persisted
//! session log to verify:
//!
//! - `prune.skip { reason: "no_candidates" }` lands when the protected-turn
//! window covers the entire history.
//! - `prune.fire` lands once enough turns + usage measurements exist for
//! the projection to actually apply.
//! - `prune.skip { reason: "no_candidates" }` lands when usage estimates are
//! unavailable or the protected-token window covers all tool results.
//! - `prune.fire` lands once enough measured history exceeds the protected-token
//! budget for the projection to actually apply.
//! - The fire metric and the immediately-following `prune.post_request`
//! metric share the same `correlation_id`, so cache_read / cache_write
//! from the LlmUsage that triggered the projection can be joined back
@@ -136,7 +136,7 @@ fn text_response_with_cache(text: &str, cache_read: u64, cache_write: u64) -> Ve
]
}
fn manifest_toml(prune_protected_turns: usize, prune_min_savings: u64) -> String {
fn manifest_toml(prune_protected_tokens: u64, prune_min_savings: u64) -> String {
format!(
r#"
[pod]
@@ -151,7 +151,7 @@ model_id = "test-model"
max_tokens = 100
[compaction]
prune_protected_turns = {prune_protected_turns}
prune_protected_tokens = {prune_protected_tokens}
prune_min_savings = {prune_min_savings}
[[scope.allow]]
@@ -192,7 +192,7 @@ async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
// Run 1 (request 0): tool_use → triggers tool execution → request 1
// on the second iteration to produce the assistant reply.
// Run 2 (request 2): plain assistant text. Prune evaluation here
// sees user1's tool_result outside the 1-protected-turn window and
// sees user1's tool_result outside the protected-token suffix and
// should fire.
let client = MockClient::new(vec![
tool_use_response("call-1", "big_tool"),
@@ -250,8 +250,8 @@ async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
"fire missing candidate_count: {fire:?}"
);
assert!(
fire.dimensions.contains_key("border_turn"),
"fire missing border_turn: {fire:?}"
fire.dimensions.contains_key("protected_start_index"),
"fire missing protected_start_index: {fire:?}"
);
assert!(fire.value.is_some(), "fire missing estimated_savings value");
let fire_id = fire
@@ -277,6 +277,36 @@ async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
assert!(post.dimensions.contains_key("history_len"));
}
#[tokio::test]
async fn prune_metrics_fire_during_single_long_task_without_multiple_user_turns() {
let client = MockClient::new(vec![
tool_use_response("call-1", "big_tool"),
tool_use_response("call-2", "big_tool"),
tool_use_response("call-3", "big_tool"),
tool_use_response("call-4", "big_tool"),
text_response_with_cache("done", 100, 20),
]);
let (mut pod, _store_tmp, _pwd_tmp) = make_pod(manifest_toml(1, 1), client, "big_tool").await;
let session_id = pod.session_id();
let segment_id = pod.segment_id();
let store = pod.store().clone();
pod.run_text("one long task").await.unwrap();
let state = session_store::restore(&store, session_id, segment_id).unwrap();
let metrics = metrics_from_extensions(&state.extensions);
let fire_count = metrics.iter().filter(|m| m.name == "prune.fire").count();
assert!(
fire_count > 0,
"single-turn tool loop should produce prune.fire once old heavy ToolResults fall outside the protected-token suffix: {metrics:?}"
);
assert!(
metrics.iter().any(|m| {
m.name == "prune.fire" && m.dimensions.contains_key("protected_start_index")
})
);
}
/// `min_savings` set high enough that candidates exist but the estimated
/// savings always fall short → the second run should record
/// `prune.skip { reason: "below_min_savings" }`.
@@ -288,7 +318,7 @@ async fn prune_metrics_record_below_min_savings_skip() {
text_response_with_cache("done", 0, 0),
]);
let (mut pod, _store_tmp, _pwd_tmp) =
make_pod(manifest_toml(1, u64::MAX), client, "big_tool").await;
make_pod(manifest_toml(1, 1_000_000), client, "big_tool").await;
let session_id = pod.session_id();
let segment_id = pod.segment_id();
let store = pod.store().clone();
@@ -405,7 +435,7 @@ async fn metric_write_failure_emits_warn_alert_and_does_not_abort_run() {
// Even with a tool registered, this run will only emit
// `prune.skip { reason: "no_candidates" }` (one user message,
// protected_turns=1 covers everything). That is enough to drive
// protected token budget covers the only user message). That is enough to drive
// the failure path: at least one metric attempts to write.
let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]);
let worker = Worker::new(client);