diff --git a/crates/worker/src/compact/state.rs b/crates/worker/src/compact/state.rs index f9aeb638..cebab0f2 100644 --- a/crates/worker/src/compact/state.rs +++ b/crates/worker/src/compact/state.rs @@ -1,109 +1,265 @@ -//! Shared state for compaction decisions. -//! -//! Holds the two configured thresholds and circuit-breaker / thrash-detection -//! flags shared between: -//! - `WorkerInterceptor` (reads `request_threshold` — the *safety net* for -//! between-requests yielding) -//! - `Worker::try_pre_run_compact` (reads `post_run_threshold` — the -//! *proactive* check before the next turn starts) -//! - `Worker::run()` / `resume()` (circuit breaker, thrash detection) -//! -//! Current occupancy (input-token count) is **not** stored here. The single -//! source of truth is `session_store::UsageRecord` (persisted per LLM call) -//! projected through `Worker::total_tokens()`. Callers pass the current -//! occupancy to `exceeds_*` at check time. +use std::sync::Mutex; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use super::telemetry::CompactFailureCategory; -const MAX_COMPACT_FAILURES: usize = 3; +/// The process-local state of automatic compaction for one logical run. +/// +/// This is deliberately not persisted: restoring a worker starts from +/// [`AutomaticCompactState::NotAttempted`]. The failure counter is likewise a +/// guard for one live worker process, not session authority. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AutomaticCompactState { + NotAttempted, + Attempted, + Completed(CompactionOutcome), +} -/// Shared mutable state for compaction decisions. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CompactionOutcome { + Succeeded, + Failed(CompactFailureCategory), + Cancelled, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AutomaticCompactTrigger { + PreRun, + RequestThreshold, +} + +/// Decision returned by an atomic threshold/attempt-state evaluation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AutomaticCompactDecision { + Continue, + Start(AutomaticCompactTrigger), + Block(AutomaticCompactBlock), +} + +/// Typed reason why a provider request may not proceed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AutomaticCompactBlock { + /// Compaction succeeded but the replacement still exceeds the request + /// safety threshold before a post-compaction request was durably recorded. + Thrash, + /// This logical run already used its automatic attempt and it failed. + Failed(CompactFailureCategory), + /// This logical run already used its automatic attempt and it was + /// cancelled. Cancellation is not counted as a compaction failure. + Cancelled, + /// Consecutive automatic failures disabled further attempts until a + /// successful manual compaction or explicit re-enable. + Disabled(CompactFailureCategory), +} + +#[derive(Debug)] +struct AutomaticCompactRuntimeState { + attempt: AutomaticCompactState, + consecutive_failures: u32, + automatic_disabled: bool, + last_failure: Option, + pending_request_block: Option, +} + +const MAX_CONSECUTIVE_AUTOMATIC_FAILURES: u32 = 2; + +/// Tracks automatic compaction thresholds and process-local loop guards. +#[derive(Debug)] pub(crate) struct CompactState { - /// Between-turns threshold (proactive). Checked before the next turn - /// starts. `None` disables the pre-run check. - post_run_threshold: Option, - /// Between-requests threshold (safety net). Checked inside a turn - /// before each LLM request. `None` disables the request check. + /// Post-run threshold. Checked before a new user run starts. + compact_threshold: Option, + /// In-request safety threshold. Checked immediately before each provider + /// request, both before and after pre-request hooks. request_threshold: Option, - /// Token budget retained verbatim at the tail after compaction. retained_tokens: u64, - /// Consecutive compact failures. At `MAX_COMPACT_FAILURES`, compaction is disabled. - consecutive_failures: AtomicUsize, - /// `true` immediately after a successful compact, cleared on next normal completion. - just_compacted: AtomicBool, - /// `true` when circuit breaker has tripped. - disabled: AtomicBool, + max_consecutive_failures: u32, + runtime: Mutex, } impl CompactState { pub(crate) fn new( - post_run_threshold: Option, + compact_threshold: Option, request_threshold: Option, retained_tokens: u64, ) -> Self { Self { - post_run_threshold, + compact_threshold, request_threshold, retained_tokens, - consecutive_failures: AtomicUsize::new(0), - just_compacted: AtomicBool::new(false), - disabled: AtomicBool::new(false), + max_consecutive_failures: MAX_CONSECUTIVE_AUTOMATIC_FAILURES, + runtime: Mutex::new(AutomaticCompactRuntimeState { + attempt: AutomaticCompactState::NotAttempted, + consecutive_failures: 0, + automatic_disabled: false, + last_failure: None, + pending_request_block: None, + }), } } - /// Configured between-requests threshold (if any). - pub(crate) fn request_threshold(&self) -> Option { - self.request_threshold - } - - /// Token budget retained verbatim at the tail after compaction. pub(crate) fn retained_tokens(&self) -> u64 { self.retained_tokens } - /// Whether compaction has been disabled by the circuit breaker. - pub(crate) fn is_disabled(&self) -> bool { - self.disabled.load(Ordering::Relaxed) - } - - /// Whether `current_tokens` exceeds the between-requests threshold. - /// Returns `false` when `request_threshold` is unset. - pub(crate) fn exceeds_request(&self, current_tokens: u64) -> bool { - self.request_threshold - .map(|t| current_tokens > t) - .unwrap_or(false) - } - - /// Whether `current_tokens` exceeds the post-run threshold. - /// Returns `false` when `post_run_threshold` is unset. - pub(crate) fn exceeds_post_run(&self, current_tokens: u64) -> bool { - self.post_run_threshold - .map(|t| current_tokens > t) - .unwrap_or(false) - } - - /// Whether a compact just completed (for thrash detection). - pub(crate) fn just_compacted(&self) -> bool { - self.just_compacted.load(Ordering::Relaxed) - } - - /// Set or clear the just_compacted flag. - pub(crate) fn set_just_compacted(&self, val: bool) { - self.just_compacted.store(val, Ordering::Relaxed); - } - - /// Record a successful compaction: reset failure counter, set just_compacted. - pub(crate) fn record_compact_success(&self) { - self.consecutive_failures.store(0, Ordering::Relaxed); - self.just_compacted.store(true, Ordering::Relaxed); - } - - /// Record a compaction failure. Disables compaction after MAX_COMPACT_FAILURES. - pub(crate) fn record_compact_failure(&self) { - let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed); - if prev + 1 >= MAX_COMPACT_FAILURES { - self.disabled.store(true, Ordering::Relaxed); + pub(crate) fn pre_run_eligible(&self, total_tokens: u64) -> bool { + if !self + .compact_threshold + .is_some_and(|threshold| total_tokens > threshold) + { + return false; } + let runtime = self.lock_runtime(); + !runtime.automatic_disabled && runtime.attempt == AutomaticCompactState::NotAttempted + } + + /// Starts a fresh logical run. Pause/resume paths must not call this. + pub(crate) fn begin_logical_run(&self) { + let mut runtime = self.lock_runtime(); + runtime.attempt = AutomaticCompactState::NotAttempted; + runtime.pending_request_block = None; + } + + /// Clears per-run state after a terminal run outcome. The consecutive + /// failure circuit breaker intentionally survives logical-run boundaries. + pub(crate) fn finish_logical_run(&self) { + let mut runtime = self.lock_runtime(); + runtime.attempt = AutomaticCompactState::NotAttempted; + runtime.pending_request_block = None; + } + + /// Atomically evaluates the pre-run threshold and claims this logical + /// run's one automatic compaction attempt when eligible. + pub(crate) fn evaluate_pre_run(&self, total_tokens: u64) -> AutomaticCompactDecision { + if !self + .compact_threshold + .is_some_and(|threshold| total_tokens > threshold) + { + return AutomaticCompactDecision::Continue; + } + self.claim_attempt(AutomaticCompactTrigger::PreRun) + } + + /// Atomically evaluates the request safety threshold and either claims the + /// run's one automatic attempt or records a typed request-block reason. + pub(crate) fn evaluate_request(&self, total_tokens: u64) -> AutomaticCompactDecision { + if !self + .request_threshold + .is_some_and(|threshold| total_tokens > threshold) + { + return AutomaticCompactDecision::Continue; + } + + self.claim_attempt(AutomaticCompactTrigger::RequestThreshold) + } + + pub(crate) fn record_request_block(&self, block: AutomaticCompactBlock) { + self.lock_runtime().pending_request_block = Some(block); + } + + /// Claims a hook-originated compaction yield under the same one-attempt + /// guard used by threshold evaluation. + pub(crate) fn claim_hook_yield(&self) -> AutomaticCompactDecision { + self.claim_attempt(AutomaticCompactTrigger::RequestThreshold) + } + + pub(crate) fn has_claimed_attempt(&self) -> bool { + self.lock_runtime().attempt == AutomaticCompactState::Attempted + } + + /// Completes the currently claimed automatic attempt exactly once. + pub(crate) fn complete_automatic(&self, outcome: CompactionOutcome) -> bool { + let mut runtime = self.lock_runtime(); + if runtime.attempt != AutomaticCompactState::Attempted { + return false; + } + + runtime.attempt = AutomaticCompactState::Completed(outcome); + match outcome { + CompactionOutcome::Succeeded => { + runtime.consecutive_failures = 0; + runtime.last_failure = None; + } + CompactionOutcome::Failed(category) => { + runtime.consecutive_failures = runtime.consecutive_failures.saturating_add(1); + runtime.last_failure = Some(category); + if runtime.consecutive_failures >= self.max_consecutive_failures { + runtime.automatic_disabled = true; + } + } + CompactionOutcome::Cancelled => {} + } + true + } + + /// Re-arms automatic compaction only after the first real provider request + /// following successful compaction has a durably committed UsageRecord. + pub(crate) fn post_compact_request_committed(&self) { + let mut runtime = self.lock_runtime(); + if runtime.attempt == AutomaticCompactState::Completed(CompactionOutcome::Succeeded) { + runtime.attempt = AutomaticCompactState::NotAttempted; + } + } + + /// A successful manual compaction is the recovery path for a disabled + /// automatic circuit breaker and does not count as an automatic attempt. + pub(crate) fn reenable_automatic(&self) { + let mut runtime = self.lock_runtime(); + runtime.attempt = AutomaticCompactState::NotAttempted; + runtime.consecutive_failures = 0; + runtime.automatic_disabled = false; + runtime.last_failure = None; + runtime.pending_request_block = None; + } + + pub(crate) fn take_pending_request_block(&self) -> Option { + self.lock_runtime().pending_request_block.take() + } + + fn claim_attempt(&self, trigger: AutomaticCompactTrigger) -> AutomaticCompactDecision { + let mut runtime = self.lock_runtime(); + if runtime.automatic_disabled { + let category = runtime + .last_failure + .expect("disabled automatic compaction must retain its failure category"); + return AutomaticCompactDecision::Block(AutomaticCompactBlock::Disabled(category)); + } + + match runtime.attempt { + AutomaticCompactState::NotAttempted => { + runtime.attempt = AutomaticCompactState::Attempted; + AutomaticCompactDecision::Start(trigger) + } + AutomaticCompactState::Attempted => AutomaticCompactDecision::Continue, + AutomaticCompactState::Completed(CompactionOutcome::Succeeded) => { + AutomaticCompactDecision::Block(AutomaticCompactBlock::Thrash) + } + AutomaticCompactState::Completed(CompactionOutcome::Failed(category)) => { + AutomaticCompactDecision::Block(AutomaticCompactBlock::Failed(category)) + } + AutomaticCompactState::Completed(CompactionOutcome::Cancelled) => { + AutomaticCompactDecision::Block(AutomaticCompactBlock::Cancelled) + } + } + } + + fn lock_runtime(&self) -> std::sync::MutexGuard<'_, AutomaticCompactRuntimeState> { + self.runtime + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + #[cfg(test)] + pub(crate) fn attempt_state(&self) -> AutomaticCompactState { + self.lock_runtime().attempt + } + + #[cfg(test)] + pub(crate) fn consecutive_failures(&self) -> u32 { + self.lock_runtime().consecutive_failures + } + + #[cfg(test)] + pub(crate) fn automatic_disabled(&self) -> bool { + self.lock_runtime().automatic_disabled } } @@ -111,81 +267,120 @@ impl CompactState { mod tests { use super::*; + const FAILURE: CompactFailureCategory = CompactFailureCategory::Storage; + #[test] - fn both_thresholds_configured() { - let state = CompactState::new(Some(80_000), Some(90_000), 8_000); - assert_eq!(state.request_threshold(), Some(90_000)); - assert_eq!(state.retained_tokens(), 8_000); + fn same_logical_run_claims_only_one_automatic_attempt() { + let state = CompactState::new(Some(10), Some(10), 2); - assert!(!state.exceeds_request(70_000)); - assert!(!state.exceeds_post_run(70_000)); - - assert!(!state.exceeds_request(85_000)); - assert!(state.exceeds_post_run(85_000)); - - assert!(state.exceeds_request(95_000)); - assert!(state.exceeds_post_run(95_000)); + assert_eq!( + state.evaluate_pre_run(11), + AutomaticCompactDecision::Start(AutomaticCompactTrigger::PreRun) + ); + assert_eq!( + state.evaluate_request(11), + AutomaticCompactDecision::Continue + ); + assert_eq!(state.attempt_state(), AutomaticCompactState::Attempted); } #[test] - fn post_run_only() { - let state = CompactState::new(Some(80_000), None, 8_000); - // request check always false when threshold is None. - assert!(!state.exceeds_request(1_000_000)); - assert!(state.exceeds_post_run(85_000)); + fn completed_attempt_blocks_retry_and_preserves_failure_category() { + let state = CompactState::new(None, Some(10), 2); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Failed(FAILURE))); + + let block = AutomaticCompactBlock::Failed(FAILURE); + assert_eq!( + state.evaluate_request(11), + AutomaticCompactDecision::Block(block) + ); + state.record_request_block(block); + assert_eq!(state.take_pending_request_block(), Some(block)); + assert_eq!(state.consecutive_failures(), 1); } #[test] - fn request_only() { - let state = CompactState::new(None, Some(90_000), 8_000); - assert!(!state.exceeds_post_run(1_000_000)); - assert!(state.exceeds_request(95_000)); + fn two_failed_automatic_attempts_disable_until_manual_success() { + let state = CompactState::new(Some(10), Some(10), 2); + for expected in 1..=2 { + state.begin_logical_run(); + assert!(matches!( + state.evaluate_pre_run(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Failed(FAILURE))); + assert_eq!(state.consecutive_failures(), expected); + } + assert!(state.automatic_disabled()); + + state.begin_logical_run(); + assert_eq!( + state.evaluate_pre_run(11), + AutomaticCompactDecision::Block(AutomaticCompactBlock::Disabled(FAILURE)) + ); + + state.reenable_automatic(); + assert!(!state.automatic_disabled()); + assert_eq!(state.consecutive_failures(), 0); + assert!(matches!( + state.evaluate_pre_run(11), + AutomaticCompactDecision::Start(_) + )); } #[test] - fn both_none_disables_all_checks() { - let state = CompactState::new(None, None, 8_000); - assert!(!state.exceeds_request(1_000_000)); - assert!(!state.exceeds_post_run(1_000_000)); + fn cancellation_consumes_run_attempt_without_counting_as_failure() { + let state = CompactState::new(None, Some(10), 2); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Cancelled)); + + assert_eq!(state.consecutive_failures(), 0); + assert!(!state.automatic_disabled()); + assert_eq!( + state.evaluate_request(11), + AutomaticCompactDecision::Block(AutomaticCompactBlock::Cancelled) + ); } #[test] - fn circuit_breaker_trips_after_max_failures() { - let state = CompactState::new(Some(80_000), Some(90_000), 8_000); - assert!(!state.is_disabled()); + fn successful_compaction_requires_committed_request_before_rearming() { + let state = CompactState::new(None, Some(10), 2); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Succeeded)); - state.record_compact_failure(); - assert!(!state.is_disabled()); - state.record_compact_failure(); - assert!(!state.is_disabled()); - state.record_compact_failure(); - assert!(state.is_disabled()); + assert_eq!( + state.evaluate_request(11), + AutomaticCompactDecision::Block(AutomaticCompactBlock::Thrash) + ); + state.post_compact_request_committed(); + assert_eq!(state.attempt_state(), AutomaticCompactState::NotAttempted); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); } #[test] - fn success_resets_failure_count() { - let state = CompactState::new(Some(80_000), Some(90_000), 8_000); - state.record_compact_failure(); - state.record_compact_failure(); - assert!(!state.is_disabled()); + fn pause_resume_preserves_attempt_while_terminal_finish_clears_it() { + let state = CompactState::new(None, Some(10), 2); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + // Pause/resume deliberately performs no state transition. + assert_eq!(state.attempt_state(), AutomaticCompactState::Attempted); - state.record_compact_success(); - assert!(state.just_compacted()); - - state.record_compact_failure(); - state.record_compact_failure(); - assert!(!state.is_disabled()); - } - - #[test] - fn just_compacted_lifecycle() { - let state = CompactState::new(Some(80_000), Some(90_000), 8_000); - assert!(!state.just_compacted()); - - state.record_compact_success(); - assert!(state.just_compacted()); - - state.set_just_compacted(false); - assert!(!state.just_compacted()); + state.finish_logical_run(); + assert_eq!(state.attempt_state(), AutomaticCompactState::NotAttempted); } } diff --git a/crates/worker/src/compact/telemetry.rs b/crates/worker/src/compact/telemetry.rs index a171468e..eee3dafd 100644 --- a/crates/worker/src/compact/telemetry.rs +++ b/crates/worker/src/compact/telemetry.rs @@ -56,7 +56,7 @@ pub(crate) enum CompactFailureCategory { } impl CompactFailureCategory { - fn as_str(self) -> &'static str { + pub(crate) fn as_str(self) -> &'static str { match self { Self::Cancelled => "cancelled", Self::SummaryMissing => "summary_missing", diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 86f26230..79bc9593 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -25,7 +25,7 @@ use arc_swap::ArcSwap; use async_trait::async_trait; use tracing::info; -use crate::compact::state::CompactState; +use crate::compact::state::{AutomaticCompactDecision, CompactState}; use crate::compact::usage_tracker::UsageTracker; use session_store::SystemItem; @@ -111,6 +111,9 @@ pub(crate) struct WorkerInterceptor { tool_calls_this_turn: AtomicUsize, } +const THRESHOLD_COMPACT_BLOCKED_DIAGNOSTIC: &str = + "automatic compaction could not make the provider request context safe"; + impl WorkerInterceptor { #[cfg(test)] pub(crate) fn new( @@ -229,27 +232,45 @@ impl WorkerInterceptor { Some(total_tokens(context, &records).tokens) } - fn request_threshold_exceeded(&self, current_tokens: Option, context: &[Item]) -> bool { - if let Some(state) = self.compact_state.as_ref() { - if !state.is_disabled() && !state.just_compacted() { - let current = current_tokens.unwrap_or(0); - if state.exceeds_request(current) { - let shape = context_shape(context); - info!( - input_tokens = current, - threshold = state.request_threshold().unwrap_or(0), - items_len = shape.items_len, - items_json_bytes = shape.items_json_bytes, - reasoning_items = shape.reasoning_items, - reasoning_encrypted_content_count = shape.reasoning_encrypted_content_count, - reasoning_encrypted_content_bytes = shape.reasoning_encrypted_content_bytes, - "Between-requests compaction threshold exceeded, yielding" - ); - return true; + fn request_compact_decision( + &self, + current_tokens: Option, + context: &[Item], + ) -> AutomaticCompactDecision { + let Some(state) = self.compact_state.as_ref() else { + return AutomaticCompactDecision::Continue; + }; + let current = current_tokens.unwrap_or(0); + let decision = state.evaluate_request(current); + if !matches!(decision, AutomaticCompactDecision::Continue) { + let shape = context_shape(context); + info!( + input_tokens = current, + ?decision, + items_len = shape.items_len, + items_json_bytes = shape.items_json_bytes, + reasoning_items = shape.reasoning_items, + reasoning_encrypted_content_count = shape.reasoning_encrypted_content_count, + reasoning_encrypted_content_bytes = shape.reasoning_encrypted_content_bytes, + "Between-requests automatic compaction decision" + ); + } + decision + } + + fn decision_action(&self, decision: AutomaticCompactDecision) -> Option { + match decision { + AutomaticCompactDecision::Continue => None, + AutomaticCompactDecision::Start(_) => Some(PreRequestAction::Yield), + AutomaticCompactDecision::Block(block) => { + if let Some(state) = &self.compact_state { + state.record_request_block(block); } + Some(PreRequestAction::Cancel( + THRESHOLD_COMPACT_BLOCKED_DIAGNOSTIC.to_string(), + )) } } - false } fn attach_prompt_provenance(&self, items: &mut [SystemItem]) { let prompts = self.prompts.load(); @@ -388,8 +409,10 @@ impl Interceptor for WorkerInterceptor { ) -> InterceptorResult { let context = context.items; let initial_tokens = self.estimated_tokens(context); - if self.request_threshold_exceeded(initial_tokens, context) { - return Ok(PreRequestAction::Yield); + if let Some(action) = + self.decision_action(self.request_compact_decision(initial_tokens, context)) + { + return Ok(action); } let info = PreRequestInfo { item_count: context.len(), @@ -423,6 +446,22 @@ impl Interceptor for WorkerInterceptor { return Ok(PreRequestAction::Cancel(reason)); } if should_yield { + if let Some(state) = &self.compact_state { + match state.claim_hook_yield() { + AutomaticCompactDecision::Start(_) => {} + AutomaticCompactDecision::Block(block) => { + state.record_request_block(block); + return Ok(PreRequestAction::Cancel( + THRESHOLD_COMPACT_BLOCKED_DIAGNOSTIC.to_string(), + )); + } + AutomaticCompactDecision::Continue => { + return Ok(PreRequestAction::Cancel( + THRESHOLD_COMPACT_BLOCKED_DIAGNOSTIC.to_string(), + )); + } + } + } return Ok(PreRequestAction::Yield); } @@ -445,16 +484,26 @@ impl Interceptor for WorkerInterceptor { }; let current_tokens = self.estimated_tokens(effective_context.as_ref()); - if self.request_threshold_exceeded(current_tokens, effective_context.as_ref()) { + let compact_decision = + self.request_compact_decision(current_tokens, effective_context.as_ref()); + if !matches!(compact_decision, AutomaticCompactDecision::Continue) { if let Err(error) = self.commit_system_items(&system_items) { return Ok(PreRequestAction::Cancel(format!( "session persistence failed: {error}" ))); } - return Ok(if appended_items.is_empty() { - PreRequestAction::Yield - } else { - PreRequestAction::YieldWith(appended_items) + return Ok(match compact_decision { + AutomaticCompactDecision::Start(_) if !appended_items.is_empty() => { + PreRequestAction::YieldWith(appended_items) + } + AutomaticCompactDecision::Start(_) => PreRequestAction::Yield, + AutomaticCompactDecision::Block(block) => { + if let Some(state) = &self.compact_state { + state.record_request_block(block); + } + PreRequestAction::Cancel(THRESHOLD_COMPACT_BLOCKED_DIAGNOSTIC.to_string()) + } + AutomaticCompactDecision::Continue => unreachable!(), }); } @@ -668,6 +717,18 @@ mod tests { Arc::new(builder.build()) } + struct YieldingPreRequestHook; + + #[async_trait] + impl Hook for YieldingPreRequestHook { + async fn call( + &self, + _info: &PreRequestContext, + ) -> Result { + Ok(HookPreRequestAction::Yield) + } + } + struct RecordingSystemItemCommitter { committed: Arc>>, } @@ -750,6 +811,36 @@ mod tests { }])) } + #[tokio::test] + async fn hook_yield_claims_attempt_before_returning_to_compaction() { + let mut builder = HookRegistryBuilder::new(); + builder.add_pre_llm_request(YieldingPreRequestHook); + let registry = Arc::new(builder.build()); + let state = Arc::new(CompactState::new(None, Some(u64::MAX), 0)); + let interceptor = WorkerInterceptor::new( + registry, + Some(Arc::clone(&state)), + Some(usage_handle_with(1, 1)), + NotifyBuffer::new(), + Arc::new(Mutex::new(Vec::new())), + test_prompts(), + None, + ); + let mut ctx = vec![Item::user_message("hello")]; + + let action = interceptor + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) + .await + .unwrap(); + + assert!(matches!(action, PreRequestAction::Yield)); + assert!(state.has_claimed_attempt()); + } + #[tokio::test] async fn pre_llm_request_yields_and_skips_hooks_when_request_threshold_exceeded() { let count = Arc::new(AtomicUsize::new(0)); @@ -761,7 +852,7 @@ mod tests { let interceptor = WorkerInterceptor::new( registry, - Some(state), + Some(Arc::clone(&state)), Some(history), NotifyBuffer::new(), Arc::new(Mutex::new(Vec::new())), @@ -779,6 +870,10 @@ mod tests { .unwrap(); assert!(matches!(action, PreRequestAction::Yield)); + assert_eq!( + state.attempt_state(), + crate::compact::state::AutomaticCompactState::Attempted + ); // Hook must not run when an internal mechanism short-circuits first. assert_eq!(count.load(Ordering::Relaxed), 0); } @@ -798,7 +893,7 @@ mod tests { let interceptor = WorkerInterceptor::new( registry, - Some(state), + Some(Arc::clone(&state)), Some(history), NotifyBuffer::new(), Arc::new(Mutex::new(Vec::new())), @@ -821,10 +916,51 @@ mod tests { PreRequestAction::YieldWith(items) => assert_eq!(items.len(), 1), other => panic!("expected YieldWith queued system item, got {other:?}"), } + assert_eq!( + state.attempt_state(), + crate::compact::state::AutomaticCompactState::Attempted + ); assert!(saw_handle.load(Ordering::Relaxed)); assert_eq!(committed.lock().expect("committed system items").len(), 1); } + #[tokio::test] + async fn successful_compaction_blocks_unsafe_request_until_usage_commit() { + let registry = Arc::new(HookRegistryBuilder::new().build()); + let state = Arc::new(CompactState::new(None, Some(10), 0)); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(crate::compact::state::CompactionOutcome::Succeeded)); + let history = usage_handle_with(1, 11); + let interceptor = WorkerInterceptor::new( + registry, + Some(Arc::clone(&state)), + Some(history), + NotifyBuffer::new(), + Arc::new(Mutex::new(Vec::new())), + test_prompts(), + None, + ); + let mut ctx = vec![Item::user_message("still too large")]; + + let action = interceptor + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) + .await + .expect("pre-request interception should succeed"); + + assert!(matches!(action, PreRequestAction::Cancel(_))); + assert_eq!( + state.take_pending_request_block(), + Some(crate::compact::state::AutomaticCompactBlock::Thrash) + ); + } + #[tokio::test] async fn pre_llm_request_counts_in_flight_usage_records() { let registry = Arc::new(HookRegistryBuilder::new().build()); @@ -843,7 +979,7 @@ mod tests { let interceptor = WorkerInterceptor::new( registry, - Some(state), + Some(Arc::clone(&state)), Some(history), NotifyBuffer::new(), Arc::new(Mutex::new(Vec::new())), @@ -875,7 +1011,7 @@ mod tests { let interceptor = WorkerInterceptor::new( registry, - Some(state), + Some(Arc::clone(&state)), Some(history), NotifyBuffer::new(), Arc::new(Mutex::new(Vec::new())), @@ -923,7 +1059,7 @@ mod tests { let history = Arc::new(Mutex::new(vec![record])); let interceptor = WorkerInterceptor::new( registry, - Some(state), + Some(Arc::clone(&state)), Some(history), NotifyBuffer::new(), Arc::new(Mutex::new(Vec::new())), @@ -957,7 +1093,7 @@ mod tests { let interceptor = WorkerInterceptor::new( registry, - Some(state), + Some(Arc::clone(&state)), Some(history), NotifyBuffer::new(), Arc::new(Mutex::new(Vec::new())), diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index febf0242..7385f514 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -39,7 +39,10 @@ use manifest::{ SharedScope, WorkerManifest, WorkerManifestConfig, }; -use crate::compact::state::CompactState; +use crate::compact::state::{ + AutomaticCompactBlock, AutomaticCompactDecision, AutomaticCompactTrigger, CompactState, + CompactionOutcome, +}; use crate::compact::telemetry::{ CompactAttempt, CompactFailureCategory, CompactMode, CompactSuccessStats, CompactThresholdPolicy, correlated_post_request_metric, new_compact_metric_correlation_id, @@ -3611,11 +3614,9 @@ impl Worker { /// defensive reasons; this is the gate for joining the memory task /// before the compact runs. fn should_pre_run_compact(&self) -> bool { - self.compact_state.as_ref().is_some_and(|s| { - !s.is_disabled() - && !s.just_compacted() - && s.exceeds_post_run(self.total_tokens().tokens) - }) + self.compact_state + .as_ref() + .is_some_and(|state| state.pre_run_eligible(self.total_tokens().tokens)) } /// Prelude shared by `run` / `run_for_notification` / `resume`. @@ -3630,8 +3631,9 @@ impl Worker { self.ensure_interceptor_installed(); self.ensure_system_prompt_materialized().await?; self.ensure_segment_head().await?; - if self.should_pre_run_compact() {} - self.try_pre_run_compact().await; + if self.should_pre_run_compact() { + self.try_pre_run_compact().await; + } Ok(()) } @@ -3901,6 +3903,10 @@ impl Worker { // context via a different entry point and never triggers this // path. self.prepare_interrupted_history_for_fresh_run()?; + self.ensure_interceptor_installed(); + if let Some(state) = &self.compact_state { + state.begin_logical_run(); + } self.prepare_for_run().await?; @@ -4322,6 +4328,10 @@ impl Worker { "run_for_notification expects a non-UserSend InvokeKind; got {kind:?}" ); self.prepare_interrupted_history_for_fresh_run()?; + self.ensure_interceptor_installed(); + if let Some(state) = &self.compact_state { + state.begin_logical_run(); + } self.prepare_for_run().await?; // IDLE → active marker for the buffered notification / worker-event @@ -4544,14 +4554,25 @@ impl Worker { tracing::warn!(error = %error, "run-committed background task start failed"); } + let request_block = self + .compact_state + .as_ref() + .and_then(|state| state.take_pending_request_block()); + if let Some(block) = request_block { + if let Some(state) = &self.compact_state { + state.finish_logical_run(); + } + return Err(automatic_compact_block_error(block)); + } + if matches!(result, EngineRunExit::Yielded) { self.last_run_interrupted = true; return self.do_compact_and_resume().await; } - if !matches!(result, EngineRunExit::Interrupted(_)) { - if let Some(ref state) = self.compact_state { - state.set_just_compacted(false); + if !matches!(result, EngineRunExit::Paused) { + if let Some(state) = &self.compact_state { + state.finish_logical_run(); } } @@ -4600,17 +4621,16 @@ impl Worker { St: Clone + 'static, { Box::pin(async move { - // Thrash detection: if we just compacted and hit the threshold again, - // something is wrong. - if let Some(ref state) = self.compact_state { - if state.just_compacted() { - state.set_just_compacted(false); - return Err(WorkerError::CompactThrash); + let state = self.compact_state.clone(); + if let Some(state) = &state { + if !state.has_claimed_attempt() { + return Err(WorkerError::AutomaticCompactState( + "automatic compaction yield had no claimed attempt".to_string(), + )); } } - let retained = self - .compact_state + let retained = state .as_ref() .map(|s| s.retained_tokens()) .unwrap_or(manifest::defaults::COMPACT_RETAINED_TOKENS); @@ -4624,8 +4644,12 @@ impl Worker { new_segment_id = %new_segment_id, "Compaction succeeded, resuming execution" ); - if let Some(ref state) = self.compact_state { - state.record_compact_success(); + if let Some(state) = &state { + let completed = state.complete_automatic(CompactionOutcome::Succeeded); + debug_assert!( + completed, + "automatic compaction must complete a claimed attempt" + ); } self.resume().await } @@ -4636,8 +4660,17 @@ impl Worker { AlertSource::Compactor, format!("mid-run compaction failed: {e}"), ); - if let Some(ref state) = self.compact_state { - state.record_compact_failure(); + if let Some(state) = &state { + let outcome = if matches!(e, WorkerError::CompactCancelled) { + CompactionOutcome::Cancelled + } else { + CompactionOutcome::Failed(compact_failure_category(&e)) + }; + let completed = state.complete_automatic(outcome); + debug_assert!( + completed, + "automatic compaction must complete a claimed attempt" + ); } Err(e) } @@ -4653,13 +4686,16 @@ impl Worker { /// Best-effort: failures are logged and surfaced, but do not abort the /// user turn that triggered the check. pub async fn try_pre_run_compact(&mut self) { - let state = match self.compact_state.as_ref() { - Some(s) if !s.is_disabled() && !s.just_compacted() => s.clone(), - _ => return, + let Some(state) = self.compact_state.clone() else { + return; }; let current_tokens = self.total_tokens().tokens; - if !state.exceeds_post_run(current_tokens) { - return; + match state.evaluate_pre_run(current_tokens) { + AutomaticCompactDecision::Start(AutomaticCompactTrigger::PreRun) => {} + AutomaticCompactDecision::Continue | AutomaticCompactDecision::Block(_) => return, + AutomaticCompactDecision::Start(AutomaticCompactTrigger::RequestThreshold) => { + unreachable!("pre-run evaluation returned request-threshold trigger") + } } let retained = state.retained_tokens(); @@ -4672,7 +4708,11 @@ impl Worker { new_segment_id = %new_segment_id, "Proactive pre-run compaction succeeded" ); - state.record_compact_success(); + let completed = state.complete_automatic(CompactionOutcome::Succeeded); + debug_assert!( + completed, + "automatic compaction must complete a claimed attempt" + ); } Err(e) => { warn!(error = %e, "Proactive pre-run compaction failed"); @@ -4681,7 +4721,16 @@ impl Worker { AlertSource::Compactor, format!("pre-run compaction failed: {e}"), ); - state.record_compact_failure(); + let outcome = if matches!(e, WorkerError::CompactCancelled) { + CompactionOutcome::Cancelled + } else { + CompactionOutcome::Failed(compact_failure_category(&e)) + }; + let completed = state.complete_automatic(outcome); + debug_assert!( + completed, + "automatic compaction must complete a claimed attempt" + ); } } } @@ -4766,12 +4815,6 @@ impl Worker { self.ensure_segment_head().await?; let state = self.compact_state.clone(); - if state.as_ref().is_some_and(|s| s.is_disabled()) { - let message = - "manual compact is disabled after repeated compaction failures".to_string(); - self.alert(AlertLevel::Warn, AlertSource::Compactor, message.clone()); - return Ok(ManualCompactResult::Skipped { message }); - } let retained = state .as_ref() @@ -4795,7 +4838,7 @@ impl Worker { Ok(new_segment_id) => { info!(new_segment_id = %new_segment_id, "Manual compaction succeeded"); if let Some(ref state) = state { - state.record_compact_success(); + state.reenable_automatic(); } Ok(ManualCompactResult::Compacted { new_segment_id }) } @@ -4806,9 +4849,6 @@ impl Worker { AlertSource::Compactor, format!("manual compaction failed: {e}"), ); - if let Some(ref state) = state { - state.record_compact_failure(); - } Err(e) } } @@ -4895,6 +4935,9 @@ impl Worker { record, post_requests, } = recorded; + let committed_post_compact_request = post_requests.iter().any(|link| { + link.metric == crate::compact::usage_tracker::PostRequestMetric::Compaction + }); self.commit_entry(LogEntry::LlmUsage { ts: segment_log::now_millis(), history_len: record.history_len, @@ -4903,6 +4946,11 @@ impl Worker { cache_write_tokens: record.cache_write_tokens, output_tokens: record.output_tokens, })?; + if committed_post_compact_request { + if let Some(state) = &self.compact_state { + state.post_compact_request_committed(); + } + } for link in post_requests { let metric = correlated_post_request_metric(link.metric, &link.correlation_id, &record); @@ -7121,6 +7169,18 @@ fn restored_flow_runtime_state( .transpose() } +fn automatic_compact_block_error(block: AutomaticCompactBlock) -> WorkerError { + match block { + AutomaticCompactBlock::Thrash => WorkerError::CompactThrash, + AutomaticCompactBlock::Failed(category) | AutomaticCompactBlock::Disabled(category) => { + WorkerError::AutomaticCompactFailed { + category: category.as_str(), + } + } + AutomaticCompactBlock::Cancelled => WorkerError::CompactCancelled, + } +} + fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory { match error { WorkerError::CompactCancelled => CompactFailureCategory::Cancelled, @@ -7202,6 +7262,12 @@ pub enum WorkerError { #[error("compaction thrash: context still exceeds threshold immediately after compact")] CompactThrash, + #[error("automatic compaction failed and the provider request remained unsafe: {category}")] + AutomaticCompactFailed { category: &'static str }, + + #[error("invalid automatic compaction state: {0}")] + AutomaticCompactState(String), + #[error("compact worker did not produce a summary (write_summary was never called)")] CompactSummaryMissing, diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 673f4184..555a9044 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -23,7 +23,7 @@ use session_store::{ }; use tokio::sync::broadcast; -use worker::{Worker, WorkerController}; +use worker::{Worker, WorkerController, WorkerError}; type TestStore = CombinedStore; @@ -992,12 +992,16 @@ async fn request_threshold_compact_publishes_runtime_progress() { // [2] compact worker closes (its final "done" response). // [3] resume() after compact makes one more LLM call. let client = MockClient::new(vec![ - text_events_with_usage("a", 1000), + text_events_with_usage("a", 100_000), write_summary_tool_use_events("call-1", "summary"), single_text_events("done"), text_events_with_usage("b", 50), ]); - let mut worker = make_worker_with_manifest(MID_TURN_MANIFEST_TOML, client).await; + let manifest = MID_TURN_MANIFEST_TOML.replace( + "compact_request_threshold = 100", + "compact_request_threshold = 50000", + ); + let mut worker = make_worker_with_manifest(&manifest, client).await; let (tx, mut rx) = broadcast::channel::(64); worker.attach_working_event_tx(tx); @@ -1040,6 +1044,31 @@ async fn request_threshold_compact_publishes_runtime_progress() { assert_eq!(post.metric.correlation_id.as_deref(), Some(correlation_id)); } +#[tokio::test] +async fn compacted_context_above_request_threshold_fails_before_provider_request() { + let client = MockClient::new(vec![ + text_events_with_usage("seed", 1000), + write_summary_tool_use_events("call-1", "still too large after compaction"), + single_text_events("done"), + single_text_events("must not be requested"), + ]); + let call_count = Arc::clone(&client.call_count); + let mut worker = make_worker_with_manifest(MID_TURN_MANIFEST_TOML, client).await; + worker.run_text("first").await.unwrap(); + + let error = worker + .run_text("second") + .await + .expect_err("unsafe compacted context must fail closed"); + + assert!(matches!(error, WorkerError::CompactThrash)); + assert_eq!( + call_count.load(Ordering::SeqCst), + 3, + "the provider must receive only the seed and compaction requests" + ); +} + #[tokio::test] async fn pre_run_compact_failure_clears_runtime_progress() { // Only the first run has a response. Compaction will run the @@ -1212,7 +1241,7 @@ async fn controller_compact_method_publishes_progress_and_clear() { single_text_events("done"), single_text_events("follow-up"), ]); - let worker = make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await; + let worker = make_worker_with_manifest(MANUAL_ONLY_MANIFEST_TOML, client).await; let runtime_tmp = tempfile::tempdir().unwrap(); let bash_output_dir = runtime_tmp.path().join("bash-output"); let (handle, _shutdown) = WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir)