feat: make tool result pruning opt-in
This commit is contained in:
@@ -577,6 +577,8 @@ pub struct PermissionConfigPartial {
|
|||||||
|
|
||||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||||
pub struct CompactionConfigPartial {
|
pub struct CompactionConfigPartial {
|
||||||
|
#[serde(default)]
|
||||||
|
pub prune_enabled: Option<bool>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub prune_protected_tokens: Option<u64>,
|
pub prune_protected_tokens: Option<u64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
@@ -912,6 +914,7 @@ impl PermissionConfigPartial {
|
|||||||
impl CompactionConfigPartial {
|
impl CompactionConfigPartial {
|
||||||
fn merge(self, upper: Self) -> Self {
|
fn merge(self, upper: Self) -> 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_protected_tokens: upper.prune_protected_tokens.or(self.prune_protected_tokens),
|
||||||
prune_min_savings: upper.prune_min_savings.or(self.prune_min_savings),
|
prune_min_savings: upper.prune_min_savings.or(self.prune_min_savings),
|
||||||
threshold: upper.threshold.or(self.threshold),
|
threshold: upper.threshold.or(self.threshold),
|
||||||
@@ -1223,6 +1226,7 @@ impl TryFrom<WorkerManifestConfig> for WorkerManifest {
|
|||||||
validate_model_paths(cm, "compaction.model.auth.file")?;
|
validate_model_paths(cm, "compaction.model.auth.file")?;
|
||||||
}
|
}
|
||||||
Ok(CompactionConfig {
|
Ok(CompactionConfig {
|
||||||
|
prune_enabled: c.prune_enabled.unwrap_or(defaults::PRUNE_ENABLED),
|
||||||
prune_protected_tokens: c
|
prune_protected_tokens: c
|
||||||
.prune_protected_tokens
|
.prune_protected_tokens
|
||||||
.unwrap_or(defaults::PRUNE_PROTECTED_TOKENS),
|
.unwrap_or(defaults::PRUNE_PROTECTED_TOKENS),
|
||||||
@@ -1831,6 +1835,7 @@ mod tests {
|
|||||||
let lower = WorkerManifestConfig {
|
let lower = WorkerManifestConfig {
|
||||||
compaction: Some(CompactionConfigPartial {
|
compaction: Some(CompactionConfigPartial {
|
||||||
threshold: Some(50_000),
|
threshold: Some(50_000),
|
||||||
|
prune_enabled: Some(false),
|
||||||
prune_protected_tokens: Some(5_000),
|
prune_protected_tokens: Some(5_000),
|
||||||
..Default::default()
|
..Default::default()
|
||||||
}),
|
}),
|
||||||
@@ -1861,7 +1866,8 @@ mod tests {
|
|||||||
let merged = lower.merge(upper);
|
let merged = lower.merge(upper);
|
||||||
let c = merged.compaction.unwrap();
|
let c = merged.compaction.unwrap();
|
||||||
assert_eq!(c.threshold, Some(80_000));
|
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));
|
assert_eq!(c.prune_protected_tokens, Some(5_000));
|
||||||
let search = merged.web.unwrap().search.unwrap();
|
let search = merged.web.unwrap().search.unwrap();
|
||||||
assert_eq!(search.timeout_secs, Some(3));
|
assert_eq!(search.timeout_secs, Some(3));
|
||||||
|
|||||||
@@ -14,6 +14,10 @@ pub const TOOL_OUTPUT_MAX_BYTES: usize = 64 * 1024;
|
|||||||
/// See [`crate::FileUploadLimits`].
|
/// See [`crate::FileUploadLimits`].
|
||||||
pub const FILE_UPLOAD_MAX_BYTES: usize = 256 * 1024;
|
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
|
/// Token budget at the history tail protected from pruning. See
|
||||||
/// [`crate::CompactionConfig::prune_protected_tokens`].
|
/// [`crate::CompactionConfig::prune_protected_tokens`].
|
||||||
pub const PRUNE_PROTECTED_TOKENS: u64 = 8000;
|
pub const PRUNE_PROTECTED_TOKENS: u64 = 8000;
|
||||||
|
|||||||
@@ -768,6 +768,11 @@ pub enum ToolPermissionAction {
|
|||||||
/// (full history summarisation). Omitting `[compaction]` disables both.
|
/// (full history summarisation). Omitting `[compaction]` disables both.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct CompactionConfig {
|
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.
|
/// Token budget at the history tail protected from pruning.
|
||||||
#[serde(default = "default_prune_protected_tokens")]
|
#[serde(default = "default_prune_protected_tokens")]
|
||||||
pub prune_protected_tokens: u64,
|
pub prune_protected_tokens: u64,
|
||||||
@@ -868,6 +873,9 @@ pub struct CompactionConfig {
|
|||||||
pub model: Option<ModelManifest>,
|
pub model: Option<ModelManifest>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn default_prune_enabled() -> bool {
|
||||||
|
defaults::PRUNE_ENABLED
|
||||||
|
}
|
||||||
fn default_prune_protected_tokens() -> u64 {
|
fn default_prune_protected_tokens() -> u64 {
|
||||||
defaults::PRUNE_PROTECTED_TOKENS
|
defaults::PRUNE_PROTECTED_TOKENS
|
||||||
}
|
}
|
||||||
@@ -914,6 +922,7 @@ fn default_result_context_max_tokens() -> u64 {
|
|||||||
impl Default for CompactionConfig {
|
impl Default for CompactionConfig {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
|
prune_enabled: default_prune_enabled(),
|
||||||
prune_protected_tokens: default_prune_protected_tokens(),
|
prune_protected_tokens: default_prune_protected_tokens(),
|
||||||
prune_min_savings: default_prune_min_savings(),
|
prune_min_savings: default_prune_min_savings(),
|
||||||
threshold: None,
|
threshold: None,
|
||||||
@@ -1460,9 +1469,11 @@ model_id = "claude-sonnet-4-20250514"
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_compaction_config() {
|
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 manifest = WorkerManifest::from_toml(&toml).unwrap();
|
||||||
let c = manifest.compaction.unwrap();
|
let c = manifest.compaction.unwrap();
|
||||||
|
assert!(!c.prune_enabled);
|
||||||
assert_eq!(c.prune_protected_tokens, 8000);
|
assert_eq!(c.prune_protected_tokens, 8000);
|
||||||
assert_eq!(c.prune_min_savings, 4096);
|
assert_eq!(c.prune_min_savings, 4096);
|
||||||
assert_eq!(c.threshold, Some(80000));
|
assert_eq!(c.threshold, Some(80000));
|
||||||
@@ -1471,6 +1482,11 @@ model_id = "claude-sonnet-4-20250514"
|
|||||||
assert_eq!(c.worker_max_turns, Some(20));
|
assert_eq!(c.worker_max_turns, Some(20));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn compaction_defaults_to_pruning_disabled() {
|
||||||
|
assert!(!CompactionConfig::default().prune_enabled);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn reject_removed_prune_protected_turns_field() {
|
fn reject_removed_prune_protected_turns_field() {
|
||||||
let toml = format!("{MINIMAL_REQUIRED}\n[compaction]\nprune_protected_turns = 3\n");
|
let toml = format!("{MINIMAL_REQUIRED}\n[compaction]\nprune_protected_turns = 3\n");
|
||||||
|
|||||||
@@ -1245,6 +1245,26 @@ mod tests {
|
|||||||
.unwrap()
|
.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");
|
let companion = resolve("companion");
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
companion.manifest.model.ref_.as_deref(),
|
companion.manifest.model.ref_.as_deref(),
|
||||||
|
|||||||
@@ -22,6 +22,13 @@ use crate::compact::token_counter::{
|
|||||||
EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl,
|
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> {
|
impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
||||||
/// Enable prune projection on the underlying Engine.
|
/// 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));
|
worker.set_prune_observer(Some(observer));
|
||||||
}
|
}
|
||||||
|
|
||||||
/// If the manifest has a `[compaction]` section, build a `PruneConfig`
|
/// If the manifest enables prune under `[compaction]`, build a
|
||||||
/// from its `prune_*` fields and call [`attach_prune`](Self::attach_prune).
|
/// `PruneConfig` from its `prune_*` fields and call
|
||||||
/// Otherwise no-op. Called from all Worker constructors so prune is
|
/// [`attach_prune`](Self::attach_prune). Otherwise no-op. Called from all
|
||||||
/// active whenever the manifest asks for it.
|
/// Worker constructors.
|
||||||
pub(crate) fn apply_prune_from_manifest(&mut self) {
|
pub(crate) fn apply_prune_from_manifest(&mut self) {
|
||||||
let Some(compaction) = self.manifest().compaction.as_ref() else {
|
let Some(compaction) = self.manifest().compaction.as_ref() else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
let config = PruneConfig {
|
let Some(config) = prune_config_from_compaction(compaction) else {
|
||||||
protected_tokens: compaction.prune_protected_tokens,
|
return;
|
||||||
min_savings: compaction.prune_min_savings,
|
|
||||||
};
|
};
|
||||||
self.attach_prune(config);
|
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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -956,6 +956,7 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
|
|||||||
.compaction
|
.compaction
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.map(|c| CompactionConfigPartial {
|
.map(|c| CompactionConfigPartial {
|
||||||
|
prune_enabled: Some(c.prune_enabled),
|
||||||
prune_protected_tokens: Some(c.prune_protected_tokens),
|
prune_protected_tokens: Some(c.prune_protected_tokens),
|
||||||
prune_min_savings: Some(c.prune_min_savings),
|
prune_min_savings: Some(c.prune_min_savings),
|
||||||
threshold: c.threshold,
|
threshold: c.threshold,
|
||||||
|
|||||||
@@ -159,6 +159,7 @@ model_id = "test-model"
|
|||||||
max_tokens = 100
|
max_tokens = 100
|
||||||
|
|
||||||
[compaction]
|
[compaction]
|
||||||
|
prune_enabled = true
|
||||||
prune_protected_tokens = {prune_protected_tokens}
|
prune_protected_tokens = {prune_protected_tokens}
|
||||||
prune_min_savings = {prune_min_savings}
|
prune_min_savings = {prune_min_savings}
|
||||||
|
|
||||||
|
|||||||
+4
-1
@@ -179,9 +179,12 @@ permission = "write"
|
|||||||
|
|
||||||
# ===== [compaction] =========================================================
|
# ===== [compaction] =========================================================
|
||||||
# コンテキスト圧縮 (Prune / Compact)。セクション省略で両方無効。
|
# コンテキスト圧縮 (Prune / Compact)。セクション省略で両方無効。
|
||||||
# セクションを書いた時点で Prune は有効化、Compact は閾値が None なら無効。
|
# `prune_enabled`でPruneを制御する。Compactは閾値がNoneなら無効。
|
||||||
# [compaction]
|
# [compaction]
|
||||||
#
|
#
|
||||||
|
# # 任意。デフォルト: false。Compactの設定を保ったままPruneだけ切替できる。
|
||||||
|
# prune_enabled = true
|
||||||
|
#
|
||||||
# # 任意。デフォルト: 8000 (`defaults::PRUNE_PROTECTED_TOKENS`)。
|
# # 任意。デフォルト: 8000 (`defaults::PRUNE_PROTECTED_TOKENS`)。
|
||||||
# # pruning から保護する末尾 token budget。turn 数ではなく usage estimate で境界を引く。
|
# # pruning から保護する末尾 token budget。turn 数ではなく usage estimate で境界を引く。
|
||||||
# prune_protected_tokens = 8000
|
# prune_protected_tokens = 8000
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ engine = {
|
|||||||
|
|
||||||
compaction = {
|
compaction = {
|
||||||
kind = "tokens";
|
kind = "tokens";
|
||||||
|
prune_enabled = false;
|
||||||
threshold = 240000;
|
threshold = 240000;
|
||||||
request_threshold = 270000;
|
request_threshold = 270000;
|
||||||
worker_context_max_tokens = 100000;
|
worker_context_max_tokens = 100000;
|
||||||
|
|||||||
Reference in New Issue
Block a user