From 85d075ce6d46d65cacff5bfba9d5aadd0c2cfc7b Mon Sep 17 00:00:00 2001 From: Ray Morgan Date: Mon, 8 Jun 2026 17:04:29 -0700 Subject: [PATCH] Eagerly compact Guardian threads between reviews --- codex-rs/core/src/guardian/review_session.rs | 21 ++ .../review_session/eager_compaction.rs | 154 +++++++++ .../review_session/eager_compaction_tests.rs | 49 +++ codex-rs/core/src/guardian/tests.rs | 291 ++++++++++++++++++ codex-rs/core/src/session/turn.rs | 41 ++- codex-rs/core/tests/suite/guardian_review.rs | 167 ++++++++++ 6 files changed, 710 insertions(+), 13 deletions(-) create mode 100644 codex-rs/core/src/guardian/review_session/eager_compaction.rs create mode 100644 codex-rs/core/src/guardian/review_session/eager_compaction_tests.rs diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 7aa007891f..d0167612c5 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -57,6 +57,10 @@ use super::prompt::build_guardian_prompt_items_with_parent_turn; use super::prompt::guardian_policy_prompt; use super::prompt::guardian_policy_prompt_with_config; +mod eager_compaction; + +use eager_compaction::GuardianEagerCompaction; + const GUARDIAN_INTERRUPT_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); #[derive(Debug)] pub(crate) enum GuardianReviewSessionOutcome { @@ -105,6 +109,7 @@ struct GuardianReviewSession { cancel_token: CancellationToken, reuse_key: GuardianReviewSessionReuseKey, review_lock: Semaphore, + eager_compaction: GuardianEagerCompaction, state: Mutex, } @@ -215,6 +220,7 @@ pub(crate) fn prompt_cache_key_override_for_review_session( impl GuardianReviewSession { async fn shutdown(&self) { self.cancel_token.cancel(); + self.wait_for_eager_compaction().await; let _ = self.codex.shutdown_and_wait().await; } @@ -408,6 +414,16 @@ impl GuardianReviewSessionManager { .await; } + if let Err(outcome) = run_before_review_deadline( + deadline, + params.external_cancel.as_ref(), + 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(_) => { @@ -435,6 +451,7 @@ impl GuardianReviewSessionManager { .await; if keep_review_session && matches!(outcome, GuardianReviewSessionOutcome::Completed(_)) { trunk.refresh_last_committed_fork_snapshot().await; + trunk.schedule_eager_compaction().await; } drop(trunk_guard); @@ -459,6 +476,7 @@ impl GuardianReviewSessionManager { codex, cancel_token: CancellationToken::new(), review_lock: Semaphore::new(/*permits*/ 1), + eager_compaction: GuardianEagerCompaction::default(), state: Mutex::new(GuardianReviewState { prior_review_count: 0, last_reviewed_transcript_cursor: None, @@ -482,6 +500,7 @@ impl GuardianReviewSessionManager { codex, cancel_token: CancellationToken::new(), review_lock: Semaphore::new(/*permits*/ 1), + eager_compaction: GuardianEagerCompaction::default(), state: Mutex::new(GuardianReviewState { prior_review_count: 0, last_reviewed_transcript_cursor: None, @@ -636,6 +655,7 @@ async fn spawn_guardian_review_session( cancel_token, reuse_key, review_lock: Semaphore::new(/*permits*/ 1), + eager_compaction: GuardianEagerCompaction::default(), state: Mutex::new(GuardianReviewState { prior_review_count, last_reviewed_transcript_cursor: initial_transcript_cursor, @@ -1115,6 +1135,7 @@ mod tests { cancel_token: CancellationToken::new(), reuse_key, review_lock: Semaphore::new(/*permits*/ 1), + eager_compaction: GuardianEagerCompaction::default(), state: Mutex::new(GuardianReviewState { prior_review_count: 0, last_reviewed_transcript_cursor: None, diff --git a/codex-rs/core/src/guardian/review_session/eager_compaction.rs b/codex-rs/core/src/guardian/review_session/eager_compaction.rs new file mode 100644 index 0000000000..3edd89f090 --- /dev/null +++ b/codex-rs/core/src/guardian/review_session/eager_compaction.rs @@ -0,0 +1,154 @@ +use std::future::Future; +use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::Mutex; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; +use tracing::warn; + +use crate::guardian::GUARDIAN_REVIEW_TIMEOUT; +use crate::session::turn; + +use super::GuardianReviewSession; + +#[cfg(test)] +#[path = "eager_compaction_tests.rs"] +mod tests; + +#[derive(Default)] +pub(super) struct GuardianEagerCompaction { + completion: Mutex>>, +} + +#[derive(Debug, Eq, PartialEq)] +enum EagerCompactionRunOutcome { + Completed, + Cancelled, + TimedOut, +} + +struct EagerCompactionRun { + completion: watch::Sender, +} + +impl EagerCompactionRun { + async fn run_bounded( + self, + cancel_token: &CancellationToken, + timeout: Duration, + maintenance: F, + ) -> EagerCompactionRunOutcome + where + F: Future, + { + let outcome = tokio::select! { + _ = cancel_token.cancelled() => EagerCompactionRunOutcome::Cancelled, + result = tokio::time::timeout(timeout, maintenance) => { + if result.is_ok() { + EagerCompactionRunOutcome::Completed + } else { + EagerCompactionRunOutcome::TimedOut + } + } + }; + self.completion.send_replace(true); + outcome + } +} + +impl GuardianEagerCompaction { + async fn begin(&self) -> Option { + let mut completion = self.completion.lock().await; + if let Some(receiver) = completion.as_ref() + && !*receiver.borrow() + && receiver.has_changed().is_ok() + { + return None; + } + + let (sender, receiver) = watch::channel(false); + *completion = Some(receiver); + Some(EagerCompactionRun { completion: sender }) + } + + async fn wait(&self) { + let Some(mut completion) = self.completion.lock().await.clone() else { + return; + }; + if *completion.borrow() { + return; + } + while completion.changed().await.is_ok() { + if *completion.borrow() { + return; + } + } + } +} + +impl GuardianReviewSession { + pub(super) async fn schedule_eager_compaction(self: &Arc) { + 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(run) = self.eager_compaction.begin().await else { + return; + }; + + let review_session = Arc::clone(self); + drop(tokio::spawn(async move { + let cancel_token = review_session.cancel_token.clone(); + let outcome = run + .run_bounded( + &cancel_token, + GUARDIAN_REVIEW_TIMEOUT, + review_session.run_eager_compaction(turn_context), + ) + .await; + if outcome == EagerCompactionRunOutcome::TimedOut { + warn!( + guardian_thread_id = %review_session.codex.session.thread_id, + "eager guardian maintenance timed out after {GUARDIAN_REVIEW_TIMEOUT:?}" + ); + } + })); + } + + pub(super) async fn wait_for_eager_compaction(&self) { + self.eager_compaction.wait().await; + } + + async fn run_eager_compaction( + self: &Arc, + turn_context: Arc, + ) { + let Ok(review_guard) = self.review_lock.acquire().await else { + return; + }; + + 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, + ) + .await; + + match compact_result { + Ok(true) => { + self.refresh_last_committed_fork_snapshot().await; + } + Ok(false) => {} + Err(err) => { + warn!( + guardian_thread_id = %self.codex.session.thread_id, + "eager guardian compaction failed: {err}" + ); + } + } + + drop(review_guard); + } +} diff --git a/codex-rs/core/src/guardian/review_session/eager_compaction_tests.rs b/codex-rs/core/src/guardian/review_session/eager_compaction_tests.rs new file mode 100644 index 0000000000..f9c86c0760 --- /dev/null +++ b/codex-rs/core/src/guardian/review_session/eager_compaction_tests.rs @@ -0,0 +1,49 @@ +use std::future; +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio_util::sync::CancellationToken; + +use super::EagerCompactionRunOutcome; +use super::GuardianEagerCompaction; + +#[tokio::test] +async fn cancellation_releases_eager_compaction_waiters() { + let eager_compaction = GuardianEagerCompaction::default(); + let run = eager_compaction.begin().await.expect("start maintenance"); + let cancel_token = CancellationToken::new(); + cancel_token.cancel(); + + let outcome = run + .run_bounded( + &cancel_token, + Duration::from_secs(/*secs*/ 1), + future::pending(), + ) + .await; + + assert_eq!(outcome, EagerCompactionRunOutcome::Cancelled); + tokio::time::timeout(Duration::from_secs(/*secs*/ 1), eager_compaction.wait()) + .await + .expect("waiter should be released after cancellation"); +} + +#[tokio::test] +async fn total_timeout_releases_eager_compaction_waiters() { + let eager_compaction = GuardianEagerCompaction::default(); + let run = eager_compaction.begin().await.expect("start maintenance"); + let cancel_token = CancellationToken::new(); + + let outcome = run + .run_bounded( + &cancel_token, + Duration::from_millis(/*millis*/ 10), + future::pending(), + ) + .await; + + assert_eq!(outcome, EagerCompactionRunOutcome::TimedOut); + tokio::time::timeout(Duration::from_secs(/*secs*/ 1), eager_compaction.wait()) + .await + .expect("waiter should be released after the total maintenance timeout"); +} diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index b19cd9d35e..9203bb308d 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -57,6 +57,7 @@ use core_test_support::context_snapshot; use core_test_support::context_snapshot::ContextSnapshotOptions; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_once; @@ -2153,6 +2154,296 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: Ok(()) } +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_provider.supports_websockets = false; + let _ = config.features.enable(Feature::RemoteCompactionV2); + 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") + }) + })) +} + +fn eager_compaction_guardian_assessment(rationale: &str) -> String { + serde_json::json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": rationale, + }) + .to_string() +} + +async fn run_initial_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, + ) + .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 +} + +#[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); + 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 (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![ + 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: 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"), + ]), + }], + ]) + .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(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(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(3)).await?; + + let committed_rollout_items = session + .guardian_review_session + .committed_fork_rollout_items_for_test() + .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 + .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(2)).await?; + + let second_outcome = run_followup_eager_compaction_review( + &session, + &turn, + "Retry after eager compaction failure.", + ) + .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(), 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; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 4beef792bb..f61f7bf0fe 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -802,22 +802,37 @@ async fn run_pre_sampling_compact( client_session: &mut ModelClientSession, ) -> CodexResult<()> { maybe_run_previous_model_inline_compact(sess, turn_context, client_session).await?; - let token_status = auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await; - // Compact if the configured auto-compaction budget or usable context window is exhausted. - if token_status.token_limit_reached { - run_auto_compact( - sess, - turn_context, - client_session, - InitialContextInjection::DoNotInject, - CompactionReason::ContextLimit, - CompactionPhase::PreTurn, - ) - .await?; - } + run_pre_turn_auto_compact_if_needed(sess, turn_context, client_session).await?; Ok(()) } +pub(crate) async fn auto_compact_needed(sess: &Session, turn_context: &TurnContext) -> bool { + auto_compact_token_status(sess, turn_context) + .await + .token_limit_reached +} + +pub(crate) async fn run_pre_turn_auto_compact_if_needed( + 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); + } + + run_auto_compact( + sess, + turn_context, + client_session, + InitialContextInjection::DoNotInject, + CompactionReason::ContextLimit, + CompactionPhase::PreTurn, + ) + .await?; + Ok(true) +} + /// Returns true only when both turns declare compaction compatibility hashes and they differ. /// A missing hash does not provide enough information to trigger compaction. fn comp_hash_changed(previous: Option<&str>, current: Option<&str>) -> bool { diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index ae87350c39..eaa67c5079 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -3,6 +3,7 @@ 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::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; @@ -12,6 +13,7 @@ use codex_protocol::user_input::UserInput; use core_test_support::fs_wait; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_sequence; @@ -30,6 +32,171 @@ use std::os::unix::fs::PermissionsExt; use std::time::Duration; use tempfile::TempDir; +#[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 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; + let sandbox_policy = SandboxPolicy::WorkspaceWrite { + writable_roots: vec![], + network_access: false, + exclude_tmpdir_env_var: true, + 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 second_justification = "Approve the second eager-compaction command."; + let first_args = json!({ + "cmd": "sleep 0.2; printf first", + "yield_time_ms": 1_000_u64, + "sandbox_permissions": SandboxPermissions::RequireEscalated, + "justification": first_justification, + }); + let second_args = json!({ + "cmd": "printf second", + "yield_time_ms": 1_000_u64, + "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 responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-parent-first-tool"), + ev_function_call( + "exec-first", + "exec_command", + &serde_json::to_string(&first_args)?, + ), + ev_completed("resp-parent-first-tool"), + ]), + 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), + ]), + sse(vec![ + ev_response_created("resp-parent-second-tool"), + ev_function_call( + "exec-second", + "exec_command", + &serde_json::to_string(&second_args)?, + ), + ev_completed("resp-parent-second-tool"), + ]), + sse(vec![ + ev_response_created("resp-guardian-second"), + ev_assistant_message("msg-guardian-second", &guardian_assessment), + ev_completed("resp-guardian-second"), + ]), + sse(vec![ + ev_response_created("resp-parent-done"), + ev_assistant_message("msg-parent-done", "done"), + ev_completed("resp-parent-done"), + ]), + ], + ) + .await; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "run both commands that require Guardian review".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + environments: Some(local_selections(test.config.cwd.clone())), + approval_policy: Some(approval_policy), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + sandbox_policy: Some(sandbox_policy), + ..Default::default() + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .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)) + .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:#?}" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_review_session_does_not_inherit_legacy_notify() -> Result<()> { skip_if_no_network!(Ok(()));