Eagerly compact Guardian threads between reviews

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<GuardianReviewState>,
 }
 
@@ -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,

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<Option<watch::Receiver<bool>>>,
+}
+
+#[derive(Debug, Eq, PartialEq)]
+enum EagerCompactionRunOutcome {
+    Completed,
+    Cancelled,
+    TimedOut,
+}
+
+struct EagerCompactionRun {
+    completion: watch::Sender<bool>,
+}
+
+impl EagerCompactionRun {
+    async fn run_bounded<F>(
+        self,
+        cancel_token: &CancellationToken,
+        timeout: Duration,
+        maintenance: F,
+    ) -> EagerCompactionRunOutcome
+    where
+        F: Future<Output = ()>,
+    {
+        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<EagerCompactionRun> {
+        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<Self>) {
+        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<Self>,
+        turn_context: Arc<crate::session::turn_context::TurnContext>,
+    ) {
+        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);
+    }
+}

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");
+}

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<TurnContext>) {
+    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<bool> {
+    let body = serde_json::from_slice::<serde_json::Value>(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<Session>,
+    turn: &Arc<TurnContext>,
+) -> (
+    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<Session>,
+    turn: &Arc<TurnContext>,
+    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(()));

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<Session>,
+    turn_context: &Arc<TurnContext>,
+    client_session: &mut ModelClientSession,
+) -> CodexResult<bool> {
+    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 {