diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 9ad9a6f965..c95f409b3f 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/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> { - 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, }; diff --git a/codex-rs/core/src/agent/control/watch.rs b/codex-rs/core/src/agent/control/watch.rs new file mode 100644 index 0000000000..590ad822e5 --- /dev/null +++ b/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 { + 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()) + } +} diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 9ea1e71131..c14218ae14 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/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 { diff --git a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs index 9f6f970d08..09efb14812 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs +++ b/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, thread_id: ThreadId, - mut status_rx: Receiver, -) -> 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, 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))) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 8a563f6358..16a1e4d9b0 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/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