Warm Guardian review sessions before first approval

This commit is contained in:
Saud
2026-06-03 14:48:10 -04:00
parent 52b359b249
commit 2d9a36327a
7 changed files with 286 additions and 69 deletions

View File

@@ -32,6 +32,7 @@ pub(crate) use approval_request::guardian_approval_request_to_json;
pub(crate) use review::guardian_rejection_message;
pub(crate) use review::guardian_timeout_message;
pub(crate) use review::is_guardian_reviewer_source;
pub(crate) use review::maybe_warm_guardian_review_session;
pub(crate) use review::new_guardian_review_id;
#[cfg(test)]
pub(crate) use review::record_guardian_denial_for_test;
@@ -170,6 +171,8 @@ use review::GuardianReviewOutcome;
#[cfg(test)]
use review::run_guardian_review_session as run_guardian_review_session_for_test;
#[cfg(test)]
use review::warm_guardian_review_session as warm_guardian_review_session_for_test;
#[cfg(test)]
use review_session::build_guardian_review_session_config as build_guardian_review_session_config_for_test;
#[cfg(test)]

View File

@@ -5,6 +5,7 @@ use codex_analytics::GuardianReviewFailureReason;
use codex_analytics::GuardianReviewTerminalStatus;
use codex_analytics::GuardianReviewTrackContext;
use codex_analytics::GuardianReviewedAction;
use codex_features::Feature;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
@@ -18,9 +19,12 @@ use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::WarningEvent;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use crate::config::Config;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use crate::turn_timing::now_unix_timestamp_ms;
@@ -159,6 +163,33 @@ pub(crate) fn routes_approval_to_guardian_with_reviewer(
) && approvals_reviewer == ApprovalsReviewer::AutoReview
}
pub(crate) fn maybe_warm_guardian_review_session(session: Arc<Session>, turn: Arc<TurnContext>) {
if !turn.features.enabled(Feature::GuardianSessionWarmup)
|| turn.session_source.is_non_root_agent()
|| !routes_approval_to_guardian(turn.as_ref())
{
return;
}
drop(tokio::spawn(async move {
let started_at = Instant::now();
let result = warm_guardian_review_session(Arc::clone(&session), turn).await;
let status = match &result {
Ok(true) => "ready",
Ok(false) => "already_initialized",
Err(_) => "failed",
};
session.services.session_telemetry.record_startup_phase(
"guardian_review_session_warmup",
started_at.elapsed(),
Some(status),
);
if let Err(err) = result {
warn!("guardian review session warmup failed: {err:#}");
}
}));
}
pub(crate) fn is_guardian_reviewer_source(
session_source: &codex_protocol::protocol::SessionSource,
) -> bool {
@@ -665,69 +696,11 @@ pub(super) async fn run_guardian_review_session(
schema: serde_json::Value,
external_cancel: Option<CancellationToken>,
) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) {
let network_proxy = session.services.network_proxy.load_full();
let live_network_config = match network_proxy.as_ref() {
Some(network_proxy) => match network_proxy.proxy().current_cfg().await {
Ok(config) => Some(config),
Err(err) => {
return (
GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(err)),
GuardianReviewAnalyticsResult::without_session(),
);
}
},
None => None,
};
let available_models = session
.services
.models_manager
.list_models(codex_models_manager::manager::RefreshStrategy::Offline)
.await;
let preferred_reasoning_effort = |supports_low: bool, fallback| {
if supports_low {
Some(codex_protocol::openai_models::ReasoningEffort::Low)
} else {
fallback
}
};
let model_override = turn.model_info.auto_review_model_override.as_deref();
let review_model_id =
model_override.unwrap_or_else(|| turn.provider.approval_review_preferred_model());
let review_model = available_models
.iter()
.find(|preset| preset.model == review_model_id);
let (guardian_model, guardian_reasoning_effort) = if let Some(preset) = review_model {
let reasoning_effort = preferred_reasoning_effort(
preset
.supported_reasoning_efforts
.iter()
.any(|effort| effort.effort == codex_protocol::openai_models::ReasoningEffort::Low),
Some(preset.default_reasoning_effort),
);
(review_model_id.to_string(), reasoning_effort)
} else {
let reasoning_effort = preferred_reasoning_effort(
turn.model_info
.supported_reasoning_levels
.iter()
.any(|preset| preset.effort == codex_protocol::openai_models::ReasoningEffort::Low),
turn.reasoning_effort
.or(turn.model_info.default_reasoning_level),
);
(
model_override
.unwrap_or(turn.model_info.slug.as_str())
.to_string(),
reasoning_effort,
)
};
let guardian_config = build_guardian_review_session_config(
turn.config.as_ref(),
live_network_config.clone(),
guardian_model.as_str(),
let ResolvedGuardianReviewSessionConfig {
guardian_config,
guardian_model,
guardian_reasoning_effort,
);
let guardian_config = match guardian_config {
} = match resolve_guardian_review_session_config(session.as_ref(), turn.as_ref()).await {
Ok(config) => config,
Err(err) => {
return (
@@ -801,6 +774,89 @@ pub(super) async fn run_guardian_review_session(
}
}
struct ResolvedGuardianReviewSessionConfig {
guardian_config: Config,
guardian_model: String,
guardian_reasoning_effort: Option<codex_protocol::openai_models::ReasoningEffort>,
}
pub(super) async fn warm_guardian_review_session(
session: Arc<Session>,
turn: Arc<TurnContext>,
) -> anyhow::Result<bool> {
let resolved_config =
resolve_guardian_review_session_config(session.as_ref(), turn.as_ref()).await?;
session
.guardian_review_session
.warm_trunk(Arc::clone(&session), turn, resolved_config.guardian_config)
.await
}
async fn resolve_guardian_review_session_config(
session: &Session,
turn: &TurnContext,
) -> anyhow::Result<ResolvedGuardianReviewSessionConfig> {
let network_proxy = session.services.network_proxy.load_full();
let live_network_config = match network_proxy.as_ref() {
Some(network_proxy) => Some(network_proxy.proxy().current_cfg().await?),
None => None,
};
let available_models = session
.services
.models_manager
.list_models(codex_models_manager::manager::RefreshStrategy::Offline)
.await;
let preferred_reasoning_effort = |supports_low: bool, fallback| {
if supports_low {
Some(codex_protocol::openai_models::ReasoningEffort::Low)
} else {
fallback
}
};
let model_override = turn.model_info.auto_review_model_override.as_deref();
let review_model_id =
model_override.unwrap_or_else(|| turn.provider.approval_review_preferred_model());
let review_model = available_models
.iter()
.find(|preset| preset.model == review_model_id);
let (guardian_model, guardian_reasoning_effort) = if let Some(preset) = review_model {
let reasoning_effort = preferred_reasoning_effort(
preset
.supported_reasoning_efforts
.iter()
.any(|effort| effort.effort == codex_protocol::openai_models::ReasoningEffort::Low),
Some(preset.default_reasoning_effort),
);
(review_model_id.to_string(), reasoning_effort)
} else {
let reasoning_effort = preferred_reasoning_effort(
turn.model_info
.supported_reasoning_levels
.iter()
.any(|preset| preset.effort == codex_protocol::openai_models::ReasoningEffort::Low),
turn.reasoning_effort
.or(turn.model_info.default_reasoning_level),
);
(
model_override
.unwrap_or(turn.model_info.slug.as_str())
.to_string(),
reasoning_effort,
)
};
let guardian_config = build_guardian_review_session_config(
turn.config.as_ref(),
live_network_config,
guardian_model.as_str(),
guardian_reasoning_effort,
)?;
Ok(ResolvedGuardianReviewSessionConfig {
guardian_config,
guardian_model,
guardian_reasoning_effort,
})
}
#[cfg(test)]
mod review_tests {
use super::*;

View File

@@ -87,6 +87,7 @@ pub(crate) struct GuardianReviewSessionManager {
struct GuardianReviewSessionState {
trunk: Option<Arc<GuardianReviewSession>>,
ephemeral_reviews: Vec<Arc<GuardianReviewSession>>,
shutdown: bool,
}
struct GuardianReviewSession {
@@ -290,6 +291,7 @@ impl GuardianReviewSessionManager {
pub(crate) async fn shutdown(&self) {
let (review_session, ephemeral_reviews) = {
let mut state = self.state.lock().await;
state.shutdown = true;
(
state.trunk.take(),
std::mem::take(&mut state.ephemeral_reviews),
@@ -303,6 +305,44 @@ impl GuardianReviewSessionManager {
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "warmup and review session trunk spawning must stay serialized"
)]
pub(crate) async fn warm_trunk(
&self,
parent_session: Arc<Session>,
parent_turn: Arc<TurnContext>,
spawn_config: Config,
) -> anyhow::Result<bool> {
let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config(&spawn_config);
let mut state = self.state.lock().await;
if state.shutdown || state.trunk.is_some() {
return Ok(false);
}
let spawn_cancel_token = CancellationToken::new();
let review_session = run_before_review_deadline_with_cancel(
tokio::time::Instant::now() + GUARDIAN_REVIEW_TIMEOUT,
/*external_cancel*/ None,
&spawn_cancel_token,
Box::pin(spawn_guardian_review_session(
&parent_session,
&parent_turn,
spawn_config,
reuse_key,
spawn_cancel_token.clone(),
/*fork_snapshot*/ None,
)),
)
.await
.map_err(|outcome| {
anyhow!("guardian review session warmup did not complete: {outcome:?}")
})??;
state.trunk = Some(Arc::new(review_session));
Ok(true)
}
#[expect(
clippy::await_holding_invalid_type,
reason = "review session selection and trunk spawning must stay serialized"
@@ -337,7 +377,8 @@ impl GuardianReviewSessionManager {
params.external_cancel.as_ref(),
&spawn_cancel_token,
Box::pin(spawn_guardian_review_session(
&params,
&params.parent_session,
&params.parent_turn,
params.spawn_config.clone(),
next_reuse_key.clone(),
spawn_cancel_token.clone(),
@@ -543,7 +584,8 @@ impl GuardianReviewSessionManager {
params.external_cancel.as_ref(),
&spawn_cancel_token,
Box::pin(spawn_guardian_review_session(
&params,
&params.parent_session,
&params.parent_turn,
fork_config,
reuse_key,
spawn_cancel_token.clone(),
@@ -582,7 +624,8 @@ impl GuardianReviewSessionManager {
}
async fn spawn_guardian_review_session(
params: &GuardianReviewSessionParams,
parent_session: &Arc<Session>,
parent_turn: &Arc<TurnContext>,
spawn_config: Config,
reuse_key: GuardianReviewSessionReuseKey,
cancel_token: CancellationToken,
@@ -598,10 +641,10 @@ async fn spawn_guardian_review_session(
};
let codex = Box::pin(run_codex_thread_interactive(
spawn_config,
params.parent_session.services.auth_manager.clone(),
params.parent_session.services.models_manager.clone(),
Arc::clone(&params.parent_session),
Arc::clone(&params.parent_turn),
parent_session.services.auth_manager.clone(),
parent_session.services.models_manager.clone(),
Arc::clone(parent_session),
Arc::clone(parent_turn),
cancel_token.clone(),
SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_string()),
initial_history,
@@ -958,6 +1001,7 @@ pub(crate) fn build_guardian_review_session_config(
Feature::Plugins,
Feature::WebSearchRequest,
Feature::WebSearchCached,
Feature::GuardianSessionWarmup,
] {
guardian_config.features.disable(feature).map_err(|err| {
anyhow::anyhow!(
@@ -1141,6 +1185,24 @@ mod tests {
}
}
#[tokio::test]
async fn warm_trunk_does_not_spawn_after_shutdown() {
let manager = GuardianReviewSessionManager::default();
manager.shutdown().await;
let params = test_review_params().await;
assert!(
!manager
.warm_trunk(
params.parent_session,
params.parent_turn,
params.spawn_config,
)
.await
.expect("warmup after shutdown should be a no-op")
);
}
#[tokio::test]
async fn guardian_review_session_config_change_invalidates_cached_session() {
let parent_config = crate::config::test_config().await;
@@ -1260,6 +1322,29 @@ mod tests {
assert!(!guardian_config.features.enabled(Feature::CodexHooks));
}
#[tokio::test]
async fn guardian_review_session_config_disables_warmup() {
let mut parent_config = crate::config::test_config().await;
parent_config
.features
.enable(Feature::GuardianSessionWarmup)
.expect("enable guardian session warmup on parent config");
let guardian_config = build_guardian_review_session_config(
&parent_config,
/*live_network_config*/ None,
"active-model",
/*reasoning_effort*/ None,
)
.expect("guardian config");
assert!(
!guardian_config
.features
.enabled(Feature::GuardianSessionWarmup)
);
}
#[tokio::test]
async fn guardian_review_session_config_disables_skill_instructions() {
let mut parent_config = crate::config::test_config().await;

View File

@@ -1575,6 +1575,62 @@ async fn build_guardian_prompt_items_includes_parent_session_id() -> anyhow::Res
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn warmed_guardian_trunk_is_reused_by_first_review() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let request_log = mount_sse_once(
&server,
sse(vec![
ev_response_created("resp-guardian"),
ev_assistant_message(
"msg-guardian",
"{\"risk_level\":\"low\",\"user_authorization\":\"high\",\"outcome\":\"allow\",\"rationale\":\"warmed guardian rationale\"}",
),
ev_completed("resp-guardian"),
]),
)
.await;
let (session, turn) = guardian_test_session_and_turn(&server).await;
assert!(warm_guardian_review_session_for_test(Arc::clone(&session), Arc::clone(&turn)).await?);
seed_guardian_parent_history(&session, &turn).await;
let outcome = run_guardian_review_session_for_test(
Arc::clone(&session),
Arc::clone(&turn),
GuardianApprovalRequest::Shell {
id: "shell-1".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("Need to push the docs fix.".to_string()),
},
/*retry_reason*/ None,
guardian_output_schema(),
/*external_cancel*/ None,
)
.await;
let (GuardianReviewOutcome::Completed(assessment), metadata) = outcome else {
panic!("expected guardian assessment");
};
assert_eq!(assessment.outcome, GuardianAssessmentOutcome::Allow);
assert!(matches!(
metadata.guardian_session_kind,
Some(codex_analytics::GuardianReviewSessionKind::TrunkReused)
));
assert_eq!(metadata.had_prior_review_context, Some(false));
assert_eq!(
request_log.requests().len(),
1,
"warmup should initialize the trunk without performing a review"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));

View File

@@ -54,6 +54,7 @@ impl SessionTask for RegularTask {
});
sess.send_event(ctx.as_ref(), event).await;
sess.set_server_reasoning_included(/*included*/ false).await;
crate::guardian::maybe_warm_guardian_review_session(Arc::clone(&sess), Arc::clone(&ctx));
let prewarmed_client_session = match sess
.consume_startup_prewarm_for_regular_turn(&cancellation_token)
.await

View File

@@ -190,6 +190,8 @@ pub enum Feature {
DefaultModeRequestUserInput,
/// Enable automatic review for approval prompts.
GuardianApproval,
/// Warm the reusable automatic-review session before its first approval prompt.
GuardianSessionWarmup,
/// Enable persisted thread goals and automatic goal continuation.
Goals,
/// Route MCP tool approval prompts through the MCP elicitation request path.
@@ -1117,6 +1119,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::Stable,
default_enabled: true,
},
FeatureSpec {
id: Feature::GuardianSessionWarmup,
key: "guardian_session_warmup",
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::Goals,
key: "goals",

View File

@@ -101,6 +101,14 @@ fn guardian_approval_is_stable_and_enabled_by_default() {
assert_eq!(Feature::GuardianApproval.default_enabled(), true);
}
#[test]
fn guardian_session_warmup_is_under_development_and_disabled_by_default() {
let spec = Feature::GuardianSessionWarmup.info();
assert_eq!(spec.stage, Stage::UnderDevelopment);
assert_eq!(Feature::GuardianSessionWarmup.default_enabled(), false);
}
#[test]
fn external_migration_is_experimental_and_disabled_by_default() {
let spec = Feature::ExternalMigration.info();