From bd282df268591a61e00c2224e2406a3fbe666fac Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Thu, 16 Apr 2026 23:39:23 -0700 Subject: [PATCH] codex: address PR review feedback (#18265) --- codex-rs/app-server/src/in_process.rs | 1 + codex-rs/app-server/src/lib.rs | 87 ++++++++++++++---------- codex-rs/app-server/src/transport/mod.rs | 43 +++++++++--- 3 files changed, 86 insertions(+), 45 deletions(-) diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 42b2db2f8a..dabcaa708f 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/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 { diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 03af181da8..8aef83a0de 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/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::::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::(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; }; diff --git a/codex-rs/app-server/src/transport/mod.rs b/codex-rs/app-server/src/transport/mod.rs index 1e7a428327..8d9c3c3921 100644 --- a/codex-rs/app-server/src/transport/mod.rs +++ b/codex-rs/app-server/src/transport/mod.rs @@ -164,12 +164,14 @@ impl OutboundConnectionState { experimental_api_enabled: Arc, opted_out_notification_methods: Arc>>, disconnect_sender: Option, + disconnect_notifier: Option>, ) -> 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, 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, ), );