diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 44fee3a71fc..3a33407cba8 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -545,6 +545,8 @@ use crate::error_code::invalid_request; use crate::filters::compute_source_filters; use crate::filters::source_kind_matches; use crate::thread_state::ConnectionCapabilities; +use crate::thread_state::ConnectionReservationAction; +use crate::thread_state::ConnectionReservationOutcome; use crate::thread_state::ThreadListenerCommand; use crate::thread_state::ThreadListenerRegistration; use crate::thread_state::ThreadState; 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 ea7ac253bab..660c76b2941 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -120,6 +120,51 @@ impl UnloadingState { } } +#[expect( + clippy::await_holding_invalid_type, + reason = "idle unload admission holds both lifecycle guards through its final rechecks" +)] +async fn claim_idle_unload_if_permitted( + unloading_state: &mut UnloadingState, + thread_list_state_permit: Arc, + pending_thread_unloads: &Arc>>, + thread_state_manager: &ThreadStateManager, + conversation_id: ThreadId, + is_running: impl std::future::Future, +) -> bool { + let Ok(permit) = thread_list_state_permit.acquire_owned().await else { + return false; + }; + let mut pending = pending_thread_unloads.lock().await; + if pending.contains(&conversation_id) { + return false; + } + if !unloading_state.should_unload_now() { + return false; + } + if is_running.await { + unloading_state.note_thread_activity_observed(); + return false; + } + if thread_state_manager + .has_connections_or_reservations(conversation_id) + .await + { + return false; + } + if !unloading_state.should_unload_now() { + return false; + } + pending.insert(conversation_id); + drop(pending); + drop(permit); + true +} + +#[cfg(test)] +#[path = "thread_lifecycle_tests.rs"] +mod tests; + pub(super) enum ThreadShutdownResult { Complete, SubmitFailed, @@ -384,15 +429,19 @@ pub(super) async fn ensure_listener_task_running( unloading_state.note_thread_activity_observed(); continue; } + if !claim_idle_unload_if_permitted( + &mut unloading_state, + thread_list_state_permit.clone(), + &pending_thread_unloads, + &thread_state_manager, + conversation_id, + async { + matches!(conversation.agent_status().await, AgentStatus::Running) + }, + ) + .await { - let mut pending_thread_unloads = pending_thread_unloads.lock().await; - if pending_thread_unloads.contains(&conversation_id) { - continue; - } - if !unloading_state.should_unload_now() { - continue; - } - pending_thread_unloads.insert(conversation_id); + continue; } unload_thread_without_subscribers( thread_manager.clone(), @@ -519,7 +568,6 @@ pub(super) async fn handle_thread_listener_command( conversation_id, conversation, codex_home, - thread_state_manager, thread_state, thread_watch_manager, outgoing, @@ -568,15 +616,10 @@ pub(super) async fn handle_thread_listener_command( } #[allow(clippy::too_many_arguments)] -#[expect( - clippy::await_holding_invalid_type, - reason = "running-thread resume subscription must be serialized against pending unloads" -)] pub(super) async fn handle_pending_thread_resume_request( conversation_id: ThreadId, conversation: &Arc, _codex_home: &Path, - thread_state_manager: &ThreadStateManager, thread_state: &Arc>, thread_watch_manager: &ThreadWatchManager, outgoing: &Arc, @@ -602,6 +645,7 @@ pub(super) async fn handle_pending_thread_resume_request( .is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress)); let request_id = pending.request_id; + let resume_handled_tx = pending.resume_handled_tx; let connection_id = request_id.connection_id; let mut thread = pending.thread_summary; if pending.include_turns { @@ -633,6 +677,7 @@ pub(super) async fn handle_pending_thread_resume_request( Ok(page) => Some(page), Err(error) => { outgoing.send_error(request_id, error).await; + let _ = resume_handled_tx.send(ConnectionReservationAction::Cancel); return; } } @@ -646,31 +691,21 @@ pub(super) async fn handle_pending_thread_resume_request( } } + if pending_thread_unloads + .lock() + .await + .contains(&conversation_id) { - let pending_thread_unloads = pending_thread_unloads.lock().await; - if pending_thread_unloads.contains(&conversation_id) { - drop(pending_thread_unloads); - outgoing - .send_error( - request_id, - invalid_request(format!( - "thread {conversation_id} is closing; retry thread/resume after the thread is closed" - )), - ) - .await; - return; - } - if !thread_state_manager - .try_add_connection_to_thread(conversation_id, connection_id) - .await - { - tracing::debug!( - thread_id = %conversation_id, - connection_id = ?connection_id, - "skipping running thread resume for closed connection" - ); - return; - } + 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; @@ -714,6 +749,7 @@ pub(super) async fn handle_pending_thread_resume_request( outgoing .send_response_with_thread_originator(request_id, response, originator) .await; + let _ = resume_handled_tx.send(ConnectionReservationAction::Commit); // Match cold resume: metadata-only resume should attach the listener without // paying the cost of turn reconstruction for historical usage replay. if let Some(token_usage_thread) = token_usage_thread { diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs new file mode 100644 index 00000000000..1ab13e80c73 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs @@ -0,0 +1,135 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn resume_reservation_prevents_idle_unload_before_listener_commit() { + let now = Instant::now(); + let (_has_subscribers_tx, has_subscribers_rx) = watch::channel(false); + let (_thread_status_tx, thread_status_rx) = watch::channel(ThreadStatus::Idle); + let mut unloading_state = UnloadingState { + delay: Duration::ZERO, + has_subscribers_rx, + has_subscribers: (false, now), + thread_status_rx, + is_active: (false, now), + }; + let pending = Arc::new(Mutex::new(HashSet::new())); + let thread_state_manager = ThreadStateManager::new(); + let thread_id = ThreadId::new(); + let connection_id = ConnectionId(1); + thread_state_manager + .connection_initialized(connection_id, ConnectionCapabilities::default()) + .await; + let permit = Arc::new(Semaphore::new(1)); + let resume_permit = permit + .clone() + .acquire_owned() + .await + .expect("permit should be open"); + let (resume_handled_tx, resume_outcome_rx) = thread_state_manager + .start_connection_reservation(thread_id, connection_id) + .await + .expect("live connection should reserve the thread"); + drop(resume_permit); + assert!( + !claim_idle_unload_if_permitted( + &mut unloading_state, + permit, + &pending, + &thread_state_manager, + thread_id, + std::future::ready(false), + ) + .await + ); + assert!(pending.lock().await.is_empty()); + + assert!( + thread_state_manager + .unsubscribe_connection_from_thread(thread_id, connection_id) + .await + ); + let (replacement_tx, replacement_outcome_rx) = thread_state_manager + .start_connection_reservation(thread_id, connection_id) + .await + .expect("replacement reservation should succeed"); + resume_handled_tx + .send(ConnectionReservationAction::Commit) + .expect("reservation owner should still be running"); + assert_eq!( + resume_outcome_rx + .await + .expect("reservation owner should report an outcome"), + ConnectionReservationOutcome::Handled, + ); + assert!( + thread_state_manager + .subscribed_connection_ids(thread_id) + .await + .is_empty() + ); + assert!( + thread_state_manager + .has_connections_or_reservations(thread_id) + .await + ); + + drop(replacement_tx); + assert_eq!( + replacement_outcome_rx + .await + .expect("reservation owner should report abandonment"), + ConnectionReservationOutcome::Abandoned, + ); + assert!( + !thread_state_manager + .has_connections_or_reservations(thread_id) + .await + ); + assert!( + thread_state_manager + .subscribed_connection_ids(thread_id) + .await + .is_empty() + ); + thread_state_manager.clear_all_listeners().await; +} + +#[tokio::test] +async fn pending_resume_prevents_disconnect_from_reporting_thread_stale() { + let thread_state_manager = ThreadStateManager::new(); + let thread_id = ThreadId::new(); + let committed_connection_id = ConnectionId(1); + let pending_connection_id = ConnectionId(2); + for connection_id in [committed_connection_id, pending_connection_id] { + thread_state_manager + .connection_initialized(connection_id, ConnectionCapabilities::default()) + .await; + } + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, committed_connection_id) + .await + ); + let (reservation_tx, reservation_outcome_rx) = thread_state_manager + .start_connection_reservation(thread_id, pending_connection_id) + .await + .expect("live connection should reserve the thread"); + + assert!( + thread_state_manager + .remove_connection(committed_connection_id) + .await + .is_empty(), + "a pending reservation must keep the thread from being reported stale" + ); + + drop(reservation_tx); + assert_eq!( + reservation_outcome_rx + .await + .expect("reservation owner should report abandonment"), + ConnectionReservationOutcome::Abandoned, + ); + thread_state_manager.clear_all_listeners().await; +} 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 84cb2803c88..9bc85ba36e8 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -371,6 +371,12 @@ pub(crate) struct ThreadRequestProcessor { enum RunningThreadResumeResult { /// The request was delegated to the loaded thread. Handled, + /// The listener owns the response after the caller releases the global + /// thread-list permit. + ListenerPending { + thread_id: ThreadId, + outcome_rx: oneshot::Receiver, + }, /// No loaded thread handled the request. /// /// The optional stored thread contains the history-bearing probe that cold @@ -2672,7 +2678,7 @@ 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 thread_list_state_permit = match self.acquire_thread_list_state_permit().await { Ok(permit) => permit, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -2689,6 +2695,23 @@ impl ThreadRequestProcessor { .await { Ok(RunningThreadResumeResult::Handled) => return Ok(()), + Ok(RunningThreadResumeResult::ListenerPending { + thread_id, + outcome_rx, + }) => { + drop(thread_list_state_permit); + if !matches!(outcome_rx.await, Ok(ConnectionReservationOutcome::Handled)) { + self.outgoing + .send_error( + request_id, + internal_error(format!( + "failed to subscribe running thread resume for thread {thread_id}: thread listener closed before handling the request" + )), + ) + .await; + } + return Ok(()); + } Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -3131,10 +3154,18 @@ impl ThreadRequestProcessor { .thread_goal_processor .pending_resume_goal_state(existing_thread.as_ref()) .await; + let Some((resume_handled_tx, outcome_rx)) = self + .thread_state_manager + .start_connection_reservation(existing_thread_id, request_id.connection_id) + .await + else { + return Ok(RunningThreadResumeResult::Handled); + }; let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse( Box::new(crate::thread_state::PendingThreadResumeRequest { request_id: request_id.clone(), + resume_handled_tx, history_items, config_snapshot, instruction_sources, @@ -3151,7 +3182,10 @@ impl ThreadRequestProcessor { "failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed" ))); } - return Ok(RunningThreadResumeResult::Handled); + return Ok(RunningThreadResumeResult::ListenerPending { + thread_id: existing_thread_id, + outcome_rx, + }); } Ok(RunningThreadResumeResult::NotRunning(None)) } diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 4a10f973031..b3ae400897e 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -32,6 +32,7 @@ type PendingInterruptQueue = Vec; pub(crate) struct PendingThreadResumeRequest { pub(crate) request_id: ConnectionRequestId, + pub(crate) resume_handled_tx: oneshot::Sender, pub(crate) history_items: Vec, pub(crate) config_snapshot: ThreadConfigSnapshot, pub(crate) instruction_sources: Vec, @@ -389,6 +390,7 @@ mod tests { struct ThreadEntry { state: Arc>, connection_ids: HashSet, + pending_connection_reservations: HashMap>, has_connections_watcher: watch::Sender, } @@ -397,6 +399,7 @@ impl Default for ThreadEntry { Self { state: Arc::new(Mutex::new(ThreadState::default())), connection_ids: HashSet::new(), + pending_connection_reservations: HashMap::new(), has_connections_watcher: watch::channel(false).0, } } @@ -417,6 +420,7 @@ struct ThreadStateManagerInner { live_connections: HashMap, threads: HashMap, thread_ids_by_connection: HashMap>, + next_connection_reservation_id: u64, } #[derive(Default)] @@ -430,6 +434,25 @@ pub(crate) struct ConnectionCapabilities { pub(crate) request_attestation: bool, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ConnectionReservationOutcome { + Handled, + Abandoned, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) enum ConnectionReservationAction { + Commit, + Cancel, +} + +#[derive(Clone, Copy)] +struct ConnectionReservation { + thread_id: ThreadId, + connection_id: ConnectionId, + reservation_id: u64, +} + #[derive(Clone, Default)] pub(crate) struct ThreadStateManager { state: Arc>, @@ -748,16 +771,26 @@ impl ThreadStateManager { return false; } + let removed = if let Some(thread_entry) = state.threads.get_mut(&thread_id) { + let removed_subscription = thread_entry.connection_ids.remove(&connection_id); + let canceled_reservation = thread_entry + .pending_connection_reservations + .remove(&connection_id) + .is_some(); + thread_entry.update_has_connections(); + removed_subscription || canceled_reservation + } else { + false + }; + if !removed { + return false; + } if let Some(thread_ids) = state.thread_ids_by_connection.get_mut(&connection_id) { thread_ids.remove(&thread_id); if thread_ids.is_empty() { state.thread_ids_by_connection.remove(&connection_id); } } - if let Some(thread_entry) = state.threads.get_mut(&thread_id) { - thread_entry.connection_ids.remove(&connection_id); - thread_entry.update_has_connections(); - } }; true @@ -773,6 +806,18 @@ impl ThreadStateManager { .is_some_and(|thread_entry| !thread_entry.connection_ids.is_empty()) } + pub(crate) async fn has_connections_or_reservations(&self, thread_id: ThreadId) -> bool { + self.state + .lock() + .await + .threads + .get(&thread_id) + .is_some_and(|thread_entry| { + !thread_entry.connection_ids.is_empty() + || !thread_entry.pending_connection_reservations.is_empty() + }) + } + pub(crate) async fn try_ensure_connection_subscribed( &self, thread_id: ThreadId, @@ -809,6 +854,7 @@ impl ThreadStateManager { Some(thread_state) } + #[cfg(test)] pub(crate) async fn try_add_connection_to_thread( &self, thread_id: ThreadId, @@ -835,6 +881,117 @@ impl ThreadStateManager { true } + pub(crate) async fn start_connection_reservation( + &self, + thread_id: ThreadId, + connection_id: ConnectionId, + ) -> Option<( + oneshot::Sender, + oneshot::Receiver, + )> { + let mut state = self.state.lock().await; + let listener_registry = self + .listener_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if listener_registry.closing || !state.live_connections.contains_key(&connection_id) { + return None; + } + let reservation_id = state.next_connection_reservation_id; + let Some(next_reservation_id) = reservation_id.checked_add(1) else { + error!("connection reservation id space exhausted"); + return None; + }; + state.next_connection_reservation_id = next_reservation_id; + state + .thread_ids_by_connection + .entry(connection_id) + .or_default() + .insert(thread_id); + let thread_entry = state.threads.entry(thread_id).or_default(); + thread_entry + .pending_connection_reservations + .entry(connection_id) + .or_default() + .insert(reservation_id); + let reservation = ConnectionReservation { + thread_id, + connection_id, + reservation_id, + }; + let (handled_tx, handled_rx) = oneshot::channel(); + let (outcome_tx, outcome_rx) = oneshot::channel(); + let manager = self.clone(); + let task = self.listener_tasks.spawn(async move { + let (action, outcome) = match handled_rx.await { + Ok(action) => (action, ConnectionReservationOutcome::Handled), + Err(_) => ( + ConnectionReservationAction::Cancel, + ConnectionReservationOutcome::Abandoned, + ), + }; + manager + .finish_connection_reservation(reservation, action) + .await; + let _ = outcome_tx.send(outcome); + }); + drop(task); + drop(listener_registry); + drop(state); + Some((handled_tx, outcome_rx)) + } + + async fn finish_connection_reservation( + &self, + reservation: ConnectionReservation, + action: ConnectionReservationAction, + ) { + let ConnectionReservation { + thread_id, + connection_id, + reservation_id, + } = reservation; + let mut state = self.state.lock().await; + let connection_is_live = state.live_connections.contains_key(&connection_id); + let Some(thread_entry) = state.threads.get_mut(&thread_id) else { + return; + }; + let remove_reservation_set = { + let Some(reservations) = thread_entry + .pending_connection_reservations + .get_mut(&connection_id) + else { + return; + }; + if !reservations.remove(&reservation_id) { + return; + } + reservations.is_empty() + }; + if remove_reservation_set { + thread_entry + .pending_connection_reservations + .remove(&connection_id); + } + let committed = matches!(action, ConnectionReservationAction::Commit) && connection_is_live; + if committed { + thread_entry.connection_ids.insert(connection_id); + } + let keep_connection_mapping = thread_entry.connection_ids.contains(&connection_id) + || thread_entry + .pending_connection_reservations + .contains_key(&connection_id); + thread_entry.update_has_connections(); + if !keep_connection_mapping + && let Some(thread_ids) = state.thread_ids_by_connection.get_mut(&connection_id) + { + thread_ids.remove(&thread_id); + if thread_ids.is_empty() { + state.thread_ids_by_connection.remove(&connection_id); + } + } + } + pub(crate) async fn remove_connection(&self, connection_id: ConnectionId) -> Vec { { let mut state = self.state.lock().await; @@ -846,16 +1003,19 @@ impl ThreadStateManager { for thread_id in &thread_ids { if let Some(thread_entry) = state.threads.get_mut(thread_id) { thread_entry.connection_ids.remove(&connection_id); + thread_entry + .pending_connection_reservations + .remove(&connection_id); thread_entry.update_has_connections(); } } thread_ids .into_iter() .filter(|thread_id| { - state - .threads - .get(thread_id) - .is_some_and(|thread_entry| thread_entry.connection_ids.is_empty()) + state.threads.get(thread_id).is_some_and(|thread_entry| { + thread_entry.connection_ids.is_empty() + && thread_entry.pending_connection_reservations.is_empty() + }) }) .collect::>() }