mirror of
https://github.com/openai/codex.git
synced 2026-10-29 19:08:43 +00:00
Prewarm guardian review trunks
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
¶ms.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(¶ms.parent_session),
|
||||
Arc::clone(¶ms.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(
|
||||
¶ms,
|
||||
Arc::clone(¶ms.parent_session),
|
||||
Arc::clone(¶ms.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(
|
||||
¶ms,
|
||||
Arc::clone(¶ms.parent_session),
|
||||
Arc::clone(¶ms.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(¶ms.parent_session),
|
||||
Arc::clone(¶ms.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(
|
||||
¶ms.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(¶ms.parent_session),
|
||||
parent_turn: Arc::clone(¶ms.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;
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user