yeti_types/plugins/

mod.rs

//! Plugin traits and registration context.
//!
//! Defines `Plugin`, `RegistrationContext`, `VectorHook`, `AiHook`,
//! `ComputedFieldHook`, `FieldMapping`, `ModelInfo`, `VectorBatcher`,
//! `EmbeddedFile`. (The former `ServiceContext` was unified into
//! `yeti_types::resource::Context` .)
//!
//! # Errors
//!
//! Result-returning methods on the traits in this file (`VectorHook`,
//! `AiHook`, `Plugin`, `ComputedFieldHook`) carry implementation-defined
//! error semantics — the trait declares the *shape* (`Result<T, String>`
//! for hot-path hook traits, `Result<T, YetiError>` for `Plugin`
//! lifecycle), but the actual conditions that produce `Err` belong to
//! each implementor.
//!
//! - `VectorHook` / `AiHook` impls in yeti-ai surface model-load,
//!   tokenizer, and inference failures as `String`. Hook implementors
//!   should document their own error contracts.
//! - `Plugin::resources` / `on_ready` / `install_event_subscriber`
//!   propagate `YetiError` from plugin setup (table registration,
//!   route conflicts, bootstrap failures).
//!
//! Per-method `# Errors` sections at the trait level would just say
//! "implementation-defined", so the lint is suppressed at the module
//! level. Each impl crate documents its own error conditions.
#![allow(clippy::missing_errors_doc)]

use async_trait::async_trait;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

// ============================================================================
// Plugin extension contracts (relocated from yeti_types::extension)
// ============================================================================
//
// One trait — `tower::Service<R>` — for every plugin surface. Composition
// via `tower::Layer<S>` / `ServiceBuilder`. See ADR-006.
//
// | Module      | Domain                          | Request                 | Response               |
// |-------------|---------------------------------|-------------------------|------------------------|
// | `request`   | HTTP-shaped per-request work    | `YetiRequest`           | `YetiResponse`         |
// | `lifecycle` | Fire-and-forget lifecycle taps  | `LifecycleEvent`        | `()`                   |
// | `auth`      | Identity → access resolution    | `ResolveRequest`        | `ResolveResponse`      |
// | `token`     | JWT mint extension              | `TokenRequest`          | `TokenResponse`        |
// | `oauth`     | OAuth callback claim mutation   | `OAuthRequest`          | `OAuthResponse`        |
// | `mcp`       | MCP tool list / call hooks      | `ListToolsRequest`      | `ListToolsResponse`    |
// | `queue`     | Durable / async job execution   | `JobRequest`            | `JobResponse`          |

pub use tower::util::{BoxCloneService, BoxCloneSyncService, BoxService, ServiceFn, service_fn};
pub use tower::{Layer, Service, ServiceBuilder, ServiceExt};

pub mod auth;
pub mod lifecycle;
pub mod mcp;
pub mod oauth;
pub mod queue;
pub mod request;
pub mod token;
pub mod trust;

// Re-export trust tier at the plugins level — commonly referenced.
pub use trust::PluginTrustTier;

// Hot-path re-exports for the most common pipeline types.
pub use lifecycle::{LifecycleEvent, TelemetryEvent, TelemetryService};
pub use mcp::{
    CallToolRequest, CallToolResponse, CallToolService, ListToolsRequest, ListToolsResponse,
    ListToolsService, ToolAnnotations, ToolDefinition,
};
pub use oauth::{OAuthRequest, OAuthResponse, OAuthService};
pub use request::{ContextService, RequestPipeline, YetiRequest, YetiResponse};
pub use token::{TokenRequest, TokenResponse, TokenService};

// ============================================================================
// ComputedFieldHook
// ============================================================================

/// Hook for computing virtual field values at read time.
#[async_trait]
pub trait ComputedFieldHook: Send + Sync {
    /// Resolve computed fields for a record.
    async fn resolve(
        &self,
        record: serde_json::Value,
        fields: &[(String, String)],
    ) -> std::result::Result<serde_json::Value, String>;
}

// ============================================================================
// Vector Embedding Support
// ============================================================================

/// Mapping from a source field to a target vector field via an embedding model.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct FieldMapping {
    /// Source field name containing text or image data
    pub source: String,
    /// Target field name for the generated embedding vector
    pub target: String,
    /// Model identifier (e.g., "BAAI/bge-small-en-v1.5")
    pub model: String,
    /// Field type: "text" (default) or "image"
    #[serde(default = "default_field_type")]
    pub field_type: String,
}

fn default_field_type() -> String {
    "text".to_owned()
}

/// Hook for automatic vector embedding generation.
///
/// Implementations must be **sync** for safety across the WIT boundary.
pub trait VectorHook: Send + Sync {
    /// Embed text/image fields in a record based on mappings.
    fn vectorize_fields(
        &self,
        record: serde_json::Value,
        mappings: &[FieldMapping],
    ) -> std::result::Result<serde_json::Value, String>;

    /// Convert a text query to a vector for search.
    fn vectorize_text(&self, text: &str, model: &str) -> std::result::Result<Vec<f32>, String>;

    /// Convert raw image bytes to a vector.
    fn vectorize_image(&self, bytes: &[u8], model: &str) -> std::result::Result<Vec<f32>, String>;

    /// Check if a model is available for use.
    fn validate_model(&self, _model: &str) -> std::result::Result<(), String> {
        Ok(())
    }

    /// Eagerly load a model into memory during startup.
    fn warmup_model(&self, _model: &str, _field_type: &str) -> std::result::Result<(), String> {
        Ok(())
    }

    /// Batch-embed multiple records.
    fn vectorize_fields_batch(
        &self,
        records: Vec<serde_json::Value>,
        mappings: &[FieldMapping],
    ) -> std::result::Result<Vec<serde_json::Value>, String> {
        records
            .into_iter()
            .map(|r| self.vectorize_fields(r, mappings))
            .collect()
    }
}

// ============================================================================
// AI Inference Support
// ============================================================================

/// Metadata about an available AI model.
#[derive(Debug, Clone, serde::Serialize)]
pub struct ModelInfo {
    /// Unique identifier
    pub id: String,
    /// Human-readable model name
    pub name: String,
    /// Parameter count label ("8B", "3B")
    pub parameters: String,
    /// Quantization level ("`Q4_K_M`", "F16")
    pub quantization: String,
    /// Whether the model is currently loaded
    pub loaded: bool,
    /// Memory usage in bytes (0 if not loaded)
    pub memory_bytes: u64,
}

/// Hook for local LLM inference.
///
/// Implementations must be **sync** for safety across the WIT boundary.
pub trait AiHook: Send + Sync {
    /// Single-turn text completion.
    fn complete(&self, prompt: &str, max_tokens: u32) -> std::result::Result<String, String>;

    /// Multi-turn chat completion.
    fn chat(
        &self,
        messages: &[(&str, &str)],
        max_tokens: u32,
    ) -> std::result::Result<String, String>;

    /// Structured JSON output.
    fn complete_json(
        &self,
        prompt: &str,
        max_tokens: u32,
    ) -> std::result::Result<serde_json::Value, String>;

    /// List available models.
    fn models(&self) -> Vec<ModelInfo>;

    /// Check if a specific model is loaded.
    fn is_loaded(&self, model: &str) -> bool;
}

// ============================================================================
// VectorContext — per-app vector capabilities (set once at startup)
// ============================================================================

/// Per-app vector capabilities. Set once at startup, never changes per-request.
///
/// Stored on `BackendManager` via `OnceLock` so any code with access to the
/// backend manager can reach vector hooks without threading them through Context.
pub struct VectorContext {
    /// Vector embedding hook (from yeti-ai plugin)
    pub hook: Arc<dyn VectorHook>,
    /// Plugin `app_id` that provided the vector hook
    pub hook_ext_id: String,
    /// Vector micro-batcher for server-side batching
    pub batcher: Option<Arc<VectorBatcher>>,
    /// Pre-parsed vector field mappings
    pub mappings: Vec<FieldMapping>,
    /// Vector embedding cache backend
    pub cache: Option<Arc<dyn crate::backend::KvBackend>>,
}

impl std::fmt::Debug for VectorContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VectorContext")
            .field("hook_ext_id", &self.hook_ext_id)
            .field("mappings", &self.mappings)
            .field("has_batcher", &self.batcher.is_some())
            .field("has_cache", &self.cache.is_some())
            // `hook` (dyn VectorHook), `cache` (dyn KvBackend) elided
            .finish_non_exhaustive()
    }
}

// ============================================================================
// VectorBatcher
// ============================================================================

/// Server-side micro-batcher for vector embeddings.
#[derive(Debug)]
pub struct VectorBatcher {
    tx: tokio::sync::mpsc::UnboundedSender<(
        serde_json::Value,
        tokio::sync::oneshot::Sender<serde_json::Value>,
    )>,
}

impl VectorBatcher {
    /// Create from a channel sender.
    #[must_use]
    pub const fn from_sender(
        tx: tokio::sync::mpsc::UnboundedSender<(
            serde_json::Value,
            tokio::sync::oneshot::Sender<serde_json::Value>,
        )>,
    ) -> Self {
        Self { tx }
    }

    /// Submit a record for embedding and wait for the result.
    pub async fn embed(&self, record: serde_json::Value) -> serde_json::Value {
        let fallback = record.clone();
        let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();
        if self.tx.send((record, reply_tx)).is_err() {
            return fallback;
        }
        reply_rx.await.unwrap_or(fallback)
    }
}

// ============================================================================
// EmbeddedFile
// ============================================================================

/// A file embedded in the binary at compile time for static web serving.
#[derive(Debug)]
pub struct EmbeddedFile {
    /// Relative path (e.g., "index.html", "assets/app.js")
    pub path: &'static str,
    /// File content bytes
    pub content: &'static [u8],
}

// ============================================================================
// Context Lookup Types
// ============================================================================

/// Closure type for looking up `PubSub` managers by table name.
///
/// `Arc<dyn Fn>` (rather than `Box<dyn Fn>`) so that the enclosing
/// `Context` can be `Clone` — required for flowing the owned `Context`
/// through Tower `Service<Context>` chains (each link clones once, then
/// consumes its clone). `Arc::clone` is a single atomic op; the closure
/// body is shared, not copied.
pub type PubSubLookupFn =
    Arc<dyn Fn(&str) -> Option<Arc<crate::pubsub::PubSubManager>> + Send + Sync>;

// ============================================================================
// TableSubscriber — host-bridged table pubsub for wasm applications
// ============================================================================

/// Subscribe to row changes on a named table from a wasm application.
///
/// # Why not just `ps.subscribe(table).await` in the guest?
///
/// The guest runs in a separate wasm sandbox with its own linear memory
/// and no visibility into the host tokio runtime — there is no host
/// reactor for guest-constructed futures to register wakers against, and
/// the guest cannot hold a host `broadcast::Receiver` whose `Sender`
/// lives host-side. Long-running subscription work has to be host-driven;
/// the guest only awaits a channel the host hands it.
///
/// # How this bridge works
///
/// The app provides a `TableSubscriber` implementation via
/// `RegistrationContext::table_subscribers`. During app-loader drain,
/// the host:
///
///   1. Calls `ps.subscribe(subscriber.table()).await` — receiver is
///      **host-allocated**.
///   2. Creates an `mpsc::UnboundedChannel` pair (the mpsc is the
///      safe-across-the-boundary primitive).
///   3. `tokio::spawn`s a forwarder task on the host runtime that loops
///      `broadcast_rx.recv().await` → `mpsc_tx.send()`. Both primitives
///      live in host tokio.
///   4. `tokio::spawn`s `subscriber.run(mpsc_rx)` on the host runtime.
///      The loop body is guest-defined, but it only awaits on the
///      host-allocated `mpsc::UnboundedReceiver` — no guest-side
///      subscribe or broadcast-receive happens.
///
/// The result is symmetric with the publish path: publish goes guest →
/// host via a WIT import; subscribe goes host → guest via an mpsc
/// receiver handed in at `run()` time.
pub trait TableSubscriber: Send + Sync + 'static {
    /// Name of the table to subscribe to (matches `@table` name in schema).
    fn table(&self) -> &str;

    /// Long-running loop. `rx` carries every row change published to the
    /// table's pubsub stream. Runs on the host tokio runtime; is
    /// host-driven on behalf of the guest.
    fn run(
        self: Box<Self>,
        rx: tokio::sync::mpsc::UnboundedReceiver<crate::resource::SubscriptionMessage>,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>;
}

// ============================================================================
// Plugin trait
// ============================================================================

/// Contract between the yeti binary and service crates.
///
/// The runtime calls methods in order:
/// `is_required()` → `schemas()` → `resources()` → `on_ready()` → `on_shutdown()`
///
/// `schemas()` is called during discovery (before backends exist).
/// `resources()` is called during loading (after backends exist).
/// Each is called exactly once.
pub trait Plugin: Send + Sync + 'static {
    /// Unique identifier (e.g., "yeti-auth").
    fn id(&self) -> &'static str;
    /// Human-readable name.
    fn name(&self) -> &'static str;
    /// Semantic version.
    fn version(&self) -> &'static str {
        "0.1.0"
    }
    /// Plugin IDs this service depends on (for topological ordering).
    fn depends_on(&self) -> &[&'static str] {
        &[]
    }
    /// Whether this service should be activated given the current config.
    fn is_required(&self, ctx: &StartupContext) -> bool {
        let _ = ctx;
        true
    }
    /// Whether this is a global plugin (loaded before user apps).
    fn is_plugin(&self) -> bool {
        false
    }
    /// Whether registration failure should abort startup.
    fn is_critical(&self) -> bool {
        false
    }
    /// Embedded `Cargo.toml` content. Discovery extracts
    /// `[package.metadata.app]` (and any plugin metadata blocks) via
    /// `yeti_sdk_host::application::cargo_manifest`.
    fn config_toml(&self) -> Option<&'static str> {
        None
    }
    /// GraphQL schema strings. Called once during discovery, before backends exist.
    /// Filesystem apps return empty — their schemas come from schema.graphql files.
    fn schemas(&self) -> Vec<&'static str> {
        vec![]
    }
    /// Register resources, hooks, web files, and providers.
    /// Called once during loading, after backends exist.
    fn resources(&self, ctx: &mut RegistrationContext) -> crate::error::Result<()> {
        let _ = ctx;
        Ok(())
    }

    /// Register in-process hook services (ADR-009).
    ///
    /// Called by the app loader between [`Self::resources`] and
    /// [`Self::on_ready`]. Static plugins that want to install
    /// services into another static plugin's hook chain (e.g.
    /// yeti-auth's OAuth/Token chains, yeti-mcp's tool/resource
    /// chains) call the per-domain hook surface directly here (e.g.
    /// `yeti_mcp::McpHooks::register_list_tools_service`), once per
    /// chain they want to extend.
    ///
    /// Slotted before `on_ready` so registrations are visible by the
    /// time any request can hit a dispatcher that reads the chain.
    /// Default-empty so existing plugins compile unchanged.
    fn register_hooks(&self) -> crate::error::Result<()> {
        Ok(())
    }

    /// Post-registration setup (bootstrap data, background tasks, etc.).
    ///
    /// Receives a `Context` with `request`-less defaults (no method/headers/body).
    /// The runtime populates `backend_manager`, `root_directory`, `app_id`,
    /// and `pubsub_lookup` before calling.
    fn on_ready(&self, ctx: &crate::resource::Context) -> crate::error::Result<()> {
        let _ = ctx;
        Ok(())
    }

    /// Build and return the service's telemetry-subscriber Service, if any.
    ///
    /// Per ADR-006 this returns a `TelemetryService` (Tower
    /// `Service<TelemetryEvent>`) instead of the legacy
    /// `Box<dyn EventSubscriber>`. The runtime creates an mpsc
    /// channel, registers the sender globally so any code can
    /// `event_sender().send(event)` to fire an event, and drives a
    /// per-event call loop on the host tokio runtime:
    ///
    /// ```ignore
    /// while let Some(ev) = rx.recv().await {
    ///     if let Ok(svc) = service.ready().await {
    ///         let _ = svc.call(ev).await;
    ///     }
    /// }
    /// ```
    ///
    /// Plugin authors stack `tower::ServiceBuilder` layers
    /// (timeout, trace, concurrency-limit) onto the Service before
    /// returning it.
    ///
    /// Default: no subscriber.
    fn install_event_subscriber(
        &self,
        ctx: &crate::resource::Context,
    ) -> crate::error::Result<Option<crate::plugins::lifecycle::TelemetryService>> {
        let _ = ctx;
        Ok(None)
    }

    /// Called during graceful shutdown.
    fn on_shutdown(&self) {}
}

/// Ordered list of service instances.
pub type PluginRegistry = Vec<Box<dyn Plugin>>;

// ============================================================================
// StartupContext
// ============================================================================

/// Read-only context passed to `Plugin::is_required()`.
#[derive(Debug)]
pub struct StartupContext<'a> {
    /// All discovered table definitions across all apps.
    pub tables: &'a [crate::schema::TableDefinition],
    /// All discovered application configs (as raw JSON for introspection).
    pub app_configs: &'a [serde_json::Value],
    /// Whether the MQTT interface is enabled (`mqtt.enabled` in
    /// yeti-config.yaml). The broker requires the auth backend, so a plugin
    /// that provides auth declares itself required when this is set — even when
    /// no app otherwise needs auth.
    pub mqtt_enabled: bool,
}

impl StartupContext<'_> {
    /// Check if any application has a given top-level config key (e.g., "auth", "kafka").
    #[must_use]
    pub fn any_app_has_config(&self, key: &str) -> bool {
        self.app_configs
            .iter()
            .any(|c| c.as_object().is_some_and(|obj| obj.contains_key(key)))
    }

    /// Check if any table has a field whose type matches the given name (case-insensitive).
    #[must_use]
    pub fn any_table_has_field_directive(&self, directive: &str) -> bool {
        self.tables.iter().any(|t| {
            t.fields
                .iter()
                .any(|f| f.field_type.eq_ignore_ascii_case(directive))
        })
    }

    /// Check if any application requires authentication (has non-empty `required_roles`).
    #[must_use]
    pub fn any_app_requires_auth(&self) -> bool {
        self.app_configs.iter().any(|c| {
            c.get("requiredRoles")
                .or_else(|| c.get("required_roles"))
                .and_then(|v| v.as_array())
                .is_some_and(|arr| !arr.is_empty())
        })
    }
}

// ============================================================================
// RegistrationContext
// ============================================================================

/// Write context passed to `Plugin::register()`.
pub struct RegistrationContext {
    /// GraphQL schema strings
    pub schemas: Vec<String>,
    /// Seed data JSON strings
    pub seed_data: Vec<String>,
    /// Resource handlers
    pub resources: Vec<crate::resource::ResourceEntry>,
    /// Embedded web files for static serving
    pub web_files: Vec<EmbeddedFile>,
    /// Authentication providers
    pub auth_providers: Vec<Arc<dyn crate::auth::AuthProvider>>,
    // Per ADR-006: `auth_hooks` (legacy `AuthHook` trait objects) is deleted.
    // yeti-auth reintroduces a plugin-side registry of
    // `Service<ResolveRequest>` impls (see `plugins::auth`).
    /// Vector embedding hooks
    pub vector_hooks: Vec<Arc<dyn VectorHook>>,
    /// AI inference hooks
    pub ai_hooks: Vec<Arc<dyn AiHook>>,
    /// Computed field hooks
    pub computed_field_hooks: Vec<Arc<dyn ComputedFieldHook>>,
    /// Pipeline services — Tower-native replacement for the legacy
    /// `RequestMiddleware` Vec. Each entry is a
    /// `BoxCloneService<Context, Context, YetiError>` (see
    /// `plugins::request::ContextService`); the router runs them in
    /// registration order via `RequestPipeline::run`.
    pub request_services: Vec<crate::plugins::request::ContextService>,
    /// Table subscribers (pubsub safe across the WIT boundary via host-created mpsc bridge)
    pub table_subscribers: Vec<Box<dyn TableSubscriber>>,
    /// Background tasks to spawn after registration
    pub background_tasks: Vec<Pin<Box<dyn Future<Output = ()> + Send>>>,
    /// Application ID
    pub app_id: String,
    /// Root directory
    pub root_directory: String,
}

impl std::fmt::Debug for RegistrationContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RegistrationContext")
            .field("app_id", &self.app_id)
            .field("root_directory", &self.root_directory)
            .field("schemas", &self.schemas.len())
            .field("seed_data", &self.seed_data.len())
            .field("resources", &self.resources.len())
            .field("web_files", &self.web_files.len())
            .field("auth_providers", &self.auth_providers.len())
            .field("vector_hooks", &self.vector_hooks.len())
            .field("ai_hooks", &self.ai_hooks.len())
            .field("computed_field_hooks", &self.computed_field_hooks.len())
            .field("request_services", &self.request_services.len())
            .field("table_subscribers", &self.table_subscribers.len())
            .field("background_tasks", &self.background_tasks.len())
            .finish_non_exhaustive()
    }
}

impl RegistrationContext {
    /// Create a new empty registration context.
    #[must_use]
    pub fn new(app_id: String, root_directory: String) -> Self {
        Self {
            schemas: Vec::new(),
            seed_data: Vec::new(),
            resources: Vec::new(),
            web_files: Vec::new(),
            auth_providers: Vec::new(),
            vector_hooks: Vec::new(),
            ai_hooks: Vec::new(),
            computed_field_hooks: Vec::new(),
            request_services: Vec::new(),
            table_subscribers: Vec::new(),
            background_tasks: Vec::new(),
            app_id,
            root_directory,
        }
    }

    /// Add a GraphQL schema string.
    pub fn add_schema(&mut self, graphql: &str) {
        self.schemas.push(graphql.to_owned());
    }
    /// Add a resource handler.
    pub fn add_resource(&mut self, entry: crate::resource::ResourceEntry) {
        self.resources.push(entry);
    }
    /// Add embedded web files.
    pub fn add_web_files(&mut self, files: Vec<EmbeddedFile>) {
        self.web_files.extend(files);
    }
    /// Add an authentication provider.
    pub fn add_auth_provider(&mut self, p: Arc<dyn crate::auth::AuthProvider>) {
        self.auth_providers.push(p);
    }
    /// Add a vector embedding hook.
    pub fn add_vector_hook(&mut self, h: Arc<dyn VectorHook>) {
        self.vector_hooks.push(h);
    }
    /// Add an AI inference hook.
    pub fn add_ai_hook(&mut self, h: Arc<dyn AiHook>) {
        self.ai_hooks.push(h);
    }
    /// Register a request-pipeline service. The service runs as part of
    /// the per-request middleware chain built by yeti-router. Plugins
    /// produce a `ContextService` via `tower::service_fn` (or by
    /// implementing `tower::Service<Context>` directly) and call
    /// `BoxCloneService::new(svc)` before passing it here.
    pub fn add_request_service(&mut self, svc: crate::plugins::request::ContextService) {
        self.request_services.push(svc);
    }
    /// Add a table subscriber. The runtime will subscribe to the named
    /// table's pubsub stream on the host tokio runtime and forward each
    /// message to the subscriber's `run()` loop via a host-allocated
    /// mpsc channel. Safe to call from wasm applications.
    pub fn add_table_subscriber(&mut self, s: Box<dyn TableSubscriber>) {
        self.table_subscribers.push(s);
    }
    /// Get the root directory path.
    #[must_use]
    pub fn root_dir(&self) -> &str {
        &self.root_directory
    }
    /// Get the application ID.
    #[must_use]
    pub fn app_id(&self) -> &str {
        &self.app_id
    }
}

// now receives `&crate::resource::Context` directly. Non-request helpers
// (`root_dir`, `pubsub`, `backend`) live on Context.