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

fix: guard automatic compaction retries

Browse Source
This commit is contained in:
Hare
2026-09-16 07:16:31 +09:00
parent ffc1933f06
commit a0f8ee51a4
5 changed files with 642 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/compact/state.rs
+334 -139
View File
@@ -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<CompactFailureCategory>,
pending_request_block: Option<AutomaticCompactBlock>,
}
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<u64>,
/// 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<u64>,
/// In-request safety threshold. Checked immediately before each provider
/// request, both before and after pre-request hooks.
request_threshold: Option<u64>,
/// 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<AutomaticCompactRuntimeState>,
}
impl CompactState {
pub(crate) fn new(
post_run_threshold: Option<u64>,
compact_threshold: Option<u64>,
request_threshold: Option<u64>,
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<u64> {
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<AutomaticCompactBlock> {
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);
}
}
crates/worker/src/compact/telemetry.rs
+1 -1
View File
@@ -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",
crates/worker/src/ipc/interceptor.rs
+168 -32
View File
@@ -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<u64>, 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<u64>,
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<PreRequestAction> {
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<SessionHistoryMetadata> for WorkerInterceptor {
) -> InterceptorResult<PreRequestAction> {
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<SessionHistoryMetadata> 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<SessionHistoryMetadata> 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<PreLlmRequest> for YieldingPreRequestHook {
async fn call(
&self,
_info: &PreRequestContext,
) -> Result<HookPreRequestAction, crate::hook::HookError> {
Ok(HookPreRequestAction::Yield)
}
}
struct RecordingSystemItemCommitter {
committed: Arc<Mutex<Vec<SystemItem>>>,
}
@@ -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())),
crates/worker/src/worker.rs
+106 -40
View File
@@ -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<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
// 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
"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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
AlertSource::Compactor,
format!("manual compaction failed: {e}"),
);
if let Some(ref state) = state {
state.record_compact_failure();
}
Err(e)
}
}
@@ -4895,6 +4935,9 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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,
crates/worker/tests/compact_events_test.rs
+33 -4
View File
@@ -23,7 +23,7 @@ use session_store::{
};
use tokio::sync::broadcast;
use worker::{Worker, WorkerController};
use worker::{Worker, WorkerController, WorkerError};
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
@@ -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::<Event>(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)