Skip the network denial grace period when remote decisions drain (#59960)

## Why

Completed remote commands wait up to 100 ms for late network denials even when all execution-owned network decision callbacks and requests have finished.

## What changed

Allow deferred network approval cleanup to finish as soon as an executor-owned remote proxy's decisions drain. Retain the grace period when a shared listener could still report an untagged denial, including when that listener starts during the command. Always register the blocked-request observer when starting the session proxy.

## Testing

Extend the paused-time regression test to cover immediate completion after remote callbacks drain and propagation of a late denial from a newly registered shared listener.

GitOrigin-RevId: 41bda6504f6d97d5626beba58cf30fb4a9a2dbdd

codex-rs/core/src/session/session.rs

@@ -1720,11 +1720,6 @@ impl Session {
             } else {
                 None
             };
-            let blocked_request_observer = config
-                .permissions
-                .network
-                .as_ref()
-                .map(|_| build_blocked_request_observer(Arc::clone(&network_approval)));
             let network_policy_decider =
                 network_policy_decider_session
                     .as_ref()
@@ -1747,7 +1742,7 @@ impl Session {
                     config.permissions.permission_profile(),
                     config.effective_local_windows_sandbox_type(),
                     network_policy_decider.as_ref().map(Arc::clone),
-                    blocked_request_observer.as_ref().map(Arc::clone),
+                    Some(build_blocked_request_observer(Arc::clone(&network_approval))),
                     managed_network_requirements_configured,
                     network_proxy_audit_metadata.clone(),
                 )

codex-rs/core/src/tools/network_approval.rs

@@ -45,6 +45,8 @@ use std::io;
 use std::sync::Arc;
 use std::sync::Mutex as SyncMutex;
 use std::sync::OnceLock as SyncOnceLock;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use tokio::sync::Mutex;
 use tokio::sync::Notify;
 use tokio::sync::OnceCell;
@@ -72,12 +74,20 @@ pub(crate) struct NetworkApprovalSpec {
     pub network_policy: Option<EnvironmentNetworkPolicy>,
 }
 
+#[derive(Clone, Debug)]
+struct NetworkApprovalProxy {
+    proxy: NetworkProxy,
+    // Only an executor-owned remote proxy can drain its decisions. A shared listener
+    // can also attribute untagged traffic to this execution, even after a config refresh.
+    shared_ingress_for_remote: Option<Arc<AtomicBool>>,
+}
+
 #[derive(Clone, Debug)]
 pub(crate) struct DeferredNetworkApproval {
     registration_id: String,
     cancellation_token: CancellationToken,
     finish_outcome: Arc<OnceCell<Option<String>>>,
-    _execution_proxy: Option<NetworkProxy>,
+    _execution_proxy: Option<NetworkApprovalProxy>,
 }
 
 impl DeferredNetworkApproval {
@@ -93,6 +103,20 @@ impl DeferredNetworkApproval {
         self.cancellation_token.is_cancelled()
     }
 
+    pub(crate) async fn remote_decisions_drained(&self) {
+        if let Some(NetworkApprovalProxy {
+            proxy,
+            shared_ingress_for_remote: Some(shared_ingress),
+        }) = &self._execution_proxy
+        {
+            proxy.wait_for_remote_policy_deciders().await;
+            if !shared_ingress.load(Ordering::Acquire) {
+                return;
+            }
+        }
+        std::future::pending().await
+    }
+
     async fn finish(&self, service: &NetworkApprovalService) -> Result<(), ToolError> {
         let outcome = self
             .finish_outcome
@@ -109,7 +133,7 @@ impl DeferredNetworkApproval {
 pub(crate) struct ActiveNetworkApproval {
     registration_id: Option<String>,
     cancellation_token: CancellationToken,
-    execution_proxy: NetworkProxy,
+    execution_proxy: NetworkApprovalProxy,
 }
 
 impl ActiveNetworkApproval {
@@ -118,7 +142,7 @@ impl ActiveNetworkApproval {
     }
 
     pub(crate) fn execution_proxy(&self) -> &NetworkProxy {
-        &self.execution_proxy
+        &self.execution_proxy.proxy
     }
 
     pub(crate) fn into_deferred(self) -> Option<DeferredNetworkApproval> {
@@ -279,6 +303,7 @@ struct PendingNetworkApprovals {
 }
 
 pub(crate) struct NetworkApprovalService {
+    shared_ingress_registered: Arc<AtomicBool>,
     // 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.
@@ -461,6 +486,7 @@ impl Drop for PendingHostApprovalOwner<'_> {
 impl Default for NetworkApprovalService {
     fn default() -> Self {
         Self {
+            shared_ingress_registered: Arc::default(),
             calls: SyncMutex::new(NetworkApprovalCallState::default()),
             pending_host_approvals: SyncMutex::default(),
             session_policy_commit_lock: Mutex::new(()),
@@ -1155,6 +1181,9 @@ impl NetworkApprovalService {
 pub(crate) fn build_blocked_request_observer(
     network_approval: Arc<NetworkApprovalService>,
 ) -> Arc<dyn BlockedRequestObserver> {
+    network_approval
+        .shared_ingress_registered
+        .store(true, Ordering::Release);
     Arc::new(move |blocked: BlockedRequest| {
         let network_approval = Arc::clone(&network_approval);
         async move {
@@ -1212,12 +1241,16 @@ pub(crate) async fn begin_network_approval(
         )
     });
     let registration_id = Uuid::new_v4().to_string();
-    let Some(execution_proxy) = execution_network_proxy(
+    let environment = environments
+        .turn_environments()
+        .find(|environment| environment.selection.environment_id == environment_id);
+    let executor_owns_proxy = network_policy.is_some()
+        && turn.config.permissions.network.is_none()
+        && environment.is_some_and(|environment| environment.environment().is_remote());
+    let Some(proxy) = execution_network_proxy(
         session,
         turn,
-        environments
-            .turn_environments()
-            .find(|environment| environment.selection.environment_id == environment_id),
+        environment,
         network,
         &environment_id,
         &permission_profile,
@@ -1249,7 +1282,11 @@ pub(crate) async fn begin_network_approval(
     Ok(Some(ActiveNetworkApproval {
         registration_id: Some(registration_id),
         cancellation_token,
-        execution_proxy,
+        execution_proxy: NetworkApprovalProxy {
+            proxy,
+            shared_ingress_for_remote: executor_owns_proxy
+                .then(|| Arc::clone(&session.services.network_approval.shared_ingress_registered)),
+        },
     }))
 }
 

codex-rs/core/src/unified_exec/process_manager.rs

@@ -353,30 +353,28 @@ async fn network_denial_message_for_session(
     }
 }
 
-async fn wait_for_late_network_denial(network_cancelled: Option<CancellationToken>) -> bool {
-    let Some(network_cancelled) = network_cancelled else {
-        return false;
-    };
-    if network_cancelled.is_cancelled() {
-        return true;
-    }
-
-    tokio::select! {
-        _ = network_cancelled.cancelled() => true,
-        _ = tokio::time::sleep(LATE_NETWORK_DENIAL_GRACE_PERIOD) => false,
-    }
+async fn wait_for_late_network_denial(network_cancelled: &CancellationToken) -> bool {
+    tokio::time::timeout(
+        LATE_NETWORK_DENIAL_GRACE_PERIOD,
+        network_cancelled.cancelled(),
+    )
+    .await
+    .is_ok()
 }
 
 async fn finish_deferred_network_approval_after_process_exit_for_session(
     session: Option<&Arc<crate::session::session::Session>>,
     deferred: Option<DeferredNetworkApproval>,
 ) -> Result<(), String> {
-    wait_for_late_network_denial(
-        deferred
-            .as_ref()
-            .map(DeferredNetworkApproval::cancellation_token),
-    )
-    .await;
+    if let Some(approval) = &deferred {
+        let network_cancelled = approval.cancellation_token();
+        // Finishing below unregisters the call before we report the outcome.
+        tokio::select! {
+            biased;
+            _ = approval.remote_decisions_drained() => {}
+            _ = wait_for_late_network_denial(&network_cancelled) => {}
+        }
+    }
     finish_deferred_network_approval_for_session(session, deferred).await
 }
 
@@ -436,7 +434,8 @@ fn terminate_process_on_network_denial(
         let denied = tokio::select! {
             _ = network_cancelled.cancelled() => true,
             _ = process_exited.cancelled() => {
-                wait_for_late_network_denial(Some(network_cancelled.clone())).await
+                // A yielded process stays registered until its next poll.
+                wait_for_late_network_denial(&network_cancelled).await
             }
         };
         if !denied {

codex-rs/core/src/unified_exec/process_manager_tests.rs

@@ -442,16 +442,121 @@ async fn network_denial_fallback_message_names_sandbox_network_proxy() {
     );
 }
 
-#[tokio::test]
-async fn late_network_denial_grace_observes_cancellation_after_exit() {
-    let cancellation = CancellationToken::new();
-    let cancellation_for_task = cancellation.clone();
-    tokio::spawn(async move {
-        tokio::time::sleep(Duration::from_millis(10)).await;
-        cancellation_for_task.cancel();
-    });
+#[tokio::test(start_paused = true)]
+async fn late_network_denial_grace_observes_cancellation_after_exit() -> anyhow::Result<()> {
+    use crate::environment_selection::EnvironmentAttachment;
+    use crate::environment_selection::TurnEnvironmentSnapshot;
+    use crate::environment_selection::TurnEnvironmentState;
+    use crate::guardian::GuardianNetworkAccessTrigger;
+    use crate::tools::network_approval::NetworkApprovalSpec;
+    use crate::tools::network_approval::begin_network_approval;
+    use crate::tools::network_approval::build_blocked_request_observer;
+    use codex_network_proxy::BlockedRequest;
+    use codex_network_proxy::BlockedRequestArgs;
+    use codex_network_proxy::EnvironmentNetworkPolicy;
+    use codex_tools::ToolName;
 
-    assert!(wait_for_late_network_denial(Some(cancellation)).await);
+    let (session, mut turn) = crate::session::tests::make_session_and_context().await;
+    let session = Arc::new(session);
+    Arc::make_mut(&mut turn.config).permissions.network = None;
+    let mut remote = turn
+        .initial_environments
+        .primary()
+        .expect("primary environment")
+        .clone();
+    let remote_id = codex_exec_server::REMOTE_ENVIRONMENT_ID;
+    let (attachment, _guard, _) = EnvironmentAttachment::new(
+        remote_id,
+        Arc::new(codex_exec_server::Environment::create_for_tests(Some(
+            "ws://127.0.0.1:1".into(),
+        ))?),
+    );
+    remote.attachment = attachment;
+    remote.selection.environment_id = remote_id.to_string();
+    remote.config_mut().permission_profile =
+        codex_protocol::models::PermissionProfileSnapshot::legacy(
+            codex_protocol::models::PermissionProfile::workspace_write(),
+        );
+    remote.config_mut().network_policy = Some(EnvironmentNetworkPolicy::from_config(
+        &codex_network_proxy::NetworkProxyConfig {
+            enabled: true,
+            ..Default::default()
+        },
+        /*managed_allowed_domains_only*/ false,
+    ));
+    let environments = TurnEnvironmentSnapshot {
+        environments: vec![TurnEnvironmentState::Ready(remote)],
+    };
+    let remote = environments.primary().expect("remote environment");
+
+    for register_shared_listener in [false, true] {
+        let approval = begin_network_approval(
+            &session,
+            &turn,
+            &environments,
+            /*managed_network_active*/ true,
+            Some(NetworkApprovalSpec {
+                network: None,
+                environment_id: remote_id.to_string(),
+                tool_name: ToolName::plain("exec_command"),
+                command: "pwd".into(),
+                trigger: GuardianNetworkAccessTrigger {
+                    call_id: "exec".into(),
+                    tool_name: "exec_command".into(),
+                    command: vec!["pwd".into()],
+                    cwd: remote.cwd().clone(),
+                    sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault,
+                    additional_permissions: None,
+                    justification: None,
+                    tty: None,
+                },
+                permission_profile: remote.permission_profile().clone(),
+                network_policy: remote.config().network_policy.clone(),
+            }),
+        )
+        .await
+        .expect("build network proxy")
+        .expect("network approval");
+        let callback = approval
+            .execution_proxy()
+            .remote_policy_decider(/*allow_local_binding*/ true)
+            .expect("remote policy callback");
+        let finish = finish_deferred_network_approval_after_process_exit_for_session(
+            Some(&session),
+            approval.into_deferred(),
+        );
+        tokio::pin!(finish);
+        assert!(futures::poll!(&mut finish).is_pending());
+
+        let observer = register_shared_listener.then(|| {
+            build_blocked_request_observer(Arc::clone(&session.services.network_approval))
+        });
+        drop(callback);
+        if let Some(observer) = observer {
+            // A shared listener started during the command and can report an untagged denial.
+            tokio::time::advance(Duration::from_millis(10)).await;
+            assert!(futures::poll!(&mut finish).is_pending());
+            let blocked = BlockedRequest::new(BlockedRequestArgs {
+                host: "blocked.example".into(),
+                protocol: "http".into(),
+                port: Some(80),
+                reason: "not_allowed".into(),
+                decision: Some("deny".into()),
+                source: Some("decider".into()),
+                client: None,
+                method: None,
+                mode: None,
+            });
+            let expected = crate::network_policy_decision::denied_network_policy_message(&blocked)
+                .expect("blocked request must describe the denied host");
+            observer.on_blocked_request(blocked).await;
+            assert_eq!(finish.await, Err(expected));
+        } else {
+            assert_eq!(futures::poll!(&mut finish), std::task::Poll::Ready(Ok(())));
+        }
+    }
+
+    Ok(())
 }
 
 #[tokio::test]

codex-rs/network-proxy/src/proxy/execution_scope.rs

@@ -17,6 +17,17 @@ impl Drop for ExecutionScope {
 }
 
 impl NetworkProxy {
+    /// Waits until every remote decision callback and in-flight request for this
+    /// execution has been dropped. Call after the executor has started, when no
+    /// further callbacks will be installed for this scope. An unscoped proxy
+    /// cannot establish this boundary, so its wait remains pending.
+    pub async fn wait_for_remote_policy_deciders(&self) {
+        match &self.execution_scope {
+            Some(scope) => scope.lifetime_tx.closed().await,
+            None => std::future::pending().await,
+        }
+    }
+
     /// Returns a proxy with execution-specific attribution and attachment policy.
     pub fn for_execution(
         &self,