fabro/crates/agent/src/subagent.rs

use crate::error::AgentError;
use crate::session::Session;
use crate::tool_registry::RegisteredTool;
use crate::tools::required_str;
use crate::types::Turn;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex};
use llm::types::ToolDefinition;
use tokio_util::sync::CancellationToken;

pub type SessionFactory = Arc<dyn Fn() -> Session + Send + Sync>;

#[derive(Debug, Clone)]
pub struct SubAgentResult {
    pub output: String,
    pub success: bool,
    pub turns_used: usize,
}

pub struct SubAgent {
    #[allow(dead_code)]
    id: String,
    #[allow(dead_code)]
    depth: usize,
    task: Option<tokio::task::JoinHandle<Result<SubAgentResult, AgentError>>>,
    followup_queue: Arc<Mutex<VecDeque<String>>>,
    cancel_token: CancellationToken,
}

#[cfg(test)]
impl SubAgent {
    #[must_use] 
    pub fn depth(&self) -> usize {
        self.depth
    }
}

pub struct SubAgentManager {
    agents: HashMap<String, SubAgent>,
    max_depth: usize,
}

impl SubAgentManager {
    #[must_use] 
    pub fn new(max_depth: usize) -> Self {
        Self {
            agents: HashMap::new(),
            max_depth,
        }
    }

    pub fn spawn(
        &mut self,
        mut session: Session,
        task_prompt: String,
        depth: usize,
    ) -> Result<String, AgentError> {
        if depth >= self.max_depth {
            return Err(AgentError::InvalidState(format!(
                "Maximum subagent depth ({}) reached",
                self.max_depth
            )));
        }

        let agent_id = uuid::Uuid::new_v4().to_string();
        let followup_queue = session.followup_queue_handle();
        let cancel_token = session.cancel_token();

        let task = tokio::spawn(async move {
            session.process_input(&task_prompt).await?;
            let turns = session.history().turns();
            let last_text = turns.iter().rev().find_map(|t| match t {
                Turn::Assistant { content, .. } => Some(content.clone()),
                _ => None,
            });
            Ok(SubAgentResult {
                output: last_text.unwrap_or_default(),
                success: true,
                turns_used: turns.len(),
            })
        });

        self.agents.insert(
            agent_id.clone(),
            SubAgent {
                id: agent_id.clone(),
                depth,
                task: Some(task),
                followup_queue,
                cancel_token,
            },
        );

        Ok(agent_id)
    }

    pub fn send_input(&self, agent_id: &str, message: &str) -> Result<(), AgentError> {
        let agent = self
            .agents
            .get(agent_id)
            .ok_or_else(|| {
                AgentError::InvalidState(format!("No agent found with id: {agent_id}"))
            })?;

        agent
            .followup_queue
            .lock()
            .expect("followup queue lock poisoned")
            .push_back(message.to_string());

        Ok(())
    }

    pub async fn wait(&mut self, agent_id: &str) -> Result<SubAgentResult, AgentError> {
        let mut agent = self
            .agents
            .remove(agent_id)
            .ok_or_else(|| {
                AgentError::InvalidState(format!("No agent found with id: {agent_id}"))
            })?;

        match agent.task.take() {
            Some(join_handle) => match join_handle.await {
                Ok(result) => result,
                Err(e) => Err(AgentError::InvalidState(format!(
                    "Agent task panicked: {e}"
                ))),
            },
            None => Err(AgentError::InvalidState(format!(
                "Agent {agent_id} has no running task"
            ))),
        }
    }

    pub fn close(&mut self, agent_id: &str) -> Result<(), AgentError> {
        let agent = self
            .agents
            .remove(agent_id)
            .ok_or_else(|| {
                AgentError::InvalidState(format!("No agent found with id: {agent_id}"))
            })?;

        agent.cancel_token.cancel();

        if let Some(join_handle) = agent.task {
            join_handle.abort();
        }

        Ok(())
    }

    #[cfg(test)]
    #[must_use] 
    pub fn get(&self, agent_id: &str) -> Option<&SubAgent> {
        self.agents.get(agent_id)
    }
}

pub fn make_spawn_agent_tool(
    manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
    session_factory: SessionFactory,
    current_depth: usize,
) -> RegisteredTool {
    RegisteredTool {
        definition: ToolDefinition {
            name: "spawn_agent".into(),
            description: "Spawn a subagent to work on a delegated task".into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "task": {
                        "type": "string",
                        "description": "The task description for the subagent"
                    },
                    "working_dir": {
                        "type": "string",
                        "description": "Working directory for the subagent"
                    },
                    "model": {
                        "type": "string",
                        "description": "Model to use for the subagent"
                    },
                    "max_turns": {
                        "type": "integer",
                        "description": "Maximum number of turns for the subagent"
                    }
                },
                "required": ["task"]
            }),
        },
        executor: Arc::new(move |args, _env, _cancel| {
            let manager = manager.clone();
            let session_factory = session_factory.clone();
            Box::pin(async move {
                let task = required_str(&args, "task")?;

                // Extract optional max_turns parameter
                let max_turns = args
                    .get("max_turns")
                    .and_then(serde_json::Value::as_u64)
                    .map(|v| v as usize);

                // Note: working_dir and model require session factory changes to wire through
                let mut session = session_factory();
                // Default subagent max_turns is 50 per spec (overridable via parameter)
                session.set_max_turns(max_turns.unwrap_or(50));
                let mut mgr = manager.lock().await;
                mgr.spawn(session, task.to_string(), current_depth)
                    .map_err(|e| e.to_string())
            })
        }),
    }
}

pub fn make_send_input_tool(
    manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
) -> RegisteredTool {
    RegisteredTool {
        definition: ToolDefinition {
            name: "send_input".into(),
            description: "Send a follow-up message to a running subagent".into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "agent_id": {
                        "type": "string",
                        "description": "The ID of the agent to send input to"
                    },
                    "message": {
                        "type": "string",
                        "description": "The message to send to the agent"
                    }
                },
                "required": ["agent_id", "message"]
            }),
        },
        executor: Arc::new(move |args, _env, _cancel| {
            let manager = manager.clone();
            Box::pin(async move {
                let agent_id = required_str(&args, "agent_id")?;
                let message = required_str(&args, "message")?;

                let mgr = manager.lock().await;
                mgr.send_input(agent_id, message)
                    .map_err(|e| e.to_string())?;
                Ok(format!("Message sent to agent {agent_id}"))
            })
        }),
    }
}

pub fn make_wait_tool(
    manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
) -> RegisteredTool {
    RegisteredTool {
        definition: ToolDefinition {
            name: "wait".into(),
            description: "Wait for a subagent to complete and return its result".into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "agent_id": {
                        "type": "string",
                        "description": "The ID of the agent to wait for"
                    }
                },
                "required": ["agent_id"]
            }),
        },
        executor: Arc::new(move |args, _env, _cancel| {
            let manager = manager.clone();
            Box::pin(async move {
                let agent_id = required_str(&args, "agent_id")?;

                let mut mgr = manager.lock().await;
                let result = mgr.wait(agent_id).await
                    .map_err(|e| e.to_string())?;
                Ok(format!(
                    "Agent completed (success: {}, turns: {})\n\n{}",
                    result.success, result.turns_used, result.output
                ))
            })
        }),
    }
}

pub fn make_close_agent_tool(
    manager: Arc<tokio::sync::Mutex<SubAgentManager>>,
) -> RegisteredTool {
    RegisteredTool {
        definition: ToolDefinition {
            name: "close_agent".into(),
            description: "Close a running subagent".into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "agent_id": {
                        "type": "string",
                        "description": "The ID of the agent to close"
                    }
                },
                "required": ["agent_id"]
            }),
        },
        executor: Arc::new(move |args, _env, _cancel| {
            let manager = manager.clone();
            Box::pin(async move {
                let agent_id = required_str(&args, "agent_id")?;

                let mut mgr = manager.lock().await;
                mgr.close(agent_id)
                    .map_err(|e| e.to_string())?;
                Ok(format!("Agent {agent_id} closed"))
            })
        }),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_support::*;

    // --- Tests ---

    #[test]
    fn manager_creation() {
        let manager = SubAgentManager::new(3);
        assert_eq!(manager.max_depth, 3);
        assert!(manager.agents.is_empty());
    }

    #[tokio::test]
    async fn spawn_creates_agent_and_returns_id() {
        let mut manager = SubAgentManager::new(3);
        let session = make_session(vec![text_response("Hello")]).await;
        let result = manager.spawn(session, "Do something".into(), 0);
        assert!(result.is_ok());
        let agent_id = result.unwrap();
        assert!(!agent_id.is_empty());
        assert!(manager.get(&agent_id).is_some());
        assert_eq!(manager.get(&agent_id).unwrap().depth(), 0);
    }

    #[tokio::test]
    async fn depth_limit_enforced() {
        let mut manager = SubAgentManager::new(2);
        let session = make_session(vec![text_response("Hello")]).await;
        let result = manager.spawn(session, "Do something".into(), 2);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Maximum subagent depth"));
    }

    #[tokio::test]
    async fn close_removes_agent() {
        let mut manager = SubAgentManager::new(3);
        let session = make_session(vec![text_response("Hello")]).await;
        let agent_id = manager.spawn(session, "Do something".into(), 0).unwrap();
        assert!(manager.get(&agent_id).is_some());

        let result = manager.close(&agent_id);
        assert!(result.is_ok());
        assert!(manager.get(&agent_id).is_none());
    }

    #[tokio::test]
    async fn send_input_nonexistent_agent_errors() {
        let manager = SubAgentManager::new(3);
        let result = manager.send_input("nonexistent-id", "hello");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("No agent found"));
    }

    #[tokio::test]
    async fn wait_nonexistent_agent_errors() {
        let mut manager = SubAgentManager::new(3);
        let result = manager.wait("nonexistent-id").await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("No agent found"));
    }

    #[tokio::test]
    async fn wait_returns_result() {
        let mut manager = SubAgentManager::new(3);
        let session =
            make_session(vec![text_response("Task completed successfully")]).await;
        let agent_id = manager.spawn(session, "Do something".into(), 0).unwrap();

        let result = manager.wait(&agent_id).await;
        assert!(result.is_ok());
        let agent_result = result.unwrap();
        assert_eq!(agent_result.output, "Task completed successfully");
        assert!(agent_result.success);
        assert!(agent_result.turns_used > 0);
        assert!(manager.get(&agent_id).is_none());
    }

    #[test]
    fn tool_definitions_correct() {
        let manager = Arc::new(tokio::sync::Mutex::new(SubAgentManager::new(3)));
        let factory: SessionFactory = Arc::new(|| {
            panic!("should not be called");
        });

        let spawn_tool = make_spawn_agent_tool(manager.clone(), factory, 0);
        assert_eq!(spawn_tool.definition.name, "spawn_agent");
        assert!(spawn_tool.definition.parameters["properties"]["task"].is_object());
        let spawn_required = spawn_tool.definition.parameters["required"]
            .as_array()
            .unwrap();
        assert!(spawn_required.contains(&serde_json::json!("task")));

        let send_tool = make_send_input_tool(manager.clone());
        assert_eq!(send_tool.definition.name, "send_input");
        assert!(send_tool.definition.parameters["properties"]["agent_id"].is_object());
        assert!(send_tool.definition.parameters["properties"]["message"].is_object());
        let send_required = send_tool.definition.parameters["required"]
            .as_array()
            .unwrap();
        assert!(send_required.contains(&serde_json::json!("agent_id")));
        assert!(send_required.contains(&serde_json::json!("message")));

        let wait_tool = make_wait_tool(manager.clone());
        assert_eq!(wait_tool.definition.name, "wait");
        assert!(wait_tool.definition.parameters["properties"]["agent_id"].is_object());
        let wait_required = wait_tool.definition.parameters["required"]
            .as_array()
            .unwrap();
        assert!(wait_required.contains(&serde_json::json!("agent_id")));

        let close_tool = make_close_agent_tool(manager);
        assert_eq!(close_tool.definition.name, "close_agent");
        assert!(close_tool.definition.parameters["properties"]["agent_id"].is_object());
        let close_required = close_tool.definition.parameters["required"]
            .as_array()
            .unwrap();
        assert!(close_required.contains(&serde_json::json!("agent_id")));
    }
}