Shut down the in-process outbound router explicitly (#35523)

## Why

Detached processor work can retain outgoing message senders, so waiting for
channel closure can keep the in-process app-server's outbound router alive
during shutdown.

## What changed

Add an explicit shutdown signal for the outbound router and prioritize it over
pending outgoing messages. Send that signal after processor tasks are stopped
so runtime shutdown no longer depends on every sender being dropped.

## Testing

Add a paused-time test that verifies the router exits and closes its channel
even when an outgoing sender is still retained.

GitOrigin-RevId: a8f88d88a0950f50339b1c56909fedb781f178d4
This commit is contained in:
Felipe Coury
2026-07-26 22:11:26 +00:00
committed by copyberry
parent 61a44880a8
commit 5a3808b77d
+53 -8
View File
@@ -383,6 +383,25 @@ pub async fn start(mut args: InProcessStartArgs) -> IoResult<InProcessClientHand
Ok(client)
}
async fn run_outbound_router(
mut outgoing_rx: mpsc::Receiver<OutgoingEnvelope>,
mut outbound_connections: HashMap<ConnectionId, OutboundConnectionState>,
mut shutdown_rx: oneshot::Receiver<()>,
) {
loop {
tokio::select! {
biased;
_ = &mut shutdown_rx => break,
envelope = outgoing_rx.recv() => {
let Some(envelope) = envelope else {
break;
};
route_outgoing_envelope(&mut outbound_connections, envelope).await;
}
}
}
}
async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
let channel_capacity = args.channel_capacity.max(1);
let installation_id = resolve_installation_id(&args.config.codex_home).await?;
@@ -390,7 +409,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
let (event_tx, event_rx) = mpsc::channel::<InProcessServerEvent>(channel_capacity);
let runtime_handle = tokio::spawn(async move {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<OutgoingEnvelope>(channel_capacity);
let (outgoing_tx, outgoing_rx) = mpsc::channel::<OutgoingEnvelope>(channel_capacity);
let auth_manager =
AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env)
.await;
@@ -418,11 +437,12 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
/*disconnect_sender*/ None,
),
);
let mut outbound_handle = tokio::spawn(async move {
while let Some(envelope) = outgoing_rx.recv().await {
route_outgoing_envelope(&mut outbound_connections, envelope).await;
}
});
let (outbound_shutdown_tx, outbound_shutdown_rx) = oneshot::channel();
let mut outbound_handle = tokio::spawn(run_outbound_router(
outgoing_rx,
outbound_connections,
outbound_shutdown_rx,
));
let processor_outgoing = Arc::clone(&outgoing_message_sender);
let config_manager = ConfigManager::new(
@@ -709,8 +729,8 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
"in-process app-server runtime is shutting down",
)))
.await;
// Drop the runtime's last sender before awaiting the router task so
// `outgoing_rx.recv()` can observe channel closure and exit cleanly.
// Detached processor work can retain outgoing senders, so channel
// closure alone cannot be used to shut down the outbound router.
drop(outgoing_message_sender);
for (_, response_tx) in pending_request_responses {
let _ = response_tx.send(Err(internal_error(
@@ -722,6 +742,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
processor_handle.abort();
let _ = processor_handle.await;
}
let _ = outbound_shutdown_tx.send(());
if let Err(_elapsed) = timeout(SHUTDOWN_TIMEOUT, &mut outbound_handle).await {
outbound_handle.abort();
let _ = outbound_handle.await;
@@ -896,6 +917,30 @@ mod tests {
.expect("in-process runtime should shutdown cleanly");
}
#[tokio::test(start_paused = true)]
async fn in_process_outbound_router_shutdown_does_not_wait_for_retained_sender() {
let (outgoing_tx, outgoing_rx) = mpsc::channel(/*buffer*/ 1);
let retained_outgoing_tx = outgoing_tx.clone();
drop(outgoing_tx);
let (shutdown_tx, shutdown_rx) = oneshot::channel();
let mut outbound_handle = tokio::spawn(run_outbound_router(
outgoing_rx,
HashMap::new(),
shutdown_rx,
));
assert!(!retained_outgoing_tx.is_closed());
shutdown_tx
.send(())
.expect("outbound router should accept explicit shutdown");
timeout(SHUTDOWN_TIMEOUT, &mut outbound_handle)
.await
.expect("outbound router should not wait for its retained sender")
.expect("outbound router should complete successfully");
assert!(retained_outgoing_tx.is_closed());
}
#[tokio::test(start_paused = true)]
async fn in_process_shutdown_waits_for_analytics_flush_budget() {
let (client_tx, mut client_rx) = mpsc::channel(/*buffer*/ 1);