From e946afdac94ec648f0ff288d14c7ee71973a670c Mon Sep 17 00:00:00 2001
From: Brent Traut <btraut@openai.com>
Date: Wed, 8 Jul 2026 02:17:37 -0700
Subject: [PATCH] fix(app-server): bind catalog subscriptions to live
 connections

---
 codex-rs/app-server/src/message_processor.rs |  4 +--
 codex-rs/app-server/src/thread_catalog.rs    | 30 ++++++++------------
 codex-rs/app-server/src/thread_state.rs      | 19 +++++++++++++
 3 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs
index 2f549fe25a..45cd0533ac 100644
--- a/codex-rs/app-server/src/message_processor.rs
+++ b/codex-rs/app-server/src/message_processor.rs
@@ -255,6 +255,7 @@ impl MessageProcessor {
         let thread_catalog_subscriptions = ThreadCatalogSubscriptions::new(
             outgoing.clone(),
             Arc::clone(&thread_store),
+            thread_state_manager.clone(),
             Arc::clone(&thread_list_state_permit),
             config.model_provider_id.clone(),
             config.cwd.clone(),
@@ -729,9 +730,6 @@ impl MessageProcessor {
             .connection_closed(connection_id)
             .await;
         self.thread_processor.connection_closed(connection_id).await;
-        self.thread_catalog_subscriptions
-            .connection_closed(connection_id)
-            .await;
     }
 
     pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver<usize> {
diff --git a/codex-rs/app-server/src/thread_catalog.rs b/codex-rs/app-server/src/thread_catalog.rs
index 437dd3eee7..3ad711ae1b 100644
--- a/codex-rs/app-server/src/thread_catalog.rs
+++ b/codex-rs/app-server/src/thread_catalog.rs
@@ -1,4 +1,3 @@
-use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::Weak;
 
@@ -9,7 +8,6 @@ use codex_thread_store::ReadThreadParams;
 use codex_thread_store::ThreadCatalogChange;
 use codex_thread_store::ThreadStore;
 use codex_utils_absolute_path::AbsolutePathBuf;
-use tokio::sync::Mutex;
 use tokio::sync::Semaphore;
 use tokio::sync::broadcast::error::RecvError;
 use tracing::warn;
@@ -17,40 +15,41 @@ use tracing::warn;
 use crate::outgoing_message::ConnectionId;
 use crate::outgoing_message::OutgoingMessageSender;
 use crate::request_processors::thread_from_stored_thread;
+use crate::thread_state::ThreadStateManager;
 
 #[derive(Clone)]
 pub(crate) struct ThreadCatalogSubscriptions {
-    connection_ids: Arc<Mutex<HashSet<ConnectionId>>>,
+    thread_state_manager: ThreadStateManager,
 }
 
 impl ThreadCatalogSubscriptions {
     pub(crate) fn new(
         outgoing: Arc<OutgoingMessageSender>,
         thread_store: Arc<dyn ThreadStore>,
+        thread_state_manager: ThreadStateManager,
         catalog_state_permit: Arc<Semaphore>,
         fallback_provider: String,
         fallback_cwd: AbsolutePathBuf,
     ) -> Self {
-        let connection_ids = Arc::new(Mutex::new(HashSet::new()));
         let changes = thread_store.subscribe_catalog_changes();
         tokio::spawn(run_catalog_listener(
             changes,
             Arc::downgrade(&thread_store),
             Arc::clone(&outgoing),
-            Arc::clone(&connection_ids),
+            thread_state_manager.clone(),
             catalog_state_permit,
             fallback_provider,
             fallback_cwd,
         ));
-        Self { connection_ids }
+        Self {
+            thread_state_manager,
+        }
     }
 
     pub(crate) async fn subscribe(&self, connection_id: ConnectionId) {
-        self.connection_ids.lock().await.insert(connection_id);
-    }
-
-    pub(crate) async fn connection_closed(&self, connection_id: ConnectionId) {
-        self.connection_ids.lock().await.remove(&connection_id);
+        self.thread_state_manager
+            .subscribe_catalog(connection_id)
+            .await;
     }
 }
 
@@ -58,7 +57,7 @@ async fn run_catalog_listener(
     mut changes: tokio::sync::broadcast::Receiver<ThreadCatalogChange>,
     thread_store: Weak<dyn ThreadStore>,
     outgoing: Arc<OutgoingMessageSender>,
-    connection_ids: Arc<Mutex<HashSet<ConnectionId>>>,
+    thread_state_manager: ThreadStateManager,
     catalog_state_permit: Arc<Semaphore>,
     fallback_provider: String,
     fallback_cwd: AbsolutePathBuf,
@@ -72,12 +71,7 @@ async fn run_catalog_listener(
             }
             Err(RecvError::Closed) => break,
         };
-        let connection_ids = connection_ids
-            .lock()
-            .await
-            .iter()
-            .copied()
-            .collect::<Vec<_>>();
+        let connection_ids = thread_state_manager.catalog_subscribers().await;
         if connection_ids.is_empty() {
             continue;
         }
diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs
index fec786373c..35a80de311 100644
--- a/codex-rs/app-server/src/thread_state.rs
+++ b/codex-rs/app-server/src/thread_state.rs
@@ -278,6 +278,7 @@ impl ThreadEntry {
 #[derive(Default)]
 struct ThreadStateManagerInner {
     live_connections: HashMap<ConnectionId, ConnectionCapabilities>,
+    catalog_subscribers: HashSet<ConnectionId>,
     threads: HashMap<ThreadId, ThreadEntry>,
     thread_ids_by_connection: HashMap<ConnectionId, HashSet<ThreadId>>,
 }
@@ -313,6 +314,23 @@ impl ThreadStateManager {
             .insert(connection_id, capabilities);
     }
 
+    pub(crate) async fn subscribe_catalog(&self, connection_id: ConnectionId) {
+        let mut state = self.state.lock().await;
+        if state.live_connections.contains_key(&connection_id) {
+            state.catalog_subscribers.insert(connection_id);
+        }
+    }
+
+    pub(crate) async fn catalog_subscribers(&self) -> Vec<ConnectionId> {
+        self.state
+            .lock()
+            .await
+            .catalog_subscribers
+            .iter()
+            .copied()
+            .collect()
+    }
+
     pub(crate) async fn first_attestation_capable_connection_for_thread(
         &self,
         thread_id: ThreadId,
@@ -543,6 +561,7 @@ impl ThreadStateManager {
         {
             let mut state = self.state.lock().await;
             state.live_connections.remove(&connection_id);
+            state.catalog_subscribers.remove(&connection_id);
             let thread_ids = state
                 .thread_ids_by_connection
                 .remove(&connection_id)