PiCloud/crates/executor-core/src/module_resolver.rs

//! `PicloudModuleResolver` — the v1.1.3 per-app Rhai module resolver.
//!
//! Replaces `DummyModuleResolver` in `Engine::build_engine`. Constructed
//! fresh per `Engine::execute` call: holds an `Arc<SdkCallCx>` so every
//! `import "<name>"` request resolves against the calling app
//! (`cx.app_id`). The script-side `name` argument carries no `app_id`
//! — that's the load-bearing cross-app isolation property.
//!
//! Three runtime invariants are enforced:
//!
//! 1. **Cross-app isolation** — `ModuleSource::lookup` is called with
//!    `&cx`; the Postgres impl scopes by `cx.app_id` (never by a
//!    script-passed argument).
//! 2. **Cycle detection** — an in-progress-imports stack rejects
//!    `A → B → A` with `ErrorInModule(... circular import detected ...)`.
//! 3. **Depth limit** — guards against deep but acyclic chains
//!    (default 8, override via `PICLOUD_MODULE_IMPORT_DEPTH_MAX`).
//!
//! Compiled modules are cached per `(app_id, name)` and invalidated by
//! `updated_at` change — no explicit pub/sub. The cache is owned by
//! `Engine` and shared across calls; only the resolver state (stack,
//! depth) is per-call.

use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};

use chrono::{DateTime, Utc};
use lru::LruCache;
use picloud_shared::{AppId, ModuleSource, ModuleSourceError, SdkCallCx, ValidatedScript};
use rhai::module_resolvers::ModuleResolver;
use rhai::{Engine as RhaiEngine, EvalAltResult, Module, Position, Shared, AST};

/// Local alias for `rhai::Shared<rhai::Module>` (rhai's `SharedRhaiModule`
/// type alias is `pub(crate)`). Resolves to `Arc<Module>` under the
/// `sync` feature that the workspace pins.
type SharedRhaiModule = Shared<Module>;

/// Cache key: `(app_id, module name)`. v1.1.3 enforces module names as
/// a conservative identifier shape (migration 0015 `scripts_module_name_shape`
/// CHECK) so the `String` here is bounded by ~64 bytes.
pub type ModuleCacheKey = (AppId, String);

/// Cache value: the freshness comparator + the compiled module Rhai
/// hands to importing scripts. Cloning the `Shared<Module>` is an Arc bump.
#[derive(Clone)]
pub struct CachedModule {
    pub updated_at: DateTime<Utc>,
    pub module: Shared<Module>,
}

/// Bounded LRU cache shared across all `Engine::execute` calls. Construct
/// once at process startup; the resolver holds an Arc into it.
pub type ModuleCache = Mutex<LruCache<ModuleCacheKey, CachedModule>>;

#[must_use]
pub fn new_module_cache(capacity: usize) -> Arc<ModuleCache> {
    // capacity 0 is nonsensical for an LRU; clamp up to 1 so the cache
    // is at least usable (callers control this via env var, and 0 means
    // "I disabled caching" — but disabling caching by accident would
    // recompile every module every call, which is a worse UX than
    // capping at 1).
    let cap = NonZeroUsize::new(capacity.max(1)).expect("max(1) is non-zero");
    Arc::new(Mutex::new(LruCache::new(cap)))
}

/// The v1.1.3 module resolver. One per `Engine::execute` call.
pub struct PicloudModuleResolver {
    /// Backend the resolver consults for `(app_id, name)`. The bridge
    /// runs Rhai's sync `resolve()` and the async `lookup()` together
    /// via `tokio::runtime::Handle::block_on(...)` — safe because
    /// `LocalExecutorClient` runs `Engine::execute` inside
    /// `spawn_blocking`, which puts us on a Tokio blocking thread
    /// that still carries a `Handle`.
    source: Arc<dyn ModuleSource>,

    /// Calling context. `cx.app_id` is the cross-app isolation
    /// boundary; the resolver passes `&cx` to every `ModuleSource`
    /// call so the backend can scope its queries.
    cx: Arc<SdkCallCx>,

    /// Compiled-module cache. Shared across executions; invalidated
    /// per-entry on `updated_at` mismatch (no explicit pub/sub).
    cache: Arc<ModuleCache>,

    /// In-progress imports stack — pushed before a `lookup`+compile,
    /// popped after. A hit on this stack while resolving means the
    /// graph contains a cycle.
    in_progress: Mutex<Vec<String>>,

    /// Current import depth. Independent of the cycle check (cycles
    /// might be short; deep acyclic graphs might fit under the cap
    /// but still warrant a guard).
    depth: Mutex<u32>,

    /// Hard ceiling on import depth. Defaults to 8; env-overridable
    /// via `PICLOUD_MODULE_IMPORT_DEPTH_MAX`. Read from `Limits` at
    /// resolver construction.
    depth_limit: u32,
}

impl PicloudModuleResolver {
    #[must_use]
    pub fn new(
        source: Arc<dyn ModuleSource>,
        cx: Arc<SdkCallCx>,
        cache: Arc<ModuleCache>,
        depth_limit: u32,
    ) -> Self {
        Self {
            source,
            cx,
            cache,
            in_progress: Mutex::new(Vec::new()),
            depth: Mutex::new(0),
            depth_limit,
        }
    }

    /// Validate `ast` as a module body: only top-level `fn` decls,
    /// `const` decls, and `import` statements are allowed. Top-level
    /// expressions (which would execute on import — a footgun for
    /// cache semantics) are rejected.
    ///
    /// `fn` declarations live in a separate slot on the AST and are
    /// not in `statements()`, so the only allowed `Stmt` variants we
    /// expect to see at top level are `Var` (when `CONSTANT` flag is
    /// set) and `Import`. Anything else triggers a `ModuleShape` error.
    fn check_module_shape(ast: &AST, name: &str) -> Result<(), String> {
        use rhai::ASTFlags;
        for stmt in ast.statements() {
            match stmt {
                rhai::Stmt::Var(_, opts, _) if opts.intersects(ASTFlags::CONSTANT) => {}
                rhai::Stmt::Import(..) => {}
                rhai::Stmt::Noop(..) => {}
                other => {
                    return Err(format!(
                        "module {name:?}: top-level {} is not allowed; \
                         modules may only contain fn declarations, \
                         const declarations, and import statements",
                        stmt_kind_label(other),
                    ));
                }
            }
        }
        Ok(())
    }

    /// Walk a compiled AST and collect the literal-path `import "<name>"`
    /// declarations. Dynamic imports (e.g. `import some_var as y;`) are
    /// skipped because the dep-graph can only track names known at
    /// compile time. Exposed via [`extract_imports`] so the manager's
    /// admin endpoints can populate the `script_imports` table from
    /// the same logic the resolver uses.
    fn extract_imports_inner(ast: &AST) -> Vec<String> {
        let mut out = Vec::new();
        for stmt in ast.statements() {
            if let rhai::Stmt::Import(boxed, _) = stmt {
                let (path_expr, _alias) = boxed.as_ref();
                if let rhai::Expr::StringConstant(s, _) = path_expr {
                    out.push(s.to_string());
                }
            }
        }
        out
    }
}

/// Compile-and-validate a candidate module body. Public so the
/// `Engine::validate_module` impl in `engine.rs` can call into it
/// without duplicating the shape check.
pub fn compile_module_ast(engine: &RhaiEngine, source: &str) -> Result<AST, String> {
    let ast = engine.compile(source).map_err(|e| e.to_string())?;
    PicloudModuleResolver::check_module_shape(&ast, "<source>")?;
    Ok(ast)
}

/// Parse `source` as an endpoint script (no module-shape check) and
/// return its declared literal-path imports. Used by
/// `Engine::validate` to populate `ValidatedScript::imports` so the
/// repo can write dep-graph edges.
pub fn extract_imports(engine: &RhaiEngine, source: &str) -> Result<ValidatedScript, String> {
    let ast = engine.compile(source).map_err(|e| e.to_string())?;
    Ok(ValidatedScript {
        imports: PicloudModuleResolver::extract_imports_inner(&ast),
    })
}

/// Parse `source` as a module script: enforce shape, then extract
/// imports. Used by `Engine::validate_module`.
pub fn validate_module_source(
    engine: &RhaiEngine,
    source: &str,
) -> Result<ValidatedScript, String> {
    let ast = compile_module_ast(engine, source)?;
    Ok(ValidatedScript {
        imports: PicloudModuleResolver::extract_imports_inner(&ast),
    })
}

fn stmt_kind_label(stmt: &rhai::Stmt) -> &'static str {
    use rhai::ASTFlags;
    match stmt {
        rhai::Stmt::Var(_, opts, _) if opts.intersects(ASTFlags::CONSTANT) => "const declaration",
        rhai::Stmt::Var(..) => "let declaration",
        rhai::Stmt::Expr(..) => "expression",
        rhai::Stmt::FnCall(..) => "function call",
        rhai::Stmt::If(..) => "if statement",
        rhai::Stmt::Switch(..) => "switch statement",
        rhai::Stmt::While(..) => "while/loop statement",
        rhai::Stmt::Do(..) => "do statement",
        rhai::Stmt::For(..) => "for statement",
        rhai::Stmt::Assignment(..) => "assignment",
        rhai::Stmt::Block(..) => "block",
        rhai::Stmt::TryCatch(..) => "try/catch",
        rhai::Stmt::Return(..) => "return/throw statement",
        rhai::Stmt::BreakLoop(..) => "break/continue",
        rhai::Stmt::Import(..) => "import statement",
        rhai::Stmt::Export(..) => "export statement",
        _ => "statement",
    }
}

impl ModuleResolver for PicloudModuleResolver {
    fn resolve(
        &self,
        engine: &RhaiEngine,
        _source: Option<&str>,
        path: &str,
        pos: Position,
    ) -> Result<SharedRhaiModule, Box<EvalAltResult>> {
        // RAII guard wraps both the depth counter and the import-stack
        // push so that any early return (cycle / depth-exceeded / DB
        // error / compile error / panic) leaves both consistent for
        // any subsequent resolve() call on this resolver instance.
        struct StackGuard<'r> {
            stack: &'r Mutex<Vec<String>>,
            depth: &'r Mutex<u32>,
            armed: bool,
        }
        impl<'r> Drop for StackGuard<'r> {
            fn drop(&mut self) {
                if !self.armed {
                    return;
                }
                if let Ok(mut s) = self.stack.lock() {
                    s.pop();
                }
                if let Ok(mut d) = self.depth.lock() {
                    *d = d.saturating_sub(1);
                }
            }
        }

        // Read-only check + atomic push under one lock pair, so a
        // sibling resolve() call on a shared resolver instance can't
        // race in between. (We don't expect parallel calls on the same
        // resolver — Rhai evaluates a single AST on one thread — but
        // grouping the operations is cheaper than reasoning about the
        // future.)
        {
            let mut depth = self.depth.lock().expect("module depth lock poisoned");
            if *depth >= self.depth_limit {
                return Err(Box::new(EvalAltResult::ErrorInModule(
                    path.to_string(),
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!(
                            "import depth limit ({}) exceeded while resolving {path:?}",
                            self.depth_limit
                        )
                        .into(),
                        pos,
                    )),
                    pos,
                )));
            }
            let mut stack = self
                .in_progress
                .lock()
                .expect("module in_progress lock poisoned");
            if stack.iter().any(|p| p == path) {
                let mut chain = stack.clone();
                chain.push(path.to_string());
                return Err(Box::new(EvalAltResult::ErrorInModule(
                    path.to_string(),
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("circular import detected: {}", chain.join(" -> ")).into(),
                        pos,
                    )),
                    pos,
                )));
            }
            stack.push(path.to_string());
            *depth += 1;
        }
        let _guard = StackGuard {
            stack: &self.in_progress,
            depth: &self.depth,
            armed: true,
        };

        // Bridge to async. The resolver always runs on a `spawn_blocking`
        // thread (see LocalExecutorClient in orchestrator-core), which
        // still carries a Tokio handle. `try_current` makes the failure
        // mode explicit when callers wire up an `Engine` from a non-
        // Tokio context (typically a test harness).
        let handle = tokio::runtime::Handle::try_current().map_err(|_| {
            Box::new(EvalAltResult::ErrorInModule(
                path.to_string(),
                Box::new(EvalAltResult::ErrorRuntime(
                    "module resolver invoked outside a Tokio runtime; \
                     wrap Engine::execute in tokio::task::spawn_blocking"
                        .into(),
                    pos,
                )),
                pos,
            ))
        })?;

        let lookup_result: Result<Option<picloud_shared::ModuleScript>, ModuleSourceError> =
            handle.block_on(self.source.lookup(&self.cx, path));

        let module_row = match lookup_result {
            Ok(Some(m)) => m,
            Ok(None) => {
                return Err(Box::new(EvalAltResult::ErrorModuleNotFound(
                    path.to_string(),
                    pos,
                )));
            }
            Err(e) => {
                return Err(Box::new(EvalAltResult::ErrorInModule(
                    path.to_string(),
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("module backend error: {e}").into(),
                        pos,
                    )),
                    pos,
                )));
            }
        };

        // Cache lookup: hit only if both key matches AND updated_at
        // matches (cache is invalidated lazily on version change).
        let cache_key = (self.cx.app_id, path.to_string());
        {
            let mut cache = self.cache.lock().expect("module cache lock poisoned");
            if let Some(cached) = cache.get(&cache_key) {
                if cached.updated_at == module_row.updated_at {
                    tracing::debug!(
                        target = "picloud::modules::cache",
                        app_id = %self.cx.app_id,
                        module = path,
                        "cache hit"
                    );
                    return Ok(cached.module.clone());
                }
                tracing::debug!(
                    target = "picloud::modules::cache",
                    app_id = %self.cx.app_id,
                    module = path,
                    "cache stale; recompiling"
                );
            } else {
                tracing::debug!(
                    target = "picloud::modules::cache",
                    app_id = %self.cx.app_id,
                    module = path,
                    "cache miss"
                );
            }
        }

        // Compile + module-shape validation. Module sources MAY have
        // already been gated at create-time (admin endpoint runs
        // `validate_module`), but we revalidate here to catch DB-direct
        // inserts that bypass the API surface.
        let ast = engine.compile(&module_row.source).map_err(|e| {
            // Wrap as an ErrorRuntime to preserve the parse message
            // text without trying to reconstruct rhai's internal
            // ParseErrorType variant (which would require matching on
            // its full variant set).
            Box::new(EvalAltResult::ErrorInModule(
                path.to_string(),
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("module {path:?} parse error: {e}").into(),
                    e.position(),
                )),
                pos,
            ))
        })?;

        if let Err(msg) = Self::check_module_shape(&ast, path) {
            return Err(Box::new(EvalAltResult::ErrorInModule(
                path.to_string(),
                Box::new(EvalAltResult::ErrorRuntime(msg.into(), pos)),
                pos,
            )));
        }

        // Rhai's eval_ast_as_new compiles the AST's body + functions
        // into a Module that the importing script consumes via
        // `path::fn(...)` calls. Recursive imports inside this module
        // are resolved through the same `engine.set_module_resolver`
        // (which is THIS resolver), so cycle/depth tracking carries
        // through naturally.
        let module = Module::eval_ast_as_new(rhai::Scope::new(), &ast, engine).map_err(|e| {
            Box::new(EvalAltResult::ErrorInModule(
                path.to_string(),
                e,
                pos,
            ))
        })?;
        let shared: SharedRhaiModule = module.into();

        // Insert (possibly evicting via LRU). Subsequent imports of
        // the same module under the same updated_at hit the cache.
        {
            let mut cache = self.cache.lock().expect("module cache lock poisoned");
            cache.put(
                cache_key,
                CachedModule {
                    updated_at: module_row.updated_at,
                    module: shared.clone(),
                },
            );
        }

        Ok(shared)
    }
}