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
2026-08-17 17:33:30 +00:00
committed by copyberry
parent 1d928cad2f
commit 9dd3d6a13e
2 changed files with 137 additions and 20 deletions
+124 -20
View File
@@ -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<Arc<MockResponsesState>>,
Json(request): Json<Value>,
@@ -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
}
+13
View File
@@ -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::<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 =
match thread_extension_init.get::<Vec<SelectedCapabilityRoot>>() {
Some(roots) => roots.as_ref().clone(),