mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user