diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs index 3c8b5d4eb5..e4051a8b2c 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -59,6 +59,7 @@ 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; +use super::mcp_tool::start_mcp_server_with_tools; const TIMEOUT: Duration = Duration::from_secs(30); const MODEL: &str = "mock-model"; @@ -147,6 +148,7 @@ struct MockResponsesState { review_outcome: ReviewOutcome, transcript_content: TranscriptContent, mcp_server_name: Option<&'static str>, + mcp_tool_sequence: Option<&'static [&'static str]>, root_worker: bool, root_user_restriction: bool, root_user_input_restriction: bool, @@ -429,7 +431,7 @@ async fn parent_response( ] } else if state.user_input_restriction && request_number == 1 { user_input_request_events() - } else if request_number < 2 + } else if request_number < state.mcp_tool_sequence.map_or(/*default*/ 2, <[_]>::len) || state.user_input_restriction && request_number == 2 || (state.root_worker || state.root_user_restriction) && request_number == 3 { @@ -450,7 +452,11 @@ async fn parent_response( responses::ev_function_call_with_namespace( &call_id, &format!("mcp__{}", state.mcp_server_name.unwrap_or(TEST_SERVER_NAME)), - TEST_TOOL_NAME, + state + .mcp_tool_sequence + .and_then(|tools| tools.get(request_number)) + .copied() + .unwrap_or(TEST_TOOL_NAME), &arguments, ), responses::ev_completed(&call_id), @@ -1606,6 +1612,103 @@ async fn guardian_v2_computer_use_only_scopes_classification_and_fast_reviews( .await } +#[test_case("node_repl", &["js", "js"]; "browser startup")] +#[test_case("cua_repl", &["js", "js"]; "computer use startup")] +#[test_case("node_repl", &["js_reset", "js", "js"]; "browser reset before execution")] +#[test_case("cua_repl", &["js_reset", "js", "js"]; "computer use reset before execution")] +#[test_case("node_repl", &["js_add_node_module_dir", "js", "js"]; "browser setup before execution")] +#[test_case("cua_repl", &["js_add_node_module_dir", "js", "js"]; "computer use setup before execution")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn first_cua_review_does_not_wait_for_initial_score( + server_name: &'static str, + tool_sequence: &'static [&'static str], +) -> Result<()> { + skip_if_no_network!(Ok(())); + let state = Arc::new(MockResponsesState { + mcp_server_name: Some(server_name), + mcp_tool_sequence: Some(tool_sequence), + ..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(&state)); + let responses_server = tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + let (mcp_url, mcp_server) = start_mcp_server_with_tools( + &["js", "js_reset", "js_add_node_module_dir"], + /*sensitive_action*/ None, + ) + .await?; + let codex_home = TempDir::new()?; + MockResponsesConfig::new(&responses_url) + .with_model(MODEL) + .with_provider_config("supports_websockets = false") + .with_approval_policy("on-request") + .with_root_config("approvals_reviewer = \"auto_review\"") + .with_extra_config(&format!( + "[mcp_servers.{server_name}]\nurl = \"{mcp_url}/mcp\"\ndefault_tools_approval_mode = \"auto\"\n\n[features.guardianv2]\nenabled = true" + )) + .enable_feature(Feature::GuardianApproval) + .write(codex_home.path())?; + let config = load_default_config_for_test(&codex_home).await; + let mut model_info = codex_core::test_support::construct_model_info_offline(MODEL, &config); + model_info.node_repl_auto_review_required = true; + write_models_cache_with_models(codex_home.path(), vec![model_info])?; + 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::default()) + .await? + .thread; + let 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(), + }], + ..Default::default() + }) + .await?; + let _: TurnStartResponse = timeout(TIMEOUT, app_server.read_response(request_id)).await??; + + // Leave the classifier pending: the first execution proceeds, the second waits for review. + let _: ItemGuardianApprovalReviewStartedNotification = timeout( + TIMEOUT, + app_server.read_notification("item/autoApprovalReview/started"), + ) + .await??; + wait_for_luna_request(&state, /*index*/ 0).await?; + wait_for_guardian_reviews(&state, /*expected*/ 1).await?; + let reviewed_call = format!("guardian-{}", tool_sequence.len() - 1); + assert!( + state + .guardian_requests + .lock() + .expect("Guardian request lock")[0] + .to_string() + .contains(&reviewed_call) + ); + assert_eq!( + state.parent_requests.load(Ordering::SeqCst), + tool_sequence.len() + ); + state.allow_guardian_review.notify_one(); + let completed: TurnCompletedNotification = + timeout(TIMEOUT, app_server.read_notification("turn/completed")).await??; + assert_eq!(completed.thread_id, thread.id); + assert_eq!(state.guardian_reviews.load(Ordering::SeqCst), 1); + app_server.shutdown_gracefully().await?; + mcp_server.abort(); + responses_server.abort(); + Ok(()) +} + #[test_case("node_repl", GuardianRisk::Low, None; "browser low risk")] #[test_case("cua_repl", GuardianRisk::Low, None; "computer use low risk")] #[test_case("node_repl", GuardianRisk::Low, Some(false); "browser low risk sensitive action false")] diff --git a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs index 6b67fb195e..210ce82351 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs @@ -1265,9 +1265,10 @@ async fn mcp_tool_call_hint_survives_mid_call_thread_read_and_resume() -> Result Ok(()) } -#[derive(Clone, Default)] +#[derive(Clone)] struct ToolAppsMcpServer { sensitive_action: Option, + tool_names: &'static [&'static str], } impl ServerHandler for ToolAppsMcpServer { @@ -1300,14 +1301,22 @@ impl ServerHandler for ToolAppsMcpServer { })) .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; - let mut tool = Tool::new( - Cow::Borrowed(TEST_TOOL_NAME), - Cow::Borrowed("Echo a message."), - Arc::new(input_schema), - ); - tool.annotations = Some(ToolAnnotations::new().read_only(true)); + let input_schema = Arc::new(input_schema); + let tools = self + .tool_names + .iter() + .map(|&name| { + let mut tool = Tool::new( + Cow::Borrowed(name), + Cow::Borrowed("Echo a message."), + Arc::clone(&input_schema), + ); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + tool + }) + .collect(); - Ok(ListToolsResult::with_all_items(vec![tool])) + Ok(ListToolsResult::with_all_items(tools)) } async fn call_tool( @@ -1315,7 +1324,10 @@ impl ServerHandler for ToolAppsMcpServer { request: CallToolRequestParams, context: RequestContext, ) -> Result { - assert_eq!(request.name.as_ref(), TEST_TOOL_NAME); + assert!(self.tool_names.contains(&request.name.as_ref())); + if matches!(request.name.as_ref(), "js_reset" | "js_add_node_module_dir") { + return Ok(CallToolResult::success(vec![ContentBlock::text("setup complete")]).into()); + } let message = request .arguments .as_ref() @@ -1349,10 +1361,14 @@ impl ServerHandler for ToolAppsMcpServer { "codex_request_type": "approval_request", "codex_approval_kind": "mcp_tool_call", "codex_strict_auto_review": true, - "tool_name": TEST_TOOL_NAME, + "tool_name": request.name, "tool_params": request.arguments, "x-codex-turn-metadata": turn_metadata, }); + if request.name == "js" { + // Match the execution approval emitted by the real Node REPL server. + approval_meta["connector_id"] = json!("node_repl"); + } if let Some(sensitive_action) = self.sensitive_action { approval_meta["codex_sensitive_action"] = json!(sensitive_action); } @@ -1491,11 +1507,23 @@ impl ServerHandler for ToolAppsMcpServer { pub(super) async fn start_mcp_server( sensitive_action: Option, +) -> Result<(String, JoinHandle<()>)> { + start_mcp_server_with_tools(&[TEST_TOOL_NAME], sensitive_action).await +} + +pub(super) async fn start_mcp_server_with_tools( + tool_names: &'static [&'static str], + sensitive_action: Option, ) -> Result<(String, JoinHandle<()>)> { let listener = TcpListener::bind("127.0.0.1:0").await?; let addr = listener.local_addr()?; let mcp_service = StreamableHttpService::new( - move || Ok(ToolAppsMcpServer { sensitive_action }), + move || { + Ok(ToolAppsMcpServer { + sensitive_action, + tool_names, + }) + }, Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), ); diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index 97d5000dab..d3a4b3a3c8 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -113,6 +113,8 @@ enum ClassificationOutcome { #[derive(Default)] struct GuardianV2ScoreProgress { latest_tool_call: AtomicUsize, + // Setup and reset calls must not consume the first JS execution allowance. + js_executions: AtomicUsize, latest_scored_tool_call: AtomicUsize, latest_failed_tool_call: AtomicUsize, // Serialize successful score publication with its authorization metadata. @@ -261,6 +263,26 @@ impl ApprovalReviewContributor for GuardianV2Extension { record_fast_decision(extension_metrics.as_deref(), "deferred", "out_of_scope"); return None; } + // The first REPL execution never waits for synchronous Guardian review. + // The async classifier still runs, and later calls use the normal policy. + if action.get("tool_name").and_then(serde_json::Value::as_str) == Some("js") + && action + .get("connector_id") + .and_then(serde_json::Value::as_str) + == Some("node_repl") + && thread_store + .get::()? + .js_executions + .load(Ordering::Acquire) + == 1 + { + record_fast_decision( + extension_metrics.as_deref(), + "approved", + "initial_cua_call", + ); + return Some(ReviewDecision::Approved); + } } else if thread_store.get::().is_some() { let manager = self.thread_manager.upgrade()?; let thread_id = ThreadId::from_string(thread_store.level_id()).ok()?; @@ -400,6 +422,14 @@ impl GuardianV2Extension { } return; } + if input.mcp_tool.is_some_and(|tool| { + let info = tool.tool_info(); + is_node_repl_backed_server(&info.server_name) && info.tool.name == "js" + }) { + score_progress + .js_executions + .fetch_add(/*val*/ 1, Ordering::Relaxed); + } let metrics = score_progress.metrics.clone(); let sampled_at = SystemTime::now(); let tool_call_index = score_progress