use super::*;use crate::session::InputQueueActivity;use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions;use crate::tools::handlers::multi_agents_spec::create_wait_agent_tool_v2;use crate::turn_timing::now_unix_timestamp_ms;use codex_tools::ToolSpec;use std::collections::HashMap;use std::time::Duration;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::plain("wait_agent")    }    fn spec(&self) -> ToolSpec {        create_wait_agent_tool_v2(self.options)    }    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 min_timeout_ms = turn.config.multi_agent_v2.min_wait_timeout_ms;        let max_timeout_ms = turn.config.multi_agent_v2.max_wait_timeout_ms;        let default_timeout_ms = turn.config.multi_agent_v2.default_wait_timeout_ms;        let timeout_ms = match args.timeout_ms {            Some(ms) if ms < min_timeout_ms => {                return Err(FunctionCallError::RespondToModel(format!(                    "timeout_ms must be at least {min_timeout_ms}"                )));            }            Some(ms) if ms > max_timeout_ms => {                return Err(FunctionCallError::RespondToModel(format!(                    "timeout_ms must be at most {max_timeout_ms}"                )));            }            Some(ms) => ms,            None => default_timeout_ms,        };        let turn_state = session            .input_queue            .turn_state_for_sub_id(&session.active_turn, &turn.sub_id)            .await;        let (mut activity_rx, pending_activity) = session            .input_queue            .subscribe_activity(turn_state.as_deref())            .await;        session            .send_event(                &turn,                CollabWaitingBeginEvent {                    started_at_ms: now_unix_timestamp_ms(),                    sender_thread_id: session.thread_id,                    receiver_thread_ids: Vec::new(),                    receiver_agents: Vec::new(),                    call_id: call_id.clone(),                }                .into(),            )            .await;        let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64);        let outcome = wait_for_activity(&mut activity_rx, pending_activity, deadline).await;        let result = WaitAgentResult::from_outcome(outcome);        session            .send_event(                &turn,                CollabWaitingEndEvent {                    sender_thread_id: session.thread_id,                    call_id,                    completed_at_ms: now_unix_timestamp_ms(),                    agent_statuses: Vec::new(),                    statuses: HashMap::new(),                }                .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)]#[serde(deny_unknown_fields)]struct WaitArgs {    timeout_ms: Option<i64>,}#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)]pub(crate) struct WaitAgentResult {    pub(crate) message: String,    pub(crate) timed_out: bool,}impl WaitAgentResult {    fn from_outcome(outcome: WaitOutcome) -> Self {        let message = match outcome {            WaitOutcome::MailboxActivity => "Wait completed.",            WaitOutcome::Steered => "Wait interrupted by new input.",            WaitOutcome::TimedOut => "Wait timed out.",        };        Self {            message: message.to_string(),            timed_out: outcome == WaitOutcome::TimedOut,        }    }}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")    }}#[derive(Clone, Copy, Debug, Eq, PartialEq)]enum WaitOutcome {    MailboxActivity,    Steered,    TimedOut,}async fn wait_for_activity(    activity_rx: &mut tokio::sync::watch::Receiver<InputQueueActivity>,    pending_activity: Option<InputQueueActivity>,    deadline: Instant,) -> WaitOutcome {    if let Some(activity) = pending_activity {        return match activity {            InputQueueActivity::Mailbox => WaitOutcome::MailboxActivity,            InputQueueActivity::Steer => WaitOutcome::Steered,        };    }    match timeout_at(deadline, activity_rx.changed()).await {        Ok(Ok(())) => match *activity_rx.borrow_and_update() {            InputQueueActivity::Mailbox => WaitOutcome::MailboxActivity,            InputQueueActivity::Steer => WaitOutcome::Steered,        },        Ok(Err(_)) | Err(_) => WaitOutcome::TimedOut,    }}