diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 90a6f36ad5..29d45451b6 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -46,6 +46,8 @@ use codex_protocol::protocol::PatchApplyBeginEvent; use codex_protocol::protocol::PatchApplyEndEvent; use codex_protocol::protocol::ReviewOutputEvent; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadRolledBackEvent; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; @@ -227,6 +229,7 @@ impl ThreadHistoryChangeAccumulator { pub struct ThreadHistoryBuilder { turns: Vec, current_turn: Option, + is_guardian_session: bool, next_item_index: i64, current_rollout_index: usize, next_rollout_index: usize, @@ -244,6 +247,7 @@ impl ThreadHistoryBuilder { Self { turns: Vec::new(), current_turn: None, + is_guardian_session: false, next_item_index: 1, current_rollout_index: 0, next_rollout_index: 0, @@ -382,11 +386,17 @@ impl ThreadHistoryBuilder { self.next_rollout_index += 1; match item { RolloutItem::EventMsg(event) => self.handle_event(event), + RolloutItem::Compacted(_) + if self.is_guardian_session && self.current_turn.is_none() => {} RolloutItem::Compacted(payload) => self.handle_compacted(payload), RolloutItem::ResponseItem(item) => self.handle_response_item(item), - RolloutItem::InterAgentCommunication(_) - | RolloutItem::TurnContext(_) - | RolloutItem::SessionMeta(_) => {} + RolloutItem::SessionMeta(meta) => { + self.is_guardian_session = matches!( + &meta.meta.source, + SessionSource::SubAgent(SubAgentSource::Other(name)) if name == "guardian" + ); + } + RolloutItem::InterAgentCommunication(_) | RolloutItem::TurnContext(_) => {} } } @@ -1572,6 +1582,8 @@ mod tests { use codex_protocol::protocol::McpInvocation; use codex_protocol::protocol::McpToolCallEndEvent; use codex_protocol::protocol::PatchApplyBeginEvent; + use codex_protocol::protocol::SessionMeta; + use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::ThreadRolledBackEvent; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; @@ -3341,6 +3353,26 @@ mod tests { ); } + #[test] + fn ignores_guardian_maintenance_compaction_outside_a_turn() { + let items = vec![ + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + source: SessionSource::SubAgent(SubAgentSource::Other("guardian".to_string())), + ..Default::default() + }, + git: None, + }), + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: None, + }), + ]; + + assert_eq!(build_turns_from_rollout_items(&items[1..]).len(), 1); + assert_eq!(build_turns_from_rollout_items(&items), Vec::new()); + } + #[test] fn reconstructs_collab_resume_end_item() { let events = vec![ diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index d0167612c5..c0b1e37c60 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -30,6 +30,7 @@ use codex_protocol::protocol::TokenUsage; use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::Semaphore; +use tokio::sync::SemaphorePermit; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -60,7 +61,9 @@ use super::prompt::guardian_policy_prompt_with_config; mod eager_compaction; use eager_compaction::GuardianEagerCompaction; +use eager_compaction::GuardianEagerCompactionOutcome; +const GUARDIAN_EAGER_COMPACTION_WAIT_TIMEOUT: Duration = Duration::from_secs(5); const GUARDIAN_INTERRUPT_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); #[derive(Debug)] pub(crate) enum GuardianReviewSessionOutcome { @@ -93,17 +96,34 @@ pub(crate) struct GuardianReviewSessionParams { pub(crate) deadline: tokio::time::Instant, } -#[derive(Default)] pub(crate) struct GuardianReviewSessionManager { state: Arc>, + eager_compaction_wait_timeout: Duration, } #[derive(Default)] struct GuardianReviewSessionState { trunk: Option>, + trunk_review_generation: u64, ephemeral_reviews: Vec>, } +impl GuardianReviewSessionState { + fn try_claim_trunk<'a>( + &mut self, + trunk: &'a Arc, + review_generation: u64, + ) -> Option<(SemaphorePermit<'a>, u64)> { + let current_trunk = self.trunk.as_ref()?; + if self.trunk_review_generation != review_generation || !Arc::ptr_eq(current_trunk, trunk) { + return None; + } + let review_guard = trunk.review_lock.try_acquire().ok()?; + self.trunk_review_generation = self.trunk_review_generation.wrapping_add(1); + Some((review_guard, self.trunk_review_generation)) + } +} + struct GuardianReviewSession { codex: Codex, cancel_token: CancellationToken, @@ -235,24 +255,25 @@ impl GuardianReviewSession { self.state.lock().await.last_committed_fork_snapshot.clone() } - async fn refresh_last_committed_fork_snapshot(&self) { - match load_rollout_items_for_fork(&self.codex.session).await { - Ok(Some(items)) if !items.is_empty() => { - let mut state = self.state.lock().await; - let prior_review_count = state.prior_review_count; - let last_reviewed_transcript_cursor = state.last_reviewed_transcript_cursor; - state.last_committed_fork_snapshot = Some(GuardianReviewForkSnapshot { - initial_history: InitialHistory::Forked(items), - prior_review_count, - last_reviewed_transcript_cursor, - }); - } - Ok(Some(_)) => {} - Ok(None) => {} + async fn refresh_last_committed_fork_snapshot(&self) -> Option { + let items = match load_rollout_items_for_fork(&self.codex.session).await { + Ok(Some(items)) if !items.is_empty() => items, + Ok(_) => return None, Err(err) => { warn!("failed to refresh guardian trunk rollout snapshot: {err}"); + return None; } - } + }; + let mut state = self.state.lock().await; + let prior_review_count = state.prior_review_count; + let last_reviewed_transcript_cursor = state.last_reviewed_transcript_cursor; + let snapshot = GuardianReviewForkSnapshot { + initial_history: InitialHistory::Forked(items), + prior_review_count, + last_reviewed_transcript_cursor, + }; + state.last_committed_fork_snapshot = Some(snapshot.clone()); + Some(snapshot) } } @@ -294,7 +315,20 @@ impl Drop for EphemeralReviewCleanup { } } +impl Default for GuardianReviewSessionManager { + fn default() -> Self { + Self::new(GUARDIAN_EAGER_COMPACTION_WAIT_TIMEOUT) + } +} + impl GuardianReviewSessionManager { + pub(crate) fn new(eager_compaction_wait_timeout: Duration) -> Self { + Self { + state: Arc::new(Mutex::new(GuardianReviewSessionState::default())), + eager_compaction_wait_timeout, + } + } + pub(crate) async fn trunk_rollout_path(&self) -> Option { let trunk = self.state.lock().await.trunk.clone()?; trunk.codex.session.ensure_rollout_materialized().await; @@ -384,7 +418,10 @@ impl GuardianReviewSessionManager { spawned_trunk = true; } - state.trunk.as_ref().cloned() + state + .trunk + .as_ref() + .map(|trunk| (Arc::clone(trunk), state.trunk_review_generation)) } Err(outcome) => { return (outcome, GuardianReviewAnalyticsResult::without_session()); @@ -395,7 +432,7 @@ impl GuardianReviewSessionManager { review_session.shutdown_in_background(); } - let Some(trunk) = trunk_candidate else { + let Some((trunk, trunk_review_generation)) = trunk_candidate else { return ( GuardianReviewSessionOutcome::Completed(Err(anyhow!( "guardian review session was not available after spawn" @@ -414,28 +451,67 @@ impl GuardianReviewSessionManager { .await; } - if let Err(outcome) = run_before_review_deadline( + let eager_compaction_guard = match run_before_review_deadline( deadline, params.external_cancel.as_ref(), - trunk.wait_for_eager_compaction(), + tokio::time::timeout( + self.eager_compaction_wait_timeout, + trunk.wait_for_eager_compaction(), + ), ) .await { - return (outcome, GuardianReviewAnalyticsResult::without_session()); - } - - let trunk_guard = match trunk.review_lock.try_acquire() { - Ok(trunk_guard) => trunk_guard, - Err(_) => { + Ok(Ok(guard)) => Some(guard), + Ok(Err(_)) => None, + Err(outcome) => { + return (outcome, GuardianReviewAnalyticsResult::without_session()); + } + }; + let eager_compaction_guard = match eager_compaction_guard { + Some(guard) if *guard == GuardianEagerCompactionOutcome::Reusable => guard, + guard => { + let review_session = self + .remove_trunk_if_current(&trunk, trunk_review_generation) + .await; + drop(guard); + if let Some(review_session) = review_session { + review_session.shutdown_in_background(); + } return Box::pin(self.run_ephemeral_review( params, next_reuse_key, deadline, - trunk.fork_snapshot().await, + /*fork_snapshot*/ None, )) .await; } }; + // Hold the maintenance latch until review ownership is decided. + let mut state = self.state.lock().await; + let trunk_pointer_is_current = state + .trunk + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &trunk)); + let Some((trunk_guard, claimed_review_generation)) = + state.try_claim_trunk(&trunk, trunk_review_generation) + else { + drop(state); + drop(eager_compaction_guard); + let fork_snapshot = if trunk_pointer_is_current { + trunk.fork_snapshot().await + } else { + None + }; + return Box::pin(self.run_ephemeral_review( + params, + next_reuse_key, + deadline, + fork_snapshot, + )) + .await; + }; + drop(state); + drop(eager_compaction_guard); let guardian_session_kind = if spawned_trunk { GuardianReviewSessionKind::TrunkNew @@ -449,16 +525,22 @@ impl GuardianReviewSessionManager { deadline, )) .await; - if keep_review_session && matches!(outcome, GuardianReviewSessionOutcome::Completed(_)) { - trunk.refresh_last_committed_fork_snapshot().await; - trunk.schedule_eager_compaction().await; + if keep_review_session + && matches!(outcome, GuardianReviewSessionOutcome::Completed(_)) + && trunk.refresh_last_committed_fork_snapshot().await.is_some() + { + trunk.schedule_eager_compaction(params).await; } - drop(trunk_guard); if keep_review_session { + drop(trunk_guard); (outcome, analytics_result) } else { - if let Some(review_session) = self.remove_trunk_if_current(&trunk).await { + let review_session = self + .remove_trunk_if_current(&trunk, claimed_review_generation) + .await; + drop(trunk_guard); + if let Some(review_session) = review_session { review_session.shutdown_in_background(); } (outcome, analytics_result) @@ -535,12 +617,14 @@ impl GuardianReviewSessionManager { async fn remove_trunk_if_current( &self, trunk: &Arc, + review_generation: u64, ) -> Option> { let mut state = self.state.lock().await; if state .trunk .as_ref() .is_some_and(|current| Arc::ptr_eq(current, trunk)) + && state.trunk_review_generation == review_generation { state.trunk.take() } else { @@ -1147,6 +1231,29 @@ mod tests { ) } + #[tokio::test] + async fn stale_generation_cannot_remove_idle_trunk() { + let (review_session, _, _) = test_review_session().await; + let trunk = Arc::new(review_session); + let manager = GuardianReviewSessionManager::default(); + let review_guard = { + let mut state = manager.state.lock().await; + state.trunk = Some(Arc::clone(&trunk)); + state + .try_claim_trunk(&trunk, /*review_generation*/ 0) + .expect("current trunk should be claimable") + .0 + }; + assert!( + manager + .remove_trunk_if_current(&trunk, /*review_generation*/ 0) + .await + .is_none() + ); + assert!(manager.state.lock().await.trunk.is_some()); + drop(review_guard); + } + fn turn_complete_event( turn_id: &str, last_agent_message: Option<&str>, diff --git a/codex-rs/core/src/guardian/review_session/eager_compaction.rs b/codex-rs/core/src/guardian/review_session/eager_compaction.rs index eac8f3939b..a86903a685 100644 --- a/codex-rs/core/src/guardian/review_session/eager_compaction.rs +++ b/codex-rs/core/src/guardian/review_session/eager_compaction.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use codex_protocol::protocol::CompactedItem; use tokio::sync::Mutex; use tokio::sync::OwnedMutexGuard; use tracing::warn; @@ -8,87 +9,176 @@ use crate::guardian::GUARDIAN_REVIEW_TIMEOUT; use crate::session::turn; use super::GuardianReviewSession; +use super::GuardianReviewSessionParams; +use super::InitialHistory; +use super::RolloutItem; +use super::spawn_guardian_review_session; -/// A held guard means eager maintenance is in flight. -#[derive(Default)] -pub(super) struct GuardianEagerCompaction { - in_flight: Arc>, +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(super) enum GuardianEagerCompactionOutcome { + #[default] + Reusable, + DiscardSession, } -impl GuardianEagerCompaction { - fn begin(&self) -> Option> { - Arc::clone(&self.in_flight).try_lock_owned().ok() - } - - async fn wait(&self) { - drop(Arc::clone(&self.in_flight).lock_owned().await); - } -} +/// The held guard serializes review ownership with eager maintenance. +pub(super) type GuardianEagerCompaction = Arc>; impl GuardianReviewSession { - pub(super) async fn schedule_eager_compaction(self: &Arc) { + pub(super) async fn schedule_eager_compaction( + self: &Arc, + params: GuardianReviewSessionParams, + ) { let turn_context = self.codex.session.new_default_turn().await; if !turn::auto_compact_needed(self.codex.session.as_ref(), turn_context.as_ref()).await { return; } - let Some(in_flight_guard) = self.eager_compaction.begin() else { - return; - }; + let mut outcome_guard = Arc::clone(&self.eager_compaction).lock_owned().await; + *outcome_guard = GuardianEagerCompactionOutcome::DiscardSession; let review_session = Arc::clone(self); drop(tokio::spawn(async move { - // Keep the latch closed through compaction and snapshot refresh. Any exit path drops it. - let _in_flight_guard = in_flight_guard; - let cancel_token = review_session.cancel_token.clone(); - let timed_out = tokio::select! { - _ = cancel_token.cancelled() => false, - result = tokio::time::timeout( - GUARDIAN_REVIEW_TIMEOUT, - review_session.run_eager_compaction(turn_context), - ) => result.is_err(), - }; - if timed_out { - warn!( - guardian_thread_id = %review_session.codex.session.thread_id, - "eager guardian maintenance timed out after {GUARDIAN_REVIEW_TIMEOUT:?}" - ); + if review_session + .run_eager_compaction(turn_context, params) + .await + { + *outcome_guard = GuardianEagerCompactionOutcome::Reusable; } })); } - pub(super) async fn wait_for_eager_compaction(&self) { - self.eager_compaction.wait().await; + pub(super) async fn wait_for_eager_compaction( + &self, + ) -> OwnedMutexGuard { + Arc::clone(&self.eager_compaction).lock_owned().await } async fn run_eager_compaction( self: &Arc, turn_context: Arc, - ) { - let Ok(review_guard) = self.review_lock.acquire().await else { - return; + params: GuardianReviewSessionParams, + ) -> bool { + let Ok(_review_guard) = self.review_lock.acquire().await else { + return false; }; - - let mut client_session = self.codex.session.services.model_client.new_session(); - let compact_result = turn::run_pre_turn_auto_compact_if_needed( - &self.codex.session, - &turn_context, - &mut client_session, + let Some(fork_snapshot) = self.fork_snapshot().await else { + return false; + }; + let InitialHistory::Forked(items) = &fork_snapshot.initial_history else { + return false; + }; + let rollout_item_count = items.len(); + let shadow_cancel = self.cancel_token.child_token(); + let mut shadow_config = params.spawn_config.clone(); + shadow_config.ephemeral = true; + let shadow = match spawn_guardian_review_session( + ¶ms, + shadow_config, + self.reuse_key.clone(), + shadow_cancel.clone(), + Some(fork_snapshot), ) - .await; - - match compact_result { - Ok(true) => { - self.refresh_last_committed_fork_snapshot().await; - } - Ok(false) => {} + .await + { + Ok(shadow) => Arc::new(shadow), Err(err) => { + warn!( + guardian_thread_id = %self.codex.session.thread_id, + "failed to spawn eager guardian compaction session: {err}" + ); + return false; + } + }; + let shadow_for_events = Arc::clone(&shadow); + let event_drain = + tokio::spawn( + async move { while shadow_for_events.codex.next_event().await.is_ok() {} }, + ); + + let shadow_turn = shadow.codex.session.new_default_turn().await; + let mut client_session = shadow.codex.session.services.model_client.new_session(); + let compact_result = tokio::select! { + _ = self.cancel_token.cancelled() => None, + result = tokio::time::timeout( + GUARDIAN_REVIEW_TIMEOUT, + turn::run_pre_turn_auto_compact( + &shadow.codex.session, + &shadow_turn, + &mut client_session, + ), + ) => match result { + Ok(result) => Some(result), + Err(_) => { + warn!( + guardian_thread_id = %self.codex.session.thread_id, + "eager guardian maintenance timed out after {GUARDIAN_REVIEW_TIMEOUT:?}" + ); + None + } + }, + }; + let reusable = match compact_result { + Some(Ok(())) => { + let history = shadow + .codex + .session + .clone_history() + .await + .raw_items() + .to_vec(); + let compacted_item = CompactedItem { + message: String::new(), + replacement_history: Some(history.clone()), + }; + self.codex + .session + .replace_compacted_history( + history, + shadow.codex.session.reference_context_item().await, + compacted_item, + ) + .await; + self.codex + .session + .recompute_token_usage(&turn_context) + .await; + let durable = self + .refresh_last_committed_fork_snapshot() + .await + .is_some_and(|snapshot| { + let InitialHistory::Forked(items) = snapshot.initial_history else { + return false; + }; + items.get(rollout_item_count..).is_some_and(|items| { + items.iter().any(|item| { + matches!( + item, + RolloutItem::Compacted(compacted) + if compacted.replacement_history.is_some() + ) + }) + }) + }); + if !durable { + warn!( + guardian_thread_id = %self.codex.session.thread_id, + "eager guardian compaction was not durably published" + ); + } + durable + } + Some(Err(err)) => { warn!( guardian_thread_id = %self.codex.session.thread_id, "eager guardian compaction failed: {err}" ); + false } - } - - drop(review_guard); + None => false, + }; + shadow_cancel.cancel(); + shadow.shutdown().await; + let _ = event_drain.await; + reusable } } diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 710cd74285..5bba919e52 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -10,6 +10,8 @@ use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::test_support; use codex_analytics::GuardianApprovalRequestSource; +use codex_analytics::GuardianReviewAnalyticsResult; +use codex_analytics::GuardianReviewSessionKind; use codex_config::ConfigLayerStack; use codex_config::FeatureRequirementsToml; use codex_config::NetworkConstraints; @@ -31,6 +33,7 @@ use codex_network_proxy::NetworkProxyConfig; use codex_protocol::ThreadId; use codex_protocol::approvals::NetworkApprovalProtocol; use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::models::ContentItem; use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; @@ -2155,307 +2158,162 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: } fn configure_guardian_eager_compaction_test(turn: &mut Arc) { - const AUTO_COMPACT_TOKEN_LIMIT: i64 = 200_000; - const COMPACTION_STREAM_RETRIES: u64 = 0; - const REQUEST_RETRIES: u64 = 0; - let mut config = (*turn.config).clone(); - config.model_auto_compact_token_limit = Some(AUTO_COMPACT_TOKEN_LIMIT); - config.model_provider.request_max_retries = Some(REQUEST_RETRIES); - config.model_provider.stream_max_retries = Some(COMPACTION_STREAM_RETRIES); + config.model_auto_compact_token_limit = Some(200_000); + config.model_auto_compact_token_limit_scope = AutoCompactTokenLimitScope::BodyAfterPrefix; + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(0); config.model_provider.supports_websockets = false; - let _ = config.features.enable(Feature::RemoteCompactionV2); + config.model_provider.name = "Test".to_string(); Arc::get_mut(turn) .expect("guardian test turn should be uniquely owned") .config = Arc::new(config); } -fn request_is_remote_compaction(request: &[u8]) -> anyhow::Result { - let body = serde_json::from_slice::(request)?; - Ok(body["input"].as_array().is_some_and(|input| { - input.iter().any(|item| { - item.get("type").and_then(serde_json::Value::as_str) == Some("compaction_trigger") - }) - })) +const EAGER_COMPACTION_GUARDIAN_ASSESSMENT: &str = r#"{"outcome":"allow"}"#; +const EAGER_COMPACTION_GUARDIAN_PREFILL_TOKENS: i64 = 10_000; +const EAGER_COMPACTION_TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30); + +fn guardian_review_sse( + id: &str, + completed: impl FnOnce(&str) -> serde_json::Value, +) -> Vec { + let response_id = format!("resp-guardian-{id}"); + let message_id = format!("msg-guardian-{id}"); + vec![StreamingSseChunk { + gate: None, + body: sse(vec![ + ev_response_created(&response_id), + ev_assistant_message(&message_id, EAGER_COMPACTION_GUARDIAN_ASSESSMENT), + completed(&response_id), + ]), + }] } -fn eager_compaction_guardian_assessment(rationale: &str) -> String { - serde_json::json!({ - "risk_level": "low", - "user_authorization": "high", - "outcome": "allow", - "rationale": rationale, +fn guardian_review_sse_with_body_tokens(id: &str, body_tokens: i64) -> Vec { + guardian_review_sse(id, |response_id| { + let mut completed = ev_completed_with_tokens( + response_id, + EAGER_COMPACTION_GUARDIAN_PREFILL_TOKENS + body_tokens, + ); + completed["response"]["usage"]["input_tokens"] = + EAGER_COMPACTION_GUARDIAN_PREFILL_TOKENS.into(); + completed["response"]["usage"]["output_tokens"] = body_tokens.into(); + completed }) - .to_string() } -async fn run_initial_eager_compaction_review( +async fn run_eager_compaction_review( session: &Arc, turn: &Arc, -) -> ( - GuardianReviewOutcome, - codex_analytics::GuardianReviewAnalyticsResult, -) { - run_guardian_review_session_for_test( - Arc::clone(session), - Arc::clone(turn), - GuardianApprovalRequest::Shell { - id: "shell-ca-540-1".to_string(), - command: vec!["git".to_string(), "status".to_string()], - cwd: test_path_buf("/repo/codex-rs/core").abs(), - sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, - additional_permissions: None, - justification: Some("Inspect repository state.".to_string()), - }, - /*retry_reason*/ None, - guardian_output_schema(), - /*external_cancel*/ None, + retry_reason: Option<&str>, +) -> anyhow::Result { + let outcome = tokio::time::timeout( + EAGER_COMPACTION_TEST_TIMEOUT, + run_guardian_review_session_for_test( + Arc::clone(session), + Arc::clone(turn), + GuardianApprovalRequest::Shell { + id: "shell-ca-540".to_string(), + command: vec!["git".to_string(), "push".to_string()], + cwd: test_path_buf("/repo/codex-rs/core").abs(), + sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, + additional_permissions: None, + justification: Some("Push the reviewed change.".to_string()), + }, + retry_reason.map(str::to_string), + guardian_output_schema(), + /*external_cancel*/ None, + ), ) - .await -} - -async fn run_followup_eager_compaction_review( - session: &Arc, - turn: &Arc, - retry_reason: &str, -) -> ( - GuardianReviewOutcome, - codex_analytics::GuardianReviewAnalyticsResult, -) { - run_guardian_review_session_for_test( - Arc::clone(session), - Arc::clone(turn), - GuardianApprovalRequest::Shell { - id: "shell-ca-540-2".to_string(), - command: vec!["git".to_string(), "push".to_string()], - cwd: test_path_buf("/repo/codex-rs/core").abs(), - sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, - additional_permissions: None, - justification: Some("Push the reviewed change.".to_string()), - }, - Some(retry_reason.to_string()), - guardian_output_schema(), - /*external_cancel*/ None, - ) - .await + .await?; + let (GuardianReviewOutcome::Completed(_), analytics_result) = outcome else { + anyhow::bail!("expected guardian assessment"); + }; + Ok(analytics_result) } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn guardian_review_waits_until_eager_compaction_snapshot_is_committed() -> anyhow::Result<()> -{ - const REQUEST_WAIT_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); - const BLOCKED_REVIEW_OBSERVATION: Duration = Duration::from_millis(/*millis*/ 100); +async fn guardian_recovers_from_eager_compaction_failure_and_timeout() -> anyhow::Result<()> { + const EAGER_COMPACTION_WAIT_TIMEOUT: Duration = Duration::from_millis(/*millis*/ 50); const FIRST_REVIEW_TOTAL_TOKENS: i64 = 500_000; - const COMPACTION_TOTAL_TOKENS: i64 = 50; - - let first_assessment = eager_compaction_guardian_assessment("first guardian rationale"); - let second_assessment = eager_compaction_guardian_assessment("second guardian rationale"); + let (failed_compaction_tx, failed_compaction_rx) = tokio::sync::oneshot::channel(); let (compaction_tx, compaction_rx) = tokio::sync::oneshot::channel(); - let (second_review_tx, second_review_rx) = tokio::sync::oneshot::channel(); let (server, _) = start_streaming_sse_server(vec![ + guardian_review_sse_with_body_tokens( + "1", + FIRST_REVIEW_TOTAL_TOKENS - EAGER_COMPACTION_GUARDIAN_PREFILL_TOKENS, + ), vec![StreamingSseChunk { - gate: None, + gate: Some(failed_compaction_rx), body: sse(vec![ - ev_response_created("resp-guardian-1"), - ev_assistant_message("msg-guardian-1", &first_assessment), - ev_completed_with_tokens("resp-guardian-1", FIRST_REVIEW_TOTAL_TOKENS), + ev_response_created("resp-eager-compact-failed"), + ev_assistant_message("msg-eager-compact-partial", "partial summary"), ]), }], + guardian_review_sse("2", ev_completed), + guardian_review_sse_with_body_tokens( + "3", + FIRST_REVIEW_TOTAL_TOKENS - EAGER_COMPACTION_GUARDIAN_PREFILL_TOKENS, + ), vec![StreamingSseChunk { gate: Some(compaction_rx), - body: sse(vec![ - ev_response_created("resp-compact-success"), - serde_json::json!({ - "type": "response.output_item.done", - "item": { - "type": "compaction", - "encrypted_content": "CA540_COMPACTED_CONTEXT", - } - }), - ev_completed_with_tokens("resp-compact-success", COMPACTION_TOTAL_TOKENS), - ]), - }], - vec![StreamingSseChunk { - gate: Some(second_review_rx), - body: sse(vec![ - ev_response_created("resp-guardian-2"), - ev_assistant_message("msg-guardian-2", &second_assessment), - ev_completed("resp-guardian-2"), - ]), + body: sse(vec![ev_response_created("resp-eager-compact-timeout")]), }], + guardian_review_sse("4", ev_completed), ]) .await; - - let (session, mut turn) = guardian_test_session_and_turn_with_base_url(server.uri()).await; + let (mut session, mut turn) = guardian_test_session_and_turn_with_base_url(server.uri()).await; + Arc::get_mut(&mut session) + .expect("guardian parent session should be uniquely owned") + .guardian_review_session = GuardianReviewSessionManager::new(EAGER_COMPACTION_WAIT_TIMEOUT); configure_guardian_eager_compaction_test(&mut turn); seed_guardian_parent_history(&session, &turn).await; - let first_outcome = tokio::time::timeout( - REQUEST_WAIT_TIMEOUT, - run_initial_eager_compaction_review(&session, &turn), - ) - .await?; - let (GuardianReviewOutcome::Completed(_), _) = first_outcome else { - panic!("expected first guardian assessment"); - }; - tokio::time::timeout( - REQUEST_WAIT_TIMEOUT, - server.wait_for_request_count(/*count*/ 2), - ) - .await?; - - let session_for_second = Arc::clone(&session); - let turn_for_second = Arc::clone(&turn); - let second_review = tokio::spawn(async move { - run_followup_eager_compaction_review( - &session_for_second, - &turn_for_second, - "Follow-up approval on the reused guardian thread.", - ) - .await - }); - - assert!( - tokio::time::timeout( - BLOCKED_REVIEW_OBSERVATION, - server.wait_for_request_count(/*count*/ 3), - ) - .await - .is_err(), - "the next guardian request must not start while eager compaction is in flight" - ); - assert!(!second_review.is_finished()); - - compaction_tx - .send(()) - .expect("eager compaction response gate should still be open"); - tokio::time::timeout( - REQUEST_WAIT_TIMEOUT, - server.wait_for_request_count(/*count*/ 3), - ) - .await?; - - let committed_rollout_items = session + run_eager_compaction_review(&session, &turn, /*retry_reason*/ None).await?; + let discarded_rollout_path = session .guardian_review_session - .committed_fork_rollout_items_for_test() + .trunk_rollout_path() .await - .expect("committed guardian fork snapshot"); - assert!( - committed_rollout_items.iter().any(|item| matches!( - item, - RolloutItem::Compacted(compacted) if compacted.replacement_history.is_some() - )), - "eager compaction must refresh the fork snapshot before releasing the next review" - ); - assert!(!second_review.is_finished()); - - second_review_tx + .expect("discarded guardian rollout path"); + let committed_rollout = tokio::fs::read(&discarded_rollout_path).await?; + failed_compaction_tx .send(()) - .expect("second guardian review gate should still be open"); - let second_outcome = tokio::time::timeout(REQUEST_WAIT_TIMEOUT, second_review).await??; - let (GuardianReviewOutcome::Completed(_), second_metadata) = second_outcome else { - panic!("expected second guardian assessment"); - }; - assert!(matches!( - second_metadata.guardian_session_kind, - Some(codex_analytics::GuardianReviewSessionKind::TrunkReused) - )); - - let requests = server.requests().await; - assert_eq!(requests.len(), 3); - assert!(request_is_remote_compaction(&requests[1])?); - assert!(!request_is_remote_compaction(&requests[2])?); - - server.shutdown().await; - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn guardian_eager_compaction_failure_falls_back_to_pre_turn_compaction() -> anyhow::Result<()> -{ - const REQUEST_WAIT_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); - const FIRST_REVIEW_TOTAL_TOKENS: i64 = 500_000; - const COMPACTION_TOTAL_TOKENS: i64 = 50; - - let first_assessment = eager_compaction_guardian_assessment("first guardian rationale"); - let second_assessment = eager_compaction_guardian_assessment("second guardian rationale"); - let (server, _) = start_streaming_sse_server(vec![ - vec![StreamingSseChunk { - gate: None, - body: sse(vec![ - ev_response_created("resp-guardian-1"), - ev_assistant_message("msg-guardian-1", &first_assessment), - ev_completed_with_tokens("resp-guardian-1", FIRST_REVIEW_TOTAL_TOKENS), - ]), - }], - vec![StreamingSseChunk { - gate: None, - body: sse(vec![ev_response_created("resp-eager-compact-failed")]), - }], - vec![StreamingSseChunk { - gate: None, - body: sse(vec![ - ev_response_created("resp-pre-turn-compact-success"), - serde_json::json!({ - "type": "response.output_item.done", - "item": { - "type": "compaction", - "encrypted_content": "CA540_RETRY_COMPACTED_CONTEXT", - } - }), - ev_completed_with_tokens("resp-pre-turn-compact-success", COMPACTION_TOTAL_TOKENS), - ]), - }], - vec![StreamingSseChunk { - gate: None, - body: sse(vec![ - ev_response_created("resp-guardian-2"), - ev_assistant_message("msg-guardian-2", &second_assessment), - ev_completed("resp-guardian-2"), - ]), - }], - ]) - .await; - - let (session, mut turn) = guardian_test_session_and_turn_with_base_url(server.uri()).await; - configure_guardian_eager_compaction_test(&mut turn); - seed_guardian_parent_history(&session, &turn).await; - - let first_outcome = tokio::time::timeout( - REQUEST_WAIT_TIMEOUT, - run_initial_eager_compaction_review(&session, &turn), - ) - .await?; - let (GuardianReviewOutcome::Completed(_), _) = first_outcome else { - panic!("expected first guardian assessment"); - }; - tokio::time::timeout( - REQUEST_WAIT_TIMEOUT, - server.wait_for_request_count(/*count*/ 2), - ) - .await?; - - let second_outcome = run_followup_eager_compaction_review( + .expect("failed compaction response gate should still be open"); + let second_metadata = run_eager_compaction_review( &session, &turn, - "Retry after eager compaction failure.", + Some("Retry after eager compaction failure."), ) - .await; - let (GuardianReviewOutcome::Completed(_), second_metadata) = second_outcome else { - panic!("expected second guardian assessment"); - }; + .await?; + assert_eq!( + tokio::fs::read(&discarded_rollout_path).await?, + committed_rollout + ); + run_eager_compaction_review(&session, &turn, Some("Start a replacement guardian trunk.")) + .await?; + tokio::time::timeout( + EAGER_COMPACTION_TEST_TIMEOUT, + server.wait_for_request_count(/*count*/ 5), + ) + .await?; + let fourth_metadata = run_eager_compaction_review( + &session, + &turn, + Some("Retry after eager compaction timeout."), + ) + .await?; assert!(matches!( - second_metadata.guardian_session_kind, - Some(codex_analytics::GuardianReviewSessionKind::TrunkReused) + ( + second_metadata.guardian_session_kind, + fourth_metadata.guardian_session_kind + ), + ( + Some(GuardianReviewSessionKind::EphemeralForked), + Some(GuardianReviewSessionKind::EphemeralForked) + ) )); - - let requests = server.requests().await; - assert_eq!(requests.len(), 4); - assert!(request_is_remote_compaction(&requests[1])?); - assert!(request_is_remote_compaction(&requests[2])?); - assert!(!request_is_remote_compaction(&requests[3])?); - - server.shutdown().await; + let _ = compaction_tx.send(()); Ok(()) } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index f61f7bf0fe..c806fbe2b0 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -802,7 +802,9 @@ async fn run_pre_sampling_compact( client_session: &mut ModelClientSession, ) -> CodexResult<()> { maybe_run_previous_model_inline_compact(sess, turn_context, client_session).await?; - run_pre_turn_auto_compact_if_needed(sess, turn_context, client_session).await?; + if auto_compact_needed(sess.as_ref(), turn_context.as_ref()).await { + run_pre_turn_auto_compact(sess, turn_context, client_session).await?; + } Ok(()) } @@ -812,15 +814,11 @@ pub(crate) async fn auto_compact_needed(sess: &Session, turn_context: &TurnConte .token_limit_reached } -pub(crate) async fn run_pre_turn_auto_compact_if_needed( +pub(crate) async fn run_pre_turn_auto_compact( sess: &Arc, turn_context: &Arc, client_session: &mut ModelClientSession, -) -> CodexResult { - if !auto_compact_needed(sess.as_ref(), turn_context.as_ref()).await { - return Ok(false); - } - +) -> CodexResult<()> { run_auto_compact( sess, turn_context, @@ -829,8 +827,7 @@ pub(crate) async fn run_pre_turn_auto_compact_if_needed( CompactionReason::ContextLimit, CompactionPhase::PreTurn, ) - .await?; - Ok(true) + .await } /// Returns true only when both turns declare compaction compatibility hashes and they differ. diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index eaa67c5079..ab7bf9f2e9 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -3,10 +3,11 @@ use anyhow::Result; use codex_core::config::Constrained; use codex_core::sandboxing::SandboxPermissions; -use codex_features::Feature; use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::GuardianAssessmentStatus; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::user_input::UserInput; @@ -29,19 +30,24 @@ use serde_json::Value; use serde_json::json; use std::fs; use std::os::unix::fs::PermissionsExt; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use tempfile::TempDir; +use wiremock::Mock; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path_regex; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<()> { skip_if_no_network!(Ok(())); skip_if_sandbox!(Ok(())); - const FIRST_REVIEW_TOTAL_TOKENS: i64 = 500_000; - const COMPACTION_TOTAL_TOKENS: i64 = 50; + const FIRST_REVIEW_TOTAL_TOKENS: i64 = 220_000; const AUTO_COMPACT_TOKEN_LIMIT: i64 = 200_000; - const REQUEST_RETRIES: u64 = 0; - const STREAM_RETRIES: u64 = 0; let server = start_mock_server().await; let approval_policy = AskForApproval::OnRequest; @@ -52,29 +58,31 @@ async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<( exclude_slash_tmp: true, }; let sandbox_policy_for_config = sandbox_policy.clone(); - let mut builder = test_codex().with_config(move |config| { - config.permissions.approval_policy = Constrained::allow_any(approval_policy); - config - .set_legacy_sandbox_policy(sandbox_policy_for_config) - .expect("set sandbox policy"); - config.model_auto_compact_token_limit = Some(AUTO_COMPACT_TOKEN_LIMIT); - config.model_provider.request_max_retries = Some(REQUEST_RETRIES); - config.model_provider.stream_max_retries = Some(STREAM_RETRIES); - config.model_provider.supports_websockets = false; - config - .features - .enable(Feature::RemoteCompactionV2) - .expect("enable remote compaction v2"); - }); - let test = builder.build(&server).await?; - - let first_justification = "Approve the first eager-compaction command."; + let test = test_codex() + .with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(approval_policy); + config + .set_legacy_sandbox_policy(sandbox_policy_for_config) + .expect("set sandbox policy"); + config.model_auto_compact_token_limit = Some(AUTO_COMPACT_TOKEN_LIMIT); + config.model_auto_compact_token_limit_scope = + AutoCompactTokenLimitScope::BodyAfterPrefix; + config.model_provider.request_max_retries = Some(u64::MAX); + config.model_provider.stream_max_retries = Some(u64::default()); + config.model_provider.supports_websockets = false; + }) + .build(&server) + .await?; + let first_tool_gate = test.cwd.path().join("release-first-tool"); let second_justification = "Approve the second eager-compaction command."; let first_args = json!({ - "cmd": "sleep 0.2; printf first", + "cmd": format!( + "while [ ! -f {} ]; do sleep 0.01; done; printf first", + shlex::try_join([first_tool_gate.to_string_lossy().as_ref()])? + ), "yield_time_ms": 1_000_u64, "sandbox_permissions": SandboxPermissions::RequireEscalated, - "justification": first_justification, + "justification": "Approve the first eager-compaction command.", }); let second_args = json!({ "cmd": "printf second", @@ -82,13 +90,10 @@ async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<( "sandbox_permissions": SandboxPermissions::RequireEscalated, "justification": second_justification, }); - let guardian_assessment = json!({ - "risk_level": "low", - "user_authorization": "high", - "outcome": "allow", - "rationale": "The command writes a bounded marker file requested by the user.", - }) - .to_string(); + let mut first_guardian_completed = + ev_completed_with_tokens("resp-guardian-first", FIRST_REVIEW_TOTAL_TOKENS); + first_guardian_completed["response"]["usage"]["input_tokens"] = 10_000.into(); + first_guardian_completed["response"]["usage"]["output_tokens"] = 210_000.into(); let responses = mount_sse_sequence( &server, vec![ @@ -103,19 +108,8 @@ async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<( ]), sse(vec![ ev_response_created("resp-guardian-first"), - ev_assistant_message("msg-guardian-first", &guardian_assessment), - ev_completed_with_tokens("resp-guardian-first", FIRST_REVIEW_TOTAL_TOKENS), - ]), - sse(vec![ - ev_response_created("resp-guardian-compact"), - json!({ - "type": "response.output_item.done", - "item": { - "type": "compaction", - "encrypted_content": "EAGER_GUARDIAN_COMPACTED_CONTEXT", - } - }), - ev_completed_with_tokens("resp-guardian-compact", COMPACTION_TOTAL_TOKENS), + ev_assistant_message("msg-guardian-first", r#"{"outcome":"allow"}"#), + first_guardian_completed, ]), sse(vec![ ev_response_created("resp-parent-second-tool"), @@ -128,7 +122,7 @@ async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<( ]), sse(vec![ ev_response_created("resp-guardian-second"), - ev_assistant_message("msg-guardian-second", &guardian_assessment), + ev_assistant_message("msg-guardian-second", r#"{"outcome":"allow"}"#), ev_completed("resp-guardian-second"), ]), sse(vec![ @@ -139,6 +133,40 @@ async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<( ], ) .await; + let compact_response_released = Arc::new(AtomicBool::new(false)); + let compact_request_count = Arc::new(AtomicUsize::new(0)); + let compact_response_released_for_mock = Arc::clone(&compact_response_released); + let compact_request_count_for_mock = Arc::clone(&compact_request_count); + Mock::given(method("POST")) + .and(path_regex(".*/responses/compact$")) + .respond_with(move |_: &wiremock::Request| { + compact_request_count_for_mock.fetch_add(1, Ordering::AcqRel); + if !compact_response_released_for_mock.load(Ordering::Acquire) { + return ResponseTemplate::new(503); + } + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_json(json!({ + "output": [{ + "type": "compaction", + "encrypted_content": "EAGER_GUARDIAN_COMPACTED_CONTEXT", + }], + })) + }) + .mount(&server) + .await; + let compact_request_count_for_release = Arc::clone(&compact_request_count); + let release_first_tool = tokio::spawn(async move { + tokio::time::timeout(Duration::from_secs(/*secs*/ 30), async { + while compact_request_count_for_release.load(Ordering::Acquire) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .map_err(|_| anyhow::anyhow!("compact request was not observed"))?; + fs::write(first_tool_gate, "release")?; + Ok::<(), anyhow::Error>(()) + }); test.codex .submit(Op::UserInput { @@ -158,41 +186,40 @@ async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<( }, }) .await?; + for _ in 0..2 { + wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::GuardianAssessment(assessment) + if assessment.status == GuardianAssessmentStatus::InProgress + ) + }) + .await; + } + assert!( + responses + .requests() + .iter() + .all(|request| !request.body_contains_text(second_justification)), + "the second Guardian request must wait for eager compaction" + ); + compact_response_released.store(true, Ordering::Release); wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) }) .await; + release_first_tool.await??; - let requests = responses.requests(); - let first_guardian_index = requests - .iter() - .position(|request| request.body_contains_text(first_justification)) - .expect("first Guardian request"); - let second_guardian_index = requests - .iter() - .position(|request| request.body_contains_text(second_justification)) + let second_guardian = responses + .requests() + .into_iter() + .find(|request| request.body_contains_text(second_justification)) .expect("second Guardian request"); - let compaction_indexes = requests - .iter() - .enumerate() - .filter_map(|(index, request)| { - request.body_json()["input"] - .as_array() - .is_some_and(|input| { - input.iter().any(|item| { - item.get("type").and_then(Value::as_str) == Some("compaction_trigger") - }) - }) - .then_some(index) - }) - .collect::>(); - - assert_eq!(compaction_indexes.len(), 1); assert!( - first_guardian_index < compaction_indexes[0] - && compaction_indexes[0] < second_guardian_index, - "expected eager compaction between Guardian reviews, requests: {requests:#?}" + second_guardian.body_contains_text("EAGER_GUARDIAN_COMPACTED_CONTEXT"), + "the second Guardian request must use the installed compacted history" ); + assert!(compact_request_count.load(Ordering::Acquire) >= 2); Ok(()) }