use super::*;use crate::agent::status::is_final;use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions;use crate::tools::handlers::multi_agents_spec::create_wait_agent_tool_v1;use crate::turn_timing::now_unix_timestamp_ms;use codex_protocol::error::CodexErr;use codex_tools::ToolSpec;use futures::FutureExt;use futures::StreamExt;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;#[derive(Default)]pub(crate) struct Handler {    options: WaitAgentTimeoutOptions,}impl Handler {    pub(crate) fn new(options: WaitAgentTimeoutOptions) -> Self {        Self { options }    }}impl ToolExecutor<ToolInvocation> for Handler {    fn tool_name(&self) -> ToolName {        ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, "wait_agent")    }    fn spec(&self) -> ToolSpec {        create_wait_agent_tool_v1(self.options)    }    fn search_info(&self) -> Option<ToolSearchInfo> {        multi_agent_tool_search_info(            "wait_agent wait agent subagent status final result complete timeout targets",            self.spec(),        )    }    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {        Box::pin(self.handle_call(invocation))    }}impl Handler {    async fn handle_call(        &self,        invocation: ToolInvocation,    ) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {        let ToolInvocation {            session,            turn,            payload,            call_id,            ..        } = invocation;        let arguments = function_arguments(payload)?;        let args: WaitArgs = parse_arguments(&arguments)?;        let receiver_thread_ids = parse_agent_id_targets(args.targets)?;        let mut receiver_agents = Vec::with_capacity(receiver_thread_ids.len());        let mut target_by_thread_id = HashMap::with_capacity(receiver_thread_ids.len());        for receiver_thread_id in &receiver_thread_ids {            let agent_metadata = session                .services                .agent_control                .get_agent_metadata(*receiver_thread_id)                .unwrap_or_default();            target_by_thread_id.insert(                *receiver_thread_id,                agent_metadata                    .agent_path                    .as_ref()                    .map(ToString::to_string)                    .unwrap_or_else(|| receiver_thread_id.to_string()),            );            receiver_agents.push(CollabAgentRef {                thread_id: *receiver_thread_id,                agent_nickname: agent_metadata.agent_nickname,                agent_role: agent_metadata.agent_role,            });        }        let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS);        let timeout_ms = match timeout_ms {            ms if ms <= 0 => {                return Err(FunctionCallError::RespondToModel(                    "timeout_ms must be greater than zero".to_owned(),                ));            }            ms => ms.clamp(MIN_WAIT_TIMEOUT_MS, MAX_WAIT_TIMEOUT_MS),        };        session            .send_event(                &turn,                CollabWaitingBeginEvent {                    started_at_ms: now_unix_timestamp_ms(),                    sender_thread_id: session.thread_id,                    receiver_thread_ids: receiver_thread_ids.clone(),                    receiver_agents: receiver_agents.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));                    }                    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.thread_id,                                call_id: call_id.clone(),                                completed_at_ms: now_unix_timestamp_ms(),                                agent_statuses: build_wait_agent_statuses(                                    &statuses,                                    &receiver_agents,                                ),                                statuses,                            }                            .into(),                        )                        .await;                    return Err(collab_agent_error(*id, err));                }            }        }        let statuses = if !initial_final_statuses.is_empty() {            initial_final_statuses        } else {            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;                    }                    Ok(Some(None)) => continue,                    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(None) | None => break,                    }                }            }            results        };        let timed_out = statuses.is_empty();        let statuses_by_id = statuses.clone().into_iter().collect::<HashMap<_, _>>();        let agent_statuses = build_wait_agent_statuses(&statuses_by_id, &receiver_agents);        let result = WaitAgentResult {            status: statuses                .into_iter()                .filter_map(|(thread_id, status)| {                    target_by_thread_id                        .get(&thread_id)                        .cloned()                        .map(|target| (target, status))                })                .collect(),            timed_out,        };        session            .send_event(                &turn,                CollabWaitingEndEvent {                    sender_thread_id: session.thread_id,                    call_id,                    completed_at_ms: now_unix_timestamp_ms(),                    agent_statuses,                    statuses: statuses_by_id,                }                .into(),            )            .await;        Ok(boxed_tool_output(result))    }}impl CoreToolRuntime for Handler {    fn matches_kind(&self, payload: &ToolPayload) -> bool {        matches!(payload, ToolPayload::Function { .. })    }}#[derive(Debug, Deserialize)]struct WaitArgs {    #[serde(default)]    targets: Vec<String>,    timeout_ms: Option<i64>,}#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)]pub(crate) struct WaitAgentResult {    pub(crate) status: HashMap<String, AgentStatus>,    pub(crate) timed_out: bool,}impl ToolOutput for WaitAgentResult {    fn log_preview(&self) -> String {        tool_output_json_text(self, "wait_agent")    }    fn success_for_logging(&self) -> bool {        true    }    fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {        tool_output_response_item(call_id, payload, self, /*success*/ None, "wait_agent")    }    fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {        tool_output_code_mode_result(self, "wait_agent")    }}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();        if is_final(&status) {            return Some((thread_id, status));        }    }}