oxios_kernel/

supervisor.rs

//! Supervisor: agent lifecycle management.
//!
//! The supervisor handles forking, executing, monitoring, and
//! terminating agent instances. It is the "init" of Oxios.
//!
//! When an agent is forked and executed, the supervisor delegates
//! the actual tool-calling loop to the [`AgentRuntime`].
//!
//! # Agent Pool & Session Persistence
//!
//! Agents are retained in an [`AgentPool`] after execution for:
//! - **Session continuation** via `Agent::continue_with()` — multi-turn
//!   conversations without re-creating the agent.
//! - **State export/import** — serialize agent conversation history to
//!   JSON for crash recovery, migration, or debugging.
//! - **Provider rate limiting** — all agents share a [`ProviderPool`] to
//!   respect per-provider RPM/concurrency limits.

use anyhow::Result;
use async_trait::async_trait;
use chrono::Utc;
use oxi_sdk::Agent;
use oxios_ouroboros::Seed;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::task::JoinHandle;

use crate::agent_runtime::AgentRuntime;
use crate::event_bus::EventBus;
use crate::resource_monitor::ResourceMonitor;
use crate::session_context::SessionContext;
use crate::types::{AgentId, AgentInfo, AgentStatus};
use oxios_ouroboros::ExecutionResult;

/// Tracks the runtime handles needed to cancel a running agent.
struct AgentHandle {
    /// Flag set on `kill()` to cooperatively signal cancellation.
    cancelled: Arc<AtomicBool>,
    /// The tokio task running the agent execution. Aborted on `kill()`.
    task: JoinHandle<Result<ExecutionResult>>,
}

/// Pool of live `Agent` instances, keyed by AgentId.
///
/// Retains agents after execution for:
/// - **State persistence** — `export_state()` serializes conversation history
///   to JSON for crash recovery, migration, or debugging.
/// - **State restoration** — `import_state()` restores a previous session.
#[derive(Default)]
pub struct AgentPool {
    agents: RwLock<HashMap<AgentId, Arc<Agent>>>,
}

impl AgentPool {
    /// Create an empty agent pool.
    pub fn new() -> Self {
        Self {
            agents: RwLock::new(HashMap::new()),
        }
    }

    /// Insert an agent into the pool.
    pub fn insert(&self, id: AgentId, agent: Arc<Agent>) {
        self.agents.write().insert(id, agent);
    }

    /// Get a pooled agent by ID.
    pub fn get(&self, id: &AgentId) -> Option<Arc<Agent>> {
        self.agents.read().get(id).cloned()
    }

    /// Remove an agent from the pool.
    pub fn remove(&self, id: &AgentId) -> Option<Arc<Agent>> {
        self.agents.write().remove(id)
    }

    /// Export an agent's state as JSON.
    ///
    /// Returns `None` if the agent is not in the pool or export fails.
    pub fn export_state(&self, id: &AgentId) -> Option<serde_json::Value> {
        self.agents
            .read()
            .get(id)
            .and_then(|agent| agent.export_state().ok())
    }

    /// Import agent state from JSON.
    ///
    /// Returns `false` if the agent is not in the pool or import fails.
    pub fn import_state(&self, id: &AgentId, state: serde_json::Value) -> bool {
        if let Some(agent) = self.agents.read().get(id) {
            agent.import_state(state).is_ok()
        } else {
            false
        }
    }

    /// Number of agents currently in the pool.
    pub fn len(&self) -> usize {
        self.agents.read().len()
    }

    /// Whether the pool is empty.
    pub fn is_empty(&self) -> bool {
        self.agents.read().is_empty()
    }
}

/// Supervisor trait for managing agent lifecycles.
#[async_trait]
pub trait Supervisor: Send + Sync {
    /// Fork a new agent from a seed specification.
    async fn fork(&self, spec: &Seed) -> Result<AgentId>;

    /// Start executing an agent.
    async fn exec(&self, id: AgentId) -> Result<()>;

    /// Fork and execute an agent with its seed, running to completion.
    /// Returns the execution result from the agent runtime.
    async fn run_with_seed(&self, id: AgentId, seed: &Seed) -> Result<ExecutionResult>;

    /// Wait for an agent to complete and return its final status.
    async fn wait(&self, id: AgentId) -> Result<AgentStatus>;

    /// Terminate an agent.
    async fn kill(&self, id: AgentId) -> Result<()>;

    /// List all known agents.
    async fn list(&self) -> Result<Vec<AgentInfo>>;
}

/// Basic in-memory supervisor implementation with AgentRuntime integration.
pub struct BasicSupervisor {
    agents: RwLock<HashMap<AgentId, AgentInfo>>,
    /// Per-agent cancellation tokens and join handles for task abortion.
    handles: RwLock<HashMap<AgentId, AgentHandle>>,
    /// Pool of live Agent instances for session continuation.
    agent_pool: AgentPool,
    event_bus: EventBus,
    runtime: Arc<AgentRuntime>,
    resource_monitor: Option<Arc<ResourceMonitor>>,
    /// Session context for proactive recall timing (RFC-020).
    /// Shared across all Seed executions within this supervisor's lifetime
    /// so that RecallTiming can track message count and topic changes.
    /// Uses tokio::sync::RwLock (not parking_lot) so the guard is Send,
    /// allowing it to be held across .await in tokio::spawn.
    session_context: Arc<tokio::sync::RwLock<SessionContext>>,
}

impl BasicSupervisor {
    /// Creates a new supervisor with the given event bus and agent runtime.
    pub fn new(event_bus: EventBus, runtime: AgentRuntime) -> Self {
        Self {
            agents: RwLock::new(HashMap::new()),
            handles: RwLock::new(HashMap::new()),
            agent_pool: AgentPool::new(),
            event_bus,
            runtime: Arc::new(runtime),
            resource_monitor: None,
            session_context: Arc::new(tokio::sync::RwLock::new(SessionContext::new())),
        }
    }

    /// Attach a resource monitor for agent count tracking.
    pub fn set_resource_monitor(&mut self, rm: Arc<ResourceMonitor>) {
        self.resource_monitor = Some(rm);
    }

    /// Update the resource monitor with current active agent count.
    fn update_agent_count(&self) {
        if let Some(ref rm) = self.resource_monitor {
            let count = self.agents.read().len();
            rm.set_active_agents(count);
        }
    }

    /// Access the agent pool for session continuation.
    pub fn pool(&self) -> &AgentPool {
        &self.agent_pool
    }
}

#[async_trait]
impl Supervisor for BasicSupervisor {
    async fn fork(&self, spec: &Seed) -> Result<AgentId> {
        let id = AgentId::new_v4();
        let info = AgentInfo {
            id,
            name: spec.goal.clone(),
            status: AgentStatus::Starting,
            created_at: Utc::now(),
            seed_id: Some(spec.id),
        };

        {
            let mut agents = self.agents.write();
            agents.insert(id, info);
        }

        self.update_agent_count();

        let _ = self
            .event_bus
            .publish(crate::event_bus::KernelEvent::AgentCreated {
                id,
                name: spec.goal.clone(),
            });

        tracing::info!(agent_id = %id, "Forked new agent from seed");
        Ok(id)
    }

    async fn exec(&self, id: AgentId) -> Result<()> {
        {
            let mut agents = self.agents.write();
            match agents.get_mut(&id) {
                Some(agent) => {
                    agent.status = AgentStatus::Running;
                }
                None => anyhow::bail!("Agent {id} not found"),
            }
        }

        self.update_agent_count();

        let _ = self
            .event_bus
            .publish(crate::event_bus::KernelEvent::AgentStarted { id });
        tracing::info!(agent_id = %id, "Agent execution started");

        Ok(())
    }

    async fn run_with_seed(&self, id: AgentId, seed: &Seed) -> Result<ExecutionResult> {
        // Mark as running.
        {
            let mut agents = self.agents.write();
            match agents.get_mut(&id) {
                Some(agent) => agent.status = AgentStatus::Running,
                None => anyhow::bail!("Agent {id} not found"),
            }
        }

        let _ = self
            .event_bus
            .publish(crate::event_bus::KernelEvent::AgentStarted { id });

        tracing::info!(agent_id = %id, seed_id = %seed.id, "Running agent task");

        // Spawn the execution as a tokio task so we can track and abort it.
        let cancelled = Arc::new(AtomicBool::new(false));
        let runtime = Arc::clone(&self.runtime);
        let seed = seed.clone();
        let cancelled_clone = cancelled.clone();

        // Share the session context so RecallTiming persists across Seeds.
        // Uses tokio::sync::RwLock so the guard is Send-safe across .await.
        let session_ctx = self.session_context.clone();

        let handle: JoinHandle<Result<ExecutionResult>> = tokio::spawn(async move {
            // Check for cancellation before starting.
            if cancelled_clone.load(Ordering::Relaxed) {
                return Ok(ExecutionResult {
                    output: "Agent cancelled before execution".into(),
                    steps_completed: 0,
                    success: false,
                    tool_calls: vec![],
                });
            }
            let mut ctx = session_ctx.write().await;
            runtime.execute(id, &seed, &mut ctx).await
        });

        // Store the handle so kill() can abort the task.
        {
            let mut handles = self.handles.write();
            handles.insert(
                id,
                AgentHandle {
                    cancelled,
                    task: handle,
                },
            );
        }

        // Await the spawned task.
        let result = {
            let agent_handle = {
                let mut handles = self.handles.write();
                handles.remove(&id)
            };
            // Guard is dropped above, safe to await.

            match agent_handle {
                Some(ah) => match ah.task.await {
                    Ok(res) => res,
                    Err(join_err) => {
                        // Task was aborted (e.g. kill()) or panicked.
                        tracing::warn!(agent_id = %id, error = %join_err, "Agent task join error");
                        Ok(ExecutionResult {
                            output: format!("Agent task aborted: {join_err}"),
                            steps_completed: 0,
                            success: false,
                            tool_calls: vec![],
                        })
                    }
                },
                None => anyhow::bail!("Agent {id} handle disappeared"),
            }
        };

        match result {
            Ok(result) => {
                tracing::info!(
                    agent_id = %id,
                    success = result.success,
                    steps = result.steps_completed,
                    "Agent task completed"
                );

                {
                    let mut agents = self.agents.write();
                    if let Some(agent) = agents.get_mut(&id) {
                        agent.status = if result.success {
                            AgentStatus::Idle
                        } else {
                            AgentStatus::Failed
                        };
                    }
                }

                let _ = self
                    .event_bus
                    .publish(crate::event_bus::KernelEvent::AgentStopped { id });
                self.update_agent_count();
                Ok(result)
            }
            Err(e) => {
                tracing::error!(agent_id = %id, error = %e, "Agent task failed");

                {
                    let mut agents = self.agents.write();
                    if let Some(agent) = agents.get_mut(&id) {
                        agent.status = AgentStatus::Failed;
                    }
                }

                let _ = self
                    .event_bus
                    .publish(crate::event_bus::KernelEvent::AgentFailed {
                        id,
                        error: e.to_string(),
                    });
                self.update_agent_count();

                Ok(ExecutionResult {
                    output: format!("Agent failed: {e}"),
                    steps_completed: 0,
                    success: false,
                    tool_calls: vec![],
                })
            }
        }
    }

    async fn wait(&self, id: AgentId) -> Result<AgentStatus> {
        let agents = self.agents.read();
        match agents.get(&id) {
            Some(info) => Ok(info.status),
            None => anyhow::bail!("Agent {id} not found"),
        }
    }

    async fn kill(&self, id: AgentId) -> Result<()> {
        // Cancel and abort the running task, if any.
        {
            let mut handles = self.handles.write();
            if let Some(agent_handle) = handles.remove(&id) {
                agent_handle.cancelled.store(true, Ordering::Relaxed);
                agent_handle.task.abort();
                tracing::info!(agent_id = %id, "Agent task aborted");
            }
        }

        {
            let mut agents = self.agents.write();
            if let Some(agent) = agents.get_mut(&id) {
                agent.status = AgentStatus::Stopped;
            } else {
                anyhow::bail!("Agent {id} not found");
            }
        }

        let _ = self
            .event_bus
            .publish(crate::event_bus::KernelEvent::AgentStopped { id });
        self.update_agent_count();
        tracing::info!(agent_id = %id, "Agent killed");
        Ok(())
    }

    async fn list(&self) -> Result<Vec<AgentInfo>> {
        let agents = self.agents.read();
        Ok(agents.values().cloned().collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::event_bus::EventBus;
    use crate::types::AgentStatus;
    // Note: imports kept for potential future test extensions.
    use oxios_ouroboros::Seed;

    // Note: MockProvider no longer needed — OxiosEngine handles provider resolution.
    // The engine resolves models internally, so tests just use OxiosEngine::new().

    /// Helper to create a real BasicSupervisor wired to a real EventBus.
    async fn make_supervisor() -> BasicSupervisor {
        let event_bus = EventBus::new(64);

        // Build a mock KernelHandle with temp dirs.
        let tmp = std::env::temp_dir().join(format!("oxios-test-{}", uuid::Uuid::new_v4()));
        let _ = std::fs::create_dir_all(&tmp);

        let state_store_2 =
            Arc::new(crate::state_store::StateStore::new(tmp.join("state")).expect("state store"));
        let state_store = state_store_2.clone();
        let state_store_for_space = state_store_2.clone();
        let memory_manager = Arc::new({
            let mut mm = crate::memory::MemoryManager::new(state_store.clone());
            mm.set_git_layer(Arc::new(
                crate::git_layer::GitLayer::new(tmp.join("git"), false).expect("git layer"),
            ));
            mm
        });

        let kernel_handle = Arc::new(crate::KernelHandle::new(
            crate::kernel_handle::StateApi::new(state_store),
            crate::kernel_handle::AgentApi::new(
                Arc::new(crate::supervisor::NoOpSupervisor),
                Arc::new(crate::budget::BudgetManager::new()),
                memory_manager.clone(),
                Some(event_bus.clone()),
            ),
            crate::kernel_handle::SecurityApi::new(
                Arc::new(parking_lot::Mutex::new(crate::auth::AuthManager::new())),
                Arc::new(oxi_sdk::observability::AuditTrail::new(100)),
                Arc::new(parking_lot::Mutex::new(
                    crate::access_manager::AccessManager::new(),
                )),
                Arc::new(
                    crate::state_store::StateStore::new(tmp.join("state2")).expect("state store 2"),
                ),
            ),
            crate::kernel_handle::PersonaApi::new(Arc::new(crate::persona::PersonaManager::new())),
            crate::kernel_handle::ExtensionApi::new(Arc::new(crate::skill::SkillManager::new(
                tmp.join("skills"),
                tmp.join("share/skills"),
            ))),
            crate::kernel_handle::McpApi::new(Arc::new(crate::mcp::McpBridge::new())),
            crate::kernel_handle::InfraApi::new(
                Arc::new(
                    crate::git_layer::GitLayer::new(tmp.join("git"), false).expect("git layer"),
                ),
                Arc::new(crate::scheduler::AgentScheduler::new(4, 60, 300)),
                Arc::new(crate::cron::CronScheduler::new(
                    Arc::new(
                        crate::state_store::StateStore::new(tmp.join("cron")).expect("cron state"),
                    ),
                    60,
                )),
                Arc::new(crate::resource_monitor::ResourceMonitor::new(60, 100)),
                EventBus::new(64),
                crate::config::OxiosConfig::default(),
                std::time::Instant::now(),
            ),
            None,
            crate::kernel_handle::ExecApi::new(
                Arc::new(parking_lot::RwLock::new(
                    crate::config::ExecConfig::default(),
                )),
                Arc::new(parking_lot::Mutex::new(
                    crate::access_manager::AccessManager::new(),
                )),
            ),
            crate::kernel_handle::A2aApi::new(Arc::new(crate::a2a::A2AProtocol::new(
                EventBus::new(64),
            ))),
            crate::kernel_handle::EngineApi::new(
                Arc::new(parking_lot::RwLock::new(
                    crate::config::OxiosConfig::default(),
                )),
                tmp.join("config.toml"),
                Arc::new(crate::kernel_handle::RoutingStats::new()),
                Arc::new(crate::engine::EngineHandle::new(Arc::new(
                    crate::OxiosEngine::new("anthropic/claude-sonnet-4-20250514"),
                ))),
            ),
            Arc::new(oxios_markdown::KnowledgeBase::new(tmp.join("knowledge")).unwrap()),
            Arc::new(
                crate::kernel_handle::KnowledgeLens::new(
                    Arc::new(oxios_markdown::KnowledgeBase::new(tmp.join("knowledge")).unwrap()),
                    memory_manager.clone(),
                )
                .unwrap(),
            ),
            crate::kernel_handle::MarketplaceApi::new(
                Arc::new(crate::skill::clawhub::ClawHubInstaller::new(
                    tmp.join("skills"),
                    tmp.join("state"),
                    None,
                )),
                Arc::new(
                    crate::skill::clawhub::ClawHubClient::new(None).expect("valid ClawHub client"),
                ),
                Arc::new(crate::skill::skills_sh::SkillsShInstaller::new(
                    tmp.join("skills"),
                    None,
                    None,
                )),
                Arc::new(
                    crate::skill::skills_sh::SkillsShClient::new(None, None)
                        .expect("valid Skills.sh client"),
                ),
            ),
            None, // calendar (not configured in test)
            None, // email (not configured in test)
        ));

        let engine = crate::OxiosEngine::new("mock/model");
        let engine_handle = Arc::new(crate::engine::EngineHandle::new(Arc::new(engine)));
        let runtime = AgentRuntime::new(engine_handle, "mock/model", kernel_handle, None);
        BasicSupervisor::new(event_bus, runtime)
    }

    /// Helper to create a minimal Seed for testing.
    fn make_seed(goal: &str) -> Seed {
        Seed {
            id: uuid::Uuid::new_v4(),
            goal: goal.to_string(),
            constraints: vec![],
            acceptance_criteria: vec![],
            ontology: vec![],
            created_at: chrono::Utc::now(),
            generation: 0,
            parent_seed_id: None,
            cspace_hint: None,
            original_request: String::new(),
            output_schema: None,
        }
    }

    #[tokio::test]
    async fn test_fork_creates_agent() {
        let supervisor = make_supervisor().await;
        let seed = make_seed("Test agent");

        let id = supervisor.fork(&seed).await.unwrap();

        let agents = supervisor.list().await.unwrap();
        assert_eq!(agents.len(), 1);
        assert_eq!(agents[0].id, id);
        assert_eq!(agents[0].name, "Test agent");
        assert_eq!(agents[0].status, AgentStatus::Starting);
        assert_eq!(agents[0].seed_id, Some(seed.id));
    }

    #[tokio::test]
    async fn test_exec_updates_status_to_running() {
        let supervisor = make_supervisor().await;
        let seed = make_seed("Running agent");

        let id = supervisor.fork(&seed).await.unwrap();
        assert_eq!(supervisor.wait(id).await.unwrap(), AgentStatus::Starting);

        supervisor.exec(id).await.unwrap();
        assert_eq!(supervisor.wait(id).await.unwrap(), AgentStatus::Running);
    }

    #[tokio::test]
    async fn test_kill_sets_stopped() {
        let supervisor = make_supervisor().await;
        let seed = make_seed("Doomed agent");

        let id = supervisor.fork(&seed).await.unwrap();
        supervisor.exec(id).await.unwrap();
        assert_eq!(supervisor.wait(id).await.unwrap(), AgentStatus::Running);

        supervisor.kill(id).await.unwrap();
        assert_eq!(supervisor.wait(id).await.unwrap(), AgentStatus::Stopped);
    }

    #[tokio::test]
    async fn test_kill_unknown_agent_returns_error() {
        let supervisor = make_supervisor().await;
        let unknown_id = uuid::Uuid::new_v4();

        let result = supervisor.kill(unknown_id).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("not found"));
    }

    #[tokio::test]
    async fn test_list_returns_all_agents() {
        let supervisor = make_supervisor().await;

        let id1 = supervisor.fork(&make_seed("Agent 1")).await.unwrap();
        let id2 = supervisor.fork(&make_seed("Agent 2")).await.unwrap();
        let id3 = supervisor.fork(&make_seed("Agent 3")).await.unwrap();

        let agents = supervisor.list().await.unwrap();
        assert_eq!(agents.len(), 3);

        let ids: std::collections::HashSet<AgentId> = agents.iter().map(|a| a.id).collect();
        assert!(ids.contains(&id1));
        assert!(ids.contains(&id2));
        assert!(ids.contains(&id3));
    }

    #[tokio::test]
    async fn test_exec_unknown_agent_returns_error() {
        let supervisor = make_supervisor().await;
        let unknown_id = uuid::Uuid::new_v4();

        let result = supervisor.exec(unknown_id).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("not found"));
    }

    #[tokio::test]
    async fn test_wait_unknown_agent_returns_error() {
        let supervisor = make_supervisor().await;
        let unknown_id = uuid::Uuid::new_v4();

        let result = supervisor.wait(unknown_id).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("not found"));
    }
}

/// A no-op supervisor used during KernelBuilder::build() to break the
/// KernelHandle → AgentRuntime → Supervisor → KernelHandle cycle.
///
/// AgentApi.supervisor is only used for list/kill operations, not during
/// tool registration, so this placeholder is safe during build time.
pub struct NoOpSupervisor;

#[async_trait::async_trait]
impl Supervisor for NoOpSupervisor {
    async fn fork(&self, _spec: &Seed) -> Result<AgentId> {
        Err(anyhow::anyhow!(
            "NoOpSupervisor: fork not available during build"
        ))
    }
    async fn exec(&self, _id: AgentId) -> Result<()> {
        Err(anyhow::anyhow!(
            "NoOpSupervisor: exec not available during build"
        ))
    }
    async fn run_with_seed(&self, _id: AgentId, _seed: &Seed) -> Result<ExecutionResult> {
        Err(anyhow::anyhow!(
            "NoOpSupervisor: run_with_seed not available during build"
        ))
    }
    async fn wait(&self, _id: AgentId) -> Result<AgentStatus> {
        Err(anyhow::anyhow!(
            "NoOpSupervisor: wait not available during build"
        ))
    }
    async fn kill(&self, _id: AgentId) -> Result<()> {
        Err(anyhow::anyhow!(
            "NoOpSupervisor: kill not available during build"
        ))
    }
    async fn list(&self) -> Result<Vec<AgentInfo>> {
        Ok(Vec::new())
    }
}