diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index fdb3f6ef0b..d32e849bfa 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -14,6 +14,7 @@ use codex_network_proxy::NetworkPolicyDecider; use codex_network_proxy::NetworkPolicyRequest; use codex_network_proxy::NetworkProtocol; use codex_network_proxy::NetworkProxy; +use codex_network_proxy::NetworkRequestDisconnect; use codex_protocol::approvals::NetworkApprovalContext; use codex_protocol::approvals::NetworkApprovalProtocol; use codex_protocol::approvals::NetworkPolicyRuleAction; @@ -270,7 +271,8 @@ struct NetworkApprovalCallState { } pub(crate) struct NetworkApprovalService { - calls: Mutex, + // Pending-request cleanup must publish a tool outcome synchronously before cancelling it. + calls: SyncMutex, // Owner cleanup runs from Drop, so this lock cannot require an async task. pending_host_approvals: SyncMutex>>, // Keep persisted session policy and the in-memory approval caches in the same order. @@ -288,6 +290,7 @@ struct PendingHostApprovalOwner<'a> { key: PendingHostApprovalKey, pending: Arc, execution_cancellation: Option, + disconnect: Option, decision_on_drop: PendingApprovalDecision, completed: bool, } @@ -304,6 +307,7 @@ impl<'a> PendingHostApprovalOwner<'a> { key, pending, execution_cancellation, + disconnect: None, decision_on_drop: PendingApprovalDecision::Deny, completed: false, } @@ -349,6 +353,18 @@ impl<'a> PendingHostApprovalOwner<'a> { impl Drop for PendingHostApprovalOwner<'_> { fn drop(&mut self) { if !self.completed { + if matches!(self.decision_on_drop, PendingApprovalDecision::Deny) + && let Some(registration_id) = self.key.execution_id.as_deref() + && let Some(elapsed) = self + .disconnect + .as_ref() + .and_then(NetworkRequestDisconnect::elapsed) + { + let elapsed_ms = elapsed.as_millis(); + self.service.record_call_outcome_if_absent(registration_id, format!( + "Network request disconnected after {elapsed_ms} ms, before approval could complete" + )); + } self.cancel_execution_if_denied(self.decision_on_drop); self.publish_and_remove(self.decision_on_drop); } @@ -358,7 +374,7 @@ impl Drop for PendingHostApprovalOwner<'_> { impl Default for NetworkApprovalService { fn default() -> Self { Self { - calls: Mutex::new(NetworkApprovalCallState::default()), + calls: SyncMutex::new(NetworkApprovalCallState::default()), pending_host_approvals: SyncMutex::new(HashMap::new()), session_policy_commit_lock: Mutex::new(()), session_approved_hosts: Mutex::new(HashSet::new()), @@ -383,7 +399,10 @@ impl NetworkApprovalService { } async fn register_call(&self, call: ActiveNetworkApprovalCall) { - let mut calls = self.calls.lock().await; + let mut calls = self + .calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); calls .active_calls .insert(call.registration_id.clone(), Arc::new(call)); @@ -394,7 +413,10 @@ impl NetworkApprovalService { } async fn resolve_single_active_call(&self) -> Option> { - let calls = self.calls.lock().await; + let calls = self + .calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); // Shared proxy requests can still arrive without an execution ID. Only pick an owner when // there is exactly one candidate; with concurrent calls, canceling one would be a guess. if calls.active_calls.len() == 1 { @@ -410,14 +432,17 @@ impl NetworkApprovalService { ) -> Option> { self.calls .lock() - .await + .unwrap_or_else(std::sync::PoisonError::into_inner) .active_calls .get(execution_id) .cloned() } async fn resolve_active_call_attribution(&self) -> ActiveNetworkApprovalAttribution { - let calls = self.calls.lock().await; + let calls = self + .calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); match calls.active_calls.len() { 0 => ActiveNetworkApprovalAttribution::None, 1 => calls.active_calls.values().next().cloned().map_or( @@ -495,12 +520,18 @@ impl NetworkApprovalService { #[cfg(test)] async fn take_call_outcome(&self, registration_id: &str) -> Option { - let mut calls = self.calls.lock().await; + let mut calls = self + .calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); calls.call_outcomes.remove(registration_id) } - async fn record_call_outcome(&self, registration_id: &str, outcome: String) { - let mut calls = self.calls.lock().await; + fn record_call_outcome(&self, registration_id: &str, outcome: String) { + 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; }; @@ -513,8 +544,30 @@ impl NetworkApprovalService { call.cancellation_token.cancel(); } + fn record_call_outcome_if_absent(&self, registration_id: &str, outcome: String) { + 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; + }; + // Cancelling an execution can disconnect its other pending requests. + // Those fallbacks must not replace the reason it was cancelled. + calls + .call_outcomes + .entry(registration_id.to_string()) + .or_insert(outcome); + + drop(calls); + call.cancellation_token.cancel(); + } + async fn remove_call(&self, registration_id: &str) -> Option { - let mut calls = self.calls.lock().await; + let mut calls = self + .calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); calls.active_calls.shift_remove(registration_id); calls.call_outcomes.remove(registration_id) } @@ -549,19 +602,7 @@ impl NetworkApprovalService { return; }; - let mut calls = self.calls.lock().await; - if calls - .call_outcomes - .contains_key(&owner_call.registration_id) - { - return; - } - calls - .call_outcomes - .insert(owner_call.registration_id.clone(), message); - - drop(calls); - owner_call.cancellation_token.cancel(); + self.record_call_outcome_if_absent(&owner_call.registration_id, message); } fn format_network_target(protocol: &str, host: &str, port: u16) -> String { @@ -632,8 +673,7 @@ impl NetworkApprovalService { let Some((turn_context, strict_auto_review)) = active_turn else { if let Some(owner_call) = owner_call.as_ref() { - self.record_call_outcome(&owner_call.registration_id, policy_denial_message) - .await; + self.record_call_outcome(&owner_call.registration_id, policy_denial_message); } return NetworkDecision::deny(REASON_NOT_ALLOWED); }; @@ -658,6 +698,7 @@ impl NetworkApprovalService { .as_ref() .map(|call| call.cancellation_token.clone()), ); + pending_owner.disconnect = request.disconnect.clone(); let permission_profile = owner_call .as_ref() @@ -671,16 +712,14 @@ impl NetworkApprovalService { }); if !permission_profile.is_some_and(permission_profile_allows_network_approval_flow) { if let Some(owner_call) = owner_call.as_ref() { - self.record_call_outcome(&owner_call.registration_id, policy_denial_message) - .await; + self.record_call_outcome(&owner_call.registration_id, policy_denial_message); } pending_owner.complete(PendingApprovalDecision::Deny); return NetworkDecision::deny(REASON_NOT_ALLOWED); } if !allows_network_approval_flow(turn_context.approval_policy()) { if let Some(owner_call) = owner_call.as_ref() { - self.record_call_outcome(&owner_call.registration_id, policy_denial_message) - .await; + self.record_call_outcome(&owner_call.registration_id, policy_denial_message); } pending_owner.complete(PendingApprovalDecision::Deny); return NetworkDecision::deny(REASON_NOT_ALLOWED); @@ -751,8 +790,7 @@ impl NetworkApprovalService { Ok(decision) => decision, Err(ToolError::Rejected(rejection)) => { if let Some(owner_call) = owner_call.as_ref() { - self.record_call_outcome(&owner_call.registration_id, rejection) - .await; + self.record_call_outcome(&owner_call.registration_id, rejection); } turn_context.session_telemetry.tool_decision( &telemetry_tool_name, @@ -781,8 +819,7 @@ impl NetworkApprovalService { } else { format!("Error while requesting approval: {err}") }; - self.record_call_outcome(&owner_call.registration_id, rejection) - .await; + self.record_call_outcome(&owner_call.registration_id, rejection); } turn_context.session_telemetry.tool_decision( &telemetry_tool_name, @@ -815,8 +852,7 @@ impl NetworkApprovalService { self.record_call_outcome( &owner_call.registration_id, policy_denial_message.clone(), - ) - .await; + ); } PendingApprovalDecision::Deny } else { @@ -829,8 +865,7 @@ impl NetworkApprovalService { self.record_call_outcome( &owner_call.registration_id, policy_denial_message.clone(), - ) - .await; + ); } PendingApprovalDecision::Deny } else { @@ -886,8 +921,7 @@ impl NetworkApprovalService { self.record_call_outcome( &owner_call.registration_id, policy_denial_message.clone(), - ) - .await; + ); } PendingApprovalDecision::Deny } @@ -928,8 +962,7 @@ impl NetworkApprovalService { self.record_call_outcome( &owner_call.registration_id, "rejected by user".to_string(), - ) - .await; + ); } { let mut approved_hosts = self.session_approved_hosts.lock().await; @@ -948,8 +981,7 @@ impl NetworkApprovalService { self.record_call_outcome( &owner_call.registration_id, "Error while requesting approval".to_string(), - ) - .await; + ); } PendingApprovalDecision::Deny } diff --git a/codex-rs/core/src/tools/network_approval_tests.rs b/codex-rs/core/src/tools/network_approval_tests.rs index 26ac89712a..f54f4d268c 100644 --- a/codex-rs/core/src/tools/network_approval_tests.rs +++ b/codex-rs/core/src/tools/network_approval_tests.rs @@ -572,9 +572,7 @@ async fn blocked_request_does_not_override_recorded_approval_outcome() { register_call_with_default_shell_trigger(&service, "registration-1").await; let rejection = "approval client unavailable"; - service - .record_call_outcome("registration-1", rejection.to_string()) - .await; + service.record_call_outcome("registration-1", rejection.to_string()); service .record_blocked_request(denied_blocked_request("example.com")) .await; @@ -594,9 +592,7 @@ async fn specific_approval_outcome_replaces_earlier_blocked_request() { service .record_blocked_request(denied_blocked_request("example.com")) .await; - service - .record_call_outcome("registration-1", rejection.to_string()) - .await; + service.record_call_outcome("registration-1", rejection.to_string()); let error = network_approval_outcome_to_result(service.take_call_outcome("registration-1").await) @@ -604,17 +600,51 @@ async fn specific_approval_outcome_replaces_earlier_blocked_request() { assert!(matches!(error, ToolError::Rejected(message) if message == rejection)); } +#[tokio::test] +async fn disconnect_fallback_preserves_earlier_approval_outcome() { + let service = NetworkApprovalService::default(); + let cancellation = register_call_with_default_shell_trigger(&service, "registration-1").await; + let denial = "approval client unavailable"; + + service.record_call_outcome("registration-1", denial.to_string()); + service.record_call_outcome_if_absent("registration-1", "network disconnected".to_string()); + + assert!(cancellation.is_cancelled()); + assert_eq!( + service.take_call_outcome("registration-1").await, + Some(denial.to_string()) + ); +} + +#[tokio::test] +async fn disconnect_fallback_cancels_execution_and_yields_to_explicit_denial() { + let service = NetworkApprovalService::default(); + let cancellation = register_call_with_default_shell_trigger(&service, "registration-1").await; + let disconnect = "network disconnected"; + let denial = "explicit approval denial"; + + service.record_call_outcome_if_absent("registration-1", disconnect.to_string()); + assert!(cancellation.is_cancelled()); + assert_eq!( + service.take_call_outcome("registration-1").await, + Some(disconnect.to_string()) + ); + + service.record_call_outcome_if_absent("registration-1", disconnect.to_string()); + service.record_call_outcome("registration-1", denial.to_string()); + assert_eq!( + service.take_call_outcome("registration-1").await, + Some(denial.to_string()) + ); +} + #[tokio::test] async fn latest_specific_approval_outcome_replaces_earlier_specific_outcome() { let service = NetworkApprovalService::default(); register_call_with_default_shell_trigger(&service, "registration-1").await; - service - .record_call_outcome("registration-1", "earlier approval rejection".to_string()) - .await; - service - .record_call_outcome("registration-1", "latest approval rejection".to_string()) - .await; + service.record_call_outcome("registration-1", "earlier approval rejection".to_string()); + service.record_call_outcome("registration-1", "latest approval rejection".to_string()); let error = network_approval_outcome_to_result(service.take_call_outcome("registration-1").await) @@ -645,9 +675,7 @@ async fn finish_call_returns_denial_and_unregisters_active_call() { let cancellation_token = register_call_with_default_shell_trigger(&service, "registration-1").await; - service - .record_call_outcome("registration-1", "network denied".to_string()) - .await; + service.record_call_outcome("registration-1", "network denied".to_string()); let err = service .finish_call("registration-1", &cancellation_token) @@ -688,9 +716,7 @@ async fn deferred_finish_reuses_denial_result_after_first_consumer() { finish_outcome: Arc::new(OnceCell::new()), _execution_proxy: None, }; - service - .record_call_outcome("registration-1", "network denied".to_string()) - .await; + service.record_call_outcome("registration-1", "network denied".to_string()); let first = deferred .finish(&service) @@ -712,9 +738,7 @@ async fn record_call_outcome_ignores_inactive_call() { register_call_with_default_shell_trigger(&service, "registration-1").await; service.unregister_call("registration-1").await; - service - .record_call_outcome("registration-1", "network denied".to_string()) - .await; + service.record_call_outcome("registration-1", "network denied".to_string()); assert!(!cancellation_token.is_cancelled()); assert_eq!(service.take_call_outcome("registration-1").await, None); diff --git a/codex-rs/core/tests/suite/network_approval.rs b/codex-rs/core/tests/suite/network_approval.rs index 54fe884139..03f4a8b650 100644 --- a/codex-rs/core/tests/suite/network_approval.rs +++ b/codex-rs/core/tests/suite/network_approval.rs @@ -69,6 +69,7 @@ use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; use tempfile::TempDir; +use test_case::test_case; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; @@ -422,6 +423,122 @@ async fn cancelled_guardian_network_review_fails_closed_without_rewriting_turn_s Ok(()) } +#[test_case("GET", "http://codex-network-test.invalid/"; "plain_http")] +#[test_case("CONNECT", "codex-network-test.invalid:443"; "connect")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg_attr( + not(target_os = "linux"), + ignore = "requires the trusted Linux proxy bridge" +)] +async fn disconnected_network_request_explains_failure_to_model( + method: &str, + target: &str, +) -> Result<()> { + skip_if_target_windows!(Ok(()), "uses the POSIX/Python network fixture"); + skip_if_host_windows!(Ok(())); + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + + let server = start_mock_server().await; + // This tests the controller-local proxy; remote disconnect forwarding is not supported yet. + let test = managed_network_unified_exec_test(&server).await?; + let call_id = "network-disconnect"; + let poll_call_id = "network-disconnect-poll"; + let command = format!( + r#"python3 - <<'PY' +import os, socket, time, urllib.parse +proxy = urllib.parse.urlparse(os.environ['HTTP_PROXY']) +sock = socket.create_connection((proxy.hostname, proxy.port), timeout=10) +sock.sendall(b'{method} {target} HTTP/1.1\r\nHost: codex-network-test.invalid\r\n\r\n') +while not os.path.exists('disconnect-now'): + time.sleep(0.01) +sock.close() +time.sleep(60) +PY"# + ); + let mut args = network_exec_args(&command); + args["environment_id"] = json!(LOCAL_ENVIRONMENT_ID); + mount_sse_once_match( + &server, + |request: &wiremock::Request| { + !is_guardian_request(request) && !request_body_contains(request, call_id) + }, + sse(vec![ + ev_function_call(call_id, "exec_command", &serde_json::to_string(&args)?), + ev_completed("parent-start"), + ]), + ) + .await; + let pending_guardian = mount_response_once_match( + &server, + is_guardian_request, + sse_response(sse(vec![ev_completed("guardian")])).set_delay(Duration::from_secs(60)), + ) + .await; + let parent_poll = mount_sse_once_match( + &server, + |request: &wiremock::Request| { + !is_guardian_request(request) + && request_body_contains(request, call_id) + && !request_body_contains(request, poll_call_id) + }, + sse(vec![ + ev_function_call( + poll_call_id, + "write_stdin", + &json!({ + "session_id": 1000, "chars": "", "yield_time_ms": 10_000, + }) + .to_string(), + ), + ev_completed("parent-poll"), + ]), + ) + .await; + let parent_final = mount_sse_once_match( + &server, + |request: &wiremock::Request| { + !is_guardian_request(request) && request_body_contains(request, poll_call_id) + }, + sse(vec![ + ev_assistant_message("done", "understood"), + ev_completed("parent-done"), + ]), + ) + .await; + submit_managed_network_turn( + &test, + "explain why the network request fails", + vec![local(test.config.cwd.clone())], + ApprovalsReviewer::AutoReview, + AskForApproval::OnRequest, + ) + .await?; + wait_for_response_request(&pending_guardian).await; + wait_for_response_request(&parent_poll).await; + fs::write(test.config.cwd.join("disconnect-now"), "close")?; + wait_for_completion_without_network_prompt(&test).await; + + let output = parent_final + .single_request() + .function_call_output_text(poll_call_id) + .context("expected model-visible disconnect output")?; + let prefix = "Network request disconnected after "; + let suffix = " ms, before approval could complete"; + let elapsed = output + .split_once(prefix) + .and_then(|(_, rest)| rest.split_once(suffix)) + .map(|(elapsed, _)| elapsed) + .with_context(|| format!("missing disconnect explanation: {output}"))?; + assert!(elapsed.parse::()? > 0); + let message = &output[output.find(prefix).context("missing disconnect prefix")?..]; + let message = &message[..prefix.len() + elapsed.len() + suffix.len()]; + insta::assert_snapshot!(message.replacen(elapsed, "", 1), @r" + Network request disconnected after ms, before approval could complete + "); + Ok(()) +} + #[tokio::test(flavor = "current_thread")] #[cfg_attr( not(target_os = "linux"), diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 49575b5b1a..1d6144f9c2 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -704,6 +704,7 @@ impl ExecServerClient { command: None, exec_policy_hint: None, execution_id: None, + disconnect: None, }); let inner = Arc::downgrade(&inner); let rpc_client = Arc::downgrade(&rpc_client); diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index ad7884707b..8d9670a41d 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/codex-rs/network-proxy/src/http_proxy.rs @@ -19,6 +19,7 @@ use crate::reasons::REASON_MITM_REQUIRED; use crate::reasons::REASON_NOT_ALLOWED; use crate::reasons::REASON_PROXY_DISABLED; use crate::reasons::REASON_UNIX_SOCKET_UNSUPPORTED; +use crate::request_disconnect::NetworkRequestDisconnect; use crate::responses::PolicyDecisionDetails; use crate::responses::blocked_header_value; use crate::responses::blocked_message_with_policy; @@ -179,6 +180,7 @@ async fn http_connect_accept( environment_id: Option, mut req: Request, ) -> Result<(Response, Request), Response> { + let started_at = Instant::now(); let app_state = req .extensions() .get::>() @@ -218,7 +220,8 @@ async fn http_connect_accept( .await); } - let request = NetworkPolicyRequest::new(NetworkPolicyRequestArgs { + let disconnect = NetworkRequestDisconnect::default(); + let mut request = NetworkPolicyRequest::new(NetworkPolicyRequestArgs { protocol: NetworkProtocol::HttpsConnect, host: host.clone(), port: authority.port, @@ -229,7 +232,14 @@ async fn http_connect_accept( exec_policy_hint: None, }); - match evaluate_host_policy(&app_state, policy_decider.as_ref(), &request).await { + request.disconnect = Some(disconnect.clone()); + match disconnect + .track_http_request( + started_at, + evaluate_host_policy(&app_state, policy_decider.as_ref(), &request), + ) + .await + { Ok(NetworkDecision::Deny { reason, source, @@ -510,6 +520,7 @@ async fn http_plain_proxy( environment_id: Option, mut req: Request, ) -> Result { + let started_at = Instant::now(); let app_state = match req.extensions().get::>().cloned() { Some(state) => state, None => { @@ -709,7 +720,8 @@ async fn http_plain_proxy( .await); } - let request = NetworkPolicyRequest::new(NetworkPolicyRequestArgs { + let disconnect = NetworkRequestDisconnect::default(); + let mut request = NetworkPolicyRequest::new(NetworkPolicyRequestArgs { protocol: NetworkProtocol::Http, host: host.clone(), port, @@ -720,7 +732,14 @@ async fn http_plain_proxy( exec_policy_hint: None, }); - match evaluate_host_policy(&app_state, policy_decider.as_ref(), &request).await { + request.disconnect = Some(disconnect.clone()); + match disconnect + .track_http_request( + started_at, + evaluate_host_policy(&app_state, policy_decider.as_ref(), &request), + ) + .await + { Ok(NetworkDecision::Deny { reason, source, diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index 99b8b7f172..6a2a3ed494 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -16,6 +16,7 @@ mod policy; mod proxy; mod reasons; mod remote_config; +mod request_disconnect; mod responses; mod runtime; mod socks5; @@ -82,6 +83,7 @@ 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_disconnect::NetworkRequestDisconnect; pub use runtime::BlockedRequest; pub use runtime::BlockedRequestArgs; pub use runtime::BlockedRequestObserver; diff --git a/codex-rs/network-proxy/src/network_policy.rs b/codex-rs/network-proxy/src/network_policy.rs index 83daa4bcb5..b5d2d2719e 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_disconnect::NetworkRequestDisconnect; use crate::runtime::HostBlockDecision; use crate::runtime::HostBlockReason; use crate::state::NetworkProxyState; @@ -106,6 +107,8 @@ pub struct NetworkPolicyRequest { pub command: Option, pub exec_policy_hint: Option, pub execution_id: Option, + /// Present only when the local HTTP transport can identify an abandoned request. + pub disconnect: Option, } pub struct NetworkPolicyRequestArgs { @@ -141,6 +144,7 @@ impl NetworkPolicyRequest { command, exec_policy_hint, execution_id: None, + disconnect: None, } } } diff --git a/codex-rs/network-proxy/src/request_disconnect.rs b/codex-rs/network-proxy/src/request_disconnect.rs new file mode 100644 index 0000000000..b56aa09a87 --- /dev/null +++ b/codex-rs/network-proxy/src/request_disconnect.rs @@ -0,0 +1,45 @@ +use std::future::Future; +use std::pin::pin; +use std::sync::Arc; +use std::sync::OnceLock; +use std::time::Duration; +use std::time::Instant; + +/// Transport metadata available while cleaning up an abandoned HTTP policy request. +/// This does not cancel the approval reviewer or change its decision. +#[derive(Clone, Debug, Default)] +pub struct NetworkRequestDisconnect(Arc>); + +impl NetworkRequestDisconnect { + pub fn elapsed(&self) -> Option { + self.0.get().copied() + } + + pub(crate) async fn track_http_request( + &self, + started_at: Instant, + future: F, + ) -> F::Output { + // Publish the cause before dropping the policy future: its cleanup may + // immediately finish the owning tool call. + let mut future = pin!(future); + let mut guard = HttpRequestGuard(Some((self, started_at))); + let result = future.as_mut().await; + guard.0 = None; + result + } +} + +struct HttpRequestGuard<'a>(Option<(&'a NetworkRequestDisconnect, Instant)>); + +impl Drop for HttpRequestGuard<'_> { + fn drop(&mut self) { + if let Some((disconnect, started_at)) = self.0 { + let _ = disconnect.0.set(started_at.elapsed()); + } + } +} + +#[cfg(test)] +#[path = "request_disconnect_tests.rs"] +mod tests; diff --git a/codex-rs/network-proxy/src/request_disconnect_tests.rs b/codex-rs/network-proxy/src/request_disconnect_tests.rs new file mode 100644 index 0000000000..f9ebc3a15e --- /dev/null +++ b/codex-rs/network-proxy/src/request_disconnect_tests.rs @@ -0,0 +1,41 @@ +use super::NetworkRequestDisconnect; +use pretty_assertions::assert_eq; +use std::future::pending; +use std::time::Duration; +use std::time::Instant; + +#[tokio::test] +async fn disconnect_is_published_before_policy_cleanup() { + struct ObserveOnDrop(NetworkRequestDisconnect); + impl Drop for ObserveOnDrop { + fn drop(&mut self) { + assert!(self.0.elapsed().is_some()); + } + } + + let disconnect = NetworkRequestDisconnect::default(); + let observer = ObserveOnDrop(disconnect.clone()); + let started_at = Instant::now(); + let decision = disconnect.track_http_request(started_at, async move { + let _observer = observer; + pending::<()>().await; + }); + assert!( + tokio::time::timeout(Duration::from_millis(1), decision) + .await + .is_err() + ); + assert!(disconnect.elapsed().expect("disconnect time") <= started_at.elapsed()); +} + +#[tokio::test] +async fn completed_policy_request_is_not_a_disconnect() { + let disconnect = NetworkRequestDisconnect::default(); + assert_eq!( + disconnect + .track_http_request(Instant::now(), async { 42 }) + .await, + 42 + ); + assert_eq!(disconnect.elapsed(), None); +}