diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 2d3ff975f3..be146d82c7 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -1245,7 +1245,7 @@ pub(crate) async fn apply_bespoke_event_handling( state.pending_rollbacks.take() }; - if let Some(request_id) = pending { + if let Some((request_id, _completion_tx)) = pending { let _thread_list_state_permit = match thread_list_state_permit.acquire().await { Ok(permit) => permit, Err(err) => { @@ -1643,7 +1643,7 @@ async fn handle_thread_rollback_failed( ) { let pending_rollback = thread_state.lock().await.pending_rollbacks.take(); - if let Some(request_id) = pending_rollback { + if let Some((request_id, _completion_tx)) = pending_rollback { outgoing .send_error(request_id, invalid_request(message)) .await; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index aff6d9ad6c..511ea54587 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -818,6 +818,7 @@ impl MessageProcessor { session_state: &ConnectionSessionState, ) { session_state.rpc_gate.close().await; + self.request_serialization_queues.discard_closed().await; self.outgoing .disconnect_user_verification_connection(connection_id) .await; 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 4758d01947..07e5145980 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -487,7 +487,10 @@ pub(super) async fn handle_thread_listener_command( listener_command: ThreadListenerCommand, ) { match listener_command { - ThreadListenerCommand::SendThreadResumeResponse(resume_request) => { + ThreadListenerCommand::SendThreadResumeResponse { + request: resume_request, + completion_tx, + } => { handle_pending_thread_resume_request( conversation_id, conversation, @@ -500,6 +503,7 @@ pub(super) async fn handle_thread_listener_command( *resume_request, ) .await; + let _ = completion_tx.send(()); } ThreadListenerCommand::EmitThreadGoalUpdated { turn_id, goal } => { outgoing 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 df4ffe73ab..c4cce3011d 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -460,7 +460,7 @@ pub(crate) struct ThreadRequestProcessor { /// Outcome of trying to satisfy a resume request from an already loaded thread. enum RunningThreadResumeResult { /// The request was delegated to the loaded thread. - Handled, + Handled(tokio::sync::oneshot::Receiver<()>), /// No loaded thread handled the request. /// /// The optional stored thread contains the history-bearing probe that cold @@ -551,17 +551,17 @@ impl ThreadRequestProcessor { client_mcp_extensions: ClientMcpExtensions, ) -> Result, JSONRPCErrorError> { let mut prepared_config = None; - while self - .thread_resume_inner( - request_id.clone(), - ¶ms, - app_server_client_name.clone(), - app_server_client_version.clone(), - client_mcp_extensions.clone(), - &mut prepared_config, - ) - .await? - .is_continue() + // Keep the resume future off the request handler's stack. + while Box::pin(self.thread_resume_inner( + request_id.clone(), + ¶ms, + app_server_client_name.clone(), + app_server_client_version.clone(), + client_mcp_extensions.clone(), + &mut prepared_config, + )) + .await? + .is_continue() {} Ok(None) } @@ -1267,7 +1267,9 @@ impl ThreadRequestProcessor { } }; self.background_tasks - .spawn(thread_start_task.instrument(request_context.span())); + .spawn(thread_start_task.instrument(request_context.span())) + .await + .map_err(|_| internal_error("thread startup task stopped before completing"))?; Ok(()) } @@ -2353,6 +2355,7 @@ impl ThreadRequestProcessor { } let request = request_id.clone(); + let (completion_tx, completion_rx) = tokio::sync::oneshot::channel(); let rollback_already_in_progress = { let thread_state = self.thread_state_manager.thread_state(thread_id).await; @@ -2360,7 +2363,7 @@ impl ThreadRequestProcessor { if thread_state.pending_rollbacks.is_some() { true } else { - thread_state.pending_rollbacks = Some(request.clone()); + thread_state.pending_rollbacks = Some((request, completion_tx)); false } }; @@ -2385,6 +2388,9 @@ impl ThreadRequestProcessor { return Err(internal_error(format!("failed to start rollback: {err}"))); } + // The listener drops the sender after queuing the response, including errors. + // Keep the RPC's drain admission alive until then, without holding thread state. + let _ = completion_rx.await; Ok(()) } @@ -3678,7 +3684,12 @@ impl ThreadRequestProcessor { ) .await { - Ok(RunningThreadResumeResult::Handled) => return Ok(ControlFlow::Break(())), + Ok(RunningThreadResumeResult::Handled(completion)) => { + // The listener may need this permit to finish the response. + drop(_thread_list_state_permit); + let _ = completion.await; + return Ok(ControlFlow::Break(())); + } Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -3801,7 +3812,11 @@ impl ThreadRequestProcessor { ) .await? { - RunningThreadResumeResult::Handled => Ok(ControlFlow::Break(())), + RunningThreadResumeResult::Handled(completion) => { + drop(_thread_list_state_permit); + let _ = completion.await; + Ok(ControlFlow::Break(())) + } RunningThreadResumeResult::NotRunning(_) => Err(invalid_request( "cannot resume an unloaded multi-agent v2 sub-agent through its parent; resume the parent first, or use thread/read to inspect it", )), @@ -4496,8 +4511,9 @@ impl ThreadRequestProcessor { }; let resume_cursor_store = paginated_resume.then(|| Arc::clone(&self.thread_store)); - let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse( - Box::new(crate::thread_state::PendingThreadResumeRequest { + let (completion_tx, completion_rx) = tokio::sync::oneshot::channel(); + let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse { + request: Box::new(crate::thread_state::PendingThreadResumeRequest { request_id: request_id.clone(), history_items, cold_resume_token_usage_turn_id, @@ -4514,13 +4530,14 @@ impl ThreadRequestProcessor { resume_cursor_store, redact_resume_payloads, }), - ); + completion_tx, + }; if listener_command_tx.send(command).is_err() { return Err(internal_error(format!( "failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed" ))); } - return Ok(RunningThreadResumeResult::Handled); + return Ok(RunningThreadResumeResult::Handled(completion_rx)); } Ok(RunningThreadResumeResult::NotRunning(None)) } diff --git a/codex-rs/app-server/src/request_serialization.rs b/codex-rs/app-server/src/request_serialization.rs index ae897271ca..a55b737460 100644 --- a/codex-rs/app-server/src/request_serialization.rs +++ b/codex-rs/app-server/src/request_serialization.rs @@ -158,6 +158,24 @@ pub(crate) struct RequestSerializationQueues { } impl RequestSerializationQueues { + /// Release requests whose connection closed while they waited in a queue. + pub(crate) async fn discard_closed(&self) { + let mut queues = self.inner.lock().await; + for queue in queues.values_mut() { + let previous_len = queue.requests.len(); + queue.requests.retain(|request| { + !request + .request + .gate + .as_ref() + .is_some_and(|gate| gate.is_closed()) + }); + if queue.requests.len() != previous_len { + queue.changed.notify_one(); + } + } + } + /// Enqueue app-owned work alongside RPCs that mutate the same serialized resource. pub(crate) async fn enqueue_background( &self, @@ -473,6 +491,7 @@ mod tests { #[tokio::test] async fn shutdown_of_live_gate_skips_already_queued_requests() { let queues = RequestSerializationQueues::default(); + let (queued_tx, mut queued_rx) = oneshot::channel::<()>(); let key = RequestSerializationQueueKey::Global("test"); let live_gate = gate(); let (tx, mut rx) = mpsc::unbounded_channel(); @@ -499,6 +518,7 @@ mod tests { key, RequestSerializationAccess::Exclusive, QueuedInitializedRequest::new(live_gate.clone(), async move { + let _queued_tx = queued_tx; tx.send(SECOND_REQUEST_VALUE) .expect("receiver should be open"); }), @@ -514,6 +534,17 @@ mod tests { Some(FIRST_REQUEST_VALUE) ); + assert_eq!( + queued_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + live_gate.close().await; + queues.discard_closed().await; + assert_eq!( + queued_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + ); + let gate_for_shutdown = Arc::clone(&live_gate); let shutdown_task = tokio::spawn(async move { gate_for_shutdown.shutdown().await; diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 3795638cdd..ba011dd73a 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -57,7 +57,10 @@ pub(crate) struct PendingThreadResumeRequest { // ThreadListenerCommand is used to perform operations in the context of the thread listener, for serialization purposes. pub(crate) enum ThreadListenerCommand { // SendThreadResumeResponse is used to resume an already running thread by sending the thread's history to the client and atomically subscribing for new updates. - SendThreadResumeResponse(Box), + SendThreadResumeResponse { + request: Box, + completion_tx: oneshot::Sender<()>, + }, // EmitThreadGoalUpdated is used to order goal updates with running-thread resume responses and goal clears. EmitThreadGoalUpdated { turn_id: Option, @@ -95,7 +98,7 @@ pub(crate) struct TurnSummary { #[derive(Default)] pub(crate) struct ThreadState { pub(crate) pending_interrupts: PendingInterruptQueue, - pub(crate) pending_rollbacks: Option, + pub(crate) pending_rollbacks: Option<(ConnectionRequestId, oneshot::Sender<()>)>, pub(crate) turn_summary: TurnSummary, pub(crate) last_terminal_turn_id: Option, /// Lets an internal runtime replacement wait until the old listener has processed Core's @@ -143,6 +146,7 @@ impl ThreadState { let _ = cancel_tx.send(()); } self.shutdown_drain_waiter = None; + self.pending_rollbacks = None; self.listener_command_tx = None; self.current_turn_history.reset(); self.listener_thread = None; @@ -247,6 +251,28 @@ mod tests { use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; + #[test] + fn clear_listener_releases_pending_rollback() { + let (completion_tx, mut completion_rx) = oneshot::channel(); + let mut state = ThreadState { + pending_rollbacks: Some(( + ConnectionRequestId { + connection_id: ConnectionId(1), + request_id: RequestId::Integer(1), + }, + completion_tx, + )), + ..Default::default() + }; + + state.clear_listener(); + + assert_eq!( + completion_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + ); + } + #[test] fn note_thread_settings_reports_only_effective_changes() { let mut state = ThreadState::default(); diff --git a/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs b/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs index 3fba0104f7..83432e883d 100644 --- a/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs +++ b/codex-rs/app-server/tests/suite/v2/server_diagnostics.rs @@ -47,24 +47,25 @@ async fn server_diagnostics_exposes_process_and_registered_thread_gauge() -> Res assert!(diagnostics.process.physical_footprint_bytes.is_some()); #[cfg(not(target_os = "macos"))] assert_eq!(diagnostics.process.physical_footprint_bytes, None); - for expected_gauge in [ - ServerDiagnosticsGauge { - name: "app.requests.in_flight".to_string(), - value: 1, - }, - ServerDiagnosticsGauge { + // Startup may still be finishing after its response was queued. + assert!(matches!( + diagnostics + .gauges + .iter() + .find(|gauge| gauge.name == "app.requests.in_flight") + .map(|gauge| gauge.value), + Some(1 | 2) + )); + assert_eq!( + diagnostics + .gauges + .iter() + .find(|gauge| gauge.name == "core.threads.live"), + Some(&ServerDiagnosticsGauge { name: "core.threads.live".to_string(), value: 1, - }, - ] { - assert_eq!( - diagnostics - .gauges - .iter() - .find(|gauge| gauge.name == expected_gauge.name), - Some(&expected_gauge) - ); - } + }) + ); Ok(()) } diff --git a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs index 8bad7219dd..48c48d8ed1 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs @@ -230,6 +230,15 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() num_turns: 1, }) .await?; + // Pipeline resume while the listener may still be handling rollback. Waiting for + // its response must not retain the thread-list permit needed by that listener. + let pipelined_resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: thread.id.clone(), + exclude_turns: true, + ..Default::default() + }) + .await?; let deprecation_notice = timeout(DEFAULT_READ_TIMEOUT, mcp.read_next_message()).await??; let JSONRPCMessage::Notification(deprecation_notice) = deprecation_notice else { panic!("thread/rollback should emit deprecationNotice before its response"); @@ -297,7 +306,10 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() other => panic!("expected user message item, got {other:?}"), } - // Resume and confirm the history is pruned. + let _: ThreadResumeResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(pipelined_resume_id)).await??; + + // Resume after rollback completes to verify the pruned history. let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id: thread.id,