openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-12 21:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Compact Guardian sessions through ephemeral shadows

Browse Source
This commit is contained in:
Ray Morgan
2026-06-17 10:35:56 -07:00
parent c45bc9f3b4
commit d8283de26d
6 changed files with 540 additions and 429 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/src/protocol/thread_history.rs
+35 -3
View File
@@ -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<Turn>,
current_turn: Option<PendingTurn>,
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![
codex-rs/core/src/guardian/review_session.rs
+139 -32
View File
@@ -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<Mutex<GuardianReviewSessionState>>,
eager_compaction_wait_timeout: Duration,
}
#[derive(Default)]
struct GuardianReviewSessionState {
trunk: Option<Arc<GuardianReviewSession>>,
trunk_review_generation: u64,
ephemeral_reviews: Vec<Arc<GuardianReviewSession>>,
}
impl GuardianReviewSessionState {
fn try_claim_trunk<'a>(
&mut self,
trunk: &'a Arc<GuardianReviewSession>,
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<GuardianReviewForkSnapshot> {
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<PathBuf> {
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<GuardianReviewSession>,
review_generation: u64,
) -> Option<Arc<GuardianReviewSession>> {
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>,
codex-rs/core/src/guardian/review_session/eager_compaction.rs
+143 -53
View File
@@ -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<Mutex<()>>,
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(super) enum GuardianEagerCompactionOutcome {
#[default]
Reusable,
DiscardSession,
}
impl GuardianEagerCompaction {
fn begin(&self) -> Option<OwnedMutexGuard<()>> {
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<Mutex<GuardianEagerCompactionOutcome>>;
impl GuardianReviewSession {
pub(super) async fn schedule_eager_compaction(self: &Arc<Self>) {
pub(super) async fn schedule_eager_compaction(
self: &Arc<Self>,
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<GuardianEagerCompactionOutcome> {
Arc::clone(&self.eager_compaction).lock_owned().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;
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(
&params,
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
}
}
codex-rs/core/src/guardian/tests.rs
+118 -260
View File
@@ -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<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_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<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")
})
}))
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<StreamingSseChunk> {
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<StreamingSseChunk> {
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<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,
retry_reason: Option<&str>,
) -> anyhow::Result<GuardianReviewAnalyticsResult> {
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<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
.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(())
}
codex-rs/core/src/session/turn.rs
+6 -9
View File
@@ -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<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);
}
) -> 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.
codex-rs/core/tests/suite/guardian_review.rs
+99 -72
View File
@@ -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::<Vec<_>>();
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(())
}