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

feat: compact worker サーキットブレーカーを占有量ベースに統一

Browse Source
This commit is contained in:
Hare
2026-05-11 00:43:16 +09:00
parent d6f27f7c45
commit 8100a5dfd1
10 changed files with 246 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/manifest/src/config.rs
+38 -1
View File
@@ -116,6 +116,8 @@ pub struct CompactionConfigPartial {
#[serde(default)]
pub compact_worker_max_input_tokens: Option<u64>,
#[serde(default)]
pub compact_worker_max_turns: Option<u32>,
#[serde(default)]
pub model: Option<ModelManifest>,
}
@@ -325,6 +327,9 @@ impl CompactionConfigPartial {
compact_worker_max_input_tokens: upper
.compact_worker_max_input_tokens
.or(self.compact_worker_max_input_tokens),
compact_worker_max_turns: upper
.compact_worker_max_turns
.or(self.compact_worker_max_turns),
model: merge_option(self.model, upper.model, ModelManifest::merge),
}
}
@@ -461,6 +466,9 @@ impl TryFrom<PodManifestConfig> for PodManifest {
compact_worker_max_input_tokens: c
.compact_worker_max_input_tokens
.unwrap_or(defaults::COMPACT_WORKER_MAX_INPUT_TOKENS),
compact_worker_max_turns: c
.compact_worker_max_turns
.or(defaults::COMPACT_WORKER_MAX_TURNS),
model: c.model,
})
})
@@ -949,6 +957,32 @@ stop_sequences = ["\n\n", "</stop>"]
);
}
#[test]
fn from_toml_accepts_compact_worker_max_turns() {
let cfg = PodManifestConfig::from_toml(
r#"
[compaction]
compact_worker_max_turns = 7
"#,
)
.unwrap();
assert_eq!(cfg.compaction.unwrap().compact_worker_max_turns, Some(7));
}
#[test]
fn try_from_compaction_defaults_compact_worker_max_turns() {
let mut cfg = minimal_valid();
cfg.compaction = Some(CompactionConfigPartial::default());
let manifest = PodManifest::try_from(cfg).unwrap();
assert_eq!(
manifest.compaction.unwrap().compact_worker_max_turns,
defaults::COMPACT_WORKER_MAX_TURNS
);
}
#[test]
fn from_toml_partial_layer_succeeds() {
// A project-layer manifest with only scope set must parse fine.
@@ -1042,7 +1076,10 @@ name = "dbg"
fn skills_directories_resolved_against_base() {
let mut cfg = minimal_valid();
cfg.skills = Some(SkillsConfig {
directories: vec![PathBuf::from(".claude/skills"), PathBuf::from("/abs/elsewhere")],
directories: vec![
PathBuf::from(".claude/skills"),
PathBuf::from("/abs/elsewhere"),
],
});
let resolved = cfg.resolve_paths(Path::new("/workspace/proj"));
let dirs = resolved.skills.as_ref().unwrap().directories.clone();
crates/manifest/src/defaults.rs
+5 -1
View File
@@ -36,12 +36,16 @@ pub const DEFAULT_INSTRUCTION: &str = "$insomnia/default";
/// [`crate::CompactionConfig::compact_auto_read_budget`].
pub const COMPACT_AUTO_READ_BUDGET: u64 = 8000;
/// Cumulative input-token cap for the compact worker's own LLM
/// Current prompt-occupancy cap for the compact worker's own LLM
/// calls. Exceeding this aborts the compact run (circuit-breaker
/// path). See
/// [`crate::CompactionConfig::compact_worker_max_input_tokens`].
pub const COMPACT_WORKER_MAX_INPUT_TOKENS: u64 = 50_000;
/// Optional maximum compact-worker tool-loop depth. `None` means unlimited.
/// See [`crate::CompactionConfig::compact_worker_max_turns`].
pub const COMPACT_WORKER_MAX_TURNS: Option<u32> = Some(20);
/// Number of recently-touched files fed to the compact worker as
/// default references.
pub const COMPACT_DEFAULT_REFERENCE_COUNT: usize = 5;
crates/manifest/src/lib.rs
+24 -2
View File
@@ -321,11 +321,16 @@ pub struct CompactionConfig {
#[serde(default = "default_compact_auto_read_budget")]
pub compact_auto_read_budget: u64,
/// Cumulative input-token cap for the compact worker's own LLM
/// calls. Exceeding this aborts the compact run.
/// Current prompt-occupancy cap for the compact worker's own LLM
/// requests. Exceeding this aborts the compact run.
#[serde(default = "default_compact_worker_max_input_tokens")]
pub compact_worker_max_input_tokens: u64,
/// Optional maximum compact-worker tool-loop depth. `None` leaves the
/// worker unlimited; the default bounds runaway short-context loops.
#[serde(default = "default_compact_worker_max_turns")]
pub compact_worker_max_turns: Option<u32>,
/// Optional model for the compactor (summary) LLM.
/// If omitted, the main model is cloned via `clone_boxed()`.
#[serde(default)]
@@ -347,6 +352,9 @@ fn default_compact_auto_read_budget() -> u64 {
fn default_compact_worker_max_input_tokens() -> u64 {
defaults::COMPACT_WORKER_MAX_INPUT_TOKENS
}
fn default_compact_worker_max_turns() -> Option<u32> {
defaults::COMPACT_WORKER_MAX_TURNS
}
impl Default for CompactionConfig {
fn default() -> Self {
@@ -358,6 +366,7 @@ impl Default for CompactionConfig {
compact_retained_tokens: default_compact_retained_tokens(),
compact_auto_read_budget: default_compact_auto_read_budget(),
compact_worker_max_input_tokens: default_compact_worker_max_input_tokens(),
compact_worker_max_turns: default_compact_worker_max_turns(),
model: None,
}
}
@@ -521,6 +530,19 @@ model_id = "claude-sonnet-4-20250514"
assert_eq!(c.compact_threshold, Some(80000));
assert_eq!(c.compact_request_threshold, None);
assert_eq!(c.compact_retained_tokens, 8000);
assert_eq!(c.compact_worker_max_turns, Some(20));
}
#[test]
fn parse_compaction_worker_max_turns() {
let toml = format!(
"{MINIMAL_REQUIRED}\n\
[compaction]\n\
compact_worker_max_turns = 7\n"
);
let manifest = PodManifest::from_toml(&toml).unwrap();
let c = manifest.compaction.unwrap();
assert_eq!(c.compact_worker_max_turns, Some(7));
}
#[test]
crates/pod/src/compact/usage_tracker.rs
+25
View File
@@ -99,6 +99,18 @@ impl UsageTracker {
});
}
/// Return a clone of the accumulated `UsageRecord`s without clearing them.
/// Used by request-time circuit breakers that need the same occupancy
/// projection as Pod persistence while the run is still active.
pub(crate) fn records(&self) -> Vec<UsageRecord> {
self.pending_records
.lock()
.unwrap()
.iter()
.map(|r| r.record.clone())
.collect()
}
/// Drain accumulated records. Called by Pod after a run completes,
/// before persisting the turn.
pub(crate) fn drain(&self) -> Vec<RecordedUsage> {
@@ -136,6 +148,19 @@ mod tests {
assert!(records[0].correlation_id.is_none());
}
#[test]
fn records_clones_without_clearing() {
let tracker = UsageTracker::new();
tracker.note_request(1);
tracker.record_usage(&make_event(10, 0, 0, 5));
let records = tracker.records();
assert_eq!(records.len(), 1);
assert_eq!(records[0].history_len, 1);
assert_eq!(records[0].input_total_tokens, 10);
assert_eq!(tracker.records().len(), 1);
}
#[test]
fn drain_clears_buffer() {
let tracker = UsageTracker::new();
crates/pod/src/compact/worker.rs
+74 -9
View File
@@ -18,7 +18,6 @@
//! compacted session's opening system messages.
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
@@ -28,6 +27,7 @@ use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use serde::Deserialize;
use tools::ScopedFs;
use crate::compact::usage_tracker::UsageTracker;
use crate::fs_view::{ReadRequirement, slice_lines};
/// Aggregated output of a compact worker run.
@@ -246,24 +246,29 @@ pub(crate) fn write_summary_tool(ctx: Arc<Mutex<CompactWorkerContext>>) -> ToolD
})
}
/// Interceptor that aborts the compact worker as soon as its cumulative
/// input-token count crosses `max_input_tokens`. Pairs with the
/// `on_usage` callback registered by `Pod::compact`, which is what
/// actually accumulates `input_so_far`.
/// Interceptor that aborts the compact worker when its current prompt
/// occupancy estimate crosses `max_input_tokens`. The estimate uses the same
/// `UsageRecord` + `llm_worker::token_counter::total_tokens` path as the main
/// Pod compaction thresholds, so prompt-cache hits are not counted cumulatively
/// across turns.
pub(crate) struct CompactWorkerInterceptor {
pub input_so_far: Arc<AtomicU64>,
pub usage_tracker: Arc<UsageTracker>,
pub max_input_tokens: u64,
}
#[async_trait]
impl Interceptor for CompactWorkerInterceptor {
async fn pre_llm_request(&self, _context: &mut Vec<Item>) -> PreRequestAction {
if self.input_so_far.load(Ordering::Relaxed) > self.max_input_tokens {
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
let records = self.usage_tracker.records();
let estimate = llm_worker::token_counter::total_tokens(context, &records);
if estimate.tokens > self.max_input_tokens {
return PreRequestAction::Cancel(format!(
"compact worker input exceeded {} tokens",
"compact worker input occupancy exceeded {} tokens",
self.max_input_tokens
));
}
self.usage_tracker.note_request(context.len());
PreRequestAction::Continue
}
}
@@ -283,6 +288,66 @@ mod tests {
ScopedFs::new(scope, tmp.to_path_buf())
}
fn make_usage(input: u64) -> llm_worker::timeline::event::UsageEvent {
llm_worker::timeline::event::UsageEvent {
input_tokens: Some(input),
output_tokens: Some(0),
total_tokens: Some(input),
cache_read_input_tokens: None,
cache_creation_input_tokens: None,
}
}
#[tokio::test]
async fn compact_worker_interceptor_uses_occupancy_not_cumulative_usage() {
let tracker = Arc::new(UsageTracker::new());
let interceptor = CompactWorkerInterceptor {
usage_tracker: tracker.clone(),
max_input_tokens: 150,
};
let mut context = vec![Item::user_message("hello")];
assert!(matches!(
interceptor.pre_llm_request(&mut context).await,
PreRequestAction::Continue
));
tracker.record_usage(&make_usage(100));
assert!(matches!(
interceptor.pre_llm_request(&mut context).await,
PreRequestAction::Continue
));
tracker.record_usage(&make_usage(100));
// Two 100-token requests would exceed a cumulative 150-token cap, but
// current occupancy is still the latest 100-token measurement.
assert!(matches!(
interceptor.pre_llm_request(&mut context).await,
PreRequestAction::Continue
));
}
#[tokio::test]
async fn compact_worker_interceptor_cancels_when_occupancy_exceeds_cap() {
let tracker = Arc::new(UsageTracker::new());
let interceptor = CompactWorkerInterceptor {
usage_tracker: tracker.clone(),
max_input_tokens: 99,
};
let mut context = vec![Item::user_message("hello")];
assert!(matches!(
interceptor.pre_llm_request(&mut context).await,
PreRequestAction::Continue
));
tracker.record_usage(&make_usage(100));
assert!(matches!(
interceptor.pre_llm_request(&mut context).await,
PreRequestAction::Cancel(message) if message.contains("occupancy")
));
}
#[tokio::test]
async fn mark_read_required_records_and_deducts_budget() {
let tmp = tempfile::TempDir::new().unwrap();
crates/pod/src/pod.rs
+21 -14
View File
@@ -12,8 +12,8 @@ use session_store::{EntryHash, PodScopeSnapshot, SessionId, SessionStartState, S
use tracing::{info, warn};
use manifest::{
Permission, PodManifest, PodManifestConfig, ResolveError, Scope, ScopeConfig, ScopeError, ScopeRule,
SharedScope, WorkerManifest,
Permission, PodManifest, PodManifestConfig, ResolveError, Scope, ScopeConfig, ScopeError,
ScopeRule, SharedScope, WorkerManifest,
};
use crate::compact::state::CompactState;
@@ -1456,8 +1456,6 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
///
/// Returns the new session ID.
pub async fn compact(&mut self, retained_tokens: u64) -> Result<SessionId, PodError> {
use std::sync::atomic::{AtomicU64, Ordering};
use crate::compact::worker::{
CompactWorkerContext, CompactWorkerInterceptor, add_reference_tool,
mark_read_required_tool, write_summary_tool,
@@ -1477,7 +1475,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
// Compaction-related knobs. Fall through to manifest defaults when
// `[compaction]` is omitted entirely.
let (auto_read_budget, compact_worker_max_input_tokens) = self
let (auto_read_budget, compact_worker_max_input_tokens, compact_worker_max_turns) = self
.manifest
.compaction
.as_ref()
@@ -1485,11 +1483,13 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
(
c.compact_auto_read_budget,
c.compact_worker_max_input_tokens,
c.compact_worker_max_turns,
)
})
.unwrap_or((
manifest::defaults::COMPACT_AUTO_READ_BUDGET,
manifest::defaults::COMPACT_WORKER_MAX_INPUT_TOKENS,
manifest::defaults::COMPACT_WORKER_MAX_TURNS,
));
// Default references: the N most-recently-touched files in the
@@ -1530,21 +1530,24 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
let mut summary_worker = Worker::new(summary_client).system_prompt(summary_system_prompt);
summary_worker.set_cache_key(Some(self.session_id.to_string()));
// Cumulative input-token meter + interceptor. The meter is bumped
// from the on_usage callback and read on every pre_llm_request.
let input_so_far = Arc::new(AtomicU64::new(0));
// Occupancy-based input-token meter + interceptor. The tracker pairs
// each pre-request history length with the following UsageEvent, then
// the interceptor projects current prompt occupancy with the same
// UsageRecord counter used by the main Pod thresholds.
let summary_usage_tracker = Arc::new(UsageTracker::new());
{
let acc = input_so_far.clone();
let tracker = summary_usage_tracker.clone();
summary_worker.on_usage(move |event| {
if let Some(tokens) = event.input_tokens {
acc.fetch_add(tokens, Ordering::Relaxed);
}
tracker.record_usage(event);
});
}
summary_worker.set_interceptor(CompactWorkerInterceptor {
input_so_far: input_so_far.clone(),
usage_tracker: summary_usage_tracker,
max_input_tokens: compact_worker_max_input_tokens,
});
if compact_worker_max_turns.is_some() {
summary_worker.set_max_turns(compact_worker_max_turns);
}
// Tools: read_file (shared scope, fresh tracker) + the three
// compact-specific tools that populate `ctx`.
@@ -3069,7 +3072,11 @@ permission = "write"
let shadows = ingest_skills(&mut registry, &manifest);
// workspace skill `alpha` should be registered (no collision).
assert!(registry.get(&memory::Slug::parse("alpha").unwrap()).is_some());
assert!(
registry
.get(&memory::Slug::parse("alpha").unwrap())
.is_some()
);
// No workflow exists to shadow `alpha`, so no shadow event for it.
assert!(shadows.iter().all(|s| s.slug.as_str() != "alpha"));
}