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
This commit is contained in:
Eric Traut
2026-10-11 15:57:06 +00:00
committed by copyberry
parent 8553558345
commit 95db2cd5d5
2 changed files with 44 additions and 18 deletions
+25 -16
View File
@@ -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]
+19 -2
View File
@@ -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");
}
}