codex: address PR review feedback (#18265)

codex-rs/app-server/src/in_process.rs

@@ -377,6 +377,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
                 Arc::clone(&outbound_experimental_api_enabled),
                 Arc::clone(&outbound_opted_out_notification_methods),
                 /*disconnect_sender*/ None,
+                /*disconnect_notifier*/ None,
             ),
         );
         let mut outbound_handle = tokio::spawn(async move {

codex-rs/app-server/src/lib.rs

@@ -29,6 +29,7 @@ use crate::transport::ConnectionState;
 use crate::transport::OutboundConnectionState;
 use crate::transport::TransportEvent;
 use crate::transport::auth::policy_from_settings;
+use crate::transport::disconnect_connection;
 use crate::transport::route_outgoing_envelope;
 use crate::transport::start_remote_control;
 use crate::transport::start_stdio_connection;
@@ -594,50 +595,62 @@ pub async fn run_main_with_transport(
 
     let outbound_handle = tokio::spawn(async move {
         let mut outbound_connections = HashMap::<ConnectionId, OutboundConnectionState>::new();
+        // Overflow workers run outside this router task. This side channel lets
+        // them remove a slow connection from routing before the transport loop's
+        // eventual ConnectionClosed event catches up.
+        let (outbound_disconnect_tx, mut outbound_disconnect_rx) =
+            mpsc::channel::<ConnectionId>(CHANNEL_CAPACITY);
         loop {
             tokio::select! {
-                    biased;
-                    event = outbound_control_rx.recv() => {
-                        let Some(event) = event else {
-                            break;
-                        };
-                        match event {
-                            OutboundControlEvent::Opened {
+                biased;
+                event = outbound_control_rx.recv() => {
+                    let Some(event) = event else {
+                        break;
+                    };
+                    match event {
+                        OutboundControlEvent::Opened {
+                            connection_id,
+                            writer,
+                            disconnect_sender,
+                            initialized,
+                            experimental_api_enabled,
+                            opted_out_notification_methods,
+                        } => {
+                            outbound_connections.insert(
                                 connection_id,
-                                writer,
-                                disconnect_sender,
-                                initialized,
-                                experimental_api_enabled,
-                                opted_out_notification_methods,
-                            } => {
-                                outbound_connections.insert(
+                                OutboundConnectionState::new(
                                     connection_id,
-                                    OutboundConnectionState::new(
-                                        connection_id,
-                                        writer,
-                                        initialized,
-                                        experimental_api_enabled,
-                                        opted_out_notification_methods,
-                                        disconnect_sender,
-                                    ),
-                                );
-                            }
-                            OutboundControlEvent::Closed { connection_id } => {
-                                outbound_connections.remove(&connection_id);
-                            }
-                            OutboundControlEvent::DisconnectAll => {
-                                info!(
-                                    "disconnecting {} outbound websocket connection(s) for graceful restart",
-                                    outbound_connections.len()
-                                );
-                                for connection_state in outbound_connections.values() {
-                                    connection_state.request_disconnect();
-                                }
-                                outbound_connections.clear();
+                                    writer,
+                                    initialized,
+                                    experimental_api_enabled,
+                                    opted_out_notification_methods,
+                                    disconnect_sender,
+                                    Some(outbound_disconnect_tx.clone()),
+                                ),
+                            );
+                        }
+                        OutboundControlEvent::Closed { connection_id } => {
+                            outbound_connections.remove(&connection_id);
+                        }
+                        OutboundControlEvent::DisconnectAll => {
+                            info!(
+                                "disconnecting {} outbound websocket connection(s) for graceful restart",
+                                outbound_connections.len()
+                            );
+                            for connection_state in outbound_connections.values() {
+                                connection_state.request_disconnect();
                             }
+                            outbound_connections.clear();
                         }
                     }
-                    envelope = outgoing_rx.recv() => {
+                }
+                connection_id = outbound_disconnect_rx.recv() => {
+                    let Some(connection_id) = connection_id else {
+                        break;
+                    };
+                    disconnect_connection(&mut outbound_connections, connection_id);
+                }
+                envelope = outgoing_rx.recv() => {
                     let Some(envelope) = envelope else {
                         break;
                     };

codex-rs/app-server/src/transport/mod.rs

@@ -164,12 +164,14 @@ impl OutboundConnectionState {
         experimental_api_enabled: Arc<AtomicBool>,
         opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
         disconnect_sender: Option<CancellationToken>,
+        disconnect_notifier: Option<mpsc::Sender<ConnectionId>>,
     ) -> Self {
         let overflow_depth = Arc::new(AtomicUsize::new(0));
         let overflow_writer = disconnect_sender.as_ref().map(|disconnect_sender| {
             let (overflow_tx, mut overflow_rx) = mpsc::channel(CHANNEL_CAPACITY);
             let writer = writer.clone();
             let disconnect_sender = disconnect_sender.clone();
+            let disconnect_notifier = disconnect_notifier.clone();
             let overflow_depth = Arc::clone(&overflow_depth);
             tokio::spawn(async move {
                 while let Some(queued_message) = overflow_rx.recv().await {
@@ -187,11 +189,22 @@ impl OutboundConnectionState {
                                 OUTBOUND_QUEUE_FULL_GRACE
                             );
                             disconnect_sender.cancel();
+                            // The websocket task will eventually report ConnectionClosed,
+                            // but notify the outbound router now so no newer messages are
+                            // routed after this timed-out one is dropped.
+                            if let Some(disconnect_notifier) = &disconnect_notifier {
+                                let _ = disconnect_notifier.send(connection_id).await;
+                            }
                             break;
                         }
                         Err(mpsc::error::SendTimeoutError::Closed(_)) => {
                             overflow_depth.fetch_sub(1, Ordering::AcqRel);
                             disconnect_sender.cancel();
+                            // Drop outbound routing state promptly even if the transport's
+                            // close event is delayed behind other incoming events.
+                            if let Some(disconnect_notifier) = &disconnect_notifier {
+                                let _ = disconnect_notifier.send(connection_id).await;
+                            }
                             break;
                         }
                     }
@@ -321,7 +334,7 @@ fn should_skip_notification_for_connection(
     }
 }
 
-fn disconnect_connection(
+pub(crate) fn disconnect_connection(
     connections: &mut HashMap<ConnectionId, OutboundConnectionState>,
     connection_id: ConnectionId,
 ) -> bool {
@@ -713,6 +726,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 opted_out_notification_methods,
                 /*disconnect_sender*/ None,
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -754,6 +768,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::from(["configWarning".to_string()]))),
                 /*disconnect_sender*/ None,
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -795,6 +810,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::new())),
                 /*disconnect_sender*/ None,
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -842,6 +858,7 @@ mod tests {
                 Arc::new(AtomicBool::new(false)),
                 Arc::new(RwLock::new(HashSet::new())),
                 /*disconnect_sender*/ None,
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -905,6 +922,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::new())),
                 /*disconnect_sender*/ None,
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -993,6 +1011,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::new())),
                 Some(disconnect_token.clone()),
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -1068,6 +1087,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::new())),
                 Some(disconnect_token.clone()),
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -1140,6 +1160,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::new())),
                 Some(disconnect_token.clone()),
+                /*disconnect_notifier*/ None,
             ),
         );
 
@@ -1199,6 +1220,7 @@ mod tests {
     async fn disconnectable_connection_requests_disconnect_after_queue_grace_expires() {
         let connection_id = ConnectionId(2);
         let (writer_tx, mut writer_rx) = mpsc::channel(1);
+        let (disconnect_notifier_tx, mut disconnect_notifier_rx) = mpsc::channel(1);
         let disconnect_token = CancellationToken::new();
 
         writer_tx
@@ -1225,6 +1247,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::new())),
                 Some(disconnect_token.clone()),
+                /*disconnect_notifier*/ Some(disconnect_notifier_tx),
             ),
         );
 
@@ -1246,16 +1269,19 @@ mod tests {
         .await;
 
         assert!(connections.contains_key(&connection_id));
-        timeout(
+        let notified_connection_id = timeout(
             OUTBOUND_QUEUE_FULL_GRACE + Duration::from_millis(100),
-            async {
-                while !disconnect_token.is_cancelled() {
-                    tokio::time::sleep(Duration::from_millis(10)).await;
-                }
-            },
+            disconnect_notifier_rx.recv(),
         )
         .await
-        .expect("full queue should request disconnect after the grace expires");
+        .expect("full queue should notify the router after the grace expires")
+        .expect("disconnect notification should contain a connection id");
+        assert_eq!(notified_connection_id, connection_id);
+        assert!(disconnect_connection(
+            &mut connections,
+            notified_connection_id
+        ));
+        assert!(!connections.contains_key(&connection_id));
         assert!(disconnect_token.is_cancelled());
         let original_message = writer_rx
             .try_recv()
@@ -1296,6 +1322,7 @@ mod tests {
                 Arc::new(AtomicBool::new(true)),
                 Arc::new(RwLock::new(HashSet::new())),
                 /*disconnect_sender*/ None,
+                /*disconnect_notifier*/ None,
             ),
         );