diff --git a/codex-rs/core/tests/suite/network_approval.rs b/codex-rs/core/tests/suite/network_approval.rs index 95fec8a8a2..c98a511f46 100644 --- a/codex-rs/core/tests/suite/network_approval.rs +++ b/codex-rs/core/tests/suite/network_approval.rs @@ -2183,31 +2183,75 @@ async fn guardian_receives_exact_trigger_for_single_network_request() -> Result< let server = start_mock_server().await; let test = managed_network_unified_exec_test(&server).await?; let command = "python3 -c \"import urllib.request; opener = urllib.request.build_opener(urllib.request.ProxyHandler()); print('OK:' + opener.open('http://1.1.1.1', timeout=10).read().decode(errors='replace'))\"".to_string(); - let responses = mount_sse_sequence( + let call_id = "exec-network-single"; + mount_sse_once_match( &server, - vec![ - sse(vec![ - ev_response_created("resp-network-single"), - ev_function_call( - "exec-network-single", - "exec_command", - &serde_json::to_string(&network_exec_args(&command))?, - ), - ev_completed("resp-network-single"), - ]), - sse(vec![ - ev_response_created("resp-network-guardian"), - ev_assistant_message("msg-network-guardian", r#"{"outcome":"deny"}"#), - ev_completed("resp-network-guardian"), - ]), - sse(vec![ - ev_response_created("resp-network-done"), - ev_assistant_message("msg-network-done", "done"), - ev_completed("resp-network-done"), - ]), - ], + move |request: &wiremock::Request| { + !is_guardian_request(request) && !request_body_contains(request, call_id) + }, + sse(vec![ + ev_response_created("resp-network-single"), + ev_function_call( + call_id, + "exec_command", + &serde_json::to_string(&network_exec_args(&command))?, + ), + ev_completed("resp-network-single"), + ]), ) .await; + let guardian = mount_sse_once_match( + &server, + is_guardian_request, + sse(vec![ + ev_response_created("resp-network-guardian"), + ev_assistant_message("msg-network-guardian", r#"{"outcome":"deny"}"#), + ev_completed("resp-network-guardian"), + ]), + ) + .await; + + // The command can yield before its network request reaches Guardian. Keep the + // parent on that process until it exits, without consuming Guardian's reply. + let _parent = wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/v1/responses")) + .and(move |request: &wiremock::Request| { + !is_guardian_request(request) && request_body_contains(request, call_id) + }) + .respond_with(move |request: &wiremock::Request| { + let body: Value = serde_json::from_slice( + &decoded_request_body(request).expect("decode parent request"), + ) + .expect("parse parent request"); + let input = body["input"].as_array().expect("parent input"); + let output = input + .iter() + .rev() + .find(|item| item["type"] == "function_call_output") + .and_then(|item| item["output"].as_str()) + .expect("network command or poll output"); + let response_id = format!("{call_id}-{}", input.len()); + let event = if let Some(session_id) = output + .lines() + .find_map(|line| line.strip_prefix("Process running with session ID ")) + { + ev_function_call( + &response_id, + "write_stdin", + &json!({ + "session_id": session_id.parse::().expect("session id"), + "chars": "", + "yield_time_ms": 1_000, + }) + .to_string(), + ) + } else { + ev_assistant_message(&response_id, "done") + }; + sse_response(sse(vec![event, ev_completed(&response_id)])) + }) + .mount_as_scoped(&server) + .await; submit_managed_network_turn( &test, @@ -2217,11 +2261,13 @@ async fn guardian_receives_exact_trigger_for_single_network_request() -> Result< AskForApproval::OnRequest, ) .await?; - wait_for_turn_complete(&test).await; + tokio::time::timeout(Duration::from_secs(30), wait_for_turn_complete(&test)) + .await + .context("timed out waiting for the network process and Guardian review to finish")?; assert_eq!( - guardian_network_triggers(&[&responses])?, - vec![("exec-network-single".to_string(), command)] + guardian_network_triggers(&[&guardian])?, + vec![(call_id.to_string(), command)] ); Ok(())