From 95db2cd5d5f37f80affb9b50ce1de70ebb4513ff Mon Sep 17 00:00:00 2001
From: Eric Traut <etraut@openai.com>
Date: Sun, 11 Oct 2026 15:40:45 +0000
Subject: [PATCH] Keep remote app-server events deliverable during shutdown
 (#53155)

## Why

Dropping the event receiver before shutdown lets an in-flight notification fail delivery and terminate the worker before it sends the WebSocket close frame.

## What changed

Keep the event receiver alive until `RemoteAppServerClient::shutdown` finishes waiting for the worker, allowing it to process the shutdown command and close the WebSocket.

## Testing

Extend regression tests to verify that an event remains deliverable after shutdown is queued and that the server receives `Message::Close(None)` with an extra notification in flight.

GitOrigin-RevId: ba14bc97ffa3453b7b59660144db6da9e688983d
---
 codex-rs/app-server-client/src/lib.rs    | 41 +++++++++++++++---------
 codex-rs/app-server-client/src/remote.rs | 21 ++++++++++--
 2 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs
index 92fe7b4662..75544173d1 100644
--- a/codex-rs/app-server-client/src/lib.rs
+++ b/codex-rs/app-server-client/src/lib.rs
@@ -1746,24 +1746,25 @@ mod tests {
 
     #[tokio::test]
     async fn remote_notifications_arrive_over_websocket() {
+        let (close_tx, close_rx) = tokio::sync::oneshot::channel();
         let websocket_url = start_test_remote_server(|mut websocket| async move {
             expect_remote_initialize(&mut websocket).await;
-            write_websocket_message(
-                &mut websocket,
-                JSONRPCMessage::Notification(
-                    serde_json::from_value(
-                        serde_json::to_value(ServerNotification::AccountUpdated(
-                            AccountUpdatedNotification {
-                                auth_mode: None,
-                                plan_type: None,
-                            },
-                        ))
-                        .expect("notification should serialize"),
-                    )
-                    .expect("notification should convert to JSON-RPC"),
-                ),
-            )
-            .await;
+            let notification = JSONRPCMessage::Notification(
+                serde_json::from_value(
+                    serde_json::to_value(ServerNotification::AccountUpdated(
+                        AccountUpdatedNotification {
+                            auth_mode: None,
+                            plan_type: None,
+                        },
+                    ))
+                    .expect("notification should serialize"),
+                )
+                .expect("notification should convert to JSON-RPC"),
+            );
+            // Only the first notification is consumed; the other can arrive during shutdown.
+            write_websocket_message(&mut websocket, notification.clone()).await;
+            write_websocket_message(&mut websocket, notification).await;
+            let _ = close_tx.send(websocket.next().await);
         })
         .await;
         let mut client = RemoteAppServerClient::connect(test_remote_connect_args(websocket_url))
@@ -1778,6 +1779,14 @@ mod tests {
         ));
 
         client.shutdown().await.expect("shutdown should complete");
+        assert_eq!(
+            close_rx
+                .await
+                .expect("server should observe shutdown")
+                .expect("server should receive a frame")
+                .expect("websocket should close cleanly"),
+            Message::Close(None),
+        );
     }
 
     #[tokio::test]
diff --git a/codex-rs/app-server-client/src/remote.rs b/codex-rs/app-server-client/src/remote.rs
index f95bee50c5..c5d9b0321f 100644
--- a/codex-rs/app-server-client/src/remote.rs
+++ b/codex-rs/app-server-client/src/remote.rs
@@ -604,7 +604,6 @@ impl RemoteAppServerClient {
             worker_handle,
         } = self;
         let mut worker_handle = worker_handle;
-        drop(event_rx);
         let (response_tx, response_rx) = oneshot::channel();
         if command_tx
             .send(RemoteClientCommand::Shutdown { response_tx })
@@ -619,6 +618,8 @@ impl RemoteAppServerClient {
             worker_handle.abort();
             let _ = worker_handle.await;
         }
+        // Keep events deliverable until the worker has closed the websocket.
+        drop(event_rx);
         Ok(())
     }
 }
@@ -1064,9 +1065,21 @@ mod tests {
     #[tokio::test]
     async fn shutdown_tolerates_worker_exit_after_command_is_queued() {
         let (command_tx, mut command_rx) = mpsc::channel(1);
-        let (_event_tx, event_rx) = mpsc::unbounded_channel::<AppServerEvent>();
+        let (event_tx, event_rx) = mpsc::unbounded_channel::<AppServerEvent>();
+        let (delivered_tx, delivered_rx) = oneshot::channel();
         let worker_handle = tokio::spawn(async move {
             let _ = command_rx.recv().await;
+            // Force an event after shutdown is queued but before this worker has exited.
+            let result = deliver_event(
+                &event_tx,
+                AppServerEvent::ServerNotification(Box::new(ServerNotification::AccountUpdated(
+                    codex_app_server_protocol::AccountUpdatedNotification {
+                        auth_mode: None,
+                        plan_type: None,
+                    },
+                ))),
+            );
+            let _ = delivered_tx.send(result);
         });
         let client = RemoteAppServerClient {
             command_tx,
@@ -1080,5 +1093,9 @@ mod tests {
             .shutdown()
             .await
             .expect("shutdown should complete when worker exits first");
+        delivered_rx
+            .await
+            .expect("worker should attempt delivery")
+            .expect("an in-flight event must not abort the worker during shutdown");
     }
 }