openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-12 21:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Report network disconnects during approval (#39284)

Browse Source 
## Why

When a local proxy request disconnects before network approval completes, the
owning tool call needs a model-visible explanation instead of remaining tied to
the abandoned request.

## What changed

- Track disconnect timing while plain HTTP and CONNECT requests await policy
  decisions.
- Cancel the owning execution and report how long the request waited when it
  disconnects before approval completes.
- Preserve an explicit approval outcome when disconnect cleanup runs afterward.

## Testing

Added unit coverage for disconnect tracking and outcome precedence, plus
end-to-end coverage for plain HTTP and CONNECT requests.

GitOrigin-RevId: b354b29bbe86f38e252fcaf529541f177480136b
This commit is contained in:
Dylan Hurd
2026-08-18 21:30:45 +00:00
committed by copyberry
parent 2fe4e06cc2
commit ba37d0c45b
9 changed files with 353 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/network_approval.rs
+75 -43
View File
@@ -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<NetworkApprovalCallState>,
// Pending-request cleanup must publish a tool outcome synchronously before cancelling it.
calls: SyncMutex<NetworkApprovalCallState>,
// Owner cleanup runs from Drop, so this lock cannot require an async task.
pending_host_approvals: SyncMutex<HashMap<PendingHostApprovalKey, Arc<PendingHostApproval>>>,
// 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<PendingHostApproval>,
execution_cancellation: Option<CancellationToken>,
disconnect: Option<NetworkRequestDisconnect>,
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<Arc<ActiveNetworkApprovalCall>> {
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<Arc<ActiveNetworkApprovalCall>> {
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<String> {
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<String> {
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
}
codex-rs/core/src/tools/network_approval_tests.rs
+45 -21
View File
@@ -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);
codex-rs/core/tests/suite/network_approval.rs
+117
View File
@@ -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::<u128>()? > 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, "<elapsed>", 1), @r"
Network request disconnected after <elapsed> ms, before approval could complete
");
Ok(())
}
#[tokio::test(flavor = "current_thread")]
#[cfg_attr(
not(target_os = "linux"),
codex-rs/exec-server/src/client_recovery.rs
+1
View File
@@ -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);
codex-rs/network-proxy/src/http_proxy.rs
+23 -4
View File
@@ -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<String>,
mut req: Request,
) -> Result<(Response, Request), Response> {
let started_at = Instant::now();
let app_state = req
.extensions()
.get::<Arc<NetworkProxyState>>()
@@ -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<String>,
mut req: Request,
) -> Result<Response, Infallible> {
let started_at = Instant::now();
let app_state = match req.extensions().get::<Arc<NetworkProxyState>>().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,
codex-rs/network-proxy/src/lib.rs
+2
View File
@@ -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;
codex-rs/network-proxy/src/network_policy.rs
+4
View File
@@ -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<String>,
pub exec_policy_hint: Option<String>,
pub execution_id: Option<String>,
/// Present only when the local HTTP transport can identify an abandoned request.
pub disconnect: Option<NetworkRequestDisconnect>,
}
pub struct NetworkPolicyRequestArgs {
@@ -141,6 +144,7 @@ impl NetworkPolicyRequest {
command,
exec_policy_hint,
execution_id: None,
disconnect: None,
}
}
}
codex-rs/network-proxy/src/request_disconnect.rs
+45
View File
@@ -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<OnceLock<Duration>>);
impl NetworkRequestDisconnect {
pub fn elapsed(&self) -> Option<Duration> {
self.0.get().copied()
}
pub(crate) async fn track_http_request<F: Future>(
&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;
codex-rs/network-proxy/src/request_disconnect_tests.rs
+41
View File
@@ -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);
}