fix: MCP leaks in app-server (openai#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

openai#16895

Author: jif-oai

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<usize> {
@@ -5591,6 +5600,66 @@ impl CodexMessageProcessor {
             .await;
     }
 
+    async fn unload_thread_without_subscribers(
+        &mut self,
+        thread_id: ThreadId,
+        thread: Arc<CodexThread>,
+    ) {
+        // 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::<ThreadId>::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::<ThreadId>::new());
 
         assert!(
             manager

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<ThreadId> {
+        {
             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::<Vec<_>>()
-        };
-
-        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"
-            );
         }
     }
 }

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<String> {
+    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::<ThreadStartResponse>(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::<Vec<_>>();
+    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::<Vec<_>>();
+    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<ThreadLoadedListResponse> {
+    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::<ThreadLoadedListResponse>(response)
+}
+
 async fn send_config_read_request(stream: &mut WsClient, id: i64) -> Result<()> {
     send_request(
         stream,