Fix app-server thread listener startup and shutdown races (#57787)

## Why A thread runtime can shut down while listener setup is awaiting I/O, and a cold resume can start before the previous listener forwards `ShutdownComplete`. Listener replacement must preserve those buffered terminal events. ## What changed - Coordinate listener installation and connection subscriptions with pending unloads, and reject closing or replaced runtimes. - Allow a newly created runtime that has already shut down to drain its queued events, while enforcing one event reader per runtime. - Track shutdown draining with a shared watch channel so cold resume and runtime replacement wait for terminal events to be processed. - Serialize idle unloading with goal edits and cold resume, and recheck thread state after acquiring the lock. GitOrigin-RevId: d05e348b0a81bddb2defa963d0cfd3fde08dbd03
2026-10-29 19:08:43 +00:00 · 2026-10-28 16:53:55 +00:00
parent c8bd6c1da9
commit 575cc5351e
10 changed files with 541 additions and 322 deletions
@@ -428,14 +428,12 @@ mod tests {
                .connection_initialized(connection_id, ConnectionCapabilities::default())
                .await;
        }
-        thread_state_manager
-            .try_ensure_connection_subscribed(
-                thread_id,
-                subscribed_connection,
-                /*experimental_raw_events*/ false,
-            )
-            .await
-            .expect("connection should be subscribed");
+        assert!(
+            thread_state_manager
+                .try_add_connection_to_thread(thread_id, subscribed_connection)
+                .await,
+            "connection should be subscribed"
+        );
        let sink = app_server_extension_event_sink(outgoing, thread_state_manager);

        sink.emit_warning(ExtensionWarning {
@@ -494,14 +492,12 @@ mod tests {
            message: "catalog was shortened".to_string(),
        });
        tokio::task::yield_now().await;
-        thread_state_manager
-            .try_ensure_connection_subscribed(
-                thread_id,
-                subscribed_connection,
-                /*experimental_raw_events*/ false,
-            )
-            .await
-            .expect("connection should be subscribed");
+        assert!(
+            thread_state_manager
+                .try_add_connection_to_thread(thread_id, subscribed_connection)
+                .await,
+            "connection should be subscribed"
+        );

        let envelope = timeout(Duration::from_secs(1), outgoing_rx.recv())
            .await
@@ -544,14 +540,12 @@ mod tests {
        thread_state_manager
            .connection_initialized(subscribed_connection, ConnectionCapabilities::default())
            .await;
-        thread_state_manager
-            .try_ensure_connection_subscribed(
-                thread_id,
-                subscribed_connection,
-                /*experimental_raw_events*/ false,
-            )
-            .await
-            .expect("connection should be subscribed");
+        assert!(
+            thread_state_manager
+                .try_add_connection_to_thread(thread_id, subscribed_connection)
+                .await,
+            "connection should be subscribed"
+        );
        let (listener_command_tx, listener_command_rx) = mpsc::unbounded_channel();
        drop(listener_command_rx);
        thread_state_manager.register_listener_command_tx(thread_id, listener_command_tx);
@@ -866,8 +866,16 @@ impl MessageProcessor {
        thread_id: ThreadId,
        connection_ids: Vec<ConnectionId>,
    ) {
+        let Ok(thread) = self
+            .thread_processor
+            .thread_manager
+            .get_thread(thread_id)
+            .await
+        else {
+            return;
+        };
        self.thread_processor
-            .try_attach_thread_listener(thread_id, connection_ids)
+            .try_attach_thread_listener(thread_id, thread, connection_ids)
            .await;
    }

@@ -1,7 +1,9 @@
 use super::*;
 use crate::extensions::send_thread_warning;
+use crate::thread_state::ListenerStateCandidate;
 use codex_app_server_protocol::ThreadQueueChangedNotification;
 use codex_core::ThreadEvictionOutcome;
+use codex_core::ThreadRegistration;
 use codex_extension_api::ThreadIdleCause;
 use codex_protocol::config_types::MultiAgentMode;
 use codex_protocol::protocol::MultiAgentVersion;
@@ -28,31 +30,23 @@ struct UnloadingState {
 }

 impl UnloadingState {
-    async fn new(
-        listener_task_context: &ListenerTaskContext,
-        thread_id: ThreadId,
+    fn new(
        delay: Duration,
-    ) -> Option<Self> {
-        let has_subscribers_rx = listener_task_context
-            .thread_state_manager
-            .subscribe_to_has_connections(thread_id)
-            .await?;
-        let thread_status_rx = listener_task_context
-            .thread_watch_manager
-            .subscribe(thread_id)
-            .await?;
+        has_subscribers_rx: watch::Receiver<bool>,
+        thread_status_rx: watch::Receiver<ThreadStatus>,
+    ) -> Self {
        let has_subscribers = (*has_subscribers_rx.borrow(), Instant::now());
        let is_active = (
            matches!(*thread_status_rx.borrow(), ThreadStatus::Active { .. }),
            Instant::now(),
        );
-        Some(Self {
+        Self {
            delay,
            has_subscribers_rx,
            has_subscribers,
            thread_status_rx,
            is_active,
-        })
+        }
    }

    fn unloading_target(&self) -> Option<Instant> {
@@ -135,10 +129,111 @@ pub(super) enum EnsureConversationListenerResult {
    ConnectionClosed,
 }

+pub(super) fn thread_closing_for_resume(thread_id: impl std::fmt::Display) -> JSONRPCErrorError {
+    invalid_request(format!(
+        "thread {thread_id} is closing; retry thread/resume after the thread is closed"
+    ))
+}
+
+/// Marks this exact runtime after listener startup, even after its app-server state is removed.
+struct ListenerStarted;
+
+enum ListenerStartup<'a> {
+    Active,
+    Subscribe {
+        connection_id: ConnectionId,
+        raw_events_enabled: bool,
+    },
+    Created {
+        connection_ids: &'a [ConnectionId],
+        raw_events_enabled: bool,
+    },
+}
+
+enum ListenerStartOutcome {
+    Started { attached: bool },
+    Stale,
+    Rejected,
+}
+
+impl ListenerStartOutcome {
+    fn into_subscription(
+        self,
+        thread_id: ThreadId,
+    ) -> Result<EnsureConversationListenerResult, JSONRPCErrorError> {
+        match self {
+            Self::Started { attached: true } => Ok(EnsureConversationListenerResult::Attached),
+            Self::Started { attached: false } => {
+                Ok(EnsureConversationListenerResult::ConnectionClosed)
+            }
+            Self::Stale => Err(invalid_request(format!(
+                "thread {thread_id} is closing; retry after the thread is closed"
+            ))),
+            Self::Rejected => Err(thread_closing_for_resume(thread_id)),
+        }
+    }
+}
+
 #[expect(
    clippy::await_holding_invalid_type,
-    reason = "listener subscription must be serialized against pending unloads"
+    reason = "listener state creation must be serialized against pending unloads"
 )]
+pub(super) async fn attach_created_thread_listener(
+    listener_task_context: ListenerTaskContext,
+    thread_id: ThreadId,
+    thread: Arc<CodexThread>,
+    connection_ids: Vec<ConnectionId>,
+) {
+    let config_snapshot = thread.config_snapshot().await;
+    let raw_events_enabled = if let Some(parent_thread_id) = config_snapshot.parent_thread_id {
+        !connection_ids.is_empty()
+            && listener_task_context
+                .thread_state_manager
+                .thread_state(parent_thread_id)
+                .await
+                .lock()
+                .await
+                .experimental_raw_events
+    } else {
+        false
+    };
+    loop {
+        let pending = listener_task_context.pending_thread_unloads.lock().await;
+        if let Some(unload) = pending.get(&thread_id) {
+            let mut completion = unload.subscribe();
+            drop(pending);
+            let _ = completion.changed().await;
+            continue;
+        }
+        if connection_ids.is_empty()
+            && (thread.multi_agent_version() != Some(MultiAgentVersion::V2)
+                || config_snapshot.ephemeral)
+        {
+            return;
+        }
+        let listener_state = listener_task_context
+            .thread_state_manager
+            .created_listener_state(thread_id)
+            .await;
+        drop(pending);
+        match start_listener(
+            listener_task_context.clone(),
+            thread_id,
+            Arc::clone(&thread),
+            listener_state,
+            ListenerStartup::Created {
+                connection_ids: &connection_ids,
+                raw_events_enabled,
+            },
+        )
+        .await
+        {
+            ListenerStartOutcome::Started { .. } | ListenerStartOutcome::Rejected => return,
+            ListenerStartOutcome::Stale => continue,
+        }
+    }
+}
+
 pub(super) async fn ensure_conversation_listener(
    listener_task_context: ListenerTaskContext,
    conversation_id: ThreadId,
@@ -157,48 +252,25 @@ pub(super) async fn ensure_conversation_listener(
            )));
        }
    };
-    let (thread_state, result) = {
-        let pending_thread_unloads = listener_task_context.pending_thread_unloads.lock().await;
-        if pending_thread_unloads.contains_key(&conversation_id) {
-            return Err(invalid_request(format!(
-                "thread {conversation_id} is closing; retry after the thread is closed"
-            )));
-        }
-        match listener_task_context
-            .thread_state_manager
-            .try_ensure_connection_subscribed(conversation_id, connection_id, raw_events_enabled)
-            .await
-        {
-            Some(thread_state) => (thread_state, EnsureConversationListenerResult::Attached),
-            None => {
-                // Startup can outlast connection cleanup; the thread still needs a
-                // listener to unload once it is idle and has no subscribers.
-                let thread_state = listener_task_context
-                    .thread_state_manager
-                    .thread_state(conversation_id)
-                    .await;
-                (
-                    thread_state,
-                    EnsureConversationListenerResult::ConnectionClosed,
-                )
-            }
-        }
-    };
-    if let Err(error) = ensure_listener_task_running(
-        listener_task_context.clone(),
+    if conversation.has_shutdown_status() {
+        return Err(thread_closing_for_resume(conversation_id));
+    }
+    let state = listener_task_context
+        .thread_state_manager
+        .thread_state(conversation_id)
+        .await;
+    start_listener(
+        listener_task_context,
        conversation_id,
        conversation,
-        thread_state,
+        ListenerStateCandidate::Existing(state),
+        ListenerStartup::Subscribe {
+            connection_id,
+            raw_events_enabled,
+        },
    )
    .await
-    {
-        let _ = listener_task_context
-            .thread_state_manager
-            .unsubscribe_connection_from_thread(conversation_id, connection_id)
-            .await;
-        return Err(error);
-    }
-    Ok(result)
+    .into_subscription(conversation_id)
 }

 pub(super) fn log_listener_attach_result(
@@ -225,30 +297,39 @@ pub(super) fn log_listener_attach_result(
    }
 }

-#[expect(
-    clippy::await_holding_invalid_type,
-    reason = "listener removal must be serialized against new subscriptions"
-)]
 pub(super) async fn ensure_listener_task_running(
    listener_task_context: ListenerTaskContext,
    conversation_id: ThreadId,
    conversation: Arc<CodexThread>,
    thread_state: Arc<Mutex<ThreadState>>,
 ) -> Result<(), JSONRPCErrorError> {
-    let (cancel_tx, mut cancel_rx) = oneshot::channel();
-    let config_snapshot = conversation.config_snapshot().await;
-    let config = conversation.config().await;
-    let Some(mut unloading_state) = UnloadingState::new(
-        &listener_task_context,
+    start_listener(
+        listener_task_context,
        conversation_id,
-        config.thread_unload_delay,
+        conversation,
+        ListenerStateCandidate::Existing(thread_state),
+        ListenerStartup::Active,
    )
    .await
-    else {
-        return Err(invalid_request(format!(
-            "thread {conversation_id} is closing; retry after the thread is closed"
-        )));
-    };
+    .into_subscription(conversation_id)
+    .map(|_| ())
+}
+
+#[expect(
+    clippy::await_holding_invalid_type,
+    reason = "listener installation and removal must be serialized against subscriptions"
+)]
+async fn start_listener(
+    listener_task_context: ListenerTaskContext,
+    conversation_id: ThreadId,
+    conversation: Arc<CodexThread>,
+    listener_state: ListenerStateCandidate,
+    startup: ListenerStartup<'_>,
+) -> ListenerStartOutcome {
+    let (cancel_tx, mut cancel_rx) = oneshot::channel();
+    let (drained_tx, drained_rx) = watch::channel(false);
+    let config_snapshot = conversation.config_snapshot().await;
+    let config = conversation.config().await;
    let environments = conversation.environment_selections().await;
    let watch_registration = listener_task_context
        .skills_watcher
@@ -261,28 +342,148 @@ pub(super) async fn ensure_listener_task_running(
    let is_v2_child = conversation.multi_agent_version() == Some(MultiAgentVersion::V2)
        && matches!(config_snapshot.session_source, SessionSource::SubAgent(_));
    let thread_settings_baseline = thread_settings_from_config_snapshot(&config_snapshot);
-    let (mut listener_command_rx, listener_generation) = {
-        let mut thread_state = thread_state.lock().await;
-        if thread_state.listener_matches(&conversation) {
-            return Ok(());
-        }
-        let (listener_command_rx, listener_generation) = thread_state.set_listener(
-            cancel_tx,
-            &conversation,
-            watch_registration,
-            thread_settings_baseline,
-            config_snapshot.history_mode,
-        );
-        let Some(listener_command_tx) = thread_state.listener_command_tx() else {
-            tracing::warn!(
-                "thread listener command sender missing immediately after listener registration"
-            );
-            return Ok(());
+    let thread_state = Arc::clone(listener_state.state());
+    let (mut listener_command_rx, listener_generation, mut unloading_state, attached) = {
+        let (created, connection_ids, raw_events_enabled) = match &startup {
+            ListenerStartup::Active => (false, &[][..], false),
+            ListenerStartup::Subscribe {
+                connection_id,
+                raw_events_enabled,
+            } => (
+                false,
+                std::slice::from_ref(connection_id),
+                *raw_events_enabled,
+            ),
+            ListenerStartup::Created {
+                connection_ids,
+                raw_events_enabled,
+            } => (true, *connection_ids, *raw_events_enabled),
        };
-        listener_task_context
+        // Publish the drain receiver before a cold resume can start a replacement.
+        let _admission_guard = if created {
+            Some(ThreadState::lock_goal_resume(&thread_state).await)
+        } else {
+            None
+        };
+        let (pending, mut state) = loop {
+            let pending = if created || !connection_ids.is_empty() {
+                let pending = listener_task_context.pending_thread_unloads.lock().await;
+                if pending.contains_key(&conversation_id) {
+                    return ListenerStartOutcome::Stale;
+                }
+                Some(pending)
+            } else {
+                None
+            };
+            let state = thread_state.lock().await;
+            if !state.listener_matches(&conversation)
+                && let Some(mut drain) = state.pending_shutdown_drain()
+            {
+                drop(state);
+                drop(pending);
+                let _ = drain.wait_for(|drained| *drained).await;
+                continue;
+            }
+            break (pending, state);
+        };
+        if created || !connection_ids.is_empty() {
+            match listener_task_context
+                .thread_manager
+                .thread_registration(conversation_id, &conversation)
+                .await
+            {
+                ThreadRegistration::Current => {}
+                // Core can remove a closed child while listener setup awaits I/O. Drain its
+                // captured runtime only if it cannot supersede a newer listener.
+                ThreadRegistration::Absent
+                    if state.cancel_tx.is_some() && !state.listener_matches(&conversation) =>
+                {
+                    return ListenerStartOutcome::Rejected;
+                }
+                ThreadRegistration::Absent if created && conversation.has_shutdown_status() => {}
+                ThreadRegistration::Absent | ThreadRegistration::Replaced => {
+                    return ListenerStartOutcome::Rejected;
+                }
+            }
+        }
+        let installed = listener_task_context
            .thread_state_manager
-            .register_listener_command_tx(conversation_id, listener_command_tx);
-        (listener_command_rx, listener_generation)
+            .subscribe_and_install_listener(
+                conversation_id,
+                &listener_state,
+                connection_ids,
+                |attached| {
+                    let listener_matches = state.listener_matches(&conversation);
+                    if !created && conversation.has_shutdown_status() {
+                        return Err(());
+                    }
+                    // Each runtime can have only one event reader. A created runtime may
+                    // have shut down before its first reader drains the queued events.
+                    if !listener_matches
+                        && !conversation
+                            .thread_extension_data()
+                            .insert_if(ListenerStarted, |existing| existing.is_none())
+                    {
+                        return Err(());
+                    }
+                    if attached && raw_events_enabled {
+                        state.set_experimental_raw_events(/*enabled*/ true);
+                    }
+                    if listener_matches {
+                        return Ok(None);
+                    }
+                    let (listener_command_rx, listener_generation) = state.set_listener(
+                        cancel_tx,
+                        drained_rx,
+                        &conversation,
+                        watch_registration,
+                        thread_settings_baseline,
+                        config_snapshot.history_mode,
+                    );
+                    let Some(listener_command_tx) = state.listener_command_tx() else {
+                        tracing::warn!(
+                            "thread listener command sender missing immediately after listener registration"
+                        );
+                        return Ok(None);
+                    };
+                    listener_task_context
+                        .thread_state_manager
+                        .register_listener_command_tx(conversation_id, listener_command_tx);
+                    Ok(Some((listener_command_rx, listener_generation)))
+                },
+            )
+            .await;
+        let ((listener_command_rx, listener_generation), has_subscribers_rx, attached) =
+            match installed {
+                Ok(Some((Some(listener), has_subscribers_rx, attached))) => {
+                    (listener, has_subscribers_rx, attached)
+                }
+                Ok(Some((None, _, attached))) => return ListenerStartOutcome::Started { attached },
+                Ok(None) => return ListenerStartOutcome::Stale,
+                Err(()) => return ListenerStartOutcome::Rejected,
+            };
+        drop(pending);
+        // A cold V2 resume can install the child's listener before its creation callback runs.
+        if created || is_v2_child {
+            listener_task_context
+                .thread_watch_manager
+                .upsert_thread(&conversation_id.to_string())
+                .await;
+        }
+        let thread_status_rx = listener_task_context
+            .thread_watch_manager
+            .subscribe(conversation_id)
+            .await;
+        (
+            listener_command_rx,
+            listener_generation,
+            UnloadingState::new(
+                config.thread_unload_delay,
+                has_subscribers_rx,
+                thread_status_rx,
+            ),
+            attached,
+        )
    };
    let ListenerTaskContext {
        outgoing,
@@ -387,19 +588,21 @@ pub(super) async fn ensure_listener_task_running(
                        thread_watch_manager.clone(),
                    )
                    .await;
-                    if is_shutdown_complete
-                        && let Some(completion_tx) = thread_state
-                            .lock()
-                            .await
-                            .take_shutdown_drain_waiter()
-                    {
-                        let _ = completion_tx.send(());
+                    if is_shutdown_complete {
+                        drained_tx.send_replace(true);
                    }
                }
-                unloading_watchers_open = unloading_state.wait_for_unloading_trigger() => {
-                    if !unloading_watchers_open {
-                        break;
+                idle_guard = async {
+                    if unloading_state.wait_for_unloading_trigger().await {
+                        // Keep serving listener commands while a cold resume attaches its subscriber.
+                        Some(ThreadState::lock_goal_resume(&thread_state).await)
+                    } else {
+                        None
                    }
+                } => {
+                    let Some(_idle_guard) = idle_guard else {
+                        break;
+                    };
                    if !unloading_state.should_unload_now() {
                        continue;
                    }
@@ -546,7 +749,7 @@ pub(super) async fn ensure_listener_task_running(
            thread_state.clear_listener();
        }
    });
-    Ok(())
+    ListenerStartOutcome::Started { attached }
 }

 pub(super) async fn wait_for_thread_shutdown(thread: &Arc<CodexThread>) -> ThreadShutdownResult {
@@ -40,7 +40,6 @@ use codex_protocol::SanitizedGitUrl;
 use codex_protocol::config_types::MultiAgentMode;
 use codex_protocol::error::CodexErrorDetails;
 use codex_protocol::mcp::ClientMcpExtensions;
-use codex_protocol::protocol::MultiAgentVersion;
 use codex_protocol::protocol::ThreadHistoryMode;
 use codex_thread_store::PersistContext;
 use std::ops::ControlFlow;
@@ -2271,41 +2270,45 @@ impl ThreadRequestProcessor {
            )));
        }
        let thread_state = self.thread_state_manager.thread_state(thread_id).await;
-        let shutdown_drain_rx = thread_state.lock().await.register_shutdown_drain_waiter();
+        let mut shutdown_drain = thread_state
+            .lock()
+            .await
+            .shutdown_drain
+            .clone()
+            .ok_or_else(|| {
+                internal_error(format!(
+                    "thread {thread_id} has no listener to drain shutdown events"
+                ))
+            })?;

        match wait_for_thread_shutdown(&thread).await {
            ThreadShutdownResult::Complete => {}
            ThreadShutdownResult::SubmitFailed => {
-                thread_state.lock().await.take_shutdown_drain_waiter();
                return Err(internal_error(format!(
                    "failed to shut down thread {thread_id} before revert"
                )));
            }
            ThreadShutdownResult::TimedOut => {
-                thread_state.lock().await.take_shutdown_drain_waiter();
                return Err(internal_error(format!(
                    "timed out shutting down thread {thread_id} before revert"
                )));
            }
        }
-        let drain_result = tokio::time::timeout(Duration::from_secs(10), shutdown_drain_rx)
-            .await
-            .map_err(|_| {
-                internal_error(format!(
-                    "timed out waiting for thread {thread_id} listener to drain shutdown events"
-                ))
-            })
-            .and_then(|result| {
-                result.map_err(|_| {
-                    internal_error(format!(
-                        "thread {thread_id} listener stopped before draining shutdown events"
-                    ))
-                })
-            });
-        if let Err(err) = drain_result {
-            thread_state.lock().await.take_shutdown_drain_waiter();
-            return Err(err);
-        }
+        tokio::time::timeout(
+            Duration::from_secs(10),
+            shutdown_drain.wait_for(|drained| *drained),
+        )
+        .await
+        .map_err(|_| {
+            internal_error(format!(
+                "timed out waiting for thread {thread_id} listener to drain shutdown events"
+            ))
+        })?
+        .map_err(|_| {
+            internal_error(format!(
+                "thread {thread_id} listener stopped before draining shutdown events"
+            ))
+        })?;
        // The listener is drained: even an interrupted turn is durable now. Inspect only the
        // suffix being reverted and remember the revision, not a collection of turn IDs.
        let receipt_to_retire = async {
@@ -3875,80 +3878,19 @@ impl ThreadRequestProcessor {
        self.thread_watch_manager.subscribe_running_turn_count()
    }

-    /// Best-effort: start the lifecycle listener and subscribe initialized connections.
-    #[expect(
-        clippy::await_holding_invalid_type,
-        reason = "listener state creation must be serialized against pending unloads"
-    )]
    pub(crate) async fn try_attach_thread_listener(
        &self,
        thread_id: ThreadId,
+        thread: Arc<CodexThread>,
        connection_ids: Vec<ConnectionId>,
    ) {
-        if connection_ids.is_empty() {
-            let result = loop {
-                let pending = self.pending_thread_unloads.lock().await;
-                if let Some(unload) = pending.get(&thread_id) {
-                    // A replacement can be published before the old listener finishes cleanup.
-                    let mut completion = unload.subscribe();
-                    drop(pending);
-                    let _ = completion.changed().await;
-                    continue;
-                }
-                let Ok(thread) = self.thread_manager.get_thread(thread_id).await else {
-                    return;
-                };
-                // Only persisted V2 children can reload after idle eviction.
-                if thread.multi_agent_version() != Some(MultiAgentVersion::V2)
-                    || thread.config().await.ephemeral
-                {
-                    return;
-                }
-                self.thread_watch_manager
-                    .upsert_thread(&thread_id.to_string())
-                    .await;
-                let thread_state = self.thread_state_manager.thread_state(thread_id).await;
-                let result = self
-                    .ensure_listener_task_running(thread_id, thread, thread_state)
-                    .await;
-                drop(pending);
-                break result;
-            };
-            if let Err(err) = result {
-                warn!(
-                    "failed to start listener for thread {thread_id}: {message}",
-                    message = err.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;
-            self.thread_watch_manager
-                .upsert_thread(&thread_id.to_string())
-                .await;
-            if let Some(parent_thread_id) = config_snapshot.parent_thread_id {
-                raw_events_enabled = self
-                    .thread_state_manager
-                    .thread_state(parent_thread_id)
-                    .await
-                    .lock()
-                    .await
-                    .experimental_raw_events;
-            }
-        }
-
-        for connection_id in connection_ids {
-            log_listener_attach_result(
-                self.ensure_conversation_listener(thread_id, connection_id, raw_events_enabled)
-                    .await,
-                thread_id,
-                connection_id,
-                "thread",
-            );
-        }
+        super::thread_lifecycle::attach_created_thread_listener(
+            self.listener_task_context(),
+            thread_id,
+            thread,
+            connection_ids,
+        )
+        .await;
    }

    async fn thread_resume_inner(
@@ -4074,11 +4016,20 @@ impl ThreadRequestProcessor {
        // Path-based resume can use an empty request thread ID. Coordinate once its real
        // identity is known; unrelated loaded threads never wait for this cold startup.
        let _goal_resume_guard = if let InitialHistory::Resumed(resumed) = &thread_history {
-            Some(
-                self.thread_state_manager
-                    .lock_goal_resume(resumed.conversation_id)
-                    .await,
-            )
+            let guard = self
+                .thread_state_manager
+                .lock_goal_resume(resumed.conversation_id)
+                .await;
+            let state = self
+                .thread_state_manager
+                .thread_state(resumed.conversation_id)
+                .await;
+            let drain = state.lock().await.pending_shutdown_drain();
+            if let Some(mut drain) = drain {
+                // Core can remove the old runtime before its listener forwards ShutdownComplete.
+                let _ = drain.wait_for(|drained| *drained).await;
+            }
+            Some(guard)
        } else {
            None
        };
@@ -1236,12 +1236,12 @@ mod thread_processor_behavior_tests {
        manager
            .connection_initialized(connection, ConnectionCapabilities::default())
            .await;
-        manager
-            .try_ensure_connection_subscribed(
-                thread_id, connection, /*experimental_raw_events*/ false,
-            )
-            .await
-            .expect("connection should be live");
+        assert!(
+            manager
+                .try_add_connection_to_thread(thread_id, connection)
+                .await,
+            "connection should be live"
+        );
        {
            let state = manager.thread_state(thread_id).await;
            let mut state = state.lock().await;
@@ -1287,22 +1287,18 @@ mod thread_processor_behavior_tests {
        manager
            .connection_initialized(connection_b, ConnectionCapabilities::default())
            .await;
-        manager
-            .try_ensure_connection_subscribed(
-                thread_id,
-                connection_a,
-                /*experimental_raw_events*/ false,
-            )
-            .await
-            .expect("connection_a should be live");
-        manager
-            .try_ensure_connection_subscribed(
-                thread_id,
-                connection_b,
-                /*experimental_raw_events*/ false,
-            )
-            .await
-            .expect("connection_b should be live");
+        assert!(
+            manager
+                .try_add_connection_to_thread(thread_id, connection_a)
+                .await,
+            "connection_a should be live"
+        );
+        assert!(
+            manager
+                .try_add_connection_to_thread(thread_id, connection_b)
+                .await,
+            "connection_b should be live"
+        );
        {
            let state = manager.thread_state(thread_id).await;
            state.lock().await.cancel_tx = Some(cancel_tx);
@@ -1336,14 +1332,12 @@ mod thread_processor_behavior_tests {
        manager
            .connection_initialized(connection_b, ConnectionCapabilities::default())
            .await;
-        manager
-            .try_ensure_connection_subscribed(
-                thread_id,
-                connection_a,
-                /*experimental_raw_events*/ false,
-            )
-            .await
-            .expect("connection_a should be live");
+        assert!(
+            manager
+                .try_add_connection_to_thread(thread_id, connection_a)
+                .await,
+            "connection_a should be live"
+        );
        let mut has_connections = manager
            .subscribe_to_has_connections(thread_id)
            .await
@@ -1412,12 +1406,9 @@ mod thread_processor_behavior_tests {
        assert_eq!(threads_to_unload, Vec::<ThreadId>::new());

        assert!(
-            manager
-                .try_ensure_connection_subscribed(
-                    thread_id, connection, /*experimental_raw_events*/ false
-                )
+            !manager
+                .try_add_connection_to_thread(thread_id, connection)
                .await
-                .is_none()
        );
        assert!(!manager.has_subscribers(thread_id).await);
        Ok(())
@@ -25,6 +25,7 @@ use codex_rollout::state_db::StateDbHandle;
 use codex_utils_path_uri::LegacyAppPathString;
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::collections::hash_map::Entry;
 use std::sync::Arc;
 use std::sync::Mutex as StdMutex;
 use std::sync::Weak;
@@ -36,7 +37,8 @@ use tracing::error;

 type PendingInterruptQueue = Vec<ConnectionRequestId>;

-// Removing an unload entry drops its sender and wakes deferred listener attachment.
+// The unload that inserts an entry owns its removal. Dropping the sender wakes deferred
+// listener attachment only after that unload finishes cleanup.
 pub(crate) type PendingThreadUnloads = Arc<Mutex<HashMap<ThreadId, watch::Sender<()>>>>;

 pub(crate) struct PendingThreadResumeRequest {
@@ -122,7 +124,7 @@ pub(crate) struct ThreadState {
    pub(crate) last_terminal_turn_id: Option<String>,
    /// Lets an internal runtime replacement wait until the old listener has processed Core's
    /// `ShutdownComplete` event before that listener is superseded.
-    shutdown_drain_waiter: Option<oneshot::Sender<()>>,
+    pub(crate) shutdown_drain: Option<watch::Receiver<bool>>,
    pub(crate) cancel_tx: Option<oneshot::Sender<()>>,
    pub(crate) experimental_raw_events: bool,
    pub(crate) listener_generation: u64,
@@ -135,6 +137,12 @@ pub(crate) struct ThreadState {
 }

 impl ThreadState {
+    /// Serializes goal edits, cold resume, and terminal cleanup for this exact state.
+    pub(crate) async fn lock_goal_resume(state: &Mutex<Self>) -> tokio::sync::OwnedMutexGuard<()> {
+        let lock = Arc::clone(&state.lock().await.goal_resume_lock);
+        lock.lock_owned().await
+    }
+
    pub(crate) fn listener_matches(&self, conversation: &Arc<CodexThread>) -> bool {
        self.listener_thread
            .as_ref()
@@ -142,9 +150,18 @@ impl ThreadState {
            .is_some_and(|existing| Arc::ptr_eq(&existing, conversation))
    }

+    /// A shut-down runtime can still have terminal events buffered for its listener.
+    pub(crate) fn pending_shutdown_drain(&self) -> Option<watch::Receiver<bool>> {
+        let drain = self.shutdown_drain.as_ref()?;
+        let thread = self.listener_thread.as_ref()?.upgrade()?;
+        (thread.has_shutdown_status() && !*drain.borrow() && drain.has_changed().is_ok())
+            .then(|| drain.clone())
+    }
+
    pub(crate) fn set_listener(
        &mut self,
        cancel_tx: oneshot::Sender<()>,
+        shutdown_drain: watch::Receiver<bool>,
        conversation: &Arc<CodexThread>,
        watch_registration: WatchRegistration,
        thread_settings_baseline: ThreadSettings,
@@ -153,6 +170,7 @@ impl ThreadState {
        if let Some(previous) = self.cancel_tx.replace(cancel_tx) {
            let _ = previous.send(());
        }
+        self.shutdown_drain = Some(shutdown_drain);
        self.listener_generation = self.listener_generation.wrapping_add(1);
        self.last_thread_settings = Some(thread_settings_baseline);
        self.history_mode = history_mode;
@@ -167,7 +185,7 @@ impl ThreadState {
        if let Some(cancel_tx) = self.cancel_tx.take() {
            let _ = cancel_tx.send(());
        }
-        self.shutdown_drain_waiter = None;
+        self.shutdown_drain = None;
        self.listener_command_tx = None;
        self.current_turn_history.reset();
        self.listener_thread = None;
@@ -205,16 +223,6 @@ impl ThreadState {
        self.current_turn_history.active_turn_metadata_snapshot()
    }

-    pub(crate) fn register_shutdown_drain_waiter(&mut self) -> oneshot::Receiver<()> {
-        let (completion_tx, completion_rx) = oneshot::channel();
-        self.shutdown_drain_waiter = Some(completion_tx);
-        completion_rx
-    }
-
-    pub(crate) fn take_shutdown_drain_waiter(&mut self) -> Option<oneshot::Sender<()>> {
-        self.shutdown_drain_waiter.take()
-    }
-
    pub(crate) fn track_current_turn_event(&mut self, event_turn_id: &str, event: &EventMsg) {
        if let EventMsg::TurnStarted(payload) = event {
            self.turn_summary.started_at = payload.started_at;
@@ -342,15 +350,19 @@ struct ThreadEntry {

 impl Default for ThreadEntry {
    fn default() -> Self {
-        Self {
-            state: Arc::new(Mutex::new(ThreadState::default())),
-            connection_ids: HashSet::new(),
-            has_connections_watcher: watch::channel(false).0,
-        }
+        Self::new(Arc::new(Mutex::new(ThreadState::default())))
    }
 }

 impl ThreadEntry {
+    fn new(state: Arc<Mutex<ThreadState>>) -> Self {
+        Self {
+            state,
+            connection_ids: HashSet::new(),
+            has_connections_watcher: watch::channel(false).0,
+        }
+    }
+
    fn update_has_connections(&self) {
        let _ = self.has_connections_watcher.send_if_modified(|current| {
            let prev = *current;
@@ -381,20 +393,41 @@ pub(crate) struct ThreadStateManager {
        Arc<StdMutex<HashMap<ThreadId, mpsc::UnboundedSender<ThreadListenerCommand>>>>,
 }

+/// A listener's state, which may not have been published to the manager yet.
+pub(crate) enum ListenerStateCandidate {
+    Existing(Arc<Mutex<ThreadState>>),
+    Unpublished(Arc<Mutex<ThreadState>>),
+}
+
+impl ListenerStateCandidate {
+    pub(crate) fn state(&self) -> &Arc<Mutex<ThreadState>> {
+        match self {
+            Self::Existing(state) | Self::Unpublished(state) => state,
+        }
+    }
+}
+
 impl ThreadStateManager {
    /// Coordinates goal edits with cold/path-based resume for just this thread.
    pub(crate) async fn lock_goal_resume(
        &self,
        thread_id: ThreadId,
    ) -> tokio::sync::OwnedMutexGuard<()> {
-        let lock = self
-            .thread_state(thread_id)
-            .await
-            .lock()
-            .await
-            .goal_resume_lock
-            .clone();
-        lock.lock_owned().await
+        loop {
+            let thread_state = self.thread_state(thread_id).await;
+            let guard = ThreadState::lock_goal_resume(&thread_state).await;
+            // Teardown may have replaced the entry while this caller waited for the lock.
+            if self
+                .state
+                .lock()
+                .await
+                .threads
+                .get(&thread_id)
+                .is_some_and(|entry| Arc::ptr_eq(&entry.state, &thread_state))
+            {
+                return guard;
+            }
+        }
    }

    pub(crate) fn new() -> Self {
@@ -464,6 +497,18 @@ impl ThreadStateManager {
        state.threads.entry(thread_id).or_default().state.clone()
    }

+    pub(crate) async fn created_listener_state(
+        &self,
+        thread_id: ThreadId,
+    ) -> ListenerStateCandidate {
+        match self.state.lock().await.threads.get(&thread_id) {
+            Some(entry) => ListenerStateCandidate::Existing(Arc::clone(&entry.state)),
+            None => {
+                ListenerStateCandidate::Unpublished(Arc::new(Mutex::new(ThreadState::default())))
+            }
+        }
+    }
+
    pub(crate) fn current_listener_command_tx(
        &self,
        thread_id: ThreadId,
@@ -578,10 +623,6 @@ impl ThreadStateManager {
    ) -> bool {
        {
            let mut state = self.state.lock().await;
-            if !state.threads.contains_key(&thread_id) {
-                return false;
-            }
-
            if !state
                .thread_ids_by_connection
                .get(&connection_id)
@@ -615,34 +656,47 @@ impl ThreadStateManager {
            .is_some_and(|thread_entry| !thread_entry.connection_ids.is_empty())
    }

-    pub(crate) async fn try_ensure_connection_subscribed(
+    /// Publishes a new state only when listener installation succeeds, or reuses a current state.
+    /// The caller must hold the candidate's state lock.
+    pub(crate) async fn subscribe_and_install_listener<T, E>(
        &self,
        thread_id: ThreadId,
-        connection_id: ConnectionId,
-        experimental_raw_events: bool,
-    ) -> Option<Arc<Mutex<ThreadState>>> {
-        let thread_state = {
-            let mut state = self.state.lock().await;
-            if !state.live_connections.contains_key(&connection_id) {
-                return None;
-            }
-            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.connection_ids.insert(connection_id);
-            thread_entry.update_has_connections();
-            thread_entry.state.clone()
-        };
-        {
-            let mut thread_state_guard = thread_state.lock().await;
-            if experimental_raw_events {
-                thread_state_guard.set_experimental_raw_events(/*enabled*/ true);
+        candidate: &ListenerStateCandidate,
+        connection_ids: &[ConnectionId],
+        install: impl FnOnce(bool) -> Result<T, E>,
+    ) -> Result<Option<(T, watch::Receiver<bool>, bool)>, E> {
+        let mut state = self.state.lock().await;
+        let ThreadStateManagerInner {
+            live_connections,
+            threads,
+            thread_ids_by_connection,
+        } = &mut *state;
+        let slot = threads.entry(thread_id);
+        match (&slot, candidate) {
+            (Entry::Occupied(entry), _) if Arc::ptr_eq(&entry.get().state, candidate.state()) => {}
+            (Entry::Vacant(_), ListenerStateCandidate::Unpublished(_)) => {}
+            _ => return Ok(None),
+        }
+        let attached = connection_ids
+            .iter()
+            .any(|connection_id| live_connections.contains_key(connection_id));
+        let installed = install(attached)?;
+        let entry = slot.or_insert_with(|| ThreadEntry::new(Arc::clone(candidate.state())));
+        for &connection_id in connection_ids {
+            if live_connections.contains_key(&connection_id) {
+                thread_ids_by_connection
+                    .entry(connection_id)
+                    .or_default()
+                    .insert(thread_id);
+                entry.connection_ids.insert(connection_id);
            }
        }
-        Some(thread_state)
+        entry.update_has_connections();
+        Ok(Some((
+            installed,
+            entry.has_connections_watcher.subscribe(),
+            attached,
+        )))
    }

    pub(crate) async fn try_add_connection_to_thread(
@@ -691,6 +745,7 @@ impl ThreadStateManager {
        }
    }

+    #[cfg(test)]
    pub(crate) async fn subscribe_to_has_connections(
        &self,
        thread_id: ThreadId,
@@ -248,11 +248,8 @@ impl ThreadWatchManager {
        }
    }

-    pub(crate) async fn subscribe(
-        &self,
-        thread_id: ThreadId,
-    ) -> Option<watch::Receiver<ThreadStatus>> {
-        Some(self.state.lock().await.subscribe(thread_id.to_string()))
+    pub(crate) async fn subscribe(&self, thread_id: ThreadId) -> watch::Receiver<ThreadStatus> {
+        self.state.lock().await.subscribe(thread_id.to_string())
    }

    async fn note_active_guard_released(
@@ -826,14 +823,8 @@ mod tests {
            .expect("interactive thread id should parse");
        let non_interactive_thread_id = ThreadId::from_string(NON_INTERACTIVE_THREAD_ID)
            .expect("non-interactive thread id should parse");
-        let mut interactive_rx = manager
-            .subscribe(interactive_thread_id)
-            .await
-            .expect("interactive status watcher should subscribe");
-        let mut non_interactive_rx = manager
-            .subscribe(non_interactive_thread_id)
-            .await
-            .expect("non-interactive status watcher should subscribe");
+        let mut interactive_rx = manager.subscribe(interactive_thread_id).await;
+        let mut non_interactive_rx = manager.subscribe(non_interactive_thread_id).await;

        manager.note_turn_started(INTERACTIVE_THREAD_ID).await;

@@ -858,6 +858,11 @@ impl CodexThread {
        self.io.agent_status().await
    }

+    /// Returns whether the runtime has published its shutdown status.
+    pub fn has_shutdown_status(&self) -> bool {
+        matches!(*self.io.agent_status.borrow(), AgentStatus::Shutdown)
+    }
+
    pub async fn list_background_terminals(&self) -> Vec<BackgroundTerminalInfo> {
        self.session.list_background_terminals().await
    }
@@ -172,6 +172,7 @@ pub use thread_manager::NewThread;
 pub use thread_manager::PromptCacheKey;
 pub use thread_manager::StartThreadOptions;
 pub use thread_manager::ThreadManager;
+pub use thread_manager::ThreadRegistration;
 pub use thread_manager::ThreadShutdownReport;
 pub use thread_manager::build_models_manager;
 pub use thread_manager::local_agent_graph_store_from_state_db;
@@ -192,6 +192,13 @@ pub struct NewThread {
    pub session_configured: SessionConfiguredEvent,
 }

+/// Describes whether a captured runtime is still registered under its thread ID.
+pub enum ThreadRegistration {
+    Current,
+    Absent,
+    Replaced,
+}
+
 // TODO(ccunningham): Add an explicit non-interrupting live-turn snapshot once
 // core can represent sampling boundaries directly instead of relying on
 // whichever items happened to be persisted mid-turn.
@@ -1008,6 +1015,19 @@ impl ThreadManager {
        self.state.get_thread(thread_id).await
    }

+    /// Checks the full registry without exposing any replacement runtime.
+    pub async fn thread_registration(
+        &self,
+        thread_id: ThreadId,
+        expected: &Arc<CodexThread>,
+    ) -> ThreadRegistration {
+        match self.state.threads.read().await.get(&thread_id) {
+            Some(current) if Arc::ptr_eq(current, expected) => ThreadRegistration::Current,
+            Some(_) => ThreadRegistration::Replaced,
+            None => ThreadRegistration::Absent,
+        }
+    }
+
    /// Board hosts may reach internal runtimes only through a matching explicit connection.
    pub(crate) async fn get_thread_for_board(
        &self,