Stream agent snapshots from local status subscriptions (#47267)

## What changed

- Return `StatusSubscription` streams of `AgentInfo` snapshots from local status subscriptions, including agent metadata and configuration.
- Emit the initial snapshot before subsequent coalesced updates. Use a weak thread reference and retain the last observed configuration so published final updates remain available after the runtime is dropped.
- Update completion notifications and `wait_agent` to consume snapshot streams, propagate subscription errors from `wait_agent`, and fall back to the latest status when a stream ends.

## Testing

Update the shutdown subscription test to verify that shutdown before the first poll preserves the initial `PendingInit` snapshot followed by `Shutdown`.

GitOrigin-RevId: beb46952cbcc9bddf65a4d6603fea939d5ace8b9

codex-rs/core/src/agent/control.rs

@@ -55,13 +55,13 @@ use codex_protocol::protocol::ThreadSource;
 use codex_protocol::user_input::UserInput;
 use codex_thread_store::LoadThreadHistoryParams;
 use codex_thread_store::ReadThreadParams;
+use futures::StreamExt;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::collections::VecDeque;
 use std::sync::Arc;
 use std::sync::OnceLock;
 use std::sync::Weak;
-use tokio::sync::watch;
 use tracing::warn;
 use uuid::Uuid;
 
@@ -83,6 +83,7 @@ mod spawn;
 mod spawn_request;
 mod target;
 mod user_authorization;
+mod watch;
 
 const MAX_ENVIRONMENT_SUBAGENTS: usize = 8;
 const MAX_ENVIRONMENT_SUBAGENT_BYTES: usize = 1_024;
@@ -426,16 +427,6 @@ impl LocalAgentControl {
         Ok(thread_ids)
     }
 
-    /// Subscribe to status updates for `agent_id`, yielding the latest value and changes.
-    pub(crate) async fn subscribe_status(
-        &self,
-        agent_id: ThreadId,
-    ) -> CodexResult<watch::Receiver<AgentStatus>> {
-        let state = self.upgrade()?;
-        let thread = state.get_thread(agent_id).await?;
-        Ok(thread.subscribe_status())
-    }
-
     pub(crate) async fn format_environment_context_subagents(
         &self,
         parent_thread_id: ThreadId,
@@ -598,16 +589,20 @@ impl LocalAgentControl {
         let control = self.clone();
         tokio::spawn(async move {
             let status = match control.subscribe_status(child_thread_id).await {
-                Ok(mut status_rx) => {
-                    let mut status = status_rx.borrow().clone();
-                    while !is_final(&status) {
-                        if status_rx.changed().await.is_err() {
-                            status = control.get_status(child_thread_id).await;
+                Ok(mut updates) => {
+                    let mut final_status = None;
+                    while let Some(Ok(snapshot)) = updates.next().await {
+                        if let Some(status) = snapshot.status()
+                            && is_final(status)
+                        {
+                            final_status = Some(status.clone());
                             break;
                         }
-                        status = status_rx.borrow().clone();
                     }
-                    status
+                    match final_status {
+                        Some(status) => status,
+                        None => control.get_status(child_thread_id).await,
+                    }
                 }
                 Err(_) => control.get_status(child_thread_id).await,
             };

codex-rs/core/src/agent/control/watch.rs

@@ -0,0 +1,61 @@
+//! Exposes local status changes as shared agent snapshots instead of Tokio receivers.
+//! Subscriptions retain no runtime handle and preserve the watch channel's coalescing.
+
+use super::LocalAgentControl;
+use crate::agent::api::AgentInfo;
+use crate::agent::api::StatusSubscription;
+use crate::agent::types::LiveAgent;
+use codex_protocol::ThreadId;
+use codex_protocol::error::Result as CodexResult;
+use futures::StreamExt;
+use futures::stream;
+use std::sync::Arc;
+
+impl LocalAgentControl {
+    pub(crate) async fn subscribe_status(
+        &self,
+        agent_id: ThreadId,
+    ) -> CodexResult<StatusSubscription> {
+        let manager = self.upgrade()?;
+        let thread = manager.get_thread(agent_id).await?;
+        // Subscribe before reading settings; mark the initial status seen only afterward.
+        let mut receiver = thread.subscribe_status();
+        let config = Box::new(thread.config_snapshot().await);
+        let initial = LiveAgent {
+            thread_id: agent_id,
+            metadata: self.get_agent_metadata(agent_id).unwrap_or_default(),
+            status: receiver.borrow_and_update().clone(),
+        };
+        let weak_thread = Arc::downgrade(&thread);
+        let changes = stream::unfold(
+            (receiver, initial.clone(), config.clone()),
+            move |(mut receiver, mut snapshot, mut config)| {
+                let weak_thread = weak_thread.clone();
+                async move {
+                    receiver.changed().await.ok()?;
+                    snapshot.status = receiver.borrow_and_update().clone();
+                    if let Some(thread) = weak_thread.upgrade() {
+                        *config = thread.config_snapshot().await;
+                    }
+                    // A final status can outlive the runtime. Keep its last observed settings
+                    // so removal from the manager does not hide an already-published update.
+                    Some((
+                        Ok(AgentInfo::Loaded {
+                            agent: snapshot.clone(),
+                            config: config.clone(),
+                        }),
+                        (receiver, snapshot, config),
+                    ))
+                }
+            },
+        );
+        Ok(stream::once(async move {
+            Ok(AgentInfo::Loaded {
+                agent: initial,
+                config,
+            })
+        })
+        .chain(changes)
+        .boxed())
+    }
+}

codex-rs/core/src/agent/control_tests.rs

@@ -709,7 +709,8 @@ async fn subscribe_status_errors_for_missing_thread() {
         .control
         .subscribe_status(thread_id)
         .await
-        .expect_err("subscribe_status should fail for missing thread");
+        .err()
+        .expect("subscribe_status should fail for missing thread");
     assert_matches!(
         err.details(),
         CodexErrorDetails::ThreadNotFound(id) if *id == thread_id
@@ -725,15 +726,25 @@ async fn subscribe_status_updates_on_shutdown() {
         .subscribe_status(thread_id)
         .await
         .expect("subscribe_status should succeed");
-    assert_eq!(status_rx.borrow().clone(), AgentStatus::PendingInit);
-
     let _ = thread
         .submit(Op::Shutdown {})
         .await
         .expect("shutdown should submit");
+    thread.wait_until_terminated().await;
 
-    let _ = status_rx.changed().await;
-    assert_eq!(status_rx.borrow().clone(), AgentStatus::Shutdown);
+    // The update arrives before the first poll, without replacing the initial snapshot.
+    let initial = status_rx
+        .next()
+        .await
+        .expect("initial snapshot")
+        .expect("snapshot");
+    assert_eq!(initial.status(), Some(&AgentStatus::PendingInit));
+    let shutdown = status_rx
+        .next()
+        .await
+        .expect("shutdown snapshot")
+        .expect("snapshot");
+    assert_eq!(shutdown.status(), Some(&AgentStatus::Shutdown));
 }
 
 #[tokio::test]
@@ -4399,11 +4410,7 @@ async fn resume_thread_subagent_restores_stored_metadata() {
         .flush_rollout()
         .await
         .expect("flush child rollout");
-    let mut status_rx = harness
-        .control
-        .subscribe_status(child_thread_id)
-        .await
-        .expect("status subscription should succeed");
+    let mut status_rx = child_thread.subscribe_status();
     if matches!(status_rx.borrow().clone(), AgentStatus::PendingInit) {
         timeout(Duration::from_secs(5), async {
             loop {

codex-rs/core/src/tools/handlers/multi_agents/wait.rs

@@ -1,8 +1,10 @@
 use super::*;
+use crate::agent::api::StatusSubscription;
 use crate::agent::status::is_final;
 use crate::session::session::Session;
 use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions;
 use crate::tools::handlers::multi_agents_spec::create_wait_agent_tool_v1;
+use codex_protocol::error::CodexErr;
 use codex_protocol::error::CodexErrorDetails;
 use codex_tools::ToolSpec;
 use futures::FutureExt;
@@ -11,7 +13,6 @@ use futures::stream::FuturesUnordered;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::sync::watch::Receiver;
 use tokio::time::Instant;
 
 use tokio::time::timeout_at;
@@ -120,13 +121,23 @@ impl Handler {
         let mut status_rxs = Vec::with_capacity(receiver_thread_ids.len());
         let mut initial_final_statuses = Vec::new();
         for id in &receiver_thread_ids {
-            match session.services.agent_control.subscribe_status(*id).await {
-                Ok(rx) => {
-                    let status = rx.borrow().clone();
+            let subscription = async {
+                let mut updates = session.services.agent_control.subscribe_status(*id).await?;
+                let initial = updates
+                    .next()
+                    .await
+                    .transpose()?
+                    .ok_or(CodexErr::InternalAgentDied)?;
+                Ok::<_, CodexErr>((initial, updates))
+            }
+            .await;
+            match subscription {
+                Ok((initial, updates)) => {
+                    let status = initial.status().cloned().unwrap_or(AgentStatus::NotFound);
                     if is_final(&status) {
                         initial_final_statuses.push((*id, status));
                     }
-                    status_rxs.push((*id, rx));
+                    status_rxs.push((*id, updates));
                 }
                 Err(err) if matches!(err.details(), CodexErrorDetails::ThreadNotFound(_)) => {
                     initial_final_statuses.push((*id, AgentStatus::NotFound));
@@ -168,19 +179,21 @@ impl Handler {
             let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64);
             loop {
                 match timeout_at(deadline, futures.next()).await {
-                    Ok(Some(Some(result))) => {
+                    Ok(Some(Ok(Some(result)))) => {
                         results.push(result);
                         break;
                     }
-                    Ok(Some(None)) => continue,
+                    Ok(Some(Ok(None))) => continue,
+                    Ok(Some(Err(err))) => return Err(err),
                     Ok(None) | Err(_) => break,
                 }
             }
             if !results.is_empty() {
                 loop {
                     match futures.next().now_or_never() {
-                        Some(Some(Some(result))) => results.push(result),
-                        Some(Some(None)) => continue,
+                        Some(Some(Ok(Some(result)))) => results.push(result),
+                        Some(Some(Ok(None))) => continue,
+                        Some(Some(Err(err))) => return Err(err),
                         Some(None) | None => break,
                     }
                 }
@@ -307,21 +320,15 @@ impl ToolOutput for WaitAgentResult {
 async fn wait_for_final_status(
     session: Arc<Session>,
     thread_id: ThreadId,
-    mut status_rx: Receiver<AgentStatus>,
-) -> Option<(ThreadId, AgentStatus)> {
-    let mut status = status_rx.borrow().clone();
-    if is_final(&status) {
-        return Some((thread_id, status));
-    }
-
-    loop {
-        if status_rx.changed().await.is_err() {
-            let latest = session.services.agent_control.get_status(thread_id).await;
-            return is_final(&latest).then_some((thread_id, latest));
-        }
-        status = status_rx.borrow().clone();
+    mut updates: StatusSubscription,
+) -> Result<Option<(ThreadId, AgentStatus)>, FunctionCallError> {
+    while let Some(snapshot) = updates.next().await {
+        let snapshot = snapshot.map_err(|err| collab_agent_error(thread_id, err))?;
+        let status = snapshot.status().cloned().unwrap_or(AgentStatus::NotFound);
         if is_final(&status) {
-            return Some((thread_id, status));
+            return Ok(Some((thread_id, status)));
         }
     }
+    let latest = session.services.agent_control.get_status(thread_id).await;
+    Ok(is_final(&latest).then_some((thread_id, latest)))
 }

codex-rs/core/src/tools/handlers/multi_agents_tests.rs

@@ -3318,11 +3318,7 @@ async fn wait_agent_returns_final_status_without_timeout() {
         .await
         .expect("start thread");
     let agent_id = thread.thread_id;
-    let mut status_rx = manager
-        .agent_control()
-        .subscribe_status(agent_id)
-        .await
-        .expect("subscribe should succeed");
+    let mut status_rx = thread.thread.subscribe_status();
 
     let _ = thread
         .thread