diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 73cd869fe8..5ca8c44a08 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -786,6 +786,7 @@ async fn network_approval_uses_published_task_authority_within_same_turn( exec_policy_hint: None, execution_id: None, disconnect: None, + cancellation: None, }, ); tokio::pin!(decision); diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index e82780c863..a3e870a92b 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -17,6 +17,8 @@ use codex_network_proxy::NetworkPolicyDecider; use codex_network_proxy::NetworkPolicyRequest; use codex_network_proxy::NetworkProtocol; use codex_network_proxy::NetworkProxy; +use codex_network_proxy::NetworkRequestCancellation; +use codex_network_proxy::NetworkRequestCancellationReason; use codex_network_proxy::NetworkRequestDisconnect; use codex_protocol::approvals::NetworkApprovalContext; use codex_protocol::approvals::NetworkApprovalProtocol; @@ -280,6 +282,7 @@ struct PendingHostApprovalOwner<'a> { pending: Arc, execution_cancellation: Option, disconnect: Option, + cancellation: Option, decision_on_drop: PendingApprovalDecision, completed: bool, } @@ -297,6 +300,7 @@ impl<'a> PendingHostApprovalOwner<'a> { pending, execution_cancellation, disconnect: None, + cancellation: None, decision_on_drop: PendingApprovalDecision::Deny, completed: false, } @@ -342,6 +346,17 @@ impl<'a> PendingHostApprovalOwner<'a> { impl Drop for PendingHostApprovalOwner<'_> { fn drop(&mut self) { if !self.completed { + if self + .cancellation + .as_ref() + .and_then(NetworkRequestCancellation::reason) + == Some(NetworkRequestCancellationReason::ProcessFinished) + { + // Normal process cleanup is not a new review failure. Any earlier + // denial is already recorded in the call outcome. + self.publish_and_remove(self.decision_on_drop); + return; + } if matches!(self.decision_on_drop, PendingApprovalDecision::Deny) && let Some(registration_id) = self.key.execution_id.as_deref() && let Some(elapsed) = self @@ -517,20 +532,27 @@ impl NetworkApprovalService { } fn record_call_outcome(&self, registration_id: &str, outcome: String) { + if let Some(cancellation_token) = self.store_call_outcome(registration_id, outcome) { + cancellation_token.cancel(); + } + } + + fn store_call_outcome( + &self, + registration_id: &str, + outcome: String, + ) -> Option { let mut calls = self .calls .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let Some(call) = calls.active_calls.get(registration_id).cloned() else { - return; - }; + let call = calls.active_calls.get(registration_id)?.clone(); // Explicit network-review outcomes replace generic blocked-request fallbacks. calls .call_outcomes .insert(registration_id.to_string(), outcome); - drop(calls); - call.cancellation_token.cancel(); + Some(call.cancellation_token.clone()) } fn record_call_outcome_if_absent(&self, registration_id: &str, outcome: String) { @@ -676,6 +698,7 @@ impl NetworkApprovalService { .map(|call| call.cancellation_token.clone()), ); pending_owner.disconnect = request.disconnect.clone(); + pending_owner.cancellation = request.cancellation.clone(); let permission_profile = owner_call .as_ref() @@ -814,6 +837,18 @@ impl NetworkApprovalService { } }; + if matches!( + &approval_decision, + ReviewDecision::NetworkPolicyAmendment { network_policy_amendment } + if network_policy_amendment.action == NetworkPolicyRuleAction::Deny + ) && let Some(owner_call) = owner_call.as_ref() + { + // Preserve the denial before waiting for the policy commit. Defer + // cancellation until after saving so it cannot interrupt persistence. + let _ = self + .store_call_outcome(&owner_call.registration_id, "rejected by user".to_string()); + } + let _session_policy_commit_guard = if matches!( &approval_decision, ReviewDecision::Approved diff --git a/codex-rs/core/tests/suite/network_approval_completion.rs b/codex-rs/core/tests/suite/network_approval_completion.rs new file mode 100644 index 0000000000..c1299b0677 --- /dev/null +++ b/codex-rs/core/tests/suite/network_approval_completion.rs @@ -0,0 +1,294 @@ +//! Exercises process cleanup while a remote network review is still pending. + +use super::PushedExecScenario; +use super::accept_initialized_exec_server; +use super::read_exec_server_json; +use super::respond_environment_info; +use super::send_exec_server_json; +use anyhow::Context; +use anyhow::Result; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_core::TurnInputRequest; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::items::CommandExecutionStatus; +use codex_protocol::items::TurnItem; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::UserInput; +use core_test_support::managed_network_requirements_loader; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::mount_response_once_match; +use core_test_support::responses::mount_sse_once_match; +use core_test_support::responses::sse; +use core_test_support::responses::sse_response; +use core_test_support::responses::start_mock_server; +use core_test_support::startup::STARTUP_TIMEOUT; +use core_test_support::test_codex::test_codex; +use core_test_support::test_codex::turn_permission_fields; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::fs; +use std::time::Duration; +use tokio::net::TcpListener; +use tokio::net::TcpStream; +use tokio::time::timeout; +use tokio_tungstenite::WebSocketStream; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn completed_remote_process_withdraws_pending_network_review() -> Result<()> { + let server = start_mock_server().await; + let listener = TcpListener::bind("127.0.0.1:0").await?; + let exec_server_url = format!("ws://{}", listener.local_addr()?); + let call_id = "pending-network-at-exit"; + let parent = mount_sse_once_match( + &server, + |request: &wiremock::Request| { + let body = review_body(request); + body.pointer("/client_metadata/x-openai-subagent") != Some(&json!("guardian")) + && !body.to_string().contains("pending-network-at-exit") + }, + sse(vec![ + ev_function_call( + call_id, + "exec_command", + &json!({"cmd": "build-site", "yield_time_ms": 10_000}).to_string(), + ), + ev_completed("start"), + ]), + ) + .await; + let pending_review = mount_response_once_match( + &server, + |request: &wiremock::Request| { + review_body(request).pointer("/client_metadata/x-openai-subagent") + == Some(&json!("guardian")) + }, + sse_response(sse(vec![ev_completed("review")])) + .set_delay(Duration::from_secs(/*secs*/ 60)), + ) + .await; + let final_response = mount_sse_once_match( + &server, + |request: &wiremock::Request| { + let body = review_body(request); + body.pointer("/client_metadata/x-openai-subagent") != Some(&json!("guardian")) + && body.to_string().contains("pending-network-at-exit") + }, + sse(vec![ + ev_assistant_message("done", "done"), + ev_completed("done"), + ]), + ) + .await; + let exec_server = tokio::spawn(async move { + let mut websocket = accept_initialized_exec_server(listener).await; + let start = loop { + let request = read_exec_server_json(&mut websocket, STARTUP_TIMEOUT).await; + if request["method"] == "process/start" { + break request; + } + respond_setup_request(&mut websocket, &request).await; + }; + assert!( + start["params"]["networkProxy"]["policyDecisionTimeoutMs"] + .as_u64() + .is_some() + ); + let process_id = &start["params"]["processId"]; + send_exec_server_json( + &mut websocket, + json!({"id": start["id"], "result": {"processId": process_id}}), + ) + .await; + send_exec_server_json(&mut websocket, json!({ + "id": 1, "method": "network/policyRequest", + "params": {"processId": process_id, "request": {"protocol": "https_connect", "host": "review.example", "port": 443}} + })).await; + timeout(Duration::from_secs(/*secs*/ 10), async { + while !pending_review.requests().iter().any(|request| { + request + .body_json() + .pointer("/client_metadata/x-openai-subagent") + == Some(&json!("guardian")) + }) { + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis(/*millis*/ 10)) => {}, + request = read_exec_server_json(&mut websocket, Duration::from_secs(/*secs*/ 10)) => { + respond_setup_request(&mut websocket, &request).await; + } + } + } + }) + .await + .expect("Guardian must start before process completion"); + for message in [ + json!({"method": "process/output", "params": {"processId": process_id, "seq": 1, "stream": "stdout", "chunk": BASE64_STANDARD.encode("build complete\n")}}), + json!({"method": "process/exited", "params": {"processId": process_id, "seq": 2, "exitCode": 0, "sandboxDenied": false}}), + json!({"method": "process/closed", "params": {"processId": process_id, "seq": 3}}), + ] { + send_exec_server_json(&mut websocket, message).await; + } + loop { + let request = + read_exec_server_json(&mut websocket, Duration::from_secs(/*secs*/ 10)).await; + match request["method"].as_str() { + Some("process/terminate") => { + send_exec_server_json( + &mut websocket, + json!({"id": request["id"], "result": {"running": false}}), + ) + .await; + break; + } + Some("process/read") => { + send_exec_server_json(&mut websocket, json!({"id": request["id"], "result": {"chunks": [], "nextSeq": 4, "exited": true, "exitCode": 0, "closed": true, "failure": null, "sandboxDenied": false}})).await; + } + None => panic!("withdrawn request must not be approved: {request}"), + _ => respond_setup_request(&mut websocket, &request).await, + } + } + }); + let test = test_codex() + .with_exec_server_url(exec_server_url) + .with_cloud_config_bundle(managed_network_requirements_loader()) + .with_pre_build_hook(|home| { + fs::write( + home.join("config.toml"), + r#"default_permissions = "workspace" +[permissions.workspace.filesystem] +":minimal" = "read" +[permissions.workspace.network] +enabled = true +mode = "full" +allow_local_binding = true +"#, + ) + .expect("managed network config"); + }) + .with_config(|config| { + config.project_doc_max_bytes = 0; + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + #[cfg(windows)] + config.set_windows_sandbox_enabled(/*value*/ true); + }) + // This test supplies its own fake executor and must not select a CI executor. + .build(&server) + .await?; + let (sandbox_policy, permission_profile) = turn_permission_fields( + test.session_configured.permission_profile.clone(), + test.config.cwd.as_path(), + ); + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "build the site".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + approval_policy: Some(AskForApproval::OnRequest), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + sandbox_policy: Some(sandbox_policy), + permission_profile, + ..Default::default() + }), + ) + .await?; + let completed = timeout(Duration::from_secs(/*secs*/ 15), async { + loop { + match test.codex.next_event().await?.msg { + EventMsg::ItemCompleted(event) => { + if let TurnItem::CommandExecution(item) = event.item { + return Ok::<_, anyhow::Error>(item); + } + } + EventMsg::ExecApprovalRequest(_) => anyhow::bail!("unexpected user approval"), + _ => {} + } + } + }) + .await + .context("command should complete")??; + assert_eq!( + ( + completed.status, + completed.exit_code, + completed.aggregated_output.as_deref() + ), + ( + CommandExecutionStatus::Completed, + Some(0), + Some("build complete\n") + ) + ); + timeout(Duration::from_secs(/*secs*/ 10), async { + while final_response.function_call_output_text(call_id).is_none() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await?; + let output = final_response + .function_call_output_text(call_id) + .context("tool output")?; + assert!(output.contains("Process exited with code 0"), "{output}"); + assert!(output.contains("build complete"), "{output}"); + assert_eq!(parent.requests().len(), 1); + timeout(Duration::from_secs(/*secs*/ 10), exec_server).await??; + Ok(()) +} + +fn review_body(request: &wiremock::Request) -> Value { + let bytes = if request + .headers + .get("content-encoding") + .is_some_and(|value| value == "zstd") + { + zstd::stream::decode_all(std::io::Cursor::new(&request.body)).expect("decode request") + } else { + request.body.clone() + }; + serde_json::from_slice(&bytes).expect("request JSON") +} + +async fn respond_setup_request(websocket: &mut WebSocketStream, request: &Value) { + match request["method"].as_str() { + Some("environment/info") => { + respond_environment_info(websocket, &request["id"], PushedExecScenario::Complete).await; + } + Some("fs/getMetadata" | "fs/readFile" | "fs/readDirectory") => { + send_exec_server_json( + websocket, + json!({ + "id": request["id"], + "error": {"code": -32004, "message": "not found"} + }), + ) + .await; + } + Some("fs/canonicalize") => { + send_exec_server_json( + websocket, + json!({ + "id": request["id"], + "result": {"path": request["params"]["path"]} + }), + ) + .await; + } + Some("fs/walk") => { + send_exec_server_json( + websocket, + json!({ + "id": request["id"], + "result": {"entries": [], "errors": [], "truncated": false} + }), + ) + .await; + } + method => panic!("unexpected setup request: {method:?}"), + } +} diff --git a/codex-rs/core/tests/suite/unified_exec_process_events.rs b/codex-rs/core/tests/suite/unified_exec_process_events.rs index 75e0b16e10..4e2210e851 100644 --- a/codex-rs/core/tests/suite/unified_exec_process_events.rs +++ b/codex-rs/core/tests/suite/unified_exec_process_events.rs @@ -1075,3 +1075,6 @@ timeout = 900 Ok(()) } + +#[path = "network_approval_completion.rs"] +mod network_approval_completion; diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index f832d9be13..d05a0c14b0 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -13,6 +13,8 @@ use arc_swap::ArcSwapOption; use codex_exec_server_protocol::JSONRPCNotification; use codex_network_proxy::NetworkPolicyDecider; use codex_network_proxy::NetworkProxyAuditMetadata; +use codex_network_proxy::NetworkRequestCancellation; +use codex_network_proxy::NetworkRequestCancellationReason; use futures::FutureExt; use futures::future::BoxFuture; use serde_json::Value; @@ -211,6 +213,7 @@ pub(crate) struct SessionState { struct NetworkPolicyState { controller: ArcSwapOption, cancelled: CancellationToken, + cancellation: NetworkRequestCancellation, audit: Option, } @@ -877,6 +880,13 @@ impl ExecServerClient { &self, process_id: &ProcessId, ) -> Result { + // A close notification may arrive before the termination response. + if let Some(session) = self.inner.get_session(process_id) { + session + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessCancelled); + } self.call_for_cleanup( EXEC_TERMINATE_METHOD, &TerminateParams { @@ -1318,6 +1328,7 @@ impl SessionState { network_policy: NetworkPolicyState { controller: ArcSwapOption::empty(), cancelled: CancellationToken::new(), + cancellation: NetworkRequestCancellation::default(), audit: None, }, } @@ -1381,6 +1392,11 @@ impl SessionState { let is_closed = matches!(&event, ExecProcessEvent::Closed { .. }); ordered_events.closed_published |= is_closed; published_closed |= is_closed; + if is_closed && ordered_events.exit_published { + self.network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessFinished); + } self.events.publish(event); } published_closed @@ -1581,6 +1597,10 @@ impl Session { } pub(crate) fn cancel_network_policy_decisions(&self) { + self.state + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessCancelled); self.state.network_policy.cancelled.cancel(); } @@ -1638,6 +1658,10 @@ impl Inner { let mut next_sessions = sessions.as_ref().clone(); next_sessions.remove(process_id); self.sessions.store(Arc::new(next_sessions)); + expected + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessCancelled); expected.network_policy.cancelled.cancel(); expected.network_policy.controller.store(None); } @@ -1678,6 +1702,10 @@ fn fail_all_sessions(inner: &Arc, message: String) { let sessions = inner.take_all_sessions(); for (_, session) in sessions { + session + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ConnectionClosed); session.network_policy.cancelled.cancel(); session.network_policy.controller.store(None); // Sessions synthesize a closed read response and emit a pushed Failed diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 8e9a8c77b2..6b4cf61df1 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -9,6 +9,8 @@ use codex_network_proxy::NetworkDecision; use codex_network_proxy::NetworkPolicyDecision; use codex_network_proxy::NetworkPolicyRequest; use codex_network_proxy::NetworkProtocol; +use codex_network_proxy::NetworkRequestCancellation; +use codex_network_proxy::NetworkRequestCancellationReason; use serde_json::Value; use tokio::sync::mpsc; use tokio::time::Instant; @@ -526,6 +528,10 @@ impl Inner { Ok(true) => self.remove_session_if(process_id, session), Ok(false) => {} Err(error) => { + session + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessCancelled); let terminated: Result = rpc_client .call_for_cleanup( EXEC_TERMINATE_METHOD, @@ -688,6 +694,10 @@ impl ExecServerClient { let process_cancelled = session .as_ref() .map(|session| session.network_policy.cancelled.clone()); + let process_cancellation = session + .as_ref() + .map(|session| session.network_policy.cancellation.clone()); + let cancellation = NetworkRequestCancellation::default(); let expected_session = session.as_ref().map(Arc::downgrade); let policy_request = (process_id_valid && host_valid).then_some(NetworkPolicyRequest { @@ -712,6 +722,7 @@ impl ExecServerClient { exec_policy_hint: None, execution_id: None, disconnect: None, + cancellation: Some(cancellation.clone()), }); let inner = Arc::downgrade(&inner); let rpc_client = Arc::downgrade(&rpc_client); @@ -721,18 +732,26 @@ impl ExecServerClient { let _request_guard = request_guard; let decision = match (controller, policy_request, process_cancelled) { (Some(controller), Some(request), Some(process_cancelled)) => { - // Core's pending-approval guard makes dropping this - // future on process removal or deadline fail closed. + // Keep the decision future outside select/timeout so its + // guard sees the cancellation cause before it is dropped. + let mut decision = controller.decider.decide(request); tokio::select! { biased; - _ = connection_cancelled.cancelled() => return, + _ = connection_cancelled.cancelled() => { + cancellation.record(NetworkRequestCancellationReason::ConnectionClosed); + return; + }, _ = process_cancelled.cancelled() => { + cancellation.record(process_cancellation.as_ref() + .and_then(NetworkRequestCancellation::reason) + .unwrap_or(NetworkRequestCancellationReason::ProcessCancelled)); NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON) } - decision = timeout( + result = timeout( controller.timeout, - controller.decider.decide(request), - ) => decision.unwrap_or_else(|_| { + &mut decision, + ) => result.unwrap_or_else(|_| { + cancellation.record(NetworkRequestCancellationReason::TimedOut); NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON) }), } diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index 5ffb23046d..b2affe0b8c 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -17,6 +17,7 @@ mod process_log_metadata; mod proxy; mod reasons; mod remote_config; +mod request_cancellation; mod request_disconnect; mod responses; mod runtime; @@ -89,6 +90,8 @@ pub use proxy::proxy_url_env_value; pub use proxy::strip_managed_proxy_env; pub use remote_config::RemoteNetworkProxyConfig; pub use remote_config::RemoteNetworkProxyLaunchConfig; +pub use request_cancellation::NetworkRequestCancellation; +pub use request_cancellation::NetworkRequestCancellationReason; pub use request_disconnect::NetworkRequestDisconnect; pub use runtime::BlockedRequest; pub use runtime::BlockedRequestArgs; diff --git a/codex-rs/network-proxy/src/network_policy.rs b/codex-rs/network-proxy/src/network_policy.rs index 1839c371e4..a47e1249b9 100644 --- a/codex-rs/network-proxy/src/network_policy.rs +++ b/codex-rs/network-proxy/src/network_policy.rs @@ -1,4 +1,5 @@ use crate::reasons::REASON_POLICY_DENIED; +use crate::request_cancellation::NetworkRequestCancellation; use crate::request_disconnect::NetworkRequestDisconnect; use crate::runtime::HostBlockDecision; use crate::runtime::HostBlockReason; @@ -109,6 +110,8 @@ pub struct NetworkPolicyRequest { pub execution_id: Option, /// Present only when the local HTTP transport can identify an abandoned request. pub disconnect: Option, + /// Controller-owned cause, published before an abandoned decision future is dropped. + pub cancellation: Option, } pub struct NetworkPolicyRequestArgs { @@ -145,6 +148,7 @@ impl NetworkPolicyRequest { exec_policy_hint, execution_id: None, disconnect: None, + cancellation: None, } } } diff --git a/codex-rs/network-proxy/src/request_cancellation.rs b/codex-rs/network-proxy/src/request_cancellation.rs new file mode 100644 index 0000000000..90d481acf1 --- /dev/null +++ b/codex-rs/network-proxy/src/request_cancellation.rs @@ -0,0 +1,32 @@ +//! Records why a policy request was withdrawn before its decision future is dropped. + +use std::sync::Arc; +use std::sync::OnceLock; + +/// The controller's first known reason for withdrawing a policy request. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum NetworkRequestCancellationReason { + /// The process exited and its output streams closed normally. + ProcessFinished, + /// The process was explicitly terminated or its handle was abandoned. + ProcessCancelled, + /// The connection to the executor was lost. + ConnectionClosed, + /// The policy decision deadline expired. + TimedOut, +} + +/// Shared, in-process metadata; recording a reason does not authorize network access. +#[derive(Clone, Debug, Default)] +pub struct NetworkRequestCancellation(Arc>); + +impl NetworkRequestCancellation { + pub fn reason(&self) -> Option { + self.0.get().copied() + } + + /// Publish before dropping the decision future. Cleanup cannot replace an earlier cause. + pub fn record(&self, reason: NetworkRequestCancellationReason) { + let _ = self.0.set(reason); + } +}