mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
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:
@@ -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]
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user