openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-24 18:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

app-server: serialize teardown admission by thread id

Browse Source
This commit is contained in:
Winston Howes
2026-07-07 04:07:01 -07:00
parent 2cf1e9dd69
commit 5368909df2
5 changed files with 208 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors.rs
+1
View File
@@ -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;
codex-rs/app-server/src/request_processors/thread_lifecycle.rs
+53 -63
View File
@@ -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<F>(
unloading_state: &mut UnloadingState,
thread_list_state_permit: Arc<Semaphore>,
@@ -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<CodexThread>),
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<EnsureConversationListenerResult, JSONRPCErrorError> {
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<Mutex<ThreadState>>,
thread_watch_manager: &ThreadWatchManager,
outgoing: &Arc<OutgoingMessageSender>,
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<Mutex<ThreadState>>,
thread_watch_manager: &ThreadWatchManager,
outgoing: &Arc<OutgoingMessageSender>,
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 {
codex-rs/app-server/src/request_processors/thread_processor.rs
+126 -31
View File
@@ -778,6 +778,20 @@ impl ThreadRequestProcessor {
})
}
pub(super) async fn acquire_thread_list_state_permit_after_unload(
&self,
thread_id: ThreadId,
) -> Result<SemaphorePermit<'_>, 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<String>,
@@ -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<watch::Receiver<bool>> {
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<ThreadUnsubscribeResponse, JSONRPCErrorError> {
let thread_id = ThreadId::from_string(&params.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<EnsureConversationListenerResult, JSONRPCErrorError> {
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<EnsureConversationListenerResult, JSONRPCErrorError> {
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<usize> {
@@ -2611,6 +2671,16 @@ impl ThreadRequestProcessor {
thread_id: ThreadId,
connection_ids: Vec<ConnectionId>,
) {
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<String>,
supports_openai_form_elicitation: bool,
) -> Result<(), JSONRPCErrorError> {
if let Ok(thread_id) = ThreadId::from_string(&params.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(&params.thread_id).ok()
} else if params.path.is_some() {
match self
.read_stored_thread_for_resume(
&params.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(&params.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));
codex-rs/app-server/src/request_processors/thread_teardown.rs
+22 -5
View File
@@ -11,7 +11,6 @@ pub(crate) struct PendingThreadUnloads {
#[derive(Default)]
struct PendingThreadUnloadsInner {
admission: Mutex<()>,
registry: StdMutex<PendingThreadUnloadRegistry>,
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<watch::Receiver<bool>> {
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<bool>)
}
}
pub(super) fn start_thread_teardown<F>(
pending: PendingThreadUnloads,
thread_id: ThreadId,
teardown: F,
) -> Option<watch::Receiver<bool>>
where
F: std::future::Future<Output = ()> + 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;
codex-rs/app-server/src/request_processors/turn_processor.rs
+6 -5
View File
@@ -960,9 +960,10 @@ impl TurnRequestProcessor {
request_id: &ConnectionRequestId,
thread_id: &str,
) -> Result<Option<(ThreadId, Arc<CodexThread>)>, 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<EnsureConversationListenerResult, JSONRPCErrorError> {
super::thread_lifecycle::ensure_conversation_listener(
super::thread_lifecycle::ensure_conversation_listener_serialized(
self.listener_task_context(),
conversation_id,
connection_id,