diff --git a/codex-rs/core/src/agent/collab_completion_warning.rs b/codex-rs/core/src/agent/collab_completion_warning.rs index b8c1bfe01c..3a63810d06 100644 --- a/codex-rs/core/src/agent/collab_completion_warning.rs +++ b/codex-rs/core/src/agent/collab_completion_warning.rs @@ -14,7 +14,14 @@ pub(crate) fn spawn_collab_completion_warning_watcher( agent_id: ThreadId, ) { tokio::spawn(async move { - if let Some(status) = wait_for_final_status(session.as_ref(), agent_id).await { + if let Some(status) = wait_for_final_status(session.as_ref(), agent_id).await + && !crate::agent::is_collab_wait_active( + session.as_ref(), + &turn_context.sub_id, + agent_id, + ) + .await + { let message = completion_warning_message(agent_id, &status); session.record_model_warning(message, &turn_context).await; } diff --git a/codex-rs/core/src/agent/collab_wait_tracking.rs b/codex-rs/core/src/agent/collab_wait_tracking.rs new file mode 100644 index 0000000000..ac751d9e99 --- /dev/null +++ b/codex-rs/core/src/agent/collab_wait_tracking.rs @@ -0,0 +1,32 @@ +use codex_protocol::ThreadId; + +use crate::codex::Session; + +pub(crate) async fn begin_collab_wait(session: &Session, turn_id: &str, agent_ids: &[ThreadId]) { + let active = session.active_turn.lock().await; + if let Some(active_turn) = active.as_ref() { + let mut state = active_turn.turn_state.lock().await; + state.begin_wait(turn_id, agent_ids); + } +} + +pub(crate) async fn end_collab_wait(session: &Session, turn_id: &str, agent_ids: &[ThreadId]) { + let active = session.active_turn.lock().await; + if let Some(active_turn) = active.as_ref() { + let mut state = active_turn.turn_state.lock().await; + state.end_wait(turn_id, agent_ids); + } +} + +pub(crate) async fn is_collab_wait_active( + session: &Session, + turn_id: &str, + agent_id: ThreadId, +) -> bool { + let active = session.active_turn.lock().await; + if let Some(active_turn) = active.as_ref() { + let state = active_turn.turn_state.lock().await; + return state.is_waiting_on(turn_id, agent_id); + } + false +} diff --git a/codex-rs/core/src/agent/mod.rs b/codex-rs/core/src/agent/mod.rs index 561438ca5d..3cc1316ca2 100644 --- a/codex-rs/core/src/agent/mod.rs +++ b/codex-rs/core/src/agent/mod.rs @@ -1,4 +1,5 @@ mod collab_completion_warning; +mod collab_wait_tracking; pub(crate) mod control; mod guards; pub(crate) mod role; @@ -6,6 +7,9 @@ pub(crate) mod status; pub(crate) use codex_protocol::protocol::AgentStatus; pub(crate) use collab_completion_warning::spawn_collab_completion_warning_watcher; +pub(crate) use collab_wait_tracking::begin_collab_wait; +pub(crate) use collab_wait_tracking::end_collab_wait; +pub(crate) use collab_wait_tracking::is_collab_wait_active; pub(crate) use control::AgentControl; pub(crate) use guards::MAX_THREAD_SPAWN_DEPTH; pub(crate) use guards::exceeds_thread_spawn_depth_limit; diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index ccc50d066b..ace7d72fdc 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -16,6 +16,7 @@ use tokio::sync::oneshot; use crate::codex::TurnContext; use crate::protocol::ReviewDecision; use crate::tasks::SessionTask; +use codex_protocol::ThreadId; /// Metadata about the currently running turn. pub(crate) struct ActiveTurn { @@ -73,6 +74,7 @@ pub(crate) struct TurnState { pending_user_input: HashMap>, pending_dynamic_tools: HashMap>, pending_input: Vec, + active_waits: HashMap>, } impl TurnState { @@ -96,6 +98,7 @@ impl TurnState { self.pending_user_input.clear(); self.pending_dynamic_tools.clear(); self.pending_input.clear(); + self.active_waits.clear(); } pub(crate) fn insert_pending_user_input( @@ -145,6 +148,43 @@ impl TurnState { pub(crate) fn has_pending_input(&self) -> bool { !self.pending_input.is_empty() } + + pub(crate) fn begin_wait(&mut self, turn_id: &str, agent_ids: &[ThreadId]) { + if agent_ids.is_empty() { + return; + } + let waits = self.active_waits.entry(turn_id.to_string()).or_default(); + for agent_id in agent_ids { + *waits.entry(*agent_id).or_default() += 1; + } + } + + pub(crate) fn end_wait(&mut self, turn_id: &str, agent_ids: &[ThreadId]) { + if agent_ids.is_empty() { + return; + } + let mut remove_turn = false; + if let Some(waits) = self.active_waits.get_mut(turn_id) { + for agent_id in agent_ids { + if let Some(count) = waits.get_mut(agent_id) { + *count = count.saturating_sub(1); + if *count == 0 { + waits.remove(agent_id); + } + } + } + remove_turn = waits.is_empty(); + } + if remove_turn { + self.active_waits.remove(turn_id); + } + } + + pub(crate) fn is_waiting_on(&self, turn_id: &str, agent_id: ThreadId) -> bool { + self.active_waits + .get(turn_id) + .is_some_and(|waits| waits.contains_key(&agent_id)) + } } impl ActiveTurn { @@ -154,3 +194,35 @@ impl ActiveTurn { ts.clear_pending(); } } + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + fn thread_id(s: &str) -> ThreadId { + ThreadId::from_string(s).expect("valid thread id") + } + + #[test] + fn wait_tracking_is_turn_scoped_and_reference_counted() { + let mut state = TurnState::default(); + let turn_a = "turn-a"; + let turn_b = "turn-b"; + let agent = thread_id("00000000-0000-7000-0000-000000000001"); + + state.begin_wait(turn_a, &[agent]); + state.begin_wait(turn_a, &[agent]); + state.begin_wait(turn_b, &[agent]); + + assert_eq!(state.is_waiting_on(turn_a, agent), true); + assert_eq!(state.is_waiting_on(turn_b, agent), true); + + state.end_wait(turn_a, &[agent]); + assert_eq!(state.is_waiting_on(turn_a, agent), true); + + state.end_wait(turn_a, &[agent]); + assert_eq!(state.is_waiting_on(turn_a, agent), false); + assert_eq!(state.is_waiting_on(turn_b, agent), true); + } +} diff --git a/codex-rs/core/src/tools/handlers/collab.rs b/codex-rs/core/src/tools/handlers/collab.rs index 4eb4fb8ab2..c848d8b7ad 100644 --- a/codex-rs/core/src/tools/handlers/collab.rs +++ b/codex-rs/core/src/tools/handlers/collab.rs @@ -333,114 +333,120 @@ mod wait { ms => ms.min(MAX_WAIT_TIMEOUT_MS), }; - session - .send_event( - &turn, - CollabWaitingBeginEvent { - sender_thread_id: session.conversation_id, - receiver_thread_ids: receiver_thread_ids.clone(), - call_id: call_id.clone(), - } - .into(), - ) - .await; - - 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(); - if is_final(&status) { - initial_final_statuses.push((*id, status)); + crate::agent::begin_collab_wait(session.as_ref(), &turn.sub_id, &receiver_thread_ids).await; + let result = async { + session + .send_event( + &turn, + CollabWaitingBeginEvent { + sender_thread_id: session.conversation_id, + receiver_thread_ids: receiver_thread_ids.clone(), + call_id: call_id.clone(), } - status_rxs.push((*id, rx)); - } - Err(CodexErr::ThreadNotFound(_)) => { - initial_final_statuses.push((*id, AgentStatus::NotFound)); - } - Err(err) => { - let mut statuses = HashMap::with_capacity(1); - statuses.insert(*id, session.services.agent_control.get_status(*id).await); - session - .send_event( - &turn, - CollabWaitingEndEvent { - sender_thread_id: session.conversation_id, - call_id: call_id.clone(), - statuses, - } - .into(), - ) - .await; - return Err(collab_agent_error(*id, err)); - } - } - } + .into(), + ) + .await; - let statuses = if !initial_final_statuses.is_empty() { - initial_final_statuses - } else { - // Wait for the first agent to reach a final status. - let mut futures = FuturesUnordered::new(); - for (id, rx) in status_rxs.into_iter() { - let session = session.clone(); - futures.push(wait_for_final_status(session, id, rx)); - } - let mut results = Vec::new(); - let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64); - loop { - match timeout_at(deadline, futures.next()).await { - Ok(Some(Some(result))) => { - results.push(result); - break; + 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(); + if is_final(&status) { + initial_final_statuses.push((*id, status)); + } + status_rxs.push((*id, rx)); + } + Err(CodexErr::ThreadNotFound(_)) => { + initial_final_statuses.push((*id, AgentStatus::NotFound)); + } + Err(err) => { + let mut statuses = HashMap::with_capacity(1); + statuses.insert(*id, session.services.agent_control.get_status(*id).await); + session + .send_event( + &turn, + CollabWaitingEndEvent { + sender_thread_id: session.conversation_id, + call_id: call_id.clone(), + statuses, + } + .into(), + ) + .await; + Err(collab_agent_error(*id, err))?; } - Ok(Some(None)) => continue, - Ok(None) | Err(_) => break, } } - if !results.is_empty() { - // Drain the unlikely last elements to prevent race. + + let statuses = if !initial_final_statuses.is_empty() { + initial_final_statuses + } else { + // Wait for the first agent to reach a final status. + let mut futures = FuturesUnordered::new(); + for (id, rx) in status_rxs.into_iter() { + let session = session.clone(); + futures.push(wait_for_final_status(session, id, rx)); + } + let mut results = Vec::new(); + let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64); loop { - match futures.next().now_or_never() { - Some(Some(Some(result))) => results.push(result), - Some(Some(None)) => continue, - Some(None) | None => break, + match timeout_at(deadline, futures.next()).await { + Ok(Some(Some(result))) => { + results.push(result); + break; + } + Ok(Some(None)) => continue, + Ok(None) | Err(_) => break, } } - } - results - }; - - // Convert payload. - let statuses_map = statuses.clone().into_iter().collect::>(); - let result = WaitResult { - status: statuses_map.clone(), - timed_out: statuses.is_empty(), - }; - - // Final event emission. - session - .send_event( - &turn, - CollabWaitingEndEvent { - sender_thread_id: session.conversation_id, - call_id, - statuses: statuses_map, + if !results.is_empty() { + // Drain the unlikely last elements to prevent race. + loop { + match futures.next().now_or_never() { + Some(Some(Some(result))) => results.push(result), + Some(Some(None)) => continue, + Some(None) | None => break, + } + } } - .into(), - ) - .await; + results + }; - let content = serde_json::to_string(&result).map_err(|err| { - FunctionCallError::Fatal(format!("failed to serialize wait result: {err}")) - })?; + // Convert payload. + let statuses_map = statuses.clone().into_iter().collect::>(); + let result = WaitResult { + status: statuses_map.clone(), + timed_out: statuses.is_empty(), + }; - Ok(ToolOutput::Function { - content, - success: None, - content_items: None, - }) + // Final event emission. + session + .send_event( + &turn, + CollabWaitingEndEvent { + sender_thread_id: session.conversation_id, + call_id, + statuses: statuses_map, + } + .into(), + ) + .await; + + let content = serde_json::to_string(&result).map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize wait result: {err}")) + })?; + + Ok(ToolOutput::Function { + content, + success: None, + content_items: None, + }) + } + .await; + crate::agent::end_collab_wait(session.as_ref(), &turn.sub_id, &receiver_thread_ids).await; + result } async fn wait_for_final_status(