diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs index d9f16091be..297637210b 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::sync::Arc; use std::sync::Mutex; use std::sync::atomic::AtomicUsize; @@ -7,6 +8,8 @@ use std::time::Duration; use anyhow::Result; use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; +use app_test_support::create_fake_rollout; +use app_test_support::rollout_path; use axum::Json; use axum::Router; use axum::extract::State; @@ -18,11 +21,18 @@ use axum::routing::get; use codex_app_server_protocol::ApprovalsReviewer; use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::ItemGuardianApprovalReviewStartedNotification; +use codex_app_server_protocol::ThreadForkParams; +use codex_app_server_protocol::ThreadForkResponse; +use codex_app_server_protocol::ThreadResumeParams; +use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput; use codex_features::Feature; +use codex_protocol::security_risk::SecurityRiskScore; +use codex_rollout::RolloutItem; +use codex_rollout::append_rollout_item_to_path; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; @@ -56,6 +66,13 @@ enum GuardianRisk { High, } +#[derive(Clone, Copy)] +enum ThreadLifecycle { + New, + Resume, + Fork, +} + async fn parent_response( State(state): State>, Json(request): Json, @@ -147,10 +164,14 @@ async fn luna_websocket( }) } -async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> { - let (luna_score, expected_guardian_reviews) = match risk { - GuardianRisk::Low => (0.25, 1), - GuardianRisk::High => (0.95, 2), +async fn guardian_v2_routes_tool_approvals( + risk: GuardianRisk, + lifecycle: ThreadLifecycle, +) -> Result<()> { + let (luna_score, expected_guardian_reviews) = match (risk, lifecycle) { + (GuardianRisk::Low, ThreadLifecycle::New) => (0.25, 1), + (GuardianRisk::Low, ThreadLifecycle::Resume | ThreadLifecycle::Fork) => (0.25, 0), + (GuardianRisk::High, _) => (0.95, 2), }; let responses_state = Arc::new(MockResponsesState { luna_score, @@ -177,18 +198,78 @@ async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> { .enable_feature(Feature::GuardianV2) .enable_feature(Feature::GuardianApproval) .write(codex_home.path())?; + let original_thread_id = match lifecycle { + ThreadLifecycle::New => None, + ThreadLifecycle::Resume | ThreadLifecycle::Fork => { + let thread_id = create_fake_rollout( + codex_home.path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + USER_CONTEXT, + Some("mock_provider"), + /*git_info*/ None, + )?; + let original_rollout = + rollout_path(codex_home.path(), "2025-01-05T12-00-00", &thread_id); + for action_risk in [0.95, 0.1] { + append_rollout_item_to_path( + &original_rollout, + &RolloutItem::SecurityRiskScore(SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), action_risk)]), + sampled_at: None, + }), + ) + .await?; + } + Some(thread_id) + } + }; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) .build_initialized_with_timeout(TIMEOUT) .await?; - let thread = app_server - .start_thread(ThreadStartParams { - approval_policy: Some(AskForApproval::OnRequest), - approvals_reviewer: Some(ApprovalsReviewer::AutoReview), - ..Default::default() - }) - .await? - .thread; + let thread = match lifecycle { + ThreadLifecycle::New => { + app_server + .start_thread(ThreadStartParams { + approval_policy: Some(AskForApproval::OnRequest), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }) + .await? + .thread + } + ThreadLifecycle::Resume => { + let original_thread_id = original_thread_id.expect("resumed thread should exist"); + let request_id = app_server + .send_thread_resume_request(ThreadResumeParams { + thread_id: original_thread_id.clone(), + approval_policy: Some(AskForApproval::OnRequest), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }) + .await?; + let resumed: ThreadResumeResponse = + timeout(TIMEOUT, app_server.read_response(request_id)).await??; + assert_eq!(resumed.thread.id, original_thread_id); + resumed.thread + } + ThreadLifecycle::Fork => { + let original_thread_id = original_thread_id.expect("forked thread should exist"); + let request_id = app_server + .send_thread_fork_request(ThreadForkParams { + thread_id: original_thread_id.clone(), + approval_policy: Some(AskForApproval::OnRequest), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }) + .await?; + let forked: ThreadForkResponse = + timeout(TIMEOUT, app_server.read_response(request_id)).await??; + assert_ne!(forked.thread.id, original_thread_id); + forked.thread + } + }; let thread_id = thread.id; let rollout = thread.path.expect("thread should be persisted"); let turn_request_id = app_server @@ -205,12 +286,14 @@ async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> { .await?; let _: TurnStartResponse = timeout(TIMEOUT, app_server.read_response(turn_request_id)).await??; - let review_started: ItemGuardianApprovalReviewStartedNotification = timeout( - TIMEOUT, - app_server.read_notification("item/autoApprovalReview/started"), - ) - .await??; - assert_eq!(review_started.thread_id, thread_id); + if matches!(lifecycle, ThreadLifecycle::New) { + let review_started: ItemGuardianApprovalReviewStartedNotification = timeout( + TIMEOUT, + app_server.read_notification("item/autoApprovalReview/started"), + ) + .await??; + assert_eq!(review_started.thread_id, thread_id); + } let luna_request = timeout(TIMEOUT, async { loop { @@ -246,6 +329,15 @@ async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> { }) }) ); + if matches!(lifecycle, ThreadLifecycle::Resume | ThreadLifecycle::Fork) { + timeout(TIMEOUT, async { + while responses_state.parent_requests.load(Ordering::SeqCst) < 3 { + tokio::task::yield_now().await; + } + }) + .await?; + assert_eq!(responses_state.guardian_reviews.load(Ordering::SeqCst), 0); + } responses_state.allow_luna.notify_one(); timeout(TIMEOUT, async { @@ -285,11 +377,23 @@ async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_v2_low_risk_actions_skip_subsequent_reviews() -> Result<()> { skip_if_no_network!(Ok(())); - guardian_v2_routes_tool_approvals(GuardianRisk::Low).await + guardian_v2_routes_tool_approvals(GuardianRisk::Low, ThreadLifecycle::New).await } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_v2_high_risk_actions_require_full_reviews() -> Result<()> { skip_if_no_network!(Ok(())); - guardian_v2_routes_tool_approvals(GuardianRisk::High).await + guardian_v2_routes_tool_approvals(GuardianRisk::High, ThreadLifecycle::New).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resumed_thread_inherits_latest_guardian_score() -> Result<()> { + skip_if_no_network!(Ok(())); + guardian_v2_routes_tool_approvals(GuardianRisk::Low, ThreadLifecycle::Resume).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn forked_thread_inherits_latest_guardian_score() -> Result<()> { + skip_if_no_network!(Ok(())); + guardian_v2_routes_tool_approvals(GuardianRisk::Low, ThreadLifecycle::Fork).await } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 1bf01294a7..5bda5821bd 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -26,6 +26,7 @@ use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnEnvironmentSelections; +use codex_protocol::security_risk::SecurityRiskScore; use codex_skills::SkillError; use std::sync::OnceLock; use tokio::sync::Semaphore; @@ -727,6 +728,18 @@ impl Session { config.current_time_reminder.as_ref(), external_time_provider, )?; + if thread_extension_init.get::().is_none() + && let Some(score) = initial_history + .get_rollout_items() + .iter() + .rev() + .find_map(|item| match item { + RolloutItem::SecurityRiskScore(score) => Some(score), + _ => None, + }) + { + thread_extension_init.insert(score.clone()); + } let selected_capability_roots = match thread_extension_init.get::>() { Some(roots) => roots.as_ref().clone(),