openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Release app-server thread runtimes after shutdown (#57788)

Browse Source 
## Why

Waiting for the event channel to close after `ShutdownComplete` retains the thread runtime because the thread owns its event sender. Remaining subscribers still need late notifications after shutdown.

## What changed

- Release runtime references after `ShutdownComplete` while keeping the listener command queue available for ordered notifications and late request resolutions.
- Preserve other subscribers when a client unsubscribes or disconnects from an unloaded thread, and serialize terminal cleanup with cold resumes.
- Return the thread-closing error when `thread/resume` encounters a terminated runtime or a listener that cannot complete the request.
- Close the status watch when removing a thread so a later load gets a fresh watch.

## Testing

Add regression coverage for listeners attached before and after shutdown, verifying runtime release, preservation of remaining subscribers, late notification delivery, and rejection of resumes against a terminated runtime. Add a test that removing and reloading a thread closes the old status watch.

GitOrigin-RevId: 6dda8f83a0fbe06c9a713e9801356753a0de172e
This commit is contained in:
Charlie Marsh
2026-10-28 16:55:57 +00:00
committed by copyberry
parent 575cc5351e
commit 7f4d53e928
5 changed files with 408 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors/thread_lifecycle.rs
+113 -22
View File
@@ -9,6 +9,10 @@ use codex_protocol::config_types::MultiAgentMode;
use codex_protocol::protocol::MultiAgentVersion;
use codex_protocol::protocol::SessionSource;
#[cfg(test)]
#[path = "thread_listener_shutdown_tests.rs"]
mod listener_shutdown_tests;
#[derive(Clone)]
pub(super) struct ListenerTaskContext {
pub(super) thread_manager: Arc<ThreadManager>,
@@ -497,20 +501,20 @@ async fn start_listener(
} = listener_task_context;
let outgoing_for_task = Arc::clone(&outgoing);
tokio::spawn(async move {
loop {
let shutdown_complete = loop {
tokio::select! {
biased;
_ = &mut cancel_rx => {
// Listener was superseded or the thread is being torn down.
break;
break false;
}
listener_command = listener_command_rx.recv() => {
let Some(listener_command) = listener_command else {
break;
break false;
};
handle_thread_listener_command(
conversation_id,
&conversation,
Some(&conversation),
codex_home.as_path(),
&thread_state_manager,
&thread_state,
@@ -526,7 +530,7 @@ async fn start_listener(
Ok(event) => event,
Err(err) => {
tracing::warn!("thread.next_event() failed with: {err}");
break;
break false;
}
};
@@ -589,7 +593,7 @@ async fn start_listener(
)
.await;
if is_shutdown_complete {
drained_tx.send_replace(true);
break true;
}
}
idle_guard = async {
@@ -601,7 +605,7 @@ async fn start_listener(
}
} => {
let Some(_idle_guard) = idle_guard else {
break;
break false;
};
if !unloading_state.should_unload_now() {
continue;
@@ -652,11 +656,9 @@ async fn start_listener(
conversation_id, &thread_state, listener_generation,
).await {
pending.remove(&conversation_id);
break;
break false;
}
}
// Release the status receiver before removing its watch entry.
drop(unloading_state);
outgoing_for_task
.cancel_requests_for_thread(conversation_id, /*error*/ None)
.await;
@@ -706,11 +708,9 @@ async fn start_listener(
conversation_id, &thread_state, listener_generation,
).await {
pending.remove(&conversation_id);
break;
break false;
}
}
// Release the status receiver before removing its watch entry.
drop(unloading_state);
outgoing_for_task
.cancel_requests_for_thread(conversation_id, /*error*/ None)
.await;
@@ -738,7 +738,97 @@ async fn start_listener(
conversation.clone(),
)
.await;
break;
break false;
}
}
};
// The thread owns its event sender, so waiting for channel closure retains
// the runtime. Keep only the notification relay after ShutdownComplete.
drop(conversation);
drop(config);
drop(thread_manager);
drop(turn_cost_worker);
if shutdown_complete {
drained_tx.send_replace(true);
let relay_command = |command| {
handle_thread_listener_command(
conversation_id,
/*conversation*/ None,
codex_home.as_path(),
&thread_state_manager,
&thread_state,
&thread_watch_manager,
&outgoing_for_task,
&pending_thread_unloads,
command,
)
};
// Preserve the existing queue for late request resolutions and ordered
// notifications until normal listener cancellation or replacement.
loop {
tokio::select! {
biased;
_ = &mut cancel_rx => break,
listener_command = listener_command_rx.recv() => {
let Some(listener_command) = listener_command else {
break;
};
relay_command(listener_command).await;
}
// Keep relaying commands while waiting for a cold resume or goal mutation.
resume_guard = async {
if unloading_state.wait_for_unloading_trigger().await {
Some(ThreadState::lock_goal_resume(&thread_state).await)
} else {
None
}
} => {
{
let Some(resume_guard) = resume_guard else {
break;
};
// A goal mutation may have queued a command between the select's
// receive poll and the lock acquisition. Deliver it before cleanup.
if !listener_command_rx.is_empty() {
continue;
}
let mut pending = pending_thread_unloads.lock().await;
if !unloading_state.should_unload_now() {
continue;
}
if let Some(existing) = pending.get(&conversation_id) {
let mut completion = existing.subscribe();
drop(pending);
drop(resume_guard);
tokio::select! {
_ = &mut cancel_rx => break,
_ = completion.changed() => continue,
}
}
if !thread_state_manager.remove_unsubscribed_listener(
conversation_id, &thread_state, listener_generation,
).await {
break;
}
// Fence this thread while cleanup awaits without blocking other threads.
pending.insert(conversation_id, watch::channel(()).0);
}
// A producer may still hold the retired sender. Deliver every command
// it queued successfully and let later sends take their fallback.
listener_command_rx.close();
while let Some(command) = listener_command_rx.recv().await {
relay_command(command).await;
}
outgoing_for_task
.cancel_requests_for_thread(conversation_id, /*error*/ None)
.await;
thread_watch_manager
.remove_thread(&conversation_id.to_string())
.await;
pending_thread_unloads.lock().await.remove(&conversation_id);
return;
}
}
}
}
@@ -830,7 +920,7 @@ pub(super) async fn unload_thread_without_subscribers(
#[allow(clippy::too_many_arguments)]
pub(super) async fn handle_thread_listener_command(
conversation_id: ThreadId,
conversation: &Arc<CodexThread>,
conversation: Option<&Arc<CodexThread>>,
codex_home: &Path,
thread_state_manager: &ThreadStateManager,
thread_state: &Arc<Mutex<ThreadState>>,
@@ -844,6 +934,10 @@ pub(super) async fn handle_thread_listener_command(
request: resume_request,
completion_tx,
} => {
let Some(conversation) = conversation else {
// Dropping the completion sender returns the closing error.
return;
};
handle_pending_thread_resume_request(
conversation_id,
conversation,
@@ -1049,15 +1143,12 @@ pub(super) async fn handle_pending_thread_resume_request(
{
let pending_thread_unloads = pending_thread_unloads.lock().await;
if pending_thread_unloads.contains_key(&conversation_id) {
if pending_thread_unloads.contains_key(&conversation_id)
|| matches!(conversation.agent_status().await, AgentStatus::Shutdown)
{
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"
)),
)
.send_error(request_id, thread_closing_for_resume(conversation_id))
.await;
return;
}
codex-rs/app-server/src/request_processors/thread_listener_shutdown_tests.rs
+230
View File
@@ -0,0 +1,230 @@
#![cfg(test)]
//! Checks runtime release and notification delivery around listener shutdown.
use super::super::thread_goal_processor::ThreadGoalRequestProcessor;
use super::super::thread_processor::ThreadRequestProcessor;
use super::super::thread_processor::ThreadResumeTarget;
use super::*;
use crate::config_manager::ConfigManager;
use crate::outgoing_message::OutgoingEnvelope;
use crate::outgoing_message::OutgoingMessage;
use crate::thread_sample::SampleResponseRegistry;
use crate::thread_state::ConnectionCapabilities;
use crate::thread_state::resolve_server_request_on_thread_listener;
use anyhow::Context;
use anyhow::Result;
use codex_analytics::AnalyticsEventsClient;
use codex_app_server_protocol::ClientResponsePayload;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadResumeParams;
use codex_app_server_protocol::ThreadUnsubscribeParams;
use codex_app_server_protocol::ThreadUnsubscribeResponse;
use codex_app_server_protocol::ThreadUnsubscribeStatus;
use codex_arg0::Arg0DispatchPaths;
use codex_goal_extension::GoalService;
use codex_login::AuthManager;
use codex_thread_store::PersistContext;
use core_test_support::test_codex::test_codex;
use pretty_assertions::assert_eq;
use std::collections::HashSet;
use test_case::test_case;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
#[test_case(true; "listener attached before shutdown")]
#[test_case(false; "listener attached after shutdown")]
#[tokio::test]
async fn listener_shutdown_releases_runtime_and_preserves_other_subscribers(
attach_before_shutdown: bool,
) -> Result<()> {
let server = wiremock::MockServer::start().await;
let test = test_codex().build_with_auto_env(&server).await?;
let thread_id = test.session_configured.thread_id;
test.thread_store
.persist_thread(thread_id, PersistContext::Standard)
.await?;
let thread = test.codex;
let runtime = Arc::downgrade(&thread);
let config = Arc::new(test.config.clone());
let auth =
AuthManager::shared_from_config(config.as_ref(), /*enable_codex_api_key_env*/ false)
.await?;
let config_manager =
ConfigManager::without_managed_config_for_tests(config.codex_home.to_path_buf());
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*buffer*/ 16);
let outgoing = Arc::new(OutgoingMessageSender::new(
outgoing_tx,
AnalyticsEventsClient::disabled(),
));
let thread_state_manager = ThreadStateManager::new();
let thread_watch_manager = ThreadWatchManager::new_with_outgoing(outgoing.clone());
let skills_watcher = SkillsWatcher::new(
test.thread_manager.skills_service(),
&config.codex_home,
outgoing.clone(),
);
let thread_goal_processor = ThreadGoalRequestProcessor::new(
test.thread_manager.clone(),
outgoing.clone(),
config.clone(),
thread_state_manager.clone(),
/*state_db*/ None,
Arc::new(GoalService::new()),
config_manager.clone(),
);
let processor = ThreadRequestProcessor::new(
auth,
test.thread_manager.clone(),
outgoing,
Arg0DispatchPaths::default(),
config,
config_manager,
test.thread_store.clone(),
Default::default(),
thread_state_manager.clone(),
thread_watch_manager,
Arc::new(Semaphore::new(/*permits*/ 1)),
thread_goal_processor,
/*state_db*/ None,
/*log_db*/ None,
skills_watcher.clone(),
/*turn_cost_worker*/ None,
Vec::new(),
Arc::new(SampleResponseRegistry::default()),
);
let connection_id = ConnectionId(1);
let disconnecting_connection_id = ConnectionId(2);
let remaining_connection_id = ConnectionId(3);
let connection_ids = vec![
connection_id,
disconnecting_connection_id,
remaining_connection_id,
];
for &connection_id in &connection_ids {
processor
.connection_initialized(connection_id, ConnectionCapabilities::default())
.await;
}
if attach_before_shutdown {
processor
.try_attach_thread_listener(thread_id, Arc::clone(&thread), connection_ids.clone())
.await;
}
thread.shutdown_and_wait().await?;
let shutdown_result = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async {
let resume = processor
.thread_resume(
ThreadResumeTarget::Client(ConnectionRequestId {
connection_id,
request_id: RequestId::Integer(1),
}),
ThreadResumeParams {
thread_id: thread_id.to_string(),
exclude_turns: true,
..Default::default()
},
/*app_server_client_name*/ None,
/*app_server_client_version*/ None,
/*client_mcp_extensions*/ Default::default(),
)
.await;
test.thread_manager.remove_thread(&thread_id).await;
if !attach_before_shutdown {
processor
.try_attach_thread_listener(thread_id, Arc::clone(&thread), connection_ids)
.await;
}
loop {
match outgoing_rx.recv().await.expect("outgoing channel closed") {
OutgoingEnvelope::Broadcast {
message: OutgoingMessage::AppServerNotification(notification),
} if matches!(
&notification.notification,
ServerNotification::ThreadStatusChanged(status)
if status.status == ThreadStatus::NotLoaded
) =>
{
break;
}
_ => {}
}
}
drop(thread);
while runtime.upgrade().is_some() {
tokio::task::yield_now().await;
}
let Some(ClientResponsePayload::ThreadUnsubscribe(unsubscribe)) = processor
.thread_unsubscribe(
&ConnectionRequestId {
connection_id,
request_id: RequestId::Integer(2),
},
ThreadUnsubscribeParams {
thread_id: thread_id.to_string(),
},
)
.await
.expect("unsubscribe failed")
else {
panic!("unsubscribe response missing");
};
assert_eq!(
unsubscribe,
ThreadUnsubscribeResponse {
status: ThreadUnsubscribeStatus::NotLoaded,
},
);
assert_eq!(
thread_state_manager
.subscribed_connection_ids(thread_id)
.await
.into_iter()
.collect::<HashSet<_>>(),
HashSet::from([disconnecting_connection_id, remaining_connection_id]),
);
processor
.connection_closed(disconnecting_connection_id)
.await;
assert_eq!(
thread_state_manager
.subscribed_connection_ids(thread_id)
.await,
vec![remaining_connection_id],
);
resolve_server_request_on_thread_listener(
&thread_state_manager.thread_state(thread_id).await,
RequestId::Integer(3),
)
.await;
loop {
if let OutgoingEnvelope::ToConnection {
connection_id: recipient,
message: OutgoingMessage::AppServerNotification(notification),
..
} = outgoing_rx.recv().await.expect("outgoing channel closed")
&& recipient == remaining_connection_id
&& matches!(
&notification.notification,
ServerNotification::ServerRequestResolved(resolved)
if resolved.request_id == RequestId::Integer(3)
)
{
break;
}
}
resume
})
.await;
thread_state_manager.clear_all_listeners().await;
skills_watcher.shutdown();
let resume = shutdown_result.context("terminal listener did not release or relay")?;
assert_eq!(
resume.expect_err("resume accepted a terminated runtime"),
invalid_request(format!(
"thread {thread_id} is closing; retry thread/resume after the thread is closed"
)),
);
Ok(())
}
codex-rs/app-server/src/request_processors/thread_processor.rs
+47 -31
View File
@@ -13,6 +13,7 @@ use super::thread_enrichment::enrich_loaded_threads;
use super::thread_fork_goal::inherit_thread_goal_snapshot;
use super::thread_input::can_accept_direct_input;
use super::thread_input::ensure_direct_input_allowed;
use super::thread_lifecycle::thread_closing_for_resume;
use super::*;
use crate::connection_rpc_gate::ConnectionRpcGate;
use crate::error_code::method_not_found;
@@ -1099,13 +1100,14 @@ impl ThreadRequestProcessor {
self.outgoing
.cancel_requests_for_thread(thread_id, /*error*/ None)
.await;
self.thread_state_manager
.remove_thread_state(thread_id)
.await;
// The state entry owns the resume lock. Remove the watch while a cold resume
// still has to wait for this teardown to finish.
self.thread_watch_manager
.remove_thread(&thread_id.to_string())
.await;
self.pending_thread_unloads.lock().await.remove(&thread_id);
self.thread_state_manager
.remove_thread_state(thread_id)
.await;
}
async fn thread_unsubscribe_response_inner(
@@ -1116,18 +1118,22 @@ impl ThreadRequestProcessor {
let thread_id = ThreadId::from_string(&params.thread_id)
.map_err(|err| invalid_request(format!("invalid thread id: {err}")))?;
if self.thread_manager.get_thread(thread_id).await.is_err() {
self.finalize_thread_teardown(thread_id).await;
return Ok(ThreadUnsubscribeResponse {
status: ThreadUnsubscribeStatus::NotLoaded,
});
};
// Terminal cleanup must not remove a replacement published by a cold resume.
let _goal_resume_guard = self.thread_state_manager.lock_goal_resume(thread_id).await;
let was_subscribed = self
.thread_state_manager
.unsubscribe_connection_from_thread(thread_id, connection_id)
.await;
if self.thread_manager.get_thread(thread_id).await.is_err() {
if !self.thread_state_manager.has_subscribers(thread_id).await {
self.finalize_thread_teardown(thread_id).await;
}
return Ok(ThreadUnsubscribeResponse {
status: ThreadUnsubscribeStatus::NotLoaded,
});
};
let status = if was_subscribed {
ThreadUnsubscribeStatus::Unsubscribed
} else {
@@ -3866,7 +3872,10 @@ impl ThreadRequestProcessor {
.await;
for thread_id in thread_ids {
if self.thread_manager.get_thread(thread_id).await.is_err() {
let _goal_resume_guard = self.thread_state_manager.lock_goal_resume(thread_id).await;
if self.thread_manager.get_thread(thread_id).await.is_err()
&& !self.thread_state_manager.has_subscribers(thread_id).await
{
// Reconcile stale app-server bookkeeping when the thread has already been
// removed from the core manager.
self.finalize_thread_teardown(thread_id).await;
@@ -3909,9 +3918,7 @@ impl ThreadRequestProcessor {
.await
.contains_key(&thread_id)
{
return Err(invalid_request(format!(
"thread {thread_id} is closing; retry thread/resume after the thread is closed"
)));
return Err(thread_closing_for_resume(thread_id));
}
if params.sandbox.is_some() && params.permissions.is_some() {
@@ -3937,7 +3944,10 @@ impl ThreadRequestProcessor {
RunningThreadResumeResult::Handled(completion) => {
// The listener may need this permit to finish the response.
drop(_thread_list_state_permit);
let _ = completion.await;
// A stopped listener can discard an accepted command before responding.
completion
.await
.map_err(|_| thread_closing_for_resume(&params.thread_id))?;
return Ok(ControlFlow::Break(()));
}
RunningThreadResumeResult::NotRunning(stored_thread) => stored_thread,
@@ -4040,10 +4050,7 @@ impl ThreadRequestProcessor {
.await
.contains_key(&resumed.conversation_id)
{
return Err(invalid_request(format!(
"thread {} is closing; retry thread/resume after the thread is closed",
resumed.conversation_id
)));
return Err(thread_closing_for_resume(resumed.conversation_id));
}
let paginated_thread_id = resume_source_thread.as_ref().and_then(|thread| {
matches!(thread.history_mode, ThreadHistoryMode::Paginated).then_some(thread.thread_id)
@@ -4104,7 +4111,9 @@ impl ThreadRequestProcessor {
{
RunningThreadResumeResult::Handled(completion) => {
drop(_thread_list_state_permit);
let _ = completion.await;
completion
.await
.map_err(|_| thread_closing_for_resume(child_thread_id))?;
Ok(ControlFlow::Break(()))
}
RunningThreadResumeResult::NotRunning(_) => Err(invalid_request(
@@ -4573,6 +4582,10 @@ impl ThreadRequestProcessor {
Some(persisted_metadata)
}
#[expect(
clippy::await_holding_invalid_type,
reason = "listener state creation must be serialized against terminal listener cleanup"
)]
#[tracing::instrument(level = "trace", skip_all)]
async fn resume_running_thread(
&self,
@@ -4740,10 +4753,17 @@ impl ThreadRequestProcessor {
Vec::new()
};
let thread_state = self
.thread_state_manager
.thread_state(existing_thread_id)
.await;
let thread_state = {
let pending = self.pending_thread_unloads.lock().await;
if pending.contains_key(&existing_thread_id)
|| matches!(existing_thread.agent_status().await, AgentStatus::Shutdown)
{
return Err(thread_closing_for_resume(existing_thread_id));
}
self.thread_state_manager
.thread_state(existing_thread_id)
.await
};
self.ensure_listener_task_running(
existing_thread_id,
existing_thread.clone(),
@@ -4776,9 +4796,7 @@ impl ThreadRequestProcessor {
thread_state.listener_command_tx()
};
let Some(listener_command_tx) = listener_command_tx else {
return Err(internal_error(format!(
"failed to enqueue running thread resume for thread {existing_thread_id}: thread listener is not running"
)));
return Err(thread_closing_for_resume(existing_thread_id));
};
let (emit_thread_goal_update, thread_goal_state_db) = self
@@ -4860,9 +4878,7 @@ impl ThreadRequestProcessor {
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 Err(thread_closing_for_resume(existing_thread_id));
}
return Ok(RunningThreadResumeResult::Handled(completion_rx));
}
codex-rs/app-server/src/thread_state.rs
+1 -2
View File
@@ -541,6 +541,7 @@ impl ThreadStateManager {
pub(crate) async fn remove_thread_state(&self, thread_id: ThreadId) {
let thread_state = {
let mut state = self.state.lock().await;
self.unregister_listener_command_tx(thread_id);
let thread_state = state
.threads
.remove(&thread_id)
@@ -551,7 +552,6 @@ impl ThreadStateManager {
});
thread_state
};
self.unregister_listener_command_tx(thread_id);
if let Some(thread_state) = thread_state {
let mut thread_state = thread_state.lock().await;
@@ -646,7 +646,6 @@ impl ThreadStateManager {
true
}
#[cfg(test)]
pub(crate) async fn has_subscribers(&self, thread_id: ThreadId) -> bool {
self.state
.lock()
codex-rs/app-server/src/thread_status.rs
+17 -6
View File
@@ -336,7 +336,9 @@ impl ThreadWatchState {
{
self.running_turn_count -= 1;
}
self.update_status_watcher(thread_id, &ThreadStatus::NotLoaded);
if let Some(sender) = self.status_watcher_by_thread_id.remove(thread_id) {
sender.send_replace(ThreadStatus::NotLoaded);
}
if previous_status.is_some() && previous_status != Some(ThreadStatus::NotLoaded) {
Some(ThreadStatusChangedNotification {
thread_id: thread_id.to_string(),
@@ -394,12 +396,7 @@ impl ThreadWatchState {
fn update_status_watcher_for_thread(&mut self, thread_id: &str) {
let status = self.loaded_status_for_thread(thread_id);
self.update_status_watcher(thread_id, &status);
}
fn update_status_watcher(&mut self, thread_id: &str, status: &ThreadStatus) {
let remove_watcher = if let Some(sender) = self.status_watcher_by_thread_id.get(thread_id) {
let status = status.clone();
let _ = sender.send_if_modified(|current| {
if *current == status {
false
@@ -814,6 +811,20 @@ mod tests {
);
}
#[tokio::test]
async fn removing_a_thread_closes_its_status_watch() {
let manager = ThreadWatchManager::new();
let thread_id = ThreadId::from_string(INTERACTIVE_THREAD_ID).unwrap();
manager.upsert_thread(INTERACTIVE_THREAD_ID).await;
let status = manager.subscribe(thread_id).await;
manager.remove_thread(INTERACTIVE_THREAD_ID).await;
manager.upsert_thread(INTERACTIVE_THREAD_ID).await;
assert_eq!(*status.borrow(), ThreadStatus::NotLoaded);
assert!(status.has_changed().is_err());
}
#[tokio::test]
async fn status_watchers_receive_only_their_thread_updates() {
let manager = ThreadWatchManager::new();