From 71e5e1ec50775bc3abf8a12edfffe5eb803971d6 Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Mon, 17 Aug 2026 16:58:14 +0000 Subject: [PATCH] Add app-server coverage for Guardian V2 approval routing (#39035) ## What changed - Add end-to-end app-server tests for MCP tool approvals with Guardian V2 and automatic review enabled. - Verify that sampling receives the thread context, uses a thread-scoped prompt cache key, and persists the resulting action risk score. - Cover both routing outcomes: low-risk actions skip later Guardian reviews, while high-risk actions continue through full review. GitOrigin-RevId: 78ba2b0560caf701ad333785c4dda4f7d2be4f2c --- .../app-server/tests/suite/v2/guardian_v2.rs | 295 ++++++++++++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + 2 files changed, 296 insertions(+) create mode 100644 codex-rs/app-server/tests/suite/v2/guardian_v2.rs diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs new file mode 100644 index 0000000000..d9f16091be --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -0,0 +1,295 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::Result; +use app_test_support::MockResponsesConfig; +use app_test_support::TestAppServer; +use axum::Json; +use axum::Router; +use axum::extract::State; +use axum::extract::ws::Message; +use axum::extract::ws::WebSocketUpgrade; +use axum::http::header; +use axum::response::IntoResponse; +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::ThreadStartParams; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::UserInput; +use codex_features::Feature; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use tempfile::TempDir; +use tokio::net::TcpListener; +use tokio::sync::Notify; +use tokio::time::timeout; + +use super::mcp_tool::TEST_SERVER_NAME; +use super::mcp_tool::TEST_TOOL_NAME; +use super::mcp_tool::start_mcp_server; + +const TIMEOUT: Duration = Duration::from_secs(30); +const USER_CONTEXT: &str = "The user authorized reading the existing project files."; + +#[derive(Default)] +struct MockResponsesState { + parent_requests: AtomicUsize, + guardian_reviews: AtomicUsize, + luna_requests: Mutex>, + allow_luna: Notify, + allow_guardian_review: Notify, + luna_score: f64, +} + +#[derive(Clone, Copy)] +enum GuardianRisk { + Low, + High, +} + +async fn parent_response( + State(state): State>, + Json(request): Json, +) -> impl IntoResponse { + let events = if request + .pointer("/client_metadata/x-openai-subagent") + .and_then(Value::as_str) + == Some("guardian") + { + let review_number = state.guardian_reviews.fetch_add(1, Ordering::SeqCst); + if review_number == 0 { + state.allow_guardian_review.notified().await; + } + vec![ + responses::ev_response_created("guardian-review"), + responses::ev_assistant_message( + "guardian-assessment", + &json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": "The requested command is safe.", + }) + .to_string(), + ), + responses::ev_completed("guardian-review"), + ] + } else { + let request_number = state.parent_requests.fetch_add(1, Ordering::SeqCst); + if request_number < 2 { + let call_id = format!("guardian-action-{request_number}"); + let arguments = json!({ "message": format!("guardian-{request_number}") }).to_string(); + vec![ + responses::ev_response_created(&call_id), + responses::ev_function_call_with_namespace( + &call_id, + &format!("mcp__{TEST_SERVER_NAME}"), + TEST_TOOL_NAME, + &arguments, + ), + responses::ev_completed(&call_id), + ] + } else { + vec![ + responses::ev_response_created("guardian-complete"), + responses::ev_assistant_message("guardian-message", "done"), + responses::ev_completed("guardian-complete"), + ] + } + }; + + ( + [(header::CONTENT_TYPE, "text/event-stream")], + responses::sse(events), + ) +} + +async fn luna_websocket( + State(state): State>, + websocket: WebSocketUpgrade, +) -> impl IntoResponse { + websocket.on_upgrade(move |mut socket| async move { + while let Some(Ok(message)) = socket.recv().await { + let Message::Text(text) = message else { + continue; + }; + let request: Value = serde_json::from_str(&text).expect("valid Luna request"); + state + .luna_requests + .lock() + .expect("Luna request lock should not be poisoned") + .push(request); + state.allow_luna.notified().await; + let score = json!({ "scores": { "action_risk": state.luna_score } }).to_string(); + for event in [ + responses::ev_response_created("luna-score"), + responses::ev_assistant_message("luna-score-message", &score), + responses::ev_completed("luna-score"), + ] { + if socket + .send(Message::Text(event.to_string().into())) + .await + .is_err() + { + break; + } + } + } + }) +} + +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), + }; + let responses_state = Arc::new(MockResponsesState { + luna_score, + ..Default::default() + }); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let responses_url = format!("http://{}", listener.local_addr()?); + let router = Router::new() + .route("/v1/responses", get(luna_websocket).post(parent_response)) + .with_state(Arc::clone(&responses_state)); + let responses_server = tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + let (mcp_server_url, mcp_server_handle) = start_mcp_server().await?; + + let codex_home = TempDir::new()?; + MockResponsesConfig::new(&responses_url) + .with_provider_config("supports_websockets = false") + .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\"" + )) + .enable_feature(Feature::GuardianV2) + .enable_feature(Feature::GuardianApproval) + .write(codex_home.path())?; + 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_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(), + input: vec![UserInput::Text { + text: USER_CONTEXT.to_owned(), + text_elements: Vec::new(), + }], + approval_policy: Some(AskForApproval::OnRequest), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }) + .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); + + let luna_request = timeout(TIMEOUT, async { + loop { + if let Some(request) = responses_state + .luna_requests + .lock() + .expect("Luna request lock should not be poisoned") + .first() + .cloned() + { + return request; + } + tokio::task::yield_now().await; + } + }) + .await?; + assert_eq!( + luna_request["prompt_cache_key"], + format!("guardian-v2:{thread_id}") + ); + assert!( + luna_request["input"] + .as_array() + .expect("Luna input should be an array") + .iter() + .any(|item| { + item["content"].as_array().is_some_and(|content| { + content.iter().any(|entry| { + entry["text"] + .as_str() + .is_some_and(|text| text.contains(USER_CONTEXT)) + }) + }) + }) + ); + + 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::(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??; + responses_state.allow_guardian_review.notify_one(); + responses_state.allow_luna.notify_one(); + timeout( + TIMEOUT, + app_server.read_stream_until_notification_message("turn/completed"), + ) + .await??; + assert_eq!( + responses_state.guardian_reviews.load(Ordering::SeqCst), + expected_guardian_reviews + ); + + mcp_server_handle.abort(); + responses_server.abort(); + Ok(()) +} + +#[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 +} + +#[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 +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 5facaf5a1e..373eb271de 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -33,6 +33,7 @@ mod external_agent_config; mod external_agent_import_sync; mod fs; mod git_attribution; +mod guardian_v2; mod hooks_list; mod host_skills; mod imagegen_extension;