fix(app-server): bind catalog subscriptions to live connections

This commit is contained in:
Brent Traut
2026-07-08 10:51:28 -07:00
parent a690532576
commit e946afdac9
3 changed files with 32 additions and 21 deletions
+1 -3
View File
@@ -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> {
+12 -18
View File
@@ -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;
}
+19
View File
@@ -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)