mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Keep Guardian v2 risk scores in memory (#39304)
## What changed - Stop writing Guardian v2 security risk scores to rollout history. - Treat resumed and forked threads as having no previous Guardian score, so their first tool approval is classified and reviewed normally. - Remove the now-unused `codex-history` dependency from the Guardian v2 extension. ## Testing - Update app-server coverage for new, resumed, and forked threads to wait for classification completion and verify the expected approval reviews. GitOrigin-RevId: 74d9caa903a1ff9d38e577b4697b5f8970f11ebb
This commit is contained in:
Generated
-1
@@ -3312,7 +3312,6 @@ dependencies = [
|
||||
"codex-core",
|
||||
"codex-extension-api",
|
||||
"codex-features",
|
||||
"codex-history",
|
||||
"codex-http-client",
|
||||
"codex-login",
|
||||
"codex-model-provider",
|
||||
|
||||
BIN
Binary file not shown.
@@ -1,4 +1,3 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
@@ -9,7 +8,6 @@ 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,6 +16,7 @@ use axum::extract::ws::WebSocketUpgrade;
|
||||
use axum::http::header;
|
||||
use axum::response::IntoResponse;
|
||||
use axum::routing::get;
|
||||
use axum::routing::post;
|
||||
use codex_app_server_protocol::ApprovalsReviewer;
|
||||
use codex_app_server_protocol::AskForApproval;
|
||||
use codex_app_server_protocol::ItemGuardianApprovalReviewStartedNotification;
|
||||
@@ -31,9 +30,6 @@ 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;
|
||||
@@ -58,6 +54,7 @@ struct MockResponsesState {
|
||||
luna_requests: Mutex<Vec<Value>>,
|
||||
allow_luna: Notify,
|
||||
allow_guardian_review: Notify,
|
||||
classification_completed: Notify,
|
||||
luna_score: f64,
|
||||
}
|
||||
|
||||
@@ -170,11 +167,10 @@ 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::Threshold, _) => (0.5, 2),
|
||||
(GuardianRisk::High, _) => (0.95, 2),
|
||||
let (luna_score, expected_guardian_reviews) = match risk {
|
||||
GuardianRisk::Low => (0.25, 1),
|
||||
GuardianRisk::Threshold => (0.5, 2),
|
||||
GuardianRisk::High => (0.95, 2),
|
||||
};
|
||||
let responses_state = Arc::new(MockResponsesState {
|
||||
luna_score,
|
||||
@@ -184,6 +180,16 @@ async fn guardian_v2_routes_tool_approvals(
|
||||
let responses_url = format!("http://{}", listener.local_addr()?);
|
||||
let router = Router::new()
|
||||
.route("/v1/responses", get(luna_websocket).post(parent_response))
|
||||
.route(
|
||||
"/metrics",
|
||||
post(
|
||||
|State(state): State<Arc<MockResponsesState>>, body: String| async move {
|
||||
if body.contains("codex.guardian_v2.classification") {
|
||||
state.classification_completed.notify_one();
|
||||
}
|
||||
},
|
||||
),
|
||||
)
|
||||
.with_state(Arc::clone(&responses_state));
|
||||
let responses_server = tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, router).await;
|
||||
@@ -196,39 +202,25 @@ async fn guardian_v2_routes_tool_approvals(
|
||||
.with_approval_policy("on-request")
|
||||
.with_root_config("approvals_reviewer = \"auto_review\"")
|
||||
.with_extra_config(&format!(
|
||||
"[mcp_servers.{TEST_SERVER_NAME}]\nurl = \"{mcp_server_url}/mcp\"\ndefault_tools_approval_mode = \"prompt\""
|
||||
"[mcp_servers.{TEST_SERVER_NAME}]\nurl = \"{mcp_server_url}/mcp\"\ndefault_tools_approval_mode = \"prompt\"\n\n[analytics]\nenabled = true\n\n[otel]\nmetrics_exporter = {{ otlp-http = {{ endpoint = \"{responses_url}/metrics\", protocol = \"json\" }} }}"
|
||||
))
|
||||
.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)
|
||||
}
|
||||
ThreadLifecycle::Resume | ThreadLifecycle::Fork => Some(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 mut app_server = TestAppServer::builder()
|
||||
.with_codex_home(codex_home.path())
|
||||
.with_env_overrides(&[("OTEL_METRIC_EXPORT_INTERVAL", Some("25"))])
|
||||
.build_initialized_with_timeout(TIMEOUT)
|
||||
.await?;
|
||||
let thread = match lifecycle {
|
||||
@@ -274,7 +266,6 @@ async fn guardian_v2_routes_tool_approvals(
|
||||
}
|
||||
};
|
||||
let thread_id = thread.id;
|
||||
let rollout = thread.path.expect("thread should be persisted");
|
||||
let turn_request_id = app_server
|
||||
.send_turn_start_request(TurnStartParams {
|
||||
thread_id: thread_id.clone(),
|
||||
@@ -289,14 +280,12 @@ async fn guardian_v2_routes_tool_approvals(
|
||||
.await?;
|
||||
let _: TurnStartResponse =
|
||||
timeout(TIMEOUT, app_server.read_response(turn_request_id)).await??;
|
||||
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 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 {
|
||||
@@ -332,34 +321,8 @@ async fn guardian_v2_routes_tool_approvals(
|
||||
})
|
||||
})
|
||||
);
|
||||
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 {
|
||||
loop {
|
||||
if tokio::fs::read_to_string(&rollout)
|
||||
.await?
|
||||
.lines()
|
||||
.filter_map(|line| serde_json::from_str::<Value>(line).ok())
|
||||
.any(|line| {
|
||||
line["type"] == "security_risk_score"
|
||||
&& line["payload"]["scores"]["action_risk"] == json!(luna_score)
|
||||
})
|
||||
{
|
||||
return Ok::<(), std::io::Error>(());
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await??;
|
||||
timeout(TIMEOUT, responses_state.classification_completed.notified()).await?;
|
||||
responses_state.allow_guardian_review.notify_one();
|
||||
responses_state.allow_luna.notify_one();
|
||||
timeout(
|
||||
@@ -425,13 +388,13 @@ async fn guardian_v2_threshold_score_requires_full_reviews() -> Result<()> {
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn resumed_thread_inherits_latest_guardian_score() -> Result<()> {
|
||||
async fn resumed_thread_starts_without_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<()> {
|
||||
async fn forked_thread_starts_without_guardian_score() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
guardian_v2_routes_tool_approvals(GuardianRisk::Low, ThreadLifecycle::Fork).await
|
||||
}
|
||||
|
||||
@@ -17,7 +17,6 @@ codex-api = { workspace = true }
|
||||
codex-core = { workspace = true }
|
||||
codex-extension-api = { workspace = true }
|
||||
codex-features = { workspace = true }
|
||||
codex-history = { workspace = true }
|
||||
codex-http-client = { workspace = true }
|
||||
codex-login = { workspace = true }
|
||||
codex-model-provider = { workspace = true }
|
||||
|
||||
@@ -27,7 +27,6 @@ use codex_extension_api::ToolName;
|
||||
use codex_extension_api::ToolPayload;
|
||||
use codex_extension_api::ToolStartInput;
|
||||
use codex_features::Feature;
|
||||
use codex_history::RolloutItem;
|
||||
use codex_login::AgentIdentityAuthPolicy;
|
||||
use codex_login::AuthManager;
|
||||
use codex_model_provider::create_model_provider;
|
||||
@@ -614,12 +613,6 @@ impl ToolLifecycleContributor for GuardianV2Extension {
|
||||
.latest_scored_tool_call
|
||||
.fetch_max(tool_call_index, Ordering::Release);
|
||||
classification_finished_at = Some(Instant::now());
|
||||
if !config.ephemeral {
|
||||
thread
|
||||
.append_rollout_items(&[RolloutItem::SecurityRiskScore(score)])
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
}
|
||||
Ok("success")
|
||||
}
|
||||
.await;
|
||||
|
||||
@@ -18,7 +18,6 @@ use codex_extension_api::ToolName;
|
||||
use codex_extension_api::ToolPayload;
|
||||
use codex_extension_api::ToolStartInput;
|
||||
use codex_features::Feature;
|
||||
use codex_history::RolloutItem;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::ExternalAuth;
|
||||
@@ -1201,21 +1200,6 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
|
||||
.await,
|
||||
None
|
||||
);
|
||||
test.codex.ensure_rollout_materialized().await;
|
||||
test.codex.flush_rollout().await?;
|
||||
let persisted_scores = test
|
||||
.codex
|
||||
.load_history(/*include_archived*/ false)
|
||||
.await?
|
||||
.items
|
||||
.into_iter()
|
||||
.filter_map(|item| match item {
|
||||
RolloutItem::SecurityRiskScore(score) => Some(score),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(persisted_scores, vec![score.as_ref().clone()]);
|
||||
|
||||
thread_store.insert(SecurityRiskScore {
|
||||
scores: BTreeMap::from([("action_risk".to_string(), 0.5)]),
|
||||
sampled_at: None,
|
||||
|
||||
@@ -5,10 +5,10 @@ use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
/// A durable, thread-owned snapshot of security risk classifier scores.
|
||||
/// A thread-owned, in-memory snapshot of security risk classifier scores.
|
||||
///
|
||||
/// These records belong to rollout history only and must not enter model-visible
|
||||
/// conversation context or user-visible thread item projections.
|
||||
/// Scores must not enter model-visible conversation context or user-visible thread
|
||||
/// item projections.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct SecurityRiskScore {
|
||||
pub scores: BTreeMap<String, f64>,
|
||||
|
||||
Reference in New Issue
Block a user