Prewarm guardian review trunks

This commit is contained in:
Joe Gershenson
2026-06-11 17:28:52 -07:00
parent b2cc01fc7c
commit 0defd45736
5 changed files with 433 additions and 48 deletions
+1
View File
@@ -33,6 +33,7 @@ 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::new_guardian_review_id;
pub(crate) use review::prewarm_guardian_review_session;
#[cfg(test)]
pub(crate) use review::record_guardian_denial_for_test;
pub(crate) use review::review_approval_request;
+75 -40
View File
@@ -23,7 +23,9 @@ use tokio::sync::oneshot;
use tokio::time::Instant;
use tokio::time::sleep_until;
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;
@@ -46,6 +48,7 @@ use super::prompt::guardian_output_schema;
use super::prompt::parse_guardian_assessment;
use super::review_session::GuardianReviewSessionOutcome;
use super::review_session::GuardianReviewSessionParams;
use super::review_session::GuardianReviewSessionPrewarmParams;
use super::review_session::build_guardian_review_session_config;
const GUARDIAN_REJECTION_INSTRUCTIONS: &str = concat!(
@@ -665,40 +668,19 @@ pub(crate) fn spawn_approval_request_review(
rx
}
/// Runs the guardian in a locked-down reusable review session.
///
/// The guardian itself should not mutate state or trigger further approvals, so
/// it is pinned to a read-only sandbox with `approval_policy = never` and
/// nonessential agent features disabled. When the cached trunk session is idle,
/// later approvals append onto that same guardian conversation to preserve a
/// stable prompt-cache key. If the trunk is already busy, the review runs in an
/// ephemeral fork from the last committed trunk rollout so parallel approvals
/// do not block each other or mutate the cached thread. The trunk is recreated
/// when the effective review-session config changes, and any future compaction
/// must continue to preserve the guardian policy as exact top-level developer
/// context. It may still reuse the parent's managed-network allowlist for
/// read-only checks, but it intentionally runs without inherited exec-policy
/// rules.
async fn run_guardian_review_session_before_deadline(
session: Arc<Session>,
turn: Arc<TurnContext>,
request: GuardianApprovalRequest,
retry_reason: Option<String>,
schema: serde_json::Value,
external_cancel: Option<CancellationToken>,
deadline: Instant,
) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) {
struct GuardianReviewSessionSetup {
spawn_config: Config,
model: String,
reasoning_effort: Option<codex_protocol::openai_models::ReasoningEffort>,
}
async fn guardian_review_session_setup(
session: &Arc<Session>,
turn: &Arc<TurnContext>,
) -> anyhow::Result<GuardianReviewSessionSetup> {
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(),
);
}
},
Some(network_proxy) => Some(network_proxy.proxy().current_cfg().await?),
None => None,
};
let available_models = session
@@ -745,14 +727,67 @@ async fn run_guardian_review_session_before_deadline(
reasoning_effort,
)
};
let guardian_config = build_guardian_review_session_config(
let spawn_config = build_guardian_review_session_config(
turn.config.as_ref(),
live_network_config.clone(),
live_network_config,
guardian_model.as_str(),
guardian_reasoning_effort.clone(),
);
let guardian_config = match guardian_config {
Ok(config) => config,
)?;
Ok(GuardianReviewSessionSetup {
spawn_config,
model: guardian_model,
reasoning_effort: guardian_reasoning_effort,
})
}
pub(crate) async fn prewarm_guardian_review_session(session: Arc<Session>, turn: Arc<TurnContext>) {
if !routes_approval_to_guardian(turn.as_ref()) {
return;
}
let setup = match guardian_review_session_setup(&session, &turn).await {
Ok(setup) => setup,
Err(err) => {
warn!("guardian trunk prewarm setup failed: {err:#}");
return;
}
};
session
.guardian_review_session
.prewarm_trunk(GuardianReviewSessionPrewarmParams {
parent_session: Arc::clone(&session),
parent_turn: turn,
spawn_config: setup.spawn_config,
})
.await;
}
/// Runs the guardian in a locked-down reusable review session.
///
/// The guardian itself should not mutate state or trigger further approvals, so
/// it is pinned to a read-only sandbox with `approval_policy = never` and
/// nonessential agent features disabled. When the cached trunk session is idle,
/// later approvals append onto that same guardian conversation to preserve a
/// stable prompt-cache key. If the trunk is already busy, the review runs in an
/// ephemeral fork from the last committed trunk rollout so parallel approvals
/// do not block each other or mutate the cached thread. The trunk is recreated
/// when the effective review-session config changes, and any future compaction
/// must continue to preserve the guardian policy as exact top-level developer
/// context. It may still reuse the parent's managed-network allowlist for
/// read-only checks, but it intentionally runs without inherited exec-policy
/// rules.
async fn run_guardian_review_session_before_deadline(
session: Arc<Session>,
turn: Arc<TurnContext>,
request: GuardianApprovalRequest,
retry_reason: Option<String>,
schema: serde_json::Value,
external_cancel: Option<CancellationToken>,
deadline: Instant,
) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) {
let setup = match guardian_review_session_setup(&session, &turn).await {
Ok(setup) => setup,
Err(err) => {
return (
GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(err)),
@@ -767,12 +802,12 @@ async fn run_guardian_review_session_before_deadline(
.run_review(GuardianReviewSessionParams {
parent_session: Arc::clone(&session),
parent_turn: turn.clone(),
spawn_config: guardian_config,
spawn_config: setup.spawn_config,
request,
retry_reason,
schema,
model: guardian_model,
reasoning_effort: guardian_reasoning_effort,
model: setup.model,
reasoning_effort: setup.reasoning_effort,
reasoning_summary: turn.reasoning_summary,
personality: turn.personality,
external_cancel,
+231 -8
View File
@@ -30,6 +30,7 @@ use serde_json::Value;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio_util::sync::CancellationToken;
use tracing::info;
use tracing::warn;
use crate::codex_delegate::run_codex_thread_interactive;
@@ -84,6 +85,12 @@ pub(crate) struct GuardianReviewSessionParams {
pub(crate) deadline: tokio::time::Instant,
}
pub(crate) struct GuardianReviewSessionPrewarmParams {
pub(crate) parent_session: Arc<Session>,
pub(crate) parent_turn: Arc<TurnContext>,
pub(crate) spawn_config: Config,
}
#[derive(Default)]
pub(crate) struct GuardianReviewSessionManager {
state: Arc<Mutex<GuardianReviewSessionState>>,
@@ -92,9 +99,15 @@ pub(crate) struct GuardianReviewSessionManager {
#[derive(Default)]
struct GuardianReviewSessionState {
trunk: Option<Arc<GuardianReviewSession>>,
prewarm: Option<GuardianReviewSessionPrewarm>,
ephemeral_reviews: Vec<Arc<GuardianReviewSession>>,
}
struct GuardianReviewSessionPrewarm {
reuse_key: GuardianReviewSessionReuseKey,
cancel_token: CancellationToken,
}
struct GuardianReviewSession {
codex: Codex,
cancel_token: CancellationToken,
@@ -297,13 +310,17 @@ impl GuardianReviewSessionManager {
}
pub(crate) async fn shutdown(&self) {
let (review_session, ephemeral_reviews) = {
let (review_session, prewarm, ephemeral_reviews) = {
let mut state = self.state.lock().await;
(
state.trunk.take(),
state.prewarm.take(),
std::mem::take(&mut state.ephemeral_reviews),
)
};
if let Some(prewarm) = prewarm {
prewarm.cancel_token.cancel();
}
if let Some(review_session) = review_session {
review_session.shutdown().await;
}
@@ -312,6 +329,124 @@ impl GuardianReviewSessionManager {
}
}
pub(super) async fn prewarm_trunk(&self, params: GuardianReviewSessionPrewarmParams) {
let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config(
&params.spawn_config,
params.parent_session.user_instructions().await,
);
let cancel_token = CancellationToken::new();
let mut stale_trunk_to_shutdown = None;
let mut stale_prewarm_to_cancel = None;
let mut should_schedule_prewarm = true;
{
let mut state = self.state.lock().await;
if let Some(prewarm) = state.prewarm.as_ref() {
if prewarm.reuse_key == reuse_key {
return;
}
stale_prewarm_to_cancel = state.prewarm.take();
}
if let Some(trunk) = state.trunk.as_ref() {
if trunk.reuse_key == reuse_key {
should_schedule_prewarm = false;
}
if should_schedule_prewarm && trunk.review_lock.try_acquire().is_ok() {
stale_trunk_to_shutdown = state.trunk.take();
} else if should_schedule_prewarm {
info!("guardian trunk prewarm skipped because a stale trunk is currently busy");
should_schedule_prewarm = false;
}
}
if should_schedule_prewarm {
state.prewarm = Some(GuardianReviewSessionPrewarm {
reuse_key: reuse_key.clone(),
cancel_token: cancel_token.clone(),
});
}
}
if let Some(prewarm) = stale_prewarm_to_cancel {
info!("cancelling stale guardian trunk prewarm before scheduling a replacement");
prewarm.cancel_token.cancel();
}
if let Some(review_session) = stale_trunk_to_shutdown {
review_session.shutdown_in_background();
}
if !should_schedule_prewarm {
return;
}
let state = Arc::clone(&self.state);
drop(tokio::spawn(async move {
let spawned_session = spawn_guardian_review_session(
Arc::clone(&params.parent_session),
Arc::clone(&params.parent_turn),
params.spawn_config,
reuse_key.clone(),
cancel_token.clone(),
/*fork_snapshot*/ None,
)
.await;
let review_session = match spawned_session {
Ok(review_session) => Arc::new(review_session),
Err(err) => {
if !cancel_token.is_cancelled() {
warn!("guardian trunk prewarm failed: {err:#}");
}
let mut state = state.lock().await;
if state
.prewarm
.as_ref()
.is_some_and(|prewarm| prewarm.reuse_key == reuse_key)
{
state.prewarm = None;
}
return;
}
};
let mut spawned_session_to_shutdown = Some(review_session);
let mut stale_trunk_to_shutdown = None;
let installed = {
let mut state = state.lock().await;
if !state
.prewarm
.as_ref()
.is_some_and(|prewarm| prewarm.reuse_key == reuse_key)
{
false
} else {
state.prewarm = None;
if let Some(trunk) = state.trunk.as_ref()
&& trunk.reuse_key != reuse_key
&& trunk.review_lock.try_acquire().is_ok()
{
stale_trunk_to_shutdown = state.trunk.take();
}
if state.trunk.is_none() {
state.trunk = spawned_session_to_shutdown.take();
true
} else {
false
}
}
};
if installed {
info!("guardian trunk prewarm installed");
}
if let Some(review_session) = stale_trunk_to_shutdown {
review_session.shutdown_in_background();
}
if let Some(review_session) = spawned_session_to_shutdown {
review_session.shutdown().await;
}
}));
}
#[expect(
clippy::await_holding_invalid_type,
reason = "review session selection and trunk spawning must stay serialized"
@@ -335,6 +470,14 @@ impl GuardianReviewSessionManager {
.await
{
Ok(mut state) => {
if let Some(prewarm) = state.prewarm.as_ref()
&& prewarm.reuse_key != next_reuse_key
{
info!("cancelling stale guardian trunk prewarm before review");
if let Some(prewarm) = state.prewarm.take() {
prewarm.cancel_token.cancel();
}
}
if let Some(trunk) = state.trunk.as_ref()
&& trunk.reuse_key != next_reuse_key
&& trunk.review_lock.try_acquire().is_ok()
@@ -343,13 +486,20 @@ impl GuardianReviewSessionManager {
}
if state.trunk.is_none() {
if let Some(prewarm) = state.prewarm.take() {
info!(
"guardian trunk prewarm was not ready before review; falling back to lazy spawn"
);
prewarm.cancel_token.cancel();
}
let spawn_cancel_token = CancellationToken::new();
let review_session = match run_before_review_deadline_with_cancel(
deadline,
params.external_cancel.as_ref(),
&spawn_cancel_token,
Box::pin(spawn_guardian_review_session(
&params,
Arc::clone(&params.parent_session),
Arc::clone(&params.parent_turn),
params.spawn_config.clone(),
next_reuse_key.clone(),
spawn_cancel_token.clone(),
@@ -559,7 +709,8 @@ impl GuardianReviewSessionManager {
params.external_cancel.as_ref(),
&spawn_cancel_token,
Box::pin(spawn_guardian_review_session(
&params,
Arc::clone(&params.parent_session),
Arc::clone(&params.parent_turn),
fork_config,
reuse_key,
spawn_cancel_token.clone(),
@@ -600,7 +751,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,
@@ -616,10 +768,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,
@@ -1072,6 +1224,7 @@ mod tests {
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnCompleteEvent;
use pretty_assertions::assert_eq;
async fn test_review_session() -> (
GuardianReviewSession,
@@ -1259,6 +1412,75 @@ mod tests {
);
}
#[tokio::test]
async fn shutdown_cancels_pending_trunk_prewarm() {
let parent_config = crate::config::test_config().await;
let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config(
&parent_config,
/*user_instructions*/ None,
);
let cancel_token = CancellationToken::new();
let manager = GuardianReviewSessionManager {
state: Arc::new(Mutex::new(GuardianReviewSessionState {
trunk: None,
prewarm: Some(GuardianReviewSessionPrewarm {
reuse_key,
cancel_token: cancel_token.clone(),
}),
ephemeral_reviews: Vec::new(),
})),
};
manager.shutdown().await;
assert!(cancel_token.is_cancelled());
assert!(manager.state.lock().await.prewarm.is_none());
}
#[tokio::test]
async fn prewarm_trunk_skips_existing_matching_trunk_and_cancels_stale_prewarm() {
let params = test_review_params().await;
let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config(
&params.spawn_config,
params.parent_session.user_instructions().await,
);
let mut stale_spawn_config = params.spawn_config.clone();
stale_spawn_config.model = Some("different-guardian-model".to_string());
let stale_reuse_key = GuardianReviewSessionReuseKey::from_spawn_config(
&stale_spawn_config,
params.parent_session.user_instructions().await,
);
let stale_cancel_token = CancellationToken::new();
let (mut review_session, _tx_event, _rx_sub) = test_review_session().await;
review_session.reuse_key = reuse_key;
let manager = GuardianReviewSessionManager {
state: Arc::new(Mutex::new(GuardianReviewSessionState {
trunk: Some(Arc::new(review_session)),
prewarm: Some(GuardianReviewSessionPrewarm {
reuse_key: stale_reuse_key,
cancel_token: stale_cancel_token.clone(),
}),
ephemeral_reviews: Vec::new(),
})),
};
manager
.prewarm_trunk(GuardianReviewSessionPrewarmParams {
parent_session: Arc::clone(&params.parent_session),
parent_turn: Arc::clone(&params.parent_turn),
spawn_config: params.spawn_config,
})
.await;
assert!(stale_cancel_token.is_cancelled());
{
let state = manager.state.lock().await;
assert!(state.prewarm.is_none());
assert!(state.trunk.is_some());
}
manager.shutdown().await;
}
#[tokio::test]
async fn guardian_review_session_compact_scope_change_invalidates_cached_session() {
let parent_config = crate::config::test_config().await;
@@ -1525,6 +1747,7 @@ mod tests {
let manager = GuardianReviewSessionManager {
state: Arc::new(Mutex::new(GuardianReviewSessionState {
trunk: Some(Arc::new(review_session)),
prewarm: None,
ephemeral_reviews: Vec::new(),
})),
};
@@ -43,6 +43,7 @@ use core_test_support::responses::sse_response;
use core_test_support::responses::start_mock_server;
use pretty_assertions::assert_eq;
use std::fs;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tempfile::tempdir;
@@ -68,6 +69,19 @@ where
}
}
async fn wait_for_guardian_trunk_rollout_path(session: &Session) -> PathBuf {
timeout(Duration::from_secs(10), async {
loop {
if let Some(path) = session.guardian_review_session.trunk_rollout_path().await {
return path;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("guardian trunk prewarm should complete")
}
#[tokio::test]
async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() {
let server = start_mock_server().await;
@@ -167,6 +181,111 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() {
assert!(guardian_request.body_contains_text("need network"));
}
#[tokio::test]
async fn prewarmed_guardian_trunk_is_reused_by_first_review() {
let server = start_mock_server().await;
let guardian_request_log = mount_sse_once(
&server,
sse(vec![
ev_response_created("resp-guardian"),
ev_assistant_message(
"msg-guardian",
&serde_json::json!({
"risk_level": "low",
"user_authorization": "high",
"outcome": "allow",
"rationale": "The request grants narrowly scoped network access for this turn.",
})
.to_string(),
),
ev_completed("resp-guardian"),
]),
)
.await;
let (mut session, mut turn_context_raw) = make_session_and_context().await;
*session.active_turn.lock().await = Some(ActiveTurn::default());
turn_context_raw
.approval_policy
.set(AskForApproval::OnRequest)
.expect("test setup should allow updating approval policy");
turn_context_raw
.features
.enable(Feature::GuardianApproval)
.expect("test setup should allow enabling guardian approvals");
let mut config = (*turn_context_raw.config).clone();
config.approvals_reviewer = ApprovalsReviewer::AutoReview;
config.model_provider.base_url = Some(format!("{}/v1", server.uri()));
let config = Arc::new(config);
let models_manager = models_manager_with_provider(
config.codex_home.to_path_buf(),
Arc::clone(&session.services.auth_manager),
config.model_provider.clone(),
);
session.services.models_manager = models_manager;
turn_context_raw.config = Arc::clone(&config);
turn_context_raw.provider = create_model_provider(
config.model_provider.clone(),
turn_context_raw.auth_manager.clone(),
);
let session = Arc::new(session);
let turn_context = Arc::new(turn_context_raw);
crate::guardian::prewarm_guardian_review_session(
Arc::clone(&session),
Arc::clone(&turn_context),
)
.await;
let prewarmed_trunk_rollout_path = wait_for_guardian_trunk_rollout_path(&session).await;
let requested_permissions = RequestPermissionProfile {
network: Some(NetworkPermissions {
enabled: Some(true),
}),
..RequestPermissionProfile::default()
};
let environment = turn_context
.environments
.primary()
.expect("primary environment")
.selection();
let response = timeout(
Duration::from_secs(45),
session.request_permissions_for_environment(
&turn_context,
"perm-call-prewarmed".to_string(),
RequestPermissionsArgs {
environment_id: None,
reason: Some("need network".to_string()),
permissions: requested_permissions.clone(),
},
environment,
CancellationToken::new(),
),
)
.await
.expect("request_permissions should not wait for a client approval");
assert_eq!(
response,
Some(RequestPermissionsResponse {
permissions: requested_permissions.clone(),
scope: PermissionGrantScope::Turn,
strict_auto_review: false,
})
);
assert_eq!(
wait_for_guardian_trunk_rollout_path(&session).await,
prewarmed_trunk_rollout_path
);
let guardian_request = guardian_request_log.single_request();
assert_eq!(guardian_request.path(), "/v1/responses");
assert!(guardian_request.body_contains_text("request_permissions"));
assert!(guardian_request.body_contains_text("need network"));
session.guardian_review_session.shutdown().await;
}
#[tokio::test]
async fn request_permissions_guardian_review_stops_when_cancelled() {
let server = start_mock_server().await;
+7
View File
@@ -54,6 +54,13 @@ impl SessionTask for RegularTask {
});
sess.send_event(ctx.as_ref(), event).await;
sess.set_server_reasoning_included(/*included*/ false).await;
if crate::guardian::routes_approval_to_guardian(ctx.as_ref()) {
let prewarm_sess = Arc::clone(&sess);
let prewarm_ctx = Arc::clone(&ctx);
drop(tokio::spawn(async move {
crate::guardian::prewarm_guardian_review_session(prewarm_sess, prewarm_ctx).await;
}));
}
let prewarmed_client_session = match sess
.consume_startup_prewarm_for_regular_turn(&cancellation_token)
.await