code_mesh_core/tool/

task.rs

//! Task tool for agent spawning and sub-task management

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};
use uuid::Uuid;

use super::{Tool, ToolContext, ToolResult, ToolError};
use crate::agent::{TaskResult, TaskStatus};

/// Task tool for agent spawning and management
#[derive(Clone)]
pub struct TaskTool {
    agent_registry: Arc<RwLock<AgentRegistry>>,
    task_queue: Arc<Mutex<TaskQueue>>,
    completed_tasks: Arc<RwLock<HashMap<String, TaskResult>>>,
}

/// Parameters for task execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskParams {
    /// Task description
    pub description: String,
    /// Optional detailed prompt for the task
    pub prompt: Option<String>,
    /// Required agent capabilities
    pub capabilities: Option<Vec<String>>,
    /// Task priority (low, medium, high, critical)
    pub priority: Option<String>,
    /// Task dependencies (task IDs that must complete first)
    pub dependencies: Option<Vec<String>>,
    /// Maximum number of agents to spawn for this task
    pub max_agents: Option<u32>,
    /// Task timeout in seconds
    pub timeout: Option<u64>,
    /// Whether to execute subtasks in parallel
    pub parallel: Option<bool>,
}

/// Agent registry for managing agent types and spawning
#[derive(Debug)]
pub struct AgentRegistry {
    /// Available agent types and their capabilities
    agent_types: HashMap<String, Vec<String>>,
    /// Maximum number of concurrent agents
    max_agents: u32,
    /// Current agent count
    current_agents: u32,
}

/// Task queue with priority scheduling
#[derive(Debug)]
pub struct TaskQueue {
    /// Pending tasks organized by priority
    pending: VecDeque<QueuedTask>,
    /// Task dependency graph
    dependencies: HashMap<String, Vec<String>>,
}

/// Queued task with metadata
#[derive(Debug, Clone)]
pub struct QueuedTask {
    /// Unique task ID
    pub id: String,
    /// Task description
    pub description: String,
    /// Detailed prompt
    pub prompt: Option<String>,
    /// Required capabilities
    pub capabilities: Vec<String>,
    /// Task priority
    pub priority: TaskPriority,
    /// Task dependencies
    pub dependencies: Vec<String>,
    /// Maximum agents to spawn
    pub max_agents: u32,
    /// Task timeout
    pub timeout: std::time::Duration,
    /// Execute in parallel
    pub parallel: bool,
    /// Task context
    pub context: Value,
}

/// Task priority levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum TaskPriority {
    Low = 0,
    Medium = 1,
    High = 2,
    Critical = 3,
}

impl TaskTool {
    /// Create a new task tool
    pub fn new() -> Self {
        let mut agent_registry = AgentRegistry {
            agent_types: HashMap::new(),
            max_agents: 10,
            current_agents: 0,
        };

        // Register default agent types with capabilities
        agent_registry.agent_types.insert(
            "researcher".to_string(),
            vec!["research".to_string(), "analysis".to_string(), "data_gathering".to_string()]
        );
        agent_registry.agent_types.insert(
            "coder".to_string(),
            vec!["programming".to_string(), "implementation".to_string(), "debugging".to_string()]
        );
        agent_registry.agent_types.insert(
            "analyst".to_string(),
            vec!["analysis".to_string(), "evaluation".to_string(), "metrics".to_string()]
        );
        agent_registry.agent_types.insert(
            "optimizer".to_string(),
            vec!["optimization".to_string(), "performance".to_string(), "efficiency".to_string()]
        );
        agent_registry.agent_types.insert(
            "coordinator".to_string(),
            vec!["coordination".to_string(), "orchestration".to_string(), "management".to_string()]
        );

        let task_queue = TaskQueue {
            pending: VecDeque::new(),
            dependencies: HashMap::new(),
        };

        Self {
            agent_registry: Arc::new(RwLock::new(agent_registry)),
            task_queue: Arc::new(Mutex::new(task_queue)),
            completed_tasks: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Queue a task for execution
    pub async fn queue_task(&self, params: TaskParams, context: Value) -> std::result::Result<String, ToolError> {
        let task_id = Uuid::new_v4().to_string();
        let priority = self.parse_priority(params.priority.as_deref().unwrap_or("medium"))?;
        
        let queued_task = QueuedTask {
            id: task_id.clone(),
            description: params.description,
            prompt: params.prompt,
            capabilities: params.capabilities.unwrap_or_default(),
            priority,
            dependencies: params.dependencies.unwrap_or_default(),
            max_agents: params.max_agents.unwrap_or(1),
            timeout: std::time::Duration::from_secs(params.timeout.unwrap_or(300)),
            parallel: params.parallel.unwrap_or(false),
            context,
        };

        let mut queue = self.task_queue.lock().await;
        
        // Add to dependency graph
        for dep in &queued_task.dependencies {
            queue.dependencies.entry(dep.clone())
                .or_insert_with(Vec::new)
                .push(task_id.clone());
        }

        // Add to queue (simple FIFO for now, can be enhanced with priority)
        queue.pending.push_back(queued_task);

        drop(queue); // Release lock

        // Try to execute the task immediately
        self.try_execute_next_task().await?;

        Ok(task_id)
    }

    /// Try to execute the next available task
    async fn try_execute_next_task(&self) -> std::result::Result<(), ToolError> {
        let next_task = {
            let mut queue = self.task_queue.lock().await;
            self.get_next_executable_task(&mut queue).await
        };

        if let Some(task) = next_task {
            self.execute_task(task).await?;
        }

        Ok(())
    }

    /// Get the next task that can be executed (dependencies met)
    async fn get_next_executable_task(&self, queue: &mut TaskQueue) -> Option<QueuedTask> {
        let mut i = 0;
        while i < queue.pending.len() {
            let task = &queue.pending[i];
            
            // Check if all dependencies are completed
            if self.are_dependencies_completed(&task.dependencies).await {
                return Some(queue.pending.remove(i).unwrap());
            }
            i += 1;
        }
        None
    }

    /// Check if all task dependencies are completed
    async fn are_dependencies_completed(&self, dependencies: &[String]) -> bool {
        let results = self.completed_tasks.read().await;
        dependencies.iter().all(|dep_id| {
            results.get(dep_id)
                .map(|result| matches!(result.status, TaskStatus::Completed))
                .unwrap_or(false)
        })
    }

    /// Execute a task by spawning an appropriate agent
    async fn execute_task(&self, task: QueuedTask) -> std::result::Result<(), ToolError> {
        let agent_type = self.find_best_agent_type(&task.capabilities).await?;
        let agent_id = self.spawn_virtual_agent(&agent_type, &task.capabilities).await?;
        
        // Execute the task (simplified mock execution)
        let result = self.execute_task_with_virtual_agent(task.clone(), &agent_id).await?;
        
        // Store result
        self.completed_tasks.write().await.insert(task.id.clone(), result);

        // Note: Removed recursive call to avoid boxing requirement
        // Future enhancement: implement proper task scheduler

        Ok(())
    }

    /// Find the best agent type for required capabilities
    async fn find_best_agent_type(&self, required_capabilities: &[String]) -> std::result::Result<String, ToolError> {
        let registry = self.agent_registry.read().await;
        
        let mut best_match = None;
        let mut best_score = 0;

        for (agent_type, capabilities) in &registry.agent_types {
            let score = required_capabilities.iter()
                .filter(|req_cap| capabilities.contains(req_cap))
                .count();
            
            if score > best_score {
                best_score = score;
                best_match = Some(agent_type.clone());
            }
        }

        best_match.ok_or_else(|| {
            ToolError::ExecutionFailed("No suitable agent type found for required capabilities".to_string())
        })
    }

    /// Spawn a virtual agent (simplified implementation)
    async fn spawn_virtual_agent(&self, agent_type: &str, _capabilities: &[String]) -> std::result::Result<String, ToolError> {
        let mut registry = self.agent_registry.write().await;
        
        if registry.current_agents >= registry.max_agents {
            return Err(ToolError::ExecutionFailed("Agent pool at maximum capacity".to_string()));
        }

        let agent_id = format!("{}_{}", agent_type, Uuid::new_v4());
        registry.current_agents += 1;
        
        Ok(agent_id)
    }

    /// Execute a task with a virtual agent (mock implementation)
    async fn execute_task_with_virtual_agent(
        &self,
        task: QueuedTask,
        agent_id: &str,
    ) -> std::result::Result<TaskResult, ToolError> {
        // Simulate task execution
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        
        let output = match agent_id.split('_').next().unwrap_or("unknown") {
            "researcher" => json!({
                "agent_type": "researcher",
                "result": format!("Research completed for: {}", task.description),
                "findings": ["Data analysis completed", "Research methodology validated"]
            }),
            "coder" => json!({
                "agent_type": "coder", 
                "result": format!("Implementation completed for: {}", task.description),
                "code_changes": ["Functions implemented", "Tests added", "Documentation updated"]
            }),
            "analyst" => json!({
                "agent_type": "analyst",
                "result": format!("Analysis completed for: {}", task.description),
                "metrics": {"performance": "good", "efficiency": "high", "quality": "excellent"}
            }),
            "optimizer" => json!({
                "agent_type": "optimizer",
                "result": format!("Optimization completed for: {}", task.description),
                "improvements": ["Performance increased by 25%", "Memory usage reduced", "Code complexity decreased"]
            }),
            "coordinator" => json!({
                "agent_type": "coordinator",
                "result": format!("Coordination completed for: {}", task.description),
                "coordination": ["Tasks synchronized", "Resources allocated", "Timeline optimized"]
            }),
            _ => json!({
                "agent_type": "generic",
                "result": format!("Task completed: {}", task.description)
            }),
        };

        Ok(TaskResult {
            task_id: task.id,
            status: TaskStatus::Completed,
            output,
            error: None,
        })
    }

    /// Parse priority string to enum
    fn parse_priority(&self, priority: &str) -> std::result::Result<TaskPriority, ToolError> {
        match priority.to_lowercase().as_str() {
            "low" => Ok(TaskPriority::Low),
            "medium" => Ok(TaskPriority::Medium),
            "high" => Ok(TaskPriority::High),
            "critical" => Ok(TaskPriority::Critical),
            _ => Err(ToolError::InvalidParameters(format!("Invalid priority: {}", priority))),
        }
    }

    /// Get task status
    pub async fn get_task_status(&self, task_id: &str) -> Option<TaskStatus> {
        // Check if task is completed
        if let Some(result) = self.completed_tasks.read().await.get(task_id) {
            return Some(result.status);
        }

        // Check if task is pending
        let queue = self.task_queue.lock().await;
        if queue.pending.iter().any(|task| task.id == task_id) {
            return Some(TaskStatus::Pending);
        }

        None
    }

    /// Get task results
    pub async fn get_task_results(&self, task_id: &str) -> Option<TaskResult> {
        self.completed_tasks.read().await.get(task_id).cloned()
    }

    /// Get agent registry status
    pub async fn get_agent_status(&self) -> Value {
        let registry = self.agent_registry.read().await;
        let queue = self.task_queue.lock().await;
        
        json!({
            "current_agents": registry.current_agents,
            "max_agents": registry.max_agents,
            "pending_tasks": queue.pending.len(),
            "agent_types": registry.agent_types.keys().collect::<Vec<_>>(),
            "completed_tasks": self.completed_tasks.read().await.len()
        })
    }

    /// List available agent types
    pub async fn list_agent_types(&self) -> Vec<String> {
        self.agent_registry.read().await.agent_types.keys().cloned().collect()
    }

    /// Get agent capabilities for a type
    pub async fn get_agent_capabilities(&self, agent_type: &str) -> Option<Vec<String>> {
        self.agent_registry.read().await.agent_types.get(agent_type).cloned()
    }
}

#[async_trait]
impl Tool for TaskTool {
    fn id(&self) -> &str {
        "task"
    }

    fn description(&self) -> &str {
        "Spawn agents and orchestrate sub-tasks with priority scheduling and dependency management"
    }

    fn parameters_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "description": {
                    "type": "string",
                    "description": "Task description"
                },
                "prompt": {
                    "type": "string",
                    "description": "Optional detailed prompt for the task"
                },
                "capabilities": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Required agent capabilities (researcher, coder, analyst, optimizer, coordinator)"
                },
                "priority": {
                    "type": "string",
                    "enum": ["low", "medium", "high", "critical"],
                    "description": "Task priority level"
                },
                "dependencies": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Task IDs that must complete before this task"
                },
                "max_agents": {
                    "type": "integer",
                    "description": "Maximum number of agents to spawn for this task"
                },
                "timeout": {
                    "type": "integer",
                    "description": "Task timeout in seconds"
                },
                "parallel": {
                    "type": "boolean",
                    "description": "Whether to execute subtasks in parallel"
                }
            },
            "required": ["description"]
        })
    }

    async fn execute(&self, args: Value, ctx: ToolContext) -> std::result::Result<ToolResult, ToolError> {
        let params: TaskParams = serde_json::from_value(args)
            .map_err(|e| ToolError::InvalidParameters(e.to_string()))?;

        let task_id = self.queue_task(params, json!({
            "session_id": ctx.session_id,
            "message_id": ctx.message_id,
            "working_directory": ctx.working_directory
        })).await?;

        Ok(ToolResult {
            title: "Task Queued".to_string(),
            metadata: json!({
                "task_id": task_id,
                "agent_status": self.get_agent_status().await
            }),
            output: format!("Task {} queued for execution with agent spawning", task_id),
        })
    }
}

impl Default for TaskTool {
    fn default() -> Self {
        Self::new()
    }
}