Restore Guardian risk scores across thread lifecycles (#39038)

## Why

Resumed and forked threads need their persisted Guardian risk state available when routing later approval requests.

## What changed

- Initialize a session's `SecurityRiskScore` from the latest score in its rollout history when no score was supplied explicitly.
- Preserve the inherited score when resuming or forking a thread so low-risk actions can skip redundant Guardian reviews.

## Testing

- Add app-server coverage for inheriting the latest Guardian score when resuming and forking threads.

GitOrigin-RevId: 95a6c803f8dae75516d1da9d66a722379a0a7ef2
This commit is contained in:
felixxia-oai authored and copyberry committed 2026-08-17 17:33:30 +00:00
1 parent 1d928cad2f
commit 9dd3d6a13e
2 files changed
+137 -20

No files matched your search

+124 -20
View File
@@ -1,3 +1,4 @@
use std::collections::BTreeMap;
use std::sync::Arc; use std::sync::Arc;
use std::sync::Mutex; use std::sync::Mutex;
use std::sync::atomic::AtomicUsize; use std::sync::atomic::AtomicUsize;
@@ -7,6 +8,8 @@ use std::time::Duration;
use anyhow::Result; use anyhow::Result;
use app_test_support::MockResponsesConfig; use app_test_support::MockResponsesConfig;
use app_test_support::TestAppServer; use app_test_support::TestAppServer;
use app_test_support::create_fake_rollout;
use app_test_support::rollout_path;
use axum::Json; use axum::Json;
use axum::Router; use axum::Router;
use axum::extract::State; use axum::extract::State;
@@ -18,11 +21,18 @@ use axum::routing::get;
use codex_app_server_protocol::ApprovalsReviewer; use codex_app_server_protocol::ApprovalsReviewer;
use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::AskForApproval;
use codex_app_server_protocol::ItemGuardianApprovalReviewStartedNotification; 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::ThreadStartParams;
use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStartResponse;
use codex_app_server_protocol::UserInput; use codex_app_server_protocol::UserInput;
use codex_features::Feature; 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::responses;
use core_test_support::skip_if_no_network; use core_test_support::skip_if_no_network;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
@@ -56,6 +66,13 @@ enum GuardianRisk {
High, High,
} }
#[derive(Clone, Copy)]
enum ThreadLifecycle {
New,
Resume,
Fork,
}
async fn parent_response( async fn parent_response(
State(state): State<Arc<MockResponsesState>>, State(state): State<Arc<MockResponsesState>>,
Json(request): Json<Value>, Json(request): Json<Value>,
@@ -147,10 +164,14 @@ async fn luna_websocket(
}) })
} }
async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> { async fn guardian_v2_routes_tool_approvals(
let (luna_score, expected_guardian_reviews) = match risk { risk: GuardianRisk,
GuardianRisk::Low => (0.25, 1), lifecycle: ThreadLifecycle,
GuardianRisk::High => (0.95, 2), ) -> 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 { let responses_state = Arc::new(MockResponsesState {
luna_score, luna_score,
@@ -177,18 +198,78 @@ async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> {
.enable_feature(Feature::GuardianV2) .enable_feature(Feature::GuardianV2)
.enable_feature(Feature::GuardianApproval) .enable_feature(Feature::GuardianApproval)
.write(codex_home.path())?; .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() let mut app_server = TestAppServer::builder()
.with_codex_home(codex_home.path()) .with_codex_home(codex_home.path())
.build_initialized_with_timeout(TIMEOUT) .build_initialized_with_timeout(TIMEOUT)
.await?; .await?;
let thread = app_server let thread = match lifecycle {
.start_thread(ThreadStartParams { ThreadLifecycle::New => {
approval_policy: Some(AskForApproval::OnRequest), app_server
approvals_reviewer: Some(ApprovalsReviewer::AutoReview), .start_thread(ThreadStartParams {
..Default::default() approval_policy: Some(AskForApproval::OnRequest),
}) approvals_reviewer: Some(ApprovalsReviewer::AutoReview),
.await? ..Default::default()
.thread; })
.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 thread_id = thread.id;
let rollout = thread.path.expect("thread should be persisted"); let rollout = thread.path.expect("thread should be persisted");
let turn_request_id = app_server let turn_request_id = app_server
@@ -205,12 +286,14 @@ async fn guardian_v2_routes_tool_approvals(risk: GuardianRisk) -> Result<()> {
.await?; .await?;
let _: TurnStartResponse = let _: TurnStartResponse =
timeout(TIMEOUT, app_server.read_response(turn_request_id)).await??; timeout(TIMEOUT, app_server.read_response(turn_request_id)).await??;
let review_started: ItemGuardianApprovalReviewStartedNotification = timeout( if matches!(lifecycle, ThreadLifecycle::New) {
TIMEOUT, let review_started: ItemGuardianApprovalReviewStartedNotification = timeout(
app_server.read_notification("item/autoApprovalReview/started"), TIMEOUT,
) app_server.read_notification("item/autoApprovalReview/started"),
.await??; )
assert_eq!(review_started.thread_id, thread_id); .await??;
assert_eq!(review_started.thread_id, thread_id);
}
let luna_request = timeout(TIMEOUT, async { let luna_request = timeout(TIMEOUT, async {
loop { 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(); responses_state.allow_luna.notify_one();
timeout(TIMEOUT, async { 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)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn guardian_v2_low_risk_actions_skip_subsequent_reviews() -> Result<()> { async fn guardian_v2_low_risk_actions_skip_subsequent_reviews() -> Result<()> {
skip_if_no_network!(Ok(())); 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)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn guardian_v2_high_risk_actions_require_full_reviews() -> Result<()> { async fn guardian_v2_high_risk_actions_require_full_reviews() -> Result<()> {
skip_if_no_network!(Ok(())); 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
} }
+13
View File
@@ -26,6 +26,7 @@ use codex_protocol::protocol::MultiAgentVersion;
use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadHistoryMode;
use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::ThreadSource;
use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::protocol::TurnEnvironmentSelections;
use codex_protocol::security_risk::SecurityRiskScore;
use codex_skills::SkillError; use codex_skills::SkillError;
use std::sync::OnceLock; use std::sync::OnceLock;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
@@ -727,6 +728,18 @@ impl Session {
config.current_time_reminder.as_ref(), config.current_time_reminder.as_ref(),
external_time_provider, external_time_provider,
)?; )?;
if thread_extension_init.get::<SecurityRiskScore>().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 = let selected_capability_roots =
match thread_extension_init.get::<Vec<SelectedCapabilityRoot>>() { match thread_extension_init.get::<Vec<SelectedCapabilityRoot>>() {
Some(roots) => roots.as_ref().clone(), Some(roots) => roots.as_ref().clone(),