astrid-capsule 0.2.0

//! Capsule trait and core types.

use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::sync::Semaphore;

use crate::context::CapsuleContext;
use crate::error::{CapsuleError, CapsuleResult};
use crate::manifest::CapsuleManifest;
use crate::tool::CapsuleTool;

/// Maximum concurrent interceptor invocations per capsule.
const MAX_CONCURRENT_INTERCEPTORS: usize = 4;

/// Unique, stable, human-readable capsule identifier.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
pub struct CapsuleId(String);

impl<'de> Deserialize<'de> for CapsuleId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        Self::new(s).map_err(serde::de::Error::custom)
    }
}

impl CapsuleId {
    pub fn new(id: impl Into<String>) -> CapsuleResult<Self> {
        let id = id.into();
        Self::validate(&id)?;
        Ok(Self(id))
    }

    #[must_use]
    pub fn from_static(id: &str) -> Self {
        Self(id.to_string())
    }

    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }

    fn validate(id: &str) -> CapsuleResult<()> {
        if id.is_empty() {
            return Err(CapsuleError::UnsupportedEntryPoint(
                "capsule id must not be empty".into(),
            ));
        }
        if !id
            .chars()
            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-')
        {
            return Err(CapsuleError::UnsupportedEntryPoint(format!(
                "capsule id must contain only lowercase alphanumeric characters and hyphens, got: {id}"
            )));
        }
        Ok(())
    }
}

impl fmt::Display for CapsuleId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.0)
    }
}

impl AsRef<str> for CapsuleId {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

/// Result of waiting for a capsule or engine to signal readiness.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReadyStatus {
    /// The capsule signaled ready (or has no background task).
    Ready,
    /// The timeout expired before readiness was signaled.
    Timeout,
    /// The capsule's run loop exited or crashed before signaling ready.
    Crashed,
}

impl ReadyStatus {
    /// Returns `true` if the status is [`ReadyStatus::Ready`].
    #[must_use]
    pub fn is_ready(self) -> bool {
        self == Self::Ready
    }
}

/// The lifecycle state of a capsule.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CapsuleState {
    Unloaded,
    Loading,
    Ready,
    Failed(String),
    Unloading,
}

/// A loaded capsule that can provide tools and integrations to the runtime.
#[async_trait]
pub trait Capsule: Send + Sync {
    /// The unique identifier for this capsule.
    fn id(&self) -> &CapsuleId;

    /// The manifest that describes this capsule.
    fn manifest(&self) -> &CapsuleManifest;

    /// Current lifecycle state.
    fn state(&self) -> CapsuleState;

    /// Load the capsule, initializing all of its execution engines.
    async fn load(&mut self, ctx: &CapsuleContext) -> CapsuleResult<()>;

    /// Unload the capsule, terminating all of its execution engines.
    async fn unload(&mut self) -> CapsuleResult<()>;

    /// Tools exposed by this capsule.
    fn tools(&self) -> &[std::sync::Arc<dyn CapsuleTool>] {
        &[] // Default implementation returning empty list
    }

    /// Extract the inbound receiver for uplink messages.
    /// This is typically called exactly once by the OS router after loading.
    fn take_inbound_rx(
        &mut self,
    ) -> Option<tokio::sync::mpsc::Receiver<astrid_core::InboundMessage>> {
        None
    }

    /// Wait for the capsule's background tasks to signal readiness.
    ///
    /// Returns [`ReadyStatus::Ready`] if all engines are ready or have no
    /// background tasks. Returns [`ReadyStatus::Timeout`] if the timeout
    /// expires, or [`ReadyStatus::Crashed`] if the run loop exited before
    /// signaling ready.
    async fn wait_ready(&self, _timeout: std::time::Duration) -> ReadyStatus {
        ReadyStatus::Ready
    }

    /// Invoke an interceptor handler by action name.
    ///
    /// Called by the event dispatcher when an IPC event matches one of
    /// this capsule's registered interceptor patterns. `action` is the
    /// handler name (e.g., `handle_user_prompt`), `payload` is the
    /// serialized IPC payload bytes.
    fn invoke_interceptor(&self, _action: &str, _payload: &[u8]) -> CapsuleResult<Vec<u8>> {
        Err(CapsuleError::NotSupported(
            "interceptors not supported".into(),
        ))
    }

    /// Probe liveness beyond what `state()` reports.
    ///
    /// Returns the current state by default. Composite capsules delegate
    /// to their engines, which can detect silently exited background tasks.
    fn check_health(&self) -> CapsuleState {
        self.state()
    }

    /// The directory this capsule was loaded from.
    ///
    /// Used by the kernel health monitor to restart crashed capsules.
    /// Returns `None` for capsules that don't have a filesystem source
    /// (e.g., test mocks).
    fn source_dir(&self) -> Option<&Path> {
        None
    }

    /// Per-capsule semaphore that bounds concurrent interceptor invocations.
    ///
    /// The event dispatcher acquires a permit before calling
    /// [`invoke_interceptor`](Self::invoke_interceptor), preventing any single
    /// capsule from spawning unbounded tasks under high event volume.
    ///
    /// # Default implementation
    ///
    /// Returns a **shared global** semaphore as a fallback. This means all
    /// capsules using the default share the same permit pool, which does NOT
    /// provide per-capsule isolation. Concrete types (e.g., `CompositeCapsule`)
    /// should override this with their own `Arc<Semaphore>`.
    fn interceptor_semaphore(&self) -> &Arc<Semaphore> {
        use std::sync::LazyLock;
        static FALLBACK: LazyLock<Arc<Semaphore>> =
            LazyLock::new(|| Arc::new(Semaphore::new(MAX_CONCURRENT_INTERCEPTORS)));
        &FALLBACK
    }
}

/// The universal, additive implementation of a Capsule.
///
/// Instead of choosing between WASM or MCP execution, the `CompositeCapsule`
/// owns a collection of `ExecutionEngine`s. When loaded, it iterates through
/// all of them, providing a unified lifecycle and security boundary for
/// everything declared in the `Capsule.toml`.
pub(crate) struct CompositeCapsule {
    id: CapsuleId,
    manifest: CapsuleManifest,
    state: CapsuleState,
    engines: Vec<Box<dyn crate::engine::ExecutionEngine>>,
    tools: Vec<Arc<dyn CapsuleTool>>,
    capsule_dir: Option<PathBuf>,
    interceptor_semaphore: Arc<Semaphore>,
}

impl CompositeCapsule {
    /// Create a new, empty Composite Capsule from a manifest.
    pub(crate) fn new(manifest: CapsuleManifest) -> CapsuleResult<Self> {
        let id = CapsuleId::new(manifest.package.name.clone())?;
        Ok(Self {
            id,
            manifest,
            state: CapsuleState::Unloaded,
            engines: Vec::new(),
            tools: Vec::new(),
            capsule_dir: None,
            interceptor_semaphore: Arc::new(Semaphore::new(MAX_CONCURRENT_INTERCEPTORS)),
        })
    }

    /// Set the source directory this capsule was loaded from.
    pub(crate) fn set_source_dir(&mut self, dir: PathBuf) {
        self.capsule_dir = Some(dir);
    }

    /// Add an execution engine (e.g., WasmEngine, McpEngine) to this capsule.
    pub(crate) fn add_engine(&mut self, engine: Box<dyn crate::engine::ExecutionEngine>) {
        self.engines.push(engine);
    }
}

#[async_trait]
impl Capsule for CompositeCapsule {
    fn id(&self) -> &CapsuleId {
        &self.id
    }

    fn manifest(&self) -> &CapsuleManifest {
        &self.manifest
    }

    fn state(&self) -> CapsuleState {
        self.state.clone()
    }

    async fn load(&mut self, ctx: &CapsuleContext) -> CapsuleResult<()> {
        self.state = CapsuleState::Loading;
        self.tools.clear();
        for engine in &mut self.engines {
            if let Err(e) = engine.load(ctx).await {
                self.state = CapsuleState::Failed(e.to_string());
                return Err(e);
            }
            self.tools.extend_from_slice(engine.tools());
        }
        self.state = CapsuleState::Ready;
        Ok(())
    }

    async fn unload(&mut self) -> CapsuleResult<()> {
        self.state = CapsuleState::Unloading;
        for engine in &mut self.engines {
            // Unload on a best-effort basis so a failing engine doesn't
            // prevent others from shutting down gracefully.
            let _ = engine.unload().await;
        }
        self.tools.clear();
        self.state = CapsuleState::Unloaded;
        Ok(())
    }

    fn tools(&self) -> &[std::sync::Arc<dyn CapsuleTool>] {
        &self.tools
    }

    async fn wait_ready(&self, timeout: std::time::Duration) -> ReadyStatus {
        let deadline = tokio::time::Instant::now() + timeout;
        for engine in &self.engines {
            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
            if remaining.is_zero() {
                return ReadyStatus::Timeout;
            }
            let status = engine.wait_ready(remaining).await;
            if !status.is_ready() {
                return status;
            }
        }
        ReadyStatus::Ready
    }

    fn take_inbound_rx(
        &mut self,
    ) -> Option<tokio::sync::mpsc::Receiver<astrid_core::InboundMessage>> {
        for engine in &mut self.engines {
            if let Some(rx) = engine.take_inbound_rx() {
                return Some(rx);
            }
        }
        None
    }

    fn invoke_interceptor(&self, action: &str, payload: &[u8]) -> CapsuleResult<Vec<u8>> {
        for engine in &self.engines {
            match engine.invoke_interceptor(action, payload) {
                Ok(result) => return Ok(result),
                // Engine doesn't support interceptors — try the next one.
                Err(CapsuleError::NotSupported(_)) => continue,
                Err(e) => return Err(e),
            }
        }
        Err(CapsuleError::NotSupported(
            "no engine supports interceptors".into(),
        ))
    }

    fn check_health(&self) -> CapsuleState {
        for engine in &self.engines {
            let health = engine.check_health();
            if let CapsuleState::Failed(_) = &health {
                return health;
            }
        }
        self.state.clone()
    }

    fn source_dir(&self) -> Option<&Path> {
        self.capsule_dir.as_deref()
    }

    fn interceptor_semaphore(&self) -> &Arc<Semaphore> {
        &self.interceptor_semaphore
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::ExecutionEngine;
    use crate::manifest::{CapabilitiesDef, PackageDef};
    use async_trait::async_trait;

    /// A mock engine that always reports healthy.
    struct HealthyEngine;

    #[async_trait]
    impl ExecutionEngine for HealthyEngine {
        async fn load(&mut self, _ctx: &crate::context::CapsuleContext) -> CapsuleResult<()> {
            Ok(())
        }
        async fn unload(&mut self) -> CapsuleResult<()> {
            Ok(())
        }
    }

    /// A mock engine that reports failed health.
    struct FailedEngine;

    #[async_trait]
    impl ExecutionEngine for FailedEngine {
        async fn load(&mut self, _ctx: &crate::context::CapsuleContext) -> CapsuleResult<()> {
            Ok(())
        }
        async fn unload(&mut self) -> CapsuleResult<()> {
            Ok(())
        }
        fn check_health(&self) -> CapsuleState {
            CapsuleState::Failed("engine crashed".into())
        }
    }

    fn test_manifest() -> CapsuleManifest {
        CapsuleManifest {
            package: PackageDef {
                name: "test-capsule".into(),
                version: "0.0.1".into(),
                description: None,
                authors: Vec::new(),
                repository: None,
                homepage: None,
                documentation: None,
                license: None,
                license_file: None,
                readme: None,
                keywords: Vec::new(),
                categories: Vec::new(),
                astrid_version: None,
                publish: None,
                include: None,
                exclude: None,
                metadata: None,
            },
            components: Vec::new(),
            dependencies: Default::default(),
            capabilities: CapabilitiesDef::default(),
            env: std::collections::HashMap::new(),
            context_files: Vec::new(),
            commands: Vec::new(),
            mcp_servers: Vec::new(),
            skills: Vec::new(),
            uplinks: Vec::new(),
            llm_providers: Vec::new(),
            interceptors: Vec::new(),
            cron_jobs: Vec::new(),
            tools: Vec::new(),
            topics: Vec::new(),
            effective_provides_cache: std::sync::OnceLock::new(),
        }
    }

    #[test]
    fn composite_check_health_all_healthy() {
        let mut capsule = CompositeCapsule::new(test_manifest()).unwrap();
        capsule.state = CapsuleState::Ready;
        capsule.add_engine(Box::new(HealthyEngine));
        capsule.add_engine(Box::new(HealthyEngine));

        assert_eq!(capsule.check_health(), CapsuleState::Ready);
    }

    #[test]
    fn composite_check_health_returns_first_failure() {
        let mut capsule = CompositeCapsule::new(test_manifest()).unwrap();
        capsule.state = CapsuleState::Ready;
        capsule.add_engine(Box::new(HealthyEngine));
        capsule.add_engine(Box::new(FailedEngine));

        assert_eq!(
            capsule.check_health(),
            CapsuleState::Failed("engine crashed".into())
        );
    }

    #[test]
    fn composite_check_health_no_engines_returns_state() {
        let mut capsule = CompositeCapsule::new(test_manifest()).unwrap();
        capsule.state = CapsuleState::Ready;

        assert_eq!(capsule.check_health(), CapsuleState::Ready);
    }

    // -- wait_ready tests --

    /// A mock engine that never signals ready (simulates slow startup).
    struct SlowEngine;

    #[async_trait]
    impl ExecutionEngine for SlowEngine {
        async fn load(&mut self, _ctx: &crate::context::CapsuleContext) -> CapsuleResult<()> {
            Ok(())
        }
        async fn unload(&mut self) -> CapsuleResult<()> {
            Ok(())
        }
        async fn wait_ready(&self, timeout: std::time::Duration) -> ReadyStatus {
            tokio::time::sleep(timeout).await;
            ReadyStatus::Timeout
        }
    }

    #[tokio::test]
    async fn composite_wait_ready_first_engine_timeout_starves_second() {
        // With a shared deadline, if the first engine consumes the entire
        // budget, the second engine gets zero time and returns Timeout
        // immediately. This test locks in the shared-deadline contract.
        let mut capsule = CompositeCapsule::new(test_manifest()).unwrap();
        capsule.add_engine(Box::new(SlowEngine));
        capsule.add_engine(Box::new(HealthyEngine));

        let status = capsule
            .wait_ready(std::time::Duration::from_millis(50))
            .await;
        assert_eq!(status, ReadyStatus::Timeout);
    }

    #[tokio::test]
    async fn composite_wait_ready_all_healthy() {
        let mut capsule = CompositeCapsule::new(test_manifest()).unwrap();
        capsule.add_engine(Box::new(HealthyEngine));
        capsule.add_engine(Box::new(HealthyEngine));

        let status = capsule
            .wait_ready(std::time::Duration::from_millis(100))
            .await;
        assert_eq!(status, ReadyStatus::Ready);
    }

    #[test]
    fn composite_interceptor_semaphore_is_bounded() {
        let capsule = CompositeCapsule::new(test_manifest()).unwrap();
        let sem = capsule.interceptor_semaphore();
        assert_eq!(sem.available_permits(), MAX_CONCURRENT_INTERCEPTORS);
    }

    #[test]
    fn trait_default_interceptor_semaphore_returns_valid_semaphore() {
        struct MinimalCapsule;
        #[async_trait]
        impl Capsule for MinimalCapsule {
            fn id(&self) -> &CapsuleId {
                unimplemented!()
            }
            fn manifest(&self) -> &CapsuleManifest {
                unimplemented!()
            }
            fn state(&self) -> CapsuleState {
                CapsuleState::Unloaded
            }
            async fn load(&mut self, _: &crate::context::CapsuleContext) -> CapsuleResult<()> {
                Ok(())
            }
            async fn unload(&mut self) -> CapsuleResult<()> {
                Ok(())
            }
        }
        let capsule = MinimalCapsule;
        let sem = capsule.interceptor_semaphore();
        assert_eq!(sem.available_permits(), MAX_CONCURRENT_INTERCEPTORS);
    }
}