agent-works 0.1.6

//! Multi-agent runtime — coordinates sub-agent lifecycle, event bridging, and
//! cancellation.
//!
//! The [`MultiAgentRuntime`] is the central coordinator. It is created once during
//! builder setup and shared via `Arc` to all 6 multi-agent tools.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use agent_base::{
    AgentBuilder, AgentResult, AgentRuntime, DenyAllApprovalHandler, Language, LlmClient,
    RunOutcome, RuntimeEvent, SessionId, Tool, UserEvent,
};
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;

use super::config::MultiAgentConfig;
use super::mailbox::{ChildMailbox, MailboxHub, MailboxResult, MailboxStatus, MailboxTask};
use super::path::AgentPath;
use super::registry::{AgentRegistry, AgentStatus};

// ---------------------------------------------------------------------------
// MultiAgentRuntime
// ---------------------------------------------------------------------------

/// Coordinates sub-agent lifecycle, event bridging, and cancellation.
///
/// Created once during builder setup and shared via `Arc` to all 6 multi-agent
/// tools. Each tool calls methods on the runtime to spawn, communicate with, or
/// close sub-agents.
pub struct MultiAgentRuntime {
    /// Agent lifecycle registry (spawn/close/query).
    registry: Mutex<AgentRegistry>,

    /// Inter-agent message hub.
    mailbox: Arc<MailboxHub>,

    /// Shared LLM client (from parent agent).
    client: Arc<dyn LlmClient>,

    /// Business tools to register on child agents (NOT the 6 multi-agent tools).
    business_tools: Vec<Arc<dyn Tool>>,

    /// Channel to the bridge task that emits events on parent's event bus.
    event_tx: Mutex<Option<tokio::sync::mpsc::UnboundedSender<RuntimeEvent>>>,

    /// Root cancellation token (propagates to all children).
    root_cancel: CancellationToken,

    /// JoinSet tracking all child agent tasks.
    join_set: Mutex<JoinSet<()>>,

    /// Per-child cancellation tokens.
    child_cancels: Mutex<HashMap<AgentPath, CancellationToken>>,

    /// Error recovery strategy (inherited from parent).
    error_recovery: Option<Arc<dyn agent_base::ToolErrorRecovery>>,

    /// Language preference.
    language: Language,
}

impl MultiAgentRuntime {
    /// Create a new multi-agent runtime.
    ///
    /// This is called internally by the builder. Tools receive an `Arc<Self>`.
    pub fn new(
        config: MultiAgentConfig,
        client: Arc<dyn LlmClient>,
        business_tools: Vec<Arc<dyn Tool>>,
        root_cancel: CancellationToken,
        error_recovery: Option<Arc<dyn agent_base::ToolErrorRecovery>>,
        language: Language,
    ) -> Self {
        Self {
            registry: Mutex::new(AgentRegistry::new(config)),
            mailbox: Arc::new(MailboxHub::new()),
            client,
            business_tools,
            event_tx: Mutex::new(None),
            root_cancel,
            join_set: Mutex::new(JoinSet::new()),
            child_cancels: Mutex::new(HashMap::new()),
            error_recovery,
            language,
        }
    }

    /// Set the event sender for bridging child events to parent.
    ///
    /// Called by the builder after creating the bridge channel.
    pub fn set_event_sender(&self, tx: tokio::sync::mpsc::UnboundedSender<RuntimeEvent>) {
        *self.event_tx.lock().unwrap() = Some(tx);
    }

    /// Spawn a child agent at the given path with a specific system prompt.
    ///
    /// This is called by the `spawn_agent` tool. It:
    /// 1. Checks spawn limits
    /// 2. Registers the agent in the registry
    /// 3. Creates a mailbox
    /// 4. Builds a child AgentRuntime
    /// 5. Spawns a tokio task for the child's event loop
    /// 6. Returns the AgentPath
    ///
    /// # Errors
    ///
    /// Returns a string error message if spawning fails (limits exceeded, etc.).
    pub async fn spawn_child(
        &self,
        name: &str,
        system_prompt: String,
        depth: i32,
        tool_count: usize,
    ) -> Result<String, String> {
        let path = AgentPath::root().join(name);

        // 1. Check limits and register
        {
            let mut registry = self.registry.lock().unwrap();
            registry.can_spawn(depth).map_err(|e| e.to_string())?;
            registry
                .register(&path, depth, tool_count)
                .map_err(|e| e.to_string())?;
        }

        // 2. Create mailbox
        let child_mailbox = self
            .mailbox
            .register(&path)
            .ok_or_else(|| "mailbox already exists".to_string())?;

        // 3. Build child AgentRuntime (roll back registry+mailbox on failure)
        let child_runtime = self.build_child_runtime(system_prompt).map_err(|e| {
            self.registry.lock().unwrap().close(&path);
            self.mailbox.unregister(&path);
            format!("failed to build child runtime: {}", e)
        })?;

        // 4. Create session for child
        let session_id = child_runtime.create_session().await;

        // 5. Create child cancellation token
        let child_cancel = self.root_cancel.child_token();
        {
            let mut cancels = self.child_cancels.lock().unwrap();
            cancels.insert(path.clone(), child_cancel.clone());
        }

        // 6. Spawn child agent event loop
        let agent_path = path.clone();
        let mailbox_for_task = self.mailbox.clone();
        let mailbox_for_close = self.mailbox.clone();
        let event_tx = self.event_tx.lock().unwrap().clone();
        let registry_agent_path = path.clone();

        self.join_set.lock().unwrap().spawn(async move {
            run_child_loop(
                child_mailbox,
                child_runtime,
                session_id,
                agent_path.clone(),
                mailbox_for_task,
                event_tx,
                child_cancel,
            )
            .await;

            // Post close notification when loop exits
            mailbox_for_close.post_result(MailboxResult {
                agent_path,
                status: MailboxStatus::Closed,
                result: None,
            });
        });

        self.registry
            .lock()
            .unwrap()
            .set_status(&registry_agent_path, AgentStatus::Idle);

        Ok(path.to_string())
    }

    /// Send a message to a child agent (no execution trigger).
    ///
    /// Called by `send_message` tool.
    pub fn send_message(&self, agent_path: &str, message: String) -> Result<bool, String> {
        let path = self.parse_path(agent_path)?;
        Ok(self.mailbox.send_message(&path, message))
    }

    /// Send a task to a child agent (triggers execution).
    ///
    /// Called by `followup_task` tool. Updates status to Running.
    pub fn send_task(
        &self,
        agent_path: &str,
        task: String,
        interrupt: bool,
    ) -> Result<bool, String> {
        let path = self.parse_path(agent_path)?;
        if !self.mailbox.contains(&path) {
            return Err("agent not found".to_string());
        }
        let sent = self.mailbox.send_task(&path, task, interrupt);
        if sent {
            self.registry
                .lock()
                .unwrap()
                .set_status(&path, AgentStatus::Running);
        }
        Ok(sent)
    }

    /// Wait for a result from any or a specific child agent.
    ///
    /// Called by `wait_agent` tool. Blocks until a result arrives or timeout.
    pub async fn wait_for_result(&self, agent_path: Option<&str>, timeout_ms: u64) -> WaitResult {
        let filter_path = match agent_path {
            Some(s) => match AgentPath::parse(s) {
                Some(p) => Some(p),
                None => {
                    return WaitResult {
                        status: "error".to_string(),
                        result: Some(format!("invalid agent path: {}", s)),
                        agent_path: None,
                        has_more: false,
                    };
                }
            },
            None => None,
        };

        let mut seq = self.mailbox.subscribe_seq();
        let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(timeout_ms);

        loop {
            // Check for existing results
            let result = match &filter_path {
                Some(path) => self.mailbox.try_recv_result(path),
                None => self.mailbox.try_recv_any(),
            };

            if let Some(r) = result {
                let has_more = self.mailbox.total_pending_results() > 0;
                let (status_str, result_text) = match r.status {
                    MailboxStatus::Ok => ("ok".to_string(), r.result),
                    MailboxStatus::Error => ("error".to_string(), r.result),
                    MailboxStatus::Closed => ("closed".to_string(), r.result),
                };
                return WaitResult {
                    status: status_str,
                    result: result_text,
                    agent_path: Some(r.agent_path.to_string()),
                    has_more,
                };
            }

            // Wait for sequence number change or timeout
            let now = tokio::time::Instant::now();
            if now >= deadline {
                return WaitResult {
                    status: "timeout".to_string(),
                    result: None,
                    agent_path: None,
                    has_more: false,
                };
            }

            let remaining = deadline - now;
            tokio::select! {
                _ = seq.changed() => {
                    // Sequence changed — loop back to check results
                    continue;
                }
                _ = tokio::time::sleep(remaining) => {
                    return WaitResult {
                        status: "timeout".to_string(),
                        result: None,
                        agent_path: None,
                        has_more: false,
                    };
                }
            }
        }
    }

    /// Close a child agent.
    ///
    /// Called by `close_agent` tool. Cancels the child's task, removes from
    /// registry, and posts a Closed result.
    pub fn close_agent(&self, agent_path: &str) -> Result<CloseResult, String> {
        let path = self.parse_path(agent_path)?;

        // Get previous status
        let previous_status = {
            let registry = self.registry.lock().unwrap();
            registry
                .get(&path)
                .map(|e| format!("{:?}", e.status).to_lowercase())
                .unwrap_or_else(|| "unknown".to_string())
        };

        // Cancel child token
        {
            let mut cancels = self.child_cancels.lock().unwrap();
            if let Some(token) = cancels.remove(&path) {
                token.cancel();
            }
        }

        // Close in registry
        let existed = { self.registry.lock().unwrap().close(&path).is_some() };

        // Unregister mailbox
        self.mailbox.unregister(&path);

        Ok(CloseResult {
            closed: existed,
            previous_status,
            message: if existed {
                "agent closed".to_string()
            } else {
                "agent not found".to_string()
            },
        })
    }

    /// List all active sub-agents.
    ///
    /// Called by `list_agents` tool.
    pub fn list_agents(&self) -> Vec<AgentInfo> {
        let registry = self.registry.lock().unwrap();
        registry
            .list()
            .into_iter()
            .map(|e| AgentInfo {
                agent_path: e.path.to_string(),
                status: format!("{:?}", e.status).to_lowercase(),
                tool_count: e.tool_count,
            })
            .collect()
    }

    /// Get the mailbox hub (for tools that need it directly).
    pub fn mailbox(&self) -> &Arc<MailboxHub> {
        &self.mailbox
    }

    /// Get reference to the registry.
    pub fn registry(&self) -> &Mutex<AgentRegistry> {
        &self.registry
    }

    /// Cancel all child agents.
    pub fn cancel_all(&self) {
        let mut cancels = self.child_cancels.lock().unwrap();
        for (_, token) in cancels.drain() {
            token.cancel();
        }
    }
}

impl Drop for MultiAgentRuntime {
    fn drop(&mut self) {
        self.cancel_all();
        // Drain any already-completed join handles to detect panics
        let mut js = self.join_set.lock().unwrap();
        while let Some(result) = js.try_join_next() {
            if let Err(e) = result
                && e.is_panic()
            {
                tracing::error!(
                    error = %e,
                    "child agent task panicked"
                );
            }
        }
    }
}

impl MultiAgentRuntime {
    fn parse_path(&self, s: &str) -> Result<AgentPath, String> {
        AgentPath::parse(s).ok_or_else(|| format!("invalid agent path: '{}'", s))
    }

    fn build_child_runtime(&self, system_prompt: String) -> AgentResult<AgentRuntime> {
        let mut builder = AgentBuilder::new(self.client.clone())
            .system_prompt(system_prompt)
            .approval_handler(Arc::new(DenyAllApprovalHandler))
            .language(self.language.clone());

        // Register business tools (NOT multi-agent tools)
        for tool in &self.business_tools {
            builder = builder.register_tool_arc(tool.clone());
        }

        if let Some(ref recovery) = self.error_recovery {
            builder = builder.error_recovery(recovery.clone());
        }

        builder.build()
    }
}

// ---------------------------------------------------------------------------
// Result types
// ---------------------------------------------------------------------------

/// Result from `wait_for_result()`.
#[derive(Clone, Debug)]
pub struct WaitResult {
    pub status: String,
    pub result: Option<String>,
    pub agent_path: Option<String>,
    pub has_more: bool,
}

/// Result from `close_agent()`.
#[derive(Clone, Debug)]
pub struct CloseResult {
    pub closed: bool,
    pub previous_status: String,
    pub message: String,
}

/// Agent info for `list_agents()`.
#[derive(Clone, Debug, serde::Serialize)]
pub struct AgentInfo {
    pub agent_path: String,
    pub status: String,
    pub tool_count: usize,
}

// ---------------------------------------------------------------------------
// Child agent event loop
// ---------------------------------------------------------------------------

/// Run the child agent's main event loop.
///
/// This function runs inside a tokio task spawned by [`MultiAgentRuntime::spawn_child`].
/// It:
/// 1. Subscribes to child agent events and bridges them to parent
/// 2. Listens for tasks from the mailbox
/// 3. Executes each task via `run_turn`
/// 4. Posts results back via the mailbox
async fn run_child_loop(
    child_mailbox: ChildMailbox,
    child_runtime: AgentRuntime,
    session_id: SessionId,
    agent_path: AgentPath,
    mailbox: Arc<MailboxHub>,
    event_tx: Option<tokio::sync::mpsc::UnboundedSender<RuntimeEvent>>,
    child_cancel: CancellationToken,
) {
    let mut task_rx = child_mailbox.task_rx;

    // Spawn event bridging: forward child events to parent as SubAgentEvent
    if let Some(tx) = event_tx {
        let mut child_events = child_runtime.subscribe_runtime_events();
        let bridge_path = agent_path.to_string();
        let bridge_cancel = child_cancel.clone();

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = bridge_cancel.cancelled() => break,
                    event = child_events.recv() => {
                        match event {
                            Ok(event) => {
                                if matches!(event, RuntimeEvent::RunFinished { .. } | RuntimeEvent::RunCancelled { .. }) {
                                    continue;
                                }
                                let _ = tx.send(RuntimeEvent::UserEvent {
                                    session_id: SessionId::new(0),
                                    event: UserEvent::SubAgentEvent {
                                        subagent: bridge_path.clone(),
                                        event: Box::new(event),
                                    },
                                    agent_id: None,
                                    trace_id: None,
                                });
                            }
                            Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                                tracing::warn!(
                                    subagent = %bridge_path,
                                    lagged = n,
                                    "child event bridge lagged"
                                );
                            }
                            Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                        }
                    }
                }
            }
        });
    }

    // Main task loop
    loop {
        tokio::select! {
            _ = child_cancel.cancelled() => {
                break;
            }
            task = task_rx.recv() => {
                match task {
                    Some(task) => {
                        let input = build_child_input(&task);
                        let result = child_runtime.run_turn_collect(
                            session_id.clone(),
                            &input,
                        ).await;

                        match result {
                            Ok((_events, outcome)) => {
                                let summary = summarize_outcome(&outcome);
                                mailbox.post_result(MailboxResult {
                                    agent_path: agent_path.clone(),
                                    status: MailboxStatus::Ok,
                                    result: Some(summary),
                                });
                            }
                            Err(e) => {
                                mailbox.post_result(MailboxResult {
                                    agent_path: agent_path.clone(),
                                    status: MailboxStatus::Error,
                                    result: Some(e.to_string()),
                                });
                            }
                        }
                    }
                    None => break, // task channel closed
                }
            }
        }
    }
}

/// Build the input text for a child agent from a mailbox task.
fn build_child_input(task: &MailboxTask) -> String {
    if task.pending_messages.is_empty() {
        task.task.clone()
    } else {
        let mut parts: Vec<String> = Vec::new();
        for msg in &task.pending_messages {
            parts.push(format!("[Message]: {}", msg));
        }
        parts.push(format!("[Task]: {}", task.task));
        parts.join("\n\n")
    }
}

/// Extract a human-readable summary from a run outcome.
fn summarize_outcome(outcome: &RunOutcome) -> String {
    match outcome {
        RunOutcome::Completed => "task completed".to_string(),
        RunOutcome::Failed { error } => format!("task failed: {}", error),
        RunOutcome::MaxTurnsExceeded { turns } => {
            format!("max turns exceeded ({} turns)", turns)
        }
        RunOutcome::Cancelled => "cancelled".to_string(),
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use agent_base::RunOutcome;

    // ── summarize_outcome ──

    #[test]
    fn test_summarize_completed() {
        let s = summarize_outcome(&RunOutcome::Completed);
        assert_eq!(s, "task completed");
    }

    #[test]
    fn test_summarize_failed() {
        let outcome = RunOutcome::Failed {
            error: "connection refused".to_string(),
        };
        let s = summarize_outcome(&outcome);
        assert_eq!(s, "task failed: connection refused");
    }

    #[test]
    fn test_summarize_max_turns() {
        let outcome = RunOutcome::MaxTurnsExceeded { turns: 42 };
        let s = summarize_outcome(&outcome);
        assert!(s.contains("max turns exceeded"));
        assert!(s.contains("42"));
    }

    #[test]
    fn test_summarize_cancelled() {
        let s = summarize_outcome(&RunOutcome::Cancelled);
        assert_eq!(s, "cancelled");
    }

    // ── build_child_input ──

    #[test]
    fn test_build_child_input_task_only() {
        let task = MailboxTask {
            task: "do work".into(),
            interrupt: true,
            pending_messages: vec![],
        };
        let out = build_child_input(&task);
        assert_eq!(out, "do work");
    }

    #[test]
    fn test_build_child_input_with_pending_messages() {
        let task = MailboxTask {
            task: "do work".into(),
            interrupt: false,
            pending_messages: vec!["context 1".into(), "context 2".into()],
        };
        let out = build_child_input(&task);
        assert!(out.contains("[Message]: context 1"));
        assert!(out.contains("[Message]: context 2"));
        assert!(out.contains("[Task]: do work"));
        // Messages come before task
        let msg_pos = out.find("[Message]:").unwrap();
        let task_pos = out.find("[Task]:").unwrap();
        assert!(msg_pos < task_pos, "messages should precede task");
    }

    #[test]
    fn test_build_child_input_single_message() {
        let task = MailboxTask {
            task: "final task".into(),
            interrupt: true,
            pending_messages: vec!["hint".into()],
        };
        let out = build_child_input(&task);
        assert_eq!(out, "[Message]: hint\n\n[Task]: final task");
    }
}