From 8d58899297deeccdebb1b162d9738eda58c46790 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 10 Apr 2026 15:31:26 +0100 Subject: [PATCH] fix: MCP leaks in app-server (#17223) The disconnect path now reuses the same teardown flow as explicit unsubscribe, and the thread-state bookkeeping consistently reports only threads that lost their last subscriber https://github.com/openai/codex/issues/16895 --- .../app-server/src/codex_message_processor.rs | 129 ++++++++++-------- codex-rs/app-server/src/thread_state.rs | 37 +---- .../suite/v2/connection_handling_websocket.rs | 108 +++++++++++++++ 3 files changed, 190 insertions(+), 84 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index f37f89662..2b2f888e5 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -3883,9 +3883,18 @@ impl CodexMessageProcessor { self.command_exec_manager .connection_closed(connection_id) .await; - self.thread_state_manager + let thread_ids_with_no_subscribers = self + .thread_state_manager .remove_connection(connection_id) .await; + for thread_id in thread_ids_with_no_subscribers { + let Ok(thread) = self.thread_manager.get_thread(thread_id).await else { + self.finalize_thread_teardown(thread_id).await; + continue; + }; + self.unload_thread_without_subscribers(thread_id, thread) + .await; + } } pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver { @@ -5591,6 +5600,66 @@ impl CodexMessageProcessor { .await; } + async fn unload_thread_without_subscribers( + &mut self, + thread_id: ThreadId, + thread: Arc, + ) { + // This connection was the last subscriber. Only now do we unload the thread. + info!("thread {thread_id} has no subscribers; shutting down"); + let should_start_unload_task = self.pending_thread_unloads.lock().await.insert(thread_id); + + // Any pending app-server -> client requests for this thread can no longer be + // answered; cancel their callbacks before shutdown/unload. + self.outgoing + .cancel_requests_for_thread(thread_id, /*error*/ None) + .await; + self.thread_state_manager + .remove_thread_state(thread_id) + .await; + + if !should_start_unload_task { + return; + } + + let outgoing = self.outgoing.clone(); + let pending_thread_unloads = self.pending_thread_unloads.clone(); + let thread_manager = self.thread_manager.clone(); + let thread_watch_manager = self.thread_watch_manager.clone(); + tokio::spawn(async move { + match Self::wait_for_thread_shutdown(&thread).await { + ThreadShutdownResult::Complete => { + if thread_manager.remove_thread(&thread_id).await.is_none() { + info!("thread {thread_id} was already removed before teardown finalized"); + thread_watch_manager + .remove_thread(&thread_id.to_string()) + .await; + pending_thread_unloads.lock().await.remove(&thread_id); + return; + } + thread_watch_manager + .remove_thread(&thread_id.to_string()) + .await; + let notification = ThreadClosedNotification { + thread_id: thread_id.to_string(), + }; + outgoing + .send_server_notification(ServerNotification::ThreadClosed(notification)) + .await; + pending_thread_unloads.lock().await.remove(&thread_id); + } + ThreadShutdownResult::SubmitFailed => { + pending_thread_unloads.lock().await.remove(&thread_id); + warn!("failed to submit Shutdown to thread {thread_id}"); + } + ThreadShutdownResult::TimedOut => { + pending_thread_unloads.lock().await.remove(&thread_id); + warn!("thread {thread_id} shutdown timed out; leaving thread loaded"); + } + } + }); + } + async fn thread_unsubscribe( &self, request_id: ConnectionRequestId, @@ -5638,58 +5707,8 @@ impl CodexMessageProcessor { } if !self.thread_state_manager.has_subscribers(thread_id).await { - // This connection was the last subscriber. Only now do we unload the thread. - info!("thread {thread_id} has no subscribers; shutting down"); - self.pending_thread_unloads.lock().await.insert(thread_id); - // Any pending app-server -> client requests for this thread can no longer be - // answered; cancel their callbacks before shutdown/unload. - self.outgoing - .cancel_requests_for_thread(thread_id, /*error*/ None) + self.unload_thread_without_subscribers(thread_id, thread) .await; - self.thread_state_manager - .remove_thread_state(thread_id) - .await; - - let outgoing = self.outgoing.clone(); - let pending_thread_unloads = self.pending_thread_unloads.clone(); - let thread_manager = self.thread_manager.clone(); - let thread_watch_manager = self.thread_watch_manager.clone(); - tokio::spawn(async move { - match Self::wait_for_thread_shutdown(&thread).await { - ThreadShutdownResult::Complete => { - if thread_manager.remove_thread(&thread_id).await.is_none() { - info!( - "thread {thread_id} was already removed before unsubscribe finalized" - ); - thread_watch_manager - .remove_thread(&thread_id.to_string()) - .await; - pending_thread_unloads.lock().await.remove(&thread_id); - return; - } - thread_watch_manager - .remove_thread(&thread_id.to_string()) - .await; - let notification = ThreadClosedNotification { - thread_id: thread_id.to_string(), - }; - outgoing - .send_server_notification(ServerNotification::ThreadClosed( - notification, - )) - .await; - pending_thread_unloads.lock().await.remove(&thread_id); - } - ThreadShutdownResult::SubmitFailed => { - pending_thread_unloads.lock().await.remove(&thread_id); - warn!("failed to submit Shutdown to thread {thread_id}"); - } - ThreadShutdownResult::TimedOut => { - pending_thread_unloads.lock().await.remove(&thread_id); - warn!("thread {thread_id} shutdown timed out; leaving thread loaded"); - } - } - }); } self.outgoing @@ -10047,7 +10066,8 @@ mod tests { state.lock().await.cancel_tx = Some(cancel_tx); } - manager.remove_connection(connection_a).await; + let threads_to_unload = manager.remove_connection(connection_a).await; + assert_eq!(threads_to_unload, Vec::::new()); assert!( tokio::time::timeout(Duration::from_millis(20), &mut cancel_rx) .await @@ -10068,7 +10088,8 @@ mod tests { let connection = ConnectionId(1); manager.connection_initialized(connection).await; - manager.remove_connection(connection).await; + let threads_to_unload = manager.remove_connection(connection).await; + assert_eq!(threads_to_unload, Vec::::new()); assert!( manager diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 0fe835fc7..805a64bfb 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -352,8 +352,8 @@ impl ThreadStateManager { true } - pub(crate) async fn remove_connection(&self, connection_id: ConnectionId) { - let thread_states = { + pub(crate) async fn remove_connection(&self, connection_id: ConnectionId) -> Vec { + { let mut state = self.state.lock().await; state.live_connections.remove(&connection_id); let thread_ids = state @@ -367,36 +367,13 @@ impl ThreadStateManager { } thread_ids .into_iter() - .map(|thread_id| { - ( - thread_id, - state - .threads - .get(&thread_id) - .is_none_or(|thread_entry| thread_entry.connection_ids.is_empty()), - state - .threads - .get(&thread_id) - .map(|thread_entry| thread_entry.state.clone()), - ) + .filter(|thread_id| { + state + .threads + .get(thread_id) + .is_some_and(|thread_entry| thread_entry.connection_ids.is_empty()) }) .collect::>() - }; - - for (thread_id, no_subscribers, thread_state) in thread_states { - if !no_subscribers { - continue; - } - let Some(thread_state) = thread_state else { - continue; - }; - let listener_generation = thread_state.lock().await.listener_generation; - tracing::debug!( - thread_id = %thread_id, - connection_id = ?connection_id, - listener_generation, - "retaining thread listener after connection disconnect left zero subscribers" - ); } } } diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index bc311b31c..30caa1376 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -2,6 +2,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::bail; use app_test_support::create_mock_responses_server_sequence_unchecked; +use app_test_support::to_response; use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_app_server_protocol::ClientInfo; @@ -12,6 +13,10 @@ use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ThreadLoadedListParams; +use codex_app_server_protocol::ThreadLoadedListResponse; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; use futures::SinkExt; use futures::StreamExt; use hmac::Hmac; @@ -332,6 +337,37 @@ async fn websocket_transport_allows_unauthenticated_non_loopback_startup_by_defa Ok(()) } +#[tokio::test] +async fn websocket_disconnect_unloads_last_subscribed_thread() -> Result<()> { + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; + + let mut ws1 = connect_websocket(bind_addr).await?; + send_initialize_request(&mut ws1, /*id*/ 1, "ws_thread_owner").await?; + read_response_for_id(&mut ws1, /*id*/ 1).await?; + + let thread_id = start_thread(&mut ws1, /*id*/ 2).await?; + assert_loaded_threads(&mut ws1, /*id*/ 3, &[thread_id.as_str()]).await?; + + ws1.close(None).await.context("failed to close websocket")?; + drop(ws1); + + let mut ws2 = connect_websocket(bind_addr).await?; + send_initialize_request(&mut ws2, /*id*/ 4, "ws_reconnect_client").await?; + read_response_for_id(&mut ws2, /*id*/ 4).await?; + + wait_for_loaded_threads(&mut ws2, /*first_id*/ 5, &[]).await?; + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + Ok(()) +} + pub(super) async fn spawn_websocket_server(codex_home: &Path) -> Result<(Child, SocketAddr)> { spawn_websocket_server_with_args(codex_home, "ws://127.0.0.1:0", &[]).await } @@ -564,6 +600,78 @@ pub(super) async fn send_initialize_request( .await } +async fn start_thread(stream: &mut WsClient, id: i64) -> Result { + send_request( + stream, + "thread/start", + id, + Some(serde_json::to_value(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + })?), + ) + .await?; + let response = read_response_for_id(stream, id).await?; + let ThreadStartResponse { thread, .. } = to_response::(response)?; + Ok(thread.id) +} + +async fn assert_loaded_threads(stream: &mut WsClient, id: i64, expected: &[&str]) -> Result<()> { + let response = request_loaded_threads(stream, id).await?; + let mut actual = response.data; + actual.sort(); + let mut expected = expected + .iter() + .map(|thread_id| (*thread_id).to_string()) + .collect::>(); + expected.sort(); + assert_eq!(actual, expected); + assert_eq!(response.next_cursor, None); + Ok(()) +} + +async fn wait_for_loaded_threads( + stream: &mut WsClient, + first_id: i64, + expected: &[&str], +) -> Result<()> { + let mut next_id = first_id; + let expected = expected + .iter() + .map(|thread_id| (*thread_id).to_string()) + .collect::>(); + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let response = request_loaded_threads(stream, next_id).await?; + next_id += 1; + let mut actual = response.data; + actual.sort(); + if actual == expected { + return Ok::<(), anyhow::Error>(()); + } + sleep(Duration::from_millis(50)).await; + } + }) + .await + .context("timed out waiting for loaded thread list")??; + Ok(()) +} + +async fn request_loaded_threads( + stream: &mut WsClient, + id: i64, +) -> Result { + send_request( + stream, + "thread/loaded/list", + id, + Some(serde_json::to_value(ThreadLoadedListParams::default())?), + ) + .await?; + let response = read_response_for_id(stream, id).await?; + to_response::(response) +} + async fn send_config_read_request(stream: &mut WsClient, id: i64) -> Result<()> { send_request( stream,