symbi_runtime/scheduler/

task_manager.rs

//! Task manager for executing and monitoring agent tasks

use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use sysinfo::{Pid, System};
use tokio::process::Command;
use tokio::sync::mpsc;
use tokio::time::timeout;

use crate::types::*;

/// Task manager for agent execution
pub struct TaskManager {
    task_timeout: Duration,
    running_tasks: Arc<RwLock<HashMap<AgentId, TaskHandle>>>,
    task_sender: mpsc::Sender<TaskCommand>,
    #[allow(dead_code)]
    system_info: Arc<RwLock<System>>,
}

impl TaskManager {
    /// Create a new task manager
    pub fn new(task_timeout: Duration) -> Self {
        let running_tasks = Arc::new(RwLock::new(HashMap::new()));
        // Bounded channel prevents unbounded memory growth under load.
        // 1024 is sufficient for most workloads; submitters block when full.
        let (task_sender, task_receiver) = mpsc::channel(1024);

        let manager = Self {
            task_timeout,
            running_tasks: running_tasks.clone(),
            task_sender,
            system_info: Arc::new(RwLock::new(System::new_all())),
        };

        // Start the task execution loop
        manager.start_task_loop(task_receiver);

        manager
    }

    /// Start a new task
    pub async fn start_task(&self, task: super::ScheduledTask) -> Result<(), SchedulerError> {
        let handle = TaskHandle::new(task.clone());
        let agent_id = task.agent_id;

        // Store the handle
        self.running_tasks.write().insert(agent_id, handle.clone());

        // Send command to start the task (try_send avoids blocking the caller;
        // returns an error if the bounded channel is full).
        self.task_sender
            .try_send(TaskCommand::Start {
                task: Box::new(task),
                handle,
            })
            .map_err(|_| SchedulerError::SchedulingFailed {
                agent_id,
                reason: "Task channel full — backpressure applied".into(),
            })?;

        Ok(())
    }

    /// Terminate a task
    pub async fn terminate_task(&self, agent_id: AgentId) -> Result<(), SchedulerError> {
        // Remove from running tasks
        let handle = self.running_tasks.write().remove(&agent_id);

        if let Some(handle) = handle {
            // Send termination command (try_send to avoid blocking).
            self.task_sender
                .try_send(TaskCommand::Terminate { agent_id, handle })
                .map_err(|_| SchedulerError::SchedulingFailed {
                    agent_id,
                    reason: "Task channel full — cannot send terminate command".into(),
                })?;
        }

        Ok(())
    }

    /// Check task health
    pub async fn check_task_health(&self, agent_id: AgentId) -> Result<TaskHealth, SchedulerError> {
        let running_tasks = self.running_tasks.read();

        if let Some(handle) = running_tasks.get(&agent_id) {
            let health = handle.get_health();

            // Check if task has exceeded timeout
            if health.uptime > self.task_timeout {
                return Err(SchedulerError::SchedulingFailed {
                    agent_id,
                    reason: "Task timeout exceeded".into(),
                });
            }

            Ok(health)
        } else {
            Err(SchedulerError::AgentNotFound { agent_id })
        }
    }

    /// Get task statistics
    pub async fn get_task_statistics(&self) -> TaskStatistics {
        let running_tasks = self.running_tasks.read();
        let total_tasks = running_tasks.len();

        let mut healthy_tasks = 0;
        let mut total_uptime = Duration::from_secs(0);
        let mut total_memory_usage = 0;

        for handle in running_tasks.values() {
            let health = handle.get_health();
            if health.is_healthy {
                healthy_tasks += 1;
            }
            total_uptime += health.uptime;
            total_memory_usage += health.memory_usage;
        }

        TaskStatistics {
            total_tasks,
            healthy_tasks,
            average_uptime: if total_tasks > 0 {
                total_uptime / total_tasks as u32
            } else {
                Duration::from_secs(0)
            },
            total_memory_usage,
        }
    }

    /// Start the task execution loop
    fn start_task_loop(&self, mut task_receiver: mpsc::Receiver<TaskCommand>) {
        let task_timeout = self.task_timeout;

        tokio::spawn(async move {
            while let Some(command) = task_receiver.recv().await {
                match command {
                    TaskCommand::Start { task, handle } => {
                        Self::execute_task(*task, handle, task_timeout).await;
                    }
                    TaskCommand::Terminate { agent_id, handle } => {
                        Self::terminate_task_execution(agent_id, handle).await;
                    }
                }
            }
        });
    }

    /// Execute a task
    async fn execute_task(task: super::ScheduledTask, handle: TaskHandle, task_timeout: Duration) {
        let agent_id = task.agent_id;

        tracing::info!("Starting execution of agent {}", agent_id);

        // Update handle status
        handle.set_status(TaskStatus::Running);

        // Execute the real agent task with proper monitoring and metrics
        let execution_result = timeout(
            task_timeout,
            Self::execute_agent_task(task.clone(), handle.clone()),
        )
        .await;

        match execution_result {
            Ok(Ok(execution_metrics)) => {
                tracing::info!("Agent {} completed successfully", agent_id);
                handle.set_status(TaskStatus::Completed);
                handle.update_execution_metrics(execution_metrics);
            }
            Ok(Err(e)) => {
                tracing::error!("Agent {} failed: {}", agent_id, e);
                handle.set_status(TaskStatus::Failed);
            }
            Err(_) => {
                tracing::error!("Agent {} timed out", agent_id);
                handle.set_status(TaskStatus::TimedOut);
            }
        }
    }

    /// Terminate task execution
    async fn terminate_task_execution(agent_id: AgentId, handle: TaskHandle) {
        tracing::info!("Terminating agent {}", agent_id);

        // If there's a process associated with this task, terminate it
        if let Some(process_id) = handle.get_process_id() {
            if let Err(e) = Self::terminate_process(process_id).await {
                tracing::warn!("Failed to terminate process {}: {}", process_id, e);
            }
        }

        handle.set_status(TaskStatus::Terminated);
    }

    /// Execute a real agent task with comprehensive monitoring
    async fn execute_agent_task(
        task: super::ScheduledTask,
        handle: TaskHandle,
    ) -> Result<ExecutionMetrics, String> {
        let _start_time = SystemTime::now();
        let agent_id = task.agent_id;

        tracing::debug!(
            "Executing agent {} with config: {:?}",
            agent_id,
            task.config
        );

        // Create execution context
        let execution_context = AgentExecutionContext::new(task.clone(), handle.clone());

        // Execute based on agent execution mode
        match task.config.execution_mode {
            ExecutionMode::Ephemeral => Self::execute_ephemeral_agent(execution_context).await,
            ExecutionMode::Persistent => Self::execute_persistent_agent(execution_context).await,
            ExecutionMode::Scheduled { interval } => {
                Self::execute_scheduled_agent(execution_context, interval).await
            }
            ExecutionMode::CronScheduled { .. } => {
                Self::execute_cron_scheduled_agent(execution_context).await
            }
            ExecutionMode::EventDriven => Self::execute_event_driven_agent(execution_context).await,
            ExecutionMode::External { .. } => {
                Err("External agents are not executed by the runtime".to_string())
            }
        }
    }

    /// Execute an ephemeral agent (runs once and terminates)
    async fn execute_ephemeral_agent(
        mut context: AgentExecutionContext,
    ) -> Result<ExecutionMetrics, String> {
        tracing::debug!("Executing ephemeral agent {}", context.task.agent_id);

        // Launch the agent process
        let process_handle = Self::launch_agent_process(&context.task).await?;
        context
            .handle
            .set_process_id(Some(process_handle.process_id));

        // Monitor the process execution
        let result = Self::monitor_process_execution(process_handle, &mut context).await;

        // Collect final metrics
        let end_time = SystemTime::now();
        let execution_time = end_time
            .duration_since(context.start_time)
            .unwrap_or_default();

        Ok(ExecutionMetrics {
            execution_time,
            memory_peak_mb: context.memory_peak_mb,
            cpu_time_ms: context.cpu_time_ms,
            exit_code: context.exit_code,
            error_count: context.error_count,
            success: result.is_ok(),
        })
    }

    /// Execute a persistent agent (long-running)
    async fn execute_persistent_agent(
        context: AgentExecutionContext,
    ) -> Result<ExecutionMetrics, String> {
        tracing::debug!("Executing persistent agent {}", context.task.agent_id);

        // Launch the agent process
        let process_handle = Self::launch_agent_process(&context.task).await?;
        context
            .handle
            .set_process_id(Some(process_handle.process_id));

        // For persistent agents, we start monitoring but don't wait for completion
        let _monitoring_handle =
            tokio::spawn(
                async move { Self::monitor_persistent_process(process_handle, context).await },
            );

        // Return immediately for persistent agents
        Ok(ExecutionMetrics {
            execution_time: Duration::from_secs(0),
            memory_peak_mb: 0,
            cpu_time_ms: 0,
            exit_code: None,
            error_count: 0,
            success: true,
        })
    }

    /// Execute a scheduled agent (single run — the CronScheduler handles recurrence).
    async fn execute_scheduled_agent(
        context: AgentExecutionContext,
        interval: Duration,
    ) -> Result<ExecutionMetrics, String> {
        tracing::debug!(
            "Executing scheduled agent {} (interval {:?}) as ephemeral run",
            context.task.agent_id,
            interval
        );
        // Execute exactly once and return metrics.
        // The external CronScheduler is responsible for computing the next
        // fire time and re-invoking.
        Self::execute_ephemeral_agent(context).await
    }

    /// Execute a cron-scheduled agent (single run triggered by CronScheduler).
    async fn execute_cron_scheduled_agent(
        context: AgentExecutionContext,
    ) -> Result<ExecutionMetrics, String> {
        tracing::debug!("Executing cron-scheduled agent {}", context.task.agent_id,);
        Self::execute_ephemeral_agent(context).await
    }

    /// Execute an event-driven agent
    async fn execute_event_driven_agent(
        context: AgentExecutionContext,
    ) -> Result<ExecutionMetrics, String> {
        tracing::debug!("Executing event-driven agent {}", context.task.agent_id);

        // Event-driven agents are similar to persistent but triggered by events
        Self::execute_persistent_agent(context).await
    }

    /// Launch an agent process
    async fn launch_agent_process(task: &super::ScheduledTask) -> Result<ProcessHandle, String> {
        let agent_id = task.agent_id;

        // Create a secure execution environment based on security tier
        let execution_env = Self::create_execution_environment(task)?;

        // Build the command to execute the agent
        let mut command = Self::build_agent_command(task, &execution_env).await?;

        tracing::debug!("Launching agent {} with command: {:?}", agent_id, command);

        // Spawn the process
        let child = command
            .spawn()
            .map_err(|e| format!("Failed to spawn agent process: {}", e))?;

        let process_id = child.id().unwrap_or(0);

        Ok(ProcessHandle {
            process_id,
            child: Arc::new(tokio::sync::Mutex::new(child)),
            start_time: SystemTime::now(),
        })
    }

    /// Create execution environment for the agent
    fn create_execution_environment(
        task: &super::ScheduledTask,
    ) -> Result<ExecutionEnvironment, String> {
        use std::env;

        // For tests, use a temporary directory that we know exists
        let working_dir = if cfg!(test) {
            env::temp_dir()
                .join(format!("agent_{}", task.agent_id))
                .to_string_lossy()
                .to_string()
        } else {
            format!("/tmp/agent_{}", task.agent_id)
        };

        Ok(ExecutionEnvironment {
            working_directory: working_dir,
            environment_variables: vec![
                ("AGENT_ID".to_string(), task.agent_id.to_string()),
                (
                    "SECURITY_TIER".to_string(),
                    format!("{:?}", task.config.security_tier),
                ),
            ],
            resource_limits: task.config.resource_limits.clone(),
        })
    }

    /// Build the command to execute the agent
    async fn build_agent_command(
        task: &super::ScheduledTask,
        env: &ExecutionEnvironment,
    ) -> Result<Command, String> {
        let mut command = Command::new("sh");

        // Create the working directory if it doesn't exist
        if let Err(e) = tokio::fs::create_dir_all(&env.working_directory).await {
            tracing::warn!(
                "Failed to create working directory {}: {}",
                env.working_directory,
                e
            );
        }

        // Set working directory
        command.current_dir(&env.working_directory);

        // Set environment variables
        for (key, value) in &env.environment_variables {
            command.env(key, value);
        }

        // Create a script that executes the agent DSL
        let script_content = if cfg!(test) {
            // Simplified script for testing that should always succeed
            format!(
                r#"echo "Test execution of agent {}" >&2
echo "DSL Source: {}" >&2
echo "Agent test execution completed" >&2"#,
                task.agent_id, task.config.dsl_source
            )
        } else {
            format!(
                r#"#!/bin/bash
set -e
echo "Executing agent {}" >&2
echo "DSL Source:" >&2
echo "{}" >&2
# In a real implementation, this would compile and execute the DSL
sleep 1
echo "Agent execution completed" >&2"#,
                task.agent_id, task.config.dsl_source
            )
        };

        command.args(["-c", &script_content]);

        Ok(command)
    }

    /// Monitor process execution
    async fn monitor_process_execution(
        process_handle: ProcessHandle,
        context: &mut AgentExecutionContext,
    ) -> Result<(), String> {
        let process_id = process_handle.process_id;

        // Start resource monitoring
        let resource_monitor = tokio::spawn(Self::monitor_process_resources(
            process_id,
            context.handle.clone(),
        ));

        // Wait for process completion
        let mut child = process_handle.child.lock().await;
        let exit_status = child
            .wait()
            .await
            .map_err(|e| format!("Failed to wait for process: {}", e))?;

        context.exit_code = exit_status.code();

        // Stop resource monitoring
        resource_monitor.abort();

        if exit_status.success() {
            Ok(())
        } else {
            Err(format!(
                "Process exited with code: {:?}",
                exit_status.code()
            ))
        }
    }

    /// Monitor persistent process
    async fn monitor_persistent_process(
        process_handle: ProcessHandle,
        context: AgentExecutionContext,
    ) -> Result<(), String> {
        let process_id = process_handle.process_id;

        // Continuous monitoring for persistent agents
        let monitor = tokio::spawn(Self::monitor_process_resources(
            process_id,
            context.handle.clone(),
        ));

        // Check if process is still running periodically
        let mut check_interval = tokio::time::interval(Duration::from_secs(30));

        loop {
            check_interval.tick().await;

            if !Self::is_process_running(process_id).await {
                tracing::warn!(
                    "Persistent agent {} process {} died",
                    context.task.agent_id,
                    process_id
                );
                break;
            }

            // Check if termination was requested
            if matches!(context.handle.get_status(), TaskStatus::Terminated) {
                break;
            }
        }

        monitor.abort();
        Ok(())
    }

    /// Monitor process resources
    async fn monitor_process_resources(process_id: u32, handle: TaskHandle) {
        let mut interval = tokio::time::interval(Duration::from_secs(5));
        let mut system = System::new();

        loop {
            interval.tick().await;

            system.refresh_process(Pid::from(process_id as usize));

            if let Some(process) = system.process(Pid::from(process_id as usize)) {
                let memory_mb = process.memory() / 1024 / 1024;
                let cpu_usage = process.cpu_usage();

                handle.update_resource_usage(memory_mb as usize, cpu_usage);

                tracing::trace!(
                    "Process {} - Memory: {}MB, CPU: {:.2}%",
                    process_id,
                    memory_mb,
                    cpu_usage
                );
            } else {
                // Process no longer exists
                break;
            }
        }
    }

    /// Check if process is still running
    async fn is_process_running(process_id: u32) -> bool {
        let mut system = System::new();
        system.refresh_process(Pid::from(process_id as usize));
        system.process(Pid::from(process_id as usize)).is_some()
    }

    /// Terminate a process
    async fn terminate_process(process_id: u32) -> Result<(), String> {
        let mut system = System::new();
        system.refresh_process(Pid::from(process_id as usize));

        if let Some(process) = system.process(Pid::from(process_id as usize)) {
            if process.kill() {
                Ok(())
            } else {
                Err("Failed to terminate process".to_string())
            }
        } else {
            Ok(()) // Process already terminated
        }
    }
}

/// Task command for the execution loop
#[derive(Debug)]
enum TaskCommand {
    Start {
        task: Box<super::ScheduledTask>,
        handle: TaskHandle,
    },
    Terminate {
        agent_id: AgentId,
        handle: TaskHandle,
    },
}

/// Handle for a running task
#[derive(Debug, Clone)]
pub struct TaskHandle {
    agent_id: AgentId,
    status: Arc<RwLock<TaskStatus>>,
    start_time: SystemTime,
    metrics: Arc<RwLock<TaskMetrics>>,
}

impl TaskHandle {
    fn new(task: super::ScheduledTask) -> Self {
        Self {
            agent_id: task.agent_id,
            status: Arc::new(RwLock::new(TaskStatus::Pending)),
            start_time: SystemTime::now(),
            metrics: Arc::new(RwLock::new(TaskMetrics::new())),
        }
    }

    fn set_status(&self, status: TaskStatus) {
        *self.status.write() = status.clone();
        self.metrics.write().update_status(&status);
    }

    fn get_status(&self) -> TaskStatus {
        self.status.read().clone()
    }

    fn set_process_id(&self, process_id: Option<u32>) {
        self.metrics.write().process_id = process_id;
    }

    fn get_process_id(&self) -> Option<u32> {
        self.metrics.read().process_id
    }

    fn update_resource_usage(&self, memory_mb: usize, cpu_usage: f32) {
        let mut metrics = self.metrics.write();
        metrics.memory_usage = memory_mb * 1024 * 1024; // Convert to bytes
        metrics.cpu_usage = cpu_usage;
        metrics.last_activity = SystemTime::now();
    }

    fn update_execution_metrics(&self, execution_metrics: ExecutionMetrics) {
        let mut metrics = self.metrics.write();
        metrics.execution_time = Some(execution_metrics.execution_time);
        metrics.memory_peak_mb = execution_metrics.memory_peak_mb;
        metrics.cpu_time_ms = execution_metrics.cpu_time_ms;
        metrics.exit_code = execution_metrics.exit_code;
        metrics.error_count = execution_metrics.error_count;
        metrics.last_activity = SystemTime::now();
    }

    fn get_health(&self) -> TaskHealth {
        let status = self.status.read().clone();
        let uptime = SystemTime::now()
            .duration_since(self.start_time)
            .unwrap_or_default();
        let metrics = self.metrics.read();

        TaskHealth {
            agent_id: self.agent_id,
            status: status.clone(),
            uptime,
            is_healthy: matches!(
                status,
                TaskStatus::Running | TaskStatus::Pending | TaskStatus::Completed
            ),
            memory_usage: metrics.memory_usage,
            cpu_usage: metrics.cpu_usage,
            last_activity: metrics.last_activity,
        }
    }
}

/// Task execution status
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TaskStatus {
    Pending,
    Running,
    Completed,
    Failed,
    TimedOut,
    Terminated,
}

/// Task health information
#[derive(Debug, Clone)]
pub struct TaskHealth {
    pub agent_id: AgentId,
    pub status: TaskStatus,
    pub uptime: Duration,
    pub is_healthy: bool,
    pub memory_usage: usize,
    pub cpu_usage: f32,
    pub last_activity: SystemTime,
}

/// Task metrics for monitoring
#[derive(Debug, Clone)]
struct TaskMetrics {
    memory_usage: usize,
    cpu_usage: f32,
    last_activity: SystemTime,
    status_changes: u32,
    process_id: Option<u32>,
    execution_time: Option<Duration>,
    memory_peak_mb: usize,
    cpu_time_ms: u64,
    exit_code: Option<i32>,
    error_count: u32,
}

impl TaskMetrics {
    fn new() -> Self {
        Self {
            memory_usage: 0,
            cpu_usage: 0.0,
            last_activity: SystemTime::now(),
            status_changes: 0,
            process_id: None,
            execution_time: None,
            memory_peak_mb: 0,
            cpu_time_ms: 0,
            exit_code: None,
            error_count: 0,
        }
    }

    fn update_status(&mut self, _status: &TaskStatus) {
        self.status_changes += 1;
        self.last_activity = SystemTime::now();
    }
}

/// Metrics collected during agent execution
#[derive(Debug, Clone)]
struct ExecutionMetrics {
    execution_time: Duration,
    memory_peak_mb: usize,
    cpu_time_ms: u64,
    exit_code: Option<i32>,
    error_count: u32,
    #[allow(dead_code)]
    success: bool,
}

/// Context for agent execution
#[derive(Debug)]
struct AgentExecutionContext {
    task: super::ScheduledTask,
    handle: TaskHandle,
    start_time: SystemTime,
    memory_peak_mb: usize,
    cpu_time_ms: u64,
    exit_code: Option<i32>,
    error_count: u32,
}

impl AgentExecutionContext {
    fn new(task: super::ScheduledTask, handle: TaskHandle) -> Self {
        Self {
            task,
            handle,
            start_time: SystemTime::now(),
            memory_peak_mb: 0,
            cpu_time_ms: 0,
            exit_code: None,
            error_count: 0,
        }
    }
}

/// Handle to a running process
#[derive(Debug)]
struct ProcessHandle {
    process_id: u32,
    child: Arc<tokio::sync::Mutex<tokio::process::Child>>,
    #[allow(dead_code)]
    start_time: SystemTime,
}

/// Execution environment configuration
#[derive(Debug, Clone)]
struct ExecutionEnvironment {
    working_directory: String,
    environment_variables: Vec<(String, String)>,
    #[allow(dead_code)]
    resource_limits: ResourceLimits,
}

/// Task statistics for monitoring
#[derive(Debug, Clone, serde::Serialize)]
pub struct TaskStatistics {
    pub total_tasks: usize,
    pub healthy_tasks: usize,
    pub average_uptime: Duration,
    pub total_memory_usage: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{AgentConfig, ExecutionMode, Priority, ResourceLimits, SecurityTier};
    use std::collections::HashMap;

    fn create_test_task() -> super::super::ScheduledTask {
        let agent_id = AgentId::new();
        let config = AgentConfig {
            id: agent_id,
            name: "test".to_string(),
            dsl_source: "test".to_string(),
            execution_mode: ExecutionMode::Ephemeral,
            security_tier: SecurityTier::Tier1,
            resource_limits: ResourceLimits::default(),
            capabilities: vec![],
            policies: vec![],
            metadata: HashMap::new(),
            priority: Priority::Normal,
        };
        super::super::ScheduledTask::new(config)
    }

    #[tokio::test]
    async fn test_task_start_and_health_check() {
        let task_manager = TaskManager::new(Duration::from_secs(60));
        let task = create_test_task();
        let agent_id = task.agent_id;

        // Start the task
        let result = task_manager.start_task(task).await;
        assert!(result.is_ok());

        // Give it a moment to start
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Check health
        let health = task_manager.check_task_health(agent_id).await;
        assert!(health.is_ok());

        let health = health.unwrap();
        assert_eq!(health.agent_id, agent_id);
        assert!(health.is_healthy);
    }

    #[tokio::test]
    async fn test_task_termination() {
        let task_manager = TaskManager::new(Duration::from_secs(60));
        let task = create_test_task();
        let agent_id = task.agent_id;

        // Start the task
        task_manager.start_task(task).await.unwrap();

        // Give it a moment to start
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Terminate the task
        let result = task_manager.terminate_task(agent_id).await;
        assert!(result.is_ok());

        // Check that health check fails for terminated task
        let health = task_manager.check_task_health(agent_id).await;
        assert!(health.is_err());
    }

    #[tokio::test]
    async fn test_task_statistics() {
        let task_manager = TaskManager::new(Duration::from_secs(60));

        // Start multiple tasks
        for _ in 0..3 {
            let task = create_test_task();
            task_manager.start_task(task).await.unwrap();
        }

        // Give them a moment to start
        tokio::time::sleep(Duration::from_millis(50)).await;

        let stats = task_manager.get_task_statistics().await;
        assert_eq!(stats.total_tasks, 3);
        assert!(stats.healthy_tasks <= 3); // Some might have completed already
    }
}