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

feat: make tool result pruning opt-in

Browse Source
This commit is contained in:
Hare
2026-09-18 19:43:29 +09:00
parent 19ee757880
commit d7917512ac
9 changed files with 97 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/manifest/src/config.rs
+7 -1
View File
@@ -577,6 +577,8 @@ pub struct PermissionConfigPartial {
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CompactionConfigPartial {
#[serde(default)]
pub prune_enabled: Option<bool>,
#[serde(default)]
pub prune_protected_tokens: Option<u64>,
#[serde(default)]
@@ -912,6 +914,7 @@ impl PermissionConfigPartial {
impl CompactionConfigPartial {
fn merge(self, upper: Self) -> Self {
Self {
prune_enabled: upper.prune_enabled.or(self.prune_enabled),
prune_protected_tokens: upper.prune_protected_tokens.or(self.prune_protected_tokens),
prune_min_savings: upper.prune_min_savings.or(self.prune_min_savings),
threshold: upper.threshold.or(self.threshold),
@@ -1223,6 +1226,7 @@ impl TryFrom<WorkerManifestConfig> for WorkerManifest {
validate_model_paths(cm, "compaction.model.auth.file")?;
}
Ok(CompactionConfig {
prune_enabled: c.prune_enabled.unwrap_or(defaults::PRUNE_ENABLED),
prune_protected_tokens: c
.prune_protected_tokens
.unwrap_or(defaults::PRUNE_PROTECTED_TOKENS),
@@ -1831,6 +1835,7 @@ mod tests {
let lower = WorkerManifestConfig {
compaction: Some(CompactionConfigPartial {
threshold: Some(50_000),
prune_enabled: Some(false),
prune_protected_tokens: Some(5_000),
..Default::default()
}),
@@ -1861,7 +1866,8 @@ mod tests {
let merged = lower.merge(upper);
let c = merged.compaction.unwrap();
assert_eq!(c.threshold, Some(80_000));
// field from lower retained when upper has None
// fields from lower retained when upper has None
assert_eq!(c.prune_enabled, Some(false));
assert_eq!(c.prune_protected_tokens, Some(5_000));
let search = merged.web.unwrap().search.unwrap();
assert_eq!(search.timeout_secs, Some(3));
crates/manifest/src/defaults.rs
+4
View File
@@ -14,6 +14,10 @@ pub const TOOL_OUTPUT_MAX_BYTES: usize = 64 * 1024;
/// See [`crate::FileUploadLimits`].
pub const FILE_UPLOAD_MAX_BYTES: usize = 256 * 1024;
/// Whether request-time tool-result pruning is enabled when `[compaction]`
/// exists. See [`crate::CompactionConfig::prune_enabled`].
pub const PRUNE_ENABLED: bool = false;
/// Token budget at the history tail protected from pruning. See
/// [`crate::CompactionConfig::prune_protected_tokens`].
pub const PRUNE_PROTECTED_TOKENS: u64 = 8000;
crates/manifest/src/lib.rs
+17 -1
View File
@@ -768,6 +768,11 @@ pub enum ToolPermissionAction {
/// (full history summarisation). Omitting `[compaction]` disables both.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactionConfig {
/// Whether request-time projection removes detail from old tool results.
/// Independent from the compaction thresholds below.
#[serde(default = "default_prune_enabled")]
pub prune_enabled: bool,
/// Token budget at the history tail protected from pruning.
#[serde(default = "default_prune_protected_tokens")]
pub prune_protected_tokens: u64,
@@ -868,6 +873,9 @@ pub struct CompactionConfig {
pub model: Option<ModelManifest>,
}
fn default_prune_enabled() -> bool {
defaults::PRUNE_ENABLED
}
fn default_prune_protected_tokens() -> u64 {
defaults::PRUNE_PROTECTED_TOKENS
}
@@ -914,6 +922,7 @@ fn default_result_context_max_tokens() -> u64 {
impl Default for CompactionConfig {
fn default() -> Self {
Self {
prune_enabled: default_prune_enabled(),
prune_protected_tokens: default_prune_protected_tokens(),
prune_min_savings: default_prune_min_savings(),
threshold: None,
@@ -1460,9 +1469,11 @@ model_id = "claude-sonnet-4-20250514"
#[test]
fn parse_compaction_config() {
let toml = format!("{MINIMAL_REQUIRED}\n[compaction]\nthreshold = 80000\n");
let toml =
format!("{MINIMAL_REQUIRED}\n[compaction]\nprune_enabled = false\nthreshold = 80000\n");
let manifest = WorkerManifest::from_toml(&toml).unwrap();
let c = manifest.compaction.unwrap();
assert!(!c.prune_enabled);
assert_eq!(c.prune_protected_tokens, 8000);
assert_eq!(c.prune_min_savings, 4096);
assert_eq!(c.threshold, Some(80000));
@@ -1471,6 +1482,11 @@ model_id = "claude-sonnet-4-20250514"
assert_eq!(c.worker_max_turns, Some(20));
}
#[test]
fn compaction_defaults_to_pruning_disabled() {
assert!(!CompactionConfig::default().prune_enabled);
}
#[test]
fn reject_removed_prune_protected_turns_field() {
let toml = format!("{MINIMAL_REQUIRED}\n[compaction]\nprune_protected_turns = 3\n");
crates/manifest/src/profile.rs
+20
View File
@@ -1245,6 +1245,26 @@ mod tests {
.unwrap()
};
for name in [
"default",
"intake",
"orchestrator",
"companion",
"coder",
"reviewer",
"memory-consolidation",
] {
let role = resolve(name);
assert_eq!(
role.manifest
.compaction
.as_ref()
.map(|compaction| compaction.prune_enabled),
Some(false),
"{name} must keep pruning disabled for the experiment"
);
}
let companion = resolve("companion");
assert_eq!(
companion.manifest.model.ref_.as_deref(),
crates/worker/src/compact/prune.rs
+42 -7
View File
@@ -22,6 +22,13 @@ use crate::compact::token_counter::{
EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl,
};
fn prune_config_from_compaction(compaction: &manifest::CompactionConfig) -> Option<PruneConfig> {
compaction.prune_enabled.then_some(PruneConfig {
protected_tokens: compaction.prune_protected_tokens,
min_savings: compaction.prune_min_savings,
})
}
impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Enable prune projection on the underlying Engine.
///
@@ -105,18 +112,46 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
worker.set_prune_observer(Some(observer));
}
/// 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 Worker constructors so prune is
/// active whenever the manifest asks for it.
/// If the manifest enables prune under `[compaction]`, build a
/// `PruneConfig` from its `prune_*` fields and call
/// [`attach_prune`](Self::attach_prune). Otherwise no-op. Called from all
/// Worker constructors.
pub(crate) fn apply_prune_from_manifest(&mut self) {
let Some(compaction) = self.manifest().compaction.as_ref() else {
return;
};
let config = PruneConfig {
protected_tokens: compaction.prune_protected_tokens,
min_savings: compaction.prune_min_savings,
let Some(config) = prune_config_from_compaction(compaction) else {
return;
};
self.attach_prune(config);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn disabled_pruning_does_not_build_prune_config() {
let compaction = manifest::CompactionConfig {
prune_enabled: false,
..Default::default()
};
assert!(prune_config_from_compaction(&compaction).is_none());
}
#[test]
fn enabled_compaction_preserves_prune_thresholds() {
let compaction = manifest::CompactionConfig {
prune_enabled: true,
prune_protected_tokens: 1234,
prune_min_savings: 5678,
..Default::default()
};
let config = prune_config_from_compaction(&compaction).unwrap();
assert_eq!(config.protected_tokens, 1234);
assert_eq!(config.min_savings, 5678);
}
}
crates/worker/src/spawn/tool.rs
+1
View File
@@ -956,6 +956,7 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
.compaction
.as_ref()
.map(|c| CompactionConfigPartial {
prune_enabled: Some(c.prune_enabled),
prune_protected_tokens: Some(c.prune_protected_tokens),
prune_min_savings: Some(c.prune_min_savings),
threshold: c.threshold,
crates/worker/tests/session_metrics_test.rs
+1
View File
@@ -159,6 +159,7 @@ model_id = "test-model"
max_tokens = 100
[compaction]
prune_enabled = true
prune_protected_tokens = {prune_protected_tokens}
prune_min_savings = {prune_min_savings}
docs/manifest.toml
+4 -1
View File
@@ -179,9 +179,12 @@ permission = "write"
# ===== [compaction] =========================================================
# コンテキスト圧縮 (Prune / Compact)。セクション省略で両方無効。
# セクションを書いた時点で Prune は有効化、Compact は閾値が None なら無効。
# `prune_enabled`でPruneを制御する。Compactは閾値がNoneなら無効。
# [compaction]
#
# # 任意。デフォルト: false。Compactの設定を保ったままPruneだけ切替できる。
# prune_enabled = true
#
# # 任意。デフォルト: 8000 (`defaults::PRUNE_PROTECTED_TOKENS`)。
# # pruning から保護する末尾 token budget。turn 数ではなく usage estimate で境界を引く。
# prune_protected_tokens = 8000
resources/profiles/base.dcdl
+1
View File
@@ -16,6 +16,7 @@ engine = {
compaction = {
kind = "tokens";
prune_enabled = false;
threshold = 240000;
request_threshold = 270000;
worker_context_max_tokens = 100000;