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

//! `PicloudModuleResolver` — the Rhai module resolver (v1.1.3; made
//! origin-aware / lexical in v1.2 Phase 4b).
//!
//! Replaces `DummyModuleResolver` in `Engine::build_engine`. Constructed
//! fresh per `Engine::execute` call.
//!
//! **Lexical resolution (§5.5).** An `import "<name>"` resolves against the
//! module set visible at the **importing script's own defining node**,
//! walking up its group chain — not the inheriting app's effective view.
//! The importing node is read from Rhai's `_source` (the importing AST's
//! source tag, which the resolver sets to each module's owner) and falls
//! back to `default_origin` (the entry script's defining node) for the
//! top-level AST. So an inherited group script's imports seal to the group
//! (a leaf can't shadow them — the trust boundary), while every SDK call
//! *inside* a module still scopes to the inheriting app via `cx.app_id`.
//!
//! Three runtime invariants are enforced:
//!
//! 1. **Lexical isolation** — `ModuleSource::resolve` is keyed by the
//!    importing node's owner (a trusted dispatch-derived value), never 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 by resolved `ScriptId` (a compiled module
//! is lexically self-contained) 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::{
    ModuleSource, ModuleSourceError, ScriptId, ScriptOwner, 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: the resolved module's `ScriptId` (Phase 4b). A compiled
/// module is lexically self-contained (its nested imports resolve from
/// its own defining node, baked in at compile time), so the *body* is a
/// pure function of `(script_id, updated_at)` — independent of which
/// script imported it. Keying by id also dedupes a shared group module
/// across every inheriting app, and keeps cross-app modules of the same
/// name distinct (distinct ids).
pub type ModuleCacheKey = ScriptId;

/// Encode a defining node as an AST `source` tag (`app:<uuid>` /
/// `group:<uuid>`). Set on a resolved module's AST so its own nested
/// `import`s carry *its* node as Rhai's `_source` — the lexical chaining
/// that seals each module's imports to where it was authored (§5.5).
fn encode_origin(owner: ScriptOwner) -> String {
    match owner {
        ScriptOwner::App(a) => format!("app:{a}"),
        ScriptOwner::Group(g) => format!("group:{g}"),
    }
}

/// Inverse of [`encode_origin`]. `None` for an unrecognised tag (e.g. a
/// source set by something other than this resolver) — the caller then
/// falls back to the resolver's `default_origin`.
fn decode_origin(s: &str) -> Option<ScriptOwner> {
    let (kind, id) = s.split_once(':')?;
    let uuid = uuid::Uuid::parse_str(id).ok()?;
    match kind {
        "app" => Some(ScriptOwner::App(uuid.into())),
        "group" => Some(ScriptOwner::Group(uuid.into())),
        _ => None,
    }
}

/// 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 to resolve a module name against an
    /// origin node's chain. The bridge runs Rhai's sync `resolve()` and the
    /// async `ModuleSource::resolve()` 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 execution boundary every SDK
    /// call inside a module scopes to; the resolver keeps it for diagnostic
    /// logging. (Module *resolution* is keyed by the importing node's owner,
    /// not `cx.app_id` — see the module docs.)
    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,

    /// Lexical origin for a top-level `import` — the defining node of the
    /// script being executed (Phase 4b, §5.5). Used when Rhai gives no
    /// `_source` (the entry AST). Nested imports inside a resolved module
    /// override this via the module AST's source (C3).
    default_origin: ScriptOwner,
}

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

    /// 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 {
    #[allow(clippy::too_many_lines)]
    fn resolve(
        &self,
        engine: &RhaiEngine,
        importer_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 Drop for StackGuard<'_> {
            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 typically runs on a
        // `spawn_blocking` thread (see LocalExecutorClient in
        // orchestrator-core), but tests may invoke `Engine::execute`
        // directly from a multi-threaded Tokio task. `try_current` +
        // `block_in_place` covers both — on a blocking thread it's a
        // no-op, on a worker thread it tells the runtime to relocate
        // other tasks. `current_thread` runtimes still panic; non-
        // Tokio contexts surface a clean Runtime error.
        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,
            ))
        })?;

        // Lexical resolution (§5.5): resolve `path` against the *importing
        // node's* chain. Rhai gives `_source` = the importing AST's source
        // tag — set to the module's defining node for a nested import, or
        // unset (`None`) for the top-level entry script, where we fall back
        // to `default_origin` (the executing script's node). An app-rooted
        // walk reaches ancestor group modules; a group-rooted walk is sealed
        // from app modules below it (the trust boundary).
        let origin = importer_source
            .and_then(decode_origin)
            .unwrap_or(self.default_origin);
        // §5.5 policy resolution: the source decides lexical (seal to `origin`)
        // vs dynamic (an extension point resolves against the inheriting app,
        // `cx.app_id`). The result is already the concrete module to bind, or a
        // terminal NoProvider/NotFound.
        let lookup_result: Result<picloud_shared::ModuleResolution, ModuleSourceError> =
            tokio::task::block_in_place(|| {
                handle.block_on(self.source.resolve_policy(origin, self.cx.app_id, path))
            });

        let module_row = match lookup_result {
            Ok(picloud_shared::ModuleResolution::Module(m)) => m,
            Ok(picloud_shared::ModuleResolution::NotFound) => {
                return Err(Box::new(EvalAltResult::ErrorModuleNotFound(
                    path.to_string(),
                    pos,
                )));
            }
            Ok(picloud_shared::ModuleResolution::NoProvider) => {
                // §5.5: an extension point with no provider for this app is a
                // hard failure (the app must supply the module or a default
                // body must exist up-chain). Distinct, actionable message.
                return Err(Box::new(EvalAltResult::ErrorInModule(
                    path.to_string(),
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!(
                            "extension point {path:?} has no provider for this app — \
                             define a module named {path:?} or a default body up-chain"
                        )
                        .into(),
                        pos,
                    )),
                    pos,
                )));
            }
            Err(e) => {
                // v1.1.4 §10a: redact the backend error before it
                // reaches a script. In public-HTTP context (principal:
                // None) the verbatim message (e.g. "connection refused")
                // leaks internal infrastructure shape. Log the original
                // at error level for operators; surface a stable generic.
                tracing::error!(
                    target = "picloud::modules",
                    app_id = %self.cx.app_id,
                    module = path,
                    error = %e,
                    "module backend error"
                );
                return Err(Box::new(EvalAltResult::ErrorInModule(
                    path.to_string(),
                    Box::new(EvalAltResult::ErrorRuntime(
                        "module backend unavailable; check server logs".into(),
                        pos,
                    )),
                    pos,
                )));
            }
        };

        // Cache lookup: hit only if both key matches AND updated_at
        // matches (cache is invalidated lazily on version change). Keyed by
        // the resolved module's id (lexically self-contained — see
        // `ModuleCacheKey`).
        let cache_key = module_row.script_id;
        {
            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 mut 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,
            ))
        })?;

        // Seal this module's own imports to its defining node (§5.5): tag
        // the AST source with the resolved module's owner so that when
        // `eval_ast_as_new` evaluates its nested `import`s, Rhai hands us
        // that tag as `_source` and we resolve them lexically from *here*,
        // not from the original importer. Fall back to the lookup origin
        // for a corrupt owner row (never reached under the DB CHECK).
        let module_owner = module_row.owner().unwrap_or(origin);
        ast.set_source(encode_origin(module_owner));

        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)
    }
}