From 5368909df2f7931a77afe9c0256e7d4ed3d77158 Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Tue, 7 Jul 2026 04:07:01 -0700 Subject: [PATCH] app-server: serialize teardown admission by thread id --- codex-rs/app-server/src/request_processors.rs | 1 + .../request_processors/thread_lifecycle.rs | 116 ++++++------- .../request_processors/thread_processor.rs | 157 ++++++++++++++---- .../src/request_processors/thread_teardown.rs | 27 ++- .../src/request_processors/turn_processor.rs | 11 +- 5 files changed, 208 insertions(+), 104 deletions(-) diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 60f52c360a6a..fbc141b8e25f 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -619,6 +619,7 @@ use self::thread_resume_redaction::*; use self::thread_summary::*; use self::thread_teardown::PendingThreadUnloadStartResult; pub(crate) use self::thread_teardown::PendingThreadUnloads; +use self::thread_teardown::start_thread_teardown; use self::thread_teardown::wait_for_thread_unload; pub(crate) use self::thread_lifecycle::populate_thread_turns_from_history; diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 3830b96e30b1..f2e625fddd7f 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -120,10 +120,6 @@ impl UnloadingState { } } -#[expect( - clippy::await_holding_invalid_type, - reason = "idle unload admission holds both lifecycle guards through its final rechecks" -)] async fn try_start_idle_thread_unload( unloading_state: &mut UnloadingState, thread_list_state_permit: Arc, @@ -139,10 +135,6 @@ where let Ok(permit) = thread_list_state_permit.acquire_owned().await else { return None; }; - let admission = pending_thread_unloads.lock_admission().await; - if pending_thread_unloads.contains(conversation_id) { - return None; - } if !unloading_state.should_unload_now() { return None; } @@ -160,7 +152,6 @@ where return None; } let start = pending_thread_unloads.try_start(conversation_id, teardown); - drop(admission); drop(permit); Some(start) } @@ -176,15 +167,11 @@ pub(super) enum ThreadShutdownResult { } pub(super) enum EnsureConversationListenerResult { - Attached, + Attached(Arc), ConnectionClosed, } -#[expect( - clippy::await_holding_invalid_type, - reason = "listener subscription must be serialized against pending unloads" -)] -pub(super) async fn ensure_conversation_listener( +pub(super) async fn ensure_conversation_listener_under_permit( listener_task_context: ListenerTaskContext, conversation_id: ThreadId, connection_id: ConnectionId, @@ -202,32 +189,25 @@ pub(super) async fn ensure_conversation_listener( ))); } }; - let thread_state = { - let _admission = listener_task_context - .pending_thread_unloads - .lock_admission() - .await; - if listener_task_context - .pending_thread_unloads - .contains(conversation_id) - { - return Err(invalid_request(format!( - "thread {conversation_id} is closing; retry after the thread is closed" - ))); - } - let Some(thread_state) = listener_task_context - .thread_state_manager - .try_ensure_connection_subscribed(conversation_id, connection_id, raw_events_enabled) - .await - else { - return Ok(EnsureConversationListenerResult::ConnectionClosed); - }; - thread_state + if listener_task_context + .pending_thread_unloads + .contains(conversation_id) + { + return Err(invalid_request(format!( + "thread {conversation_id} is closing; retry after the thread is closed" + ))); + } + let Some(thread_state) = listener_task_context + .thread_state_manager + .try_ensure_connection_subscribed(conversation_id, connection_id, raw_events_enabled) + .await + else { + return Ok(EnsureConversationListenerResult::ConnectionClosed); }; if let Err(error) = ensure_listener_task_running( listener_task_context.clone(), conversation_id, - conversation, + Arc::clone(&conversation), Arc::clone(&thread_state), ) .await @@ -242,7 +222,40 @@ pub(super) async fn ensure_conversation_listener( .await; return Err(error); } - Ok(EnsureConversationListenerResult::Attached) + Ok(EnsureConversationListenerResult::Attached(conversation)) +} + +pub(super) async fn ensure_conversation_listener_serialized( + listener_task_context: ListenerTaskContext, + conversation_id: ThreadId, + connection_id: ConnectionId, + raw_events_enabled: bool, +) -> Result { + loop { + let permit = listener_task_context + .thread_list_state_permit + .clone() + .acquire_owned() + .await + .map_err(|err| { + internal_error(format!("failed to acquire thread list state permit: {err}")) + })?; + let pending = listener_task_context + .pending_thread_unloads + .subscribe(conversation_id); + if let Some(pending) = pending { + drop(permit); + wait_for_thread_unload(pending).await; + continue; + } + return ensure_conversation_listener_under_permit( + listener_task_context, + conversation_id, + connection_id, + raw_events_enabled, + ) + .await; + } } pub(super) fn log_listener_attach_result( @@ -252,7 +265,7 @@ pub(super) fn log_listener_attach_result( thread_kind: &'static str, ) { match result { - Ok(EnsureConversationListenerResult::Attached) => {} + Ok(EnsureConversationListenerResult::Attached(_)) => {} Ok(EnsureConversationListenerResult::ConnectionClosed) => { tracing::debug!( thread_id = %thread_id, @@ -271,7 +284,7 @@ pub(super) fn log_listener_attach_result( #[expect( clippy::await_holding_invalid_type, - reason = "listener registration must be serialized against pending unloads" + reason = "listener registration holds current thread state through the exact registry swap" )] pub(super) async fn ensure_listener_task_running( listener_task_context: ListenerTaskContext, @@ -315,7 +328,6 @@ pub(super) async fn ensure_listener_task_running( .. } = listener_task_context; - let pending_thread_unloads_guard = pending_thread_unloads.lock_admission().await; if pending_thread_unloads.contains(conversation_id) { return Err(invalid_request(format!( "thread {conversation_id} is closing; retry after the thread is closed" @@ -371,7 +383,6 @@ pub(super) async fn ensure_listener_task_running( &thread_state, &thread_watch_manager, &outgoing_for_task, - &pending_thread_unloads, listener_command, ) .await; @@ -490,7 +501,6 @@ pub(super) async fn ensure_listener_task_running( let _ = thread_state_guard.clear_listener(); } drop(thread_state_guard); - drop(pending_thread_unloads_guard); match registration { ThreadListenerRegistration::Registered(retired) => { if let Some(retired) = retired { @@ -569,7 +579,6 @@ pub(super) async fn handle_thread_listener_command( thread_state: &Arc>, thread_watch_manager: &ThreadWatchManager, outgoing: &Arc, - pending_thread_unloads: &PendingThreadUnloads, listener_command: ThreadListenerCommand, ) { match listener_command { @@ -581,7 +590,6 @@ pub(super) async fn handle_thread_listener_command( thread_state, thread_watch_manager, outgoing, - pending_thread_unloads, *resume_request, ) .await; @@ -633,7 +641,6 @@ pub(super) async fn handle_pending_thread_resume_request( thread_state: &Arc>, thread_watch_manager: &ThreadWatchManager, outgoing: &Arc, - pending_thread_unloads: &PendingThreadUnloads, pending: crate::thread_state::PendingThreadResumeRequest, ) { let active_turn = { @@ -701,23 +708,6 @@ pub(super) async fn handle_pending_thread_resume_request( } } - let unload_pending = { - let _admission = pending_thread_unloads.lock_admission().await; - pending_thread_unloads.contains(conversation_id) - }; - if unload_pending { - outgoing - .send_error( - request_id, - invalid_request(format!( - "thread {conversation_id} is closing; retry thread/resume after the thread is closed" - )), - ) - .await; - let _ = resume_handled_tx.send(ConnectionReservationAction::Cancel); - return; - } - let config_snapshot = pending.config_snapshot; let cwd = config_snapshot.cwd().clone(); let ThreadConfigSnapshot { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 6834d95e497e..7b93c1215734 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -778,6 +778,20 @@ impl ThreadRequestProcessor { }) } + pub(super) async fn acquire_thread_list_state_permit_after_unload( + &self, + thread_id: ThreadId, + ) -> Result, JSONRPCErrorError> { + loop { + let permit = self.acquire_thread_list_state_permit().await?; + let Some(pending) = self.pending_thread_unloads.subscribe(thread_id) else { + return Ok(permit); + }; + drop(permit); + wait_for_thread_unload(pending).await; + } + } + async fn set_app_server_client_info( thread: &CodexThread, app_server_client_name: Option, @@ -797,14 +811,27 @@ impl ThreadRequestProcessor { .map_err(|err| internal_error(format!("failed to set app server client info: {err}"))) } - async fn finalize_thread_teardown(&self, thread_id: ThreadId) { + fn start_finalize_thread_teardown_under_permit( + &self, + thread_id: ThreadId, + ) -> Option> { + let processor = self.clone(); + start_thread_teardown(self.pending_thread_unloads.clone(), thread_id, async move { + processor.finalize_thread_teardown_owner(thread_id).await + }) + } + + async fn finalize_thread_teardown_owner(&self, thread_id: ThreadId) { self.outgoing .cancel_requests_for_thread(thread_id, /*error*/ None) .await; - let _ = self + if let Some(listener) = self .thread_state_manager .remove_thread_state(thread_id) - .await; + .await + { + listener.wait_until_closed().await; + } self.thread_watch_manager .remove_thread(&thread_id.to_string()) .await; @@ -817,9 +844,16 @@ impl ThreadRequestProcessor { ) -> Result { let thread_id = ThreadId::from_string(¶ms.thread_id) .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; + let thread_list_state_permit = self + .acquire_thread_list_state_permit_after_unload(thread_id) + .await?; if self.thread_manager.get_thread(thread_id).await.is_err() { - self.finalize_thread_teardown(thread_id).await; + let completed = self.start_finalize_thread_teardown_under_permit(thread_id); + drop(thread_list_state_permit); + if let Some(completed) = completed { + wait_for_thread_unload(completed).await; + } return Ok(ThreadUnsubscribeResponse { status: ThreadUnsubscribeStatus::NotLoaded, }); @@ -858,7 +892,7 @@ impl ThreadRequestProcessor { } } } - self.finalize_thread_teardown(thread_id).await; + self.finalize_thread_teardown_owner(thread_id).await; } fn listener_task_context(&self) -> ListenerTaskContext { @@ -881,7 +915,22 @@ impl ThreadRequestProcessor { connection_id: ConnectionId, raw_events_enabled: bool, ) -> Result { - super::thread_lifecycle::ensure_conversation_listener( + super::thread_lifecycle::ensure_conversation_listener_serialized( + self.listener_task_context(), + conversation_id, + connection_id, + raw_events_enabled, + ) + .await + } + + async fn ensure_conversation_listener_under_permit( + &self, + conversation_id: ThreadId, + connection_id: ConnectionId, + raw_events_enabled: bool, + ) -> Result { + super::thread_lifecycle::ensure_conversation_listener_under_permit( self.listener_task_context(), conversation_id, connection_id, @@ -1260,7 +1309,7 @@ impl ThreadRequestProcessor { // Auto-attach a thread listener when starting a thread. log_listener_attach_result( - super::thread_lifecycle::ensure_conversation_listener( + super::thread_lifecycle::ensure_conversation_listener_serialized( listener_task_context.clone(), thread_id, request_id.connection_id, @@ -2591,14 +2640,25 @@ impl ThreadRequestProcessor { .thread_state_manager .remove_connection(connection_id) .await; + let Ok(thread_list_state_permit) = self.thread_list_state_permit.acquire().await else { + return; + }; + let mut pending_cleanups = Vec::new(); for thread_id in thread_ids { if self.thread_manager.get_thread(thread_id).await.is_err() { // Reconcile stale app-server bookkeeping when the thread has already been // removed from the core manager. - self.finalize_thread_teardown(thread_id).await; + if let Some(completed) = self.start_finalize_thread_teardown_under_permit(thread_id) + { + pending_cleanups.push(completed); + } } } + drop(thread_list_state_permit); + for completed in pending_cleanups { + wait_for_thread_unload(completed).await; + } } pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver { @@ -2611,6 +2671,16 @@ impl ThreadRequestProcessor { thread_id: ThreadId, connection_ids: Vec, ) { + let _thread_list_state_permit = match self.acquire_thread_list_state_permit().await { + Ok(permit) => permit, + Err(error) => { + warn!( + "failed to serialize automatic listener attachment for thread {thread_id}: {}", + error.message + ); + return; + } + }; let mut raw_events_enabled = false; if let Ok(thread) = self.thread_manager.get_thread(thread_id).await { let config_snapshot = thread.config_snapshot().await; @@ -2634,8 +2704,12 @@ impl ThreadRequestProcessor { for connection_id in connection_ids { log_listener_attach_result( - self.ensure_conversation_listener(thread_id, connection_id, raw_events_enabled) - .await, + self.ensure_conversation_listener_under_permit( + thread_id, + connection_id, + raw_events_enabled, + ) + .await, thread_id, connection_id, "thread", @@ -2651,24 +2725,6 @@ impl ThreadRequestProcessor { app_server_client_version: Option, supports_openai_form_elicitation: bool, ) -> Result<(), JSONRPCErrorError> { - if let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) { - let unload_pending = { - let _admission = self.pending_thread_unloads.lock_admission().await; - self.pending_thread_unloads.contains(thread_id) - }; - if unload_pending { - self.outgoing - .send_error( - request_id, - invalid_request(format!( - "thread {thread_id} is closing; retry thread/resume after the thread is closed" - )), - ) - .await; - return Ok(()); - } - } - if params.sandbox.is_some() && params.permissions.is_some() { self.outgoing .send_error( @@ -2681,7 +2737,45 @@ impl ThreadRequestProcessor { let redact_resume_payloads = should_redact_thread_resume_payloads(app_server_client_name.as_deref()); - let thread_list_state_permit = match self.acquire_thread_list_state_permit().await { + let admission_thread_id = if params.history.is_some() { + ThreadId::from_string(¶ms.thread_id).ok() + } else if params.path.is_some() { + match self + .read_stored_thread_for_resume( + ¶ms.thread_id, + params.path.as_ref(), + /*include_history*/ false, + ) + .await + { + Ok(stored_thread) => Some(stored_thread.thread_id), + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return Ok(()); + } + } + } else { + match ThreadId::from_string(¶ms.thread_id) { + Ok(thread_id) => Some(thread_id), + Err(err) => { + self.outgoing + .send_error( + request_id, + invalid_request(format!("invalid session id: {err}")), + ) + .await; + return Ok(()); + } + } + }; + let permit_result = match admission_thread_id { + Some(thread_id) => { + self.acquire_thread_list_state_permit_after_unload(thread_id) + .await + } + None => self.acquire_thread_list_state_permit().await, + }; + let thread_list_state_permit = match permit_result { Ok(permit) => permit, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -2841,7 +2935,7 @@ impl ThreadRequestProcessor { }; // Auto-attach a thread listener when resuming a thread. log_listener_attach_result( - self.ensure_conversation_listener( + self.ensure_conversation_listener_under_permit( thread_id, request_id.connection_id, /*raw_events_enabled*/ false, @@ -3084,7 +3178,8 @@ impl ThreadRequestProcessor { match wait_for_thread_shutdown(&existing_thread).await { ThreadShutdownResult::Complete => { self.thread_manager.remove_thread(&existing_thread_id).await; - self.finalize_thread_teardown(existing_thread_id).await; + self.finalize_thread_teardown_owner(existing_thread_id) + .await; // Shutdown can flush newer rollout items, so reload the // stored thread before starting the replacement session. return Ok(RunningThreadResumeResult::NotRunning(None)); diff --git a/codex-rs/app-server/src/request_processors/thread_teardown.rs b/codex-rs/app-server/src/request_processors/thread_teardown.rs index fb3c95ceba7f..e7fa235e22ee 100644 --- a/codex-rs/app-server/src/request_processors/thread_teardown.rs +++ b/codex-rs/app-server/src/request_processors/thread_teardown.rs @@ -11,7 +11,6 @@ pub(crate) struct PendingThreadUnloads { #[derive(Default)] struct PendingThreadUnloadsInner { - admission: Mutex<()>, registry: StdMutex, tasks: TaskTracker, } @@ -64,10 +63,6 @@ impl Drop for PendingThreadUnloadOwner { } impl PendingThreadUnloads { - pub(super) async fn lock_admission(&self) -> tokio::sync::MutexGuard<'_, ()> { - self.inner.admission.lock().await - } - fn lock_registry(&self) -> std::sync::MutexGuard<'_, PendingThreadUnloadRegistry> { self.inner .registry @@ -79,6 +74,13 @@ impl PendingThreadUnloads { self.lock_registry().entries.contains_key(&thread_id) } + pub(super) fn subscribe(&self, thread_id: ThreadId) -> Option> { + self.lock_registry() + .entries + .get(&thread_id) + .map(watch::Sender::subscribe) + } + pub(super) fn try_claim(&self, thread_id: ThreadId) -> PendingThreadUnloadClaimResult { let mut registry = self.lock_registry(); if registry.closing { @@ -164,6 +166,21 @@ pub(super) async fn wait_for_thread_unload(mut completed: watch::Receiver) } } +pub(super) fn start_thread_teardown( + pending: PendingThreadUnloads, + thread_id: ThreadId, + teardown: F, +) -> Option> +where + F: std::future::Future + Send + 'static, +{ + match pending.try_start(thread_id, teardown) { + PendingThreadUnloadStartResult::Started(completed) + | PendingThreadUnloadStartResult::Pending(completed) => Some(completed), + PendingThreadUnloadStartResult::Closing => None, + } +} + #[cfg(test)] #[path = "thread_teardown_tests.rs"] mod tests; diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index 63d590f4b086..266ea7aa4945 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -960,9 +960,10 @@ impl TurnRequestProcessor { request_id: &ConnectionRequestId, thread_id: &str, ) -> Result)>, JSONRPCErrorError> { - let (thread_id, thread) = self.load_thread(thread_id).await?; + let thread_id = ThreadId::from_string(thread_id) + .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; - match self + let thread = match self .ensure_conversation_listener( thread_id, request_id.connection_id, @@ -970,12 +971,12 @@ impl TurnRequestProcessor { ) .await { - Ok(EnsureConversationListenerResult::Attached) => {} + Ok(EnsureConversationListenerResult::Attached(thread)) => thread, Ok(EnsureConversationListenerResult::ConnectionClosed) => { return Ok(None); } Err(error) => return Err(error), - } + }; if !thread.enabled(Feature::RealtimeConversation) { return Err(invalid_request(format!( @@ -1428,7 +1429,7 @@ impl TurnRequestProcessor { connection_id: ConnectionId, raw_events_enabled: bool, ) -> Result { - super::thread_lifecycle::ensure_conversation_listener( + super::thread_lifecycle::ensure_conversation_listener_serialized( self.listener_task_context(), conversation_id, connection_id,