diff --git a/codex-rs/core/src/subagents/manager.rs b/codex-rs/core/src/subagents/manager.rs index fd6ac560bb..1fd1e38d24 100644 --- a/codex-rs/core/src/subagents/manager.rs +++ b/codex-rs/core/src/subagents/manager.rs @@ -1,8 +1,12 @@ use std::collections::HashMap; use std::collections::VecDeque; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -94,8 +98,16 @@ pub struct SubagentManager { struct WatchdogHandle { cancel: tokio_util::sync::CancellationToken, - #[allow(dead_code)] - task: tokio::task::JoinHandle<()>, +} + +struct NeverFuture; + +impl Future for NeverFuture { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + Poll::Pending + } } #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] @@ -1149,7 +1161,7 @@ impl SubagentManager { let cancel_child = cancel.clone(); let manager = self.clone(); let msg_clone = message.clone(); - let task = tokio::spawn(async move { + let _task = tokio::spawn(async move { loop { tokio::select! { _ = cancel_child.cancelled() => break, @@ -1214,7 +1226,7 @@ impl SubagentManager { .watchdogs .write() .await - .insert(key, WatchdogHandle { cancel, task }) + .insert(key, WatchdogHandle { cancel }) .is_some(); Ok(if replaced { WatchdogAction::Replaced @@ -1379,6 +1391,55 @@ impl SubagentManager { } } + fn timeout_future(duration: Option) -> Pin + Send>> { + match duration { + Some(duration) => Box::pin(tokio::time::sleep(duration)), + None => Box::pin(NeverFuture), + } + } + + async fn record_completion_and_metadata( + &self, + session_id: &ConversationId, + completion: &SubagentCompletion, + ) -> Result { + { + let mut completions = self.completions.write().await; + completions.insert(*session_id, completion.clone()); + } + + let desired_status = status_from_completion(completion); + let mut metadata = self + .registry + .get(session_id) + .await + .ok_or(SubagentManagerError::NotFound)?; + + if metadata.status != desired_status { + self.update_status_and_emit(session_id, desired_status) + .await; + metadata.status = desired_status; + } + + Ok(metadata) + } + + async fn completion_result( + &self, + session_id: &ConversationId, + completion: SubagentCompletion, + ) -> Result { + let metadata = self + .record_completion_and_metadata(session_id, &completion) + .await?; + + Ok(AwaitInboxResult { + metadata, + completion: Some(completion), + messages: Vec::new(), + }) + } + pub async fn await_completion( &self, session_id: &ConversationId, @@ -1388,25 +1449,9 @@ impl SubagentManager { let completions = self.completions.read().await; completions.get(session_id).cloned() } { - let desired_status = status_from_completion(&completion); let metadata = self - .registry - .get(session_id) - .await - .ok_or(SubagentManagerError::NotFound)?; - if metadata.status != desired_status { - self.update_status_and_emit(session_id, desired_status) - .await; - let metadata = self - .registry - .get(session_id) - .await - .ok_or(SubagentManagerError::NotFound)?; - return Ok(AwaitResult { - metadata, - completion, - }); - } + .record_completion_and_metadata(session_id, &completion) + .await?; return Ok(AwaitResult { metadata, completion, @@ -1425,10 +1470,8 @@ impl SubagentManager { completions.get(session_id).cloned() } { let metadata = self - .registry - .get(session_id) - .await - .ok_or(SubagentManagerError::NotFound)?; + .record_completion_and_metadata(session_id, &completion) + .await?; return Ok(AwaitResult { metadata, completion, @@ -1442,10 +1485,8 @@ impl SubagentManager { if let Some(completion) = current_completion(&receiver) { let metadata = self - .registry - .get(session_id) - .await - .ok_or(SubagentManagerError::NotFound)?; + .record_completion_and_metadata(session_id, &completion) + .await?; return Ok(AwaitResult { metadata, completion, @@ -1497,24 +1538,9 @@ impl SubagentManager { .ok_or(SubagentManagerError::NotFound)? }; - { - let mut completions = self.completions.write().await; - completions.insert(*session_id, completion.clone()); - } - - let desired_status = status_from_completion(&completion); - if let Some(current) = self.registry.get(session_id).await - && current.status != desired_status - { - self.update_status_and_emit(session_id, desired_status) - .await; - } - let metadata = self - .registry - .get(session_id) - .await - .ok_or(SubagentManagerError::NotFound)?; + .record_completion_and_metadata(session_id, &completion) + .await?; Ok(AwaitResult { metadata, completion, @@ -1567,8 +1593,14 @@ impl SubagentManager { let mut completion_opt = current_completion(&receiver); if completion_opt.is_some() || !messages.is_empty() { if let Some(ref completion) = completion_opt { - self.ensure_completion_recorded(session_id, completion) + let metadata = self + .record_completion_and_metadata(session_id, completion) .await?; + return Ok(AwaitInboxResult { + metadata, + completion: completion_opt, + messages, + }); } let metadata = self .registry @@ -1624,8 +1656,14 @@ impl SubagentManager { // Finalize completion state if we observed a terminal result. if let Some(ref completion) = completion_opt { - self.ensure_completion_recorded(session_id, completion) + let metadata = self + .record_completion_and_metadata(session_id, completion) .await?; + return Ok(AwaitInboxResult { + metadata, + completion: completion_opt, + messages, + }); } let metadata = self