submilli-engine 0.3.0

//! Outbound `@mcp/<server>` transport dispatch.
//!
//! Every `@mcp/<server>.<tool>(...)` call lowers (in codegen) to **one** host
//! import — `submilli:mcp.call(server, tool, argsJson) -> unknown`. The host
//! converts the transport's bounded JSON tree directly to guest objects; normal
//! runtime casts validate declared tool return types.
//!
//! This module provides that host fn. It runs the `mcp.<server>` capability check
//! (one per server, with the tool in the filter context) *before* dispatch (a
//! denial throws before any network bytes leave), then
//! calls the embedder-provided [`McpTransport`] — the actual JSON-RPC `tools/call`,
//! auth injection, and refresh-on-`401` live in the embedder (submilli-server),
//! keeping this crate free of its async/HTTP dependencies.

use std::future::Future;
use std::pin::Pin;

use wasmtime::{FuncType, HeapType, Linker, RefType, Val, ValType};

use crate::runtime::StoreData;
use crate::runtime::call_log::{Payload, Side, record_payload};
use crate::runtime::fuel;
use crate::runtime::host::{intrinsic_string_type, read_string_arg, register_host_fn_async};
use crate::stdlib::shared::check_security_call;
use crate::{PackageDeclaration, Param, Span, Type, ValueKind, ValueSymbol};

/// The single internal host module every MCP call dispatches through.
pub const MCP_MODULE_NAME: &str = "submilli:mcp";
pub const MCP_MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024;

mod response;
pub use response::McpResponse;

/// Why an outbound MCP `tools/call` failed. The host fn maps each to a catchable
/// script error (see [`mcp_error_to_throw`]).
#[derive(Debug)]
pub enum McpCallError {
    /// The tool ran but its response exceeded the host's bounded output buffer.
    ResponseTooLarge,
    /// Internal initialization, allocation or serialization failure; never a guest exception.
    Internal { message: &'static str },
    /// OAuth authentication failed. A provider rejection may be transient;
    /// surfaces to the script as a catchable `McpAuthExpiredError`.
    AuthExpired,
    /// The server reported a tool-level or JSON-RPC error.
    Mcp { message: String },
    /// A non-2xx HTTP response from the server.
    Upstream { status: u16, body: String },
    /// The request never completed (connection / transport failure).
    Transport(String),
    /// Refused by this host's own configuration before anything was sent: unavailable
    /// credentials, or the network policy. Reads as a transport error to the script.
    Local(String),
}

impl McpCallError {
    /// The failure as the call log keeps it: a stable kind, a message, and the HTTP
    /// status of an upstream refusal.
    fn record(&self) -> serde_json::Value {
        let (kind, message, status) = match self {
            McpCallError::ResponseTooLarge => ("response-too-large", String::new(), None),
            McpCallError::Internal { message } => ("internal", (*message).to_owned(), None),
            McpCallError::AuthExpired => ("auth-expired", String::new(), None),
            McpCallError::Mcp { message } => ("mcp", message.clone(), None),
            McpCallError::Upstream { status, body } => ("upstream", body.clone(), Some(*status)),
            McpCallError::Transport(detail) => ("transport", detail.clone(), None),
            McpCallError::Local(detail) => ("local", detail.clone(), None),
        };
        serde_json::json!({ "kind": kind, "message": message, "status": status })
    }
}

/// Completed transport work, retained even when the tool result is an error.
pub struct McpOutcome {
    pub result: Result<McpResponse, McpCallError>,
    pub received_bytes: u64,
    pub parsed_bytes: u64,
}

/// The embedder-provided outbound MCP transport: performs the JSON-RPC
/// `tools/call` (resolving auth, refreshing on `401`, reusing connections) and
/// returns a bounded JSON tree. Transport decoders must bound response bytes
/// and nesting before allocating recursive JSON values.
///
/// Object-safe boxed-future trait (rather than `async fn`) so the interpreter
/// stays free of the embedder's async stack while still `await`-ing the call from
/// within the async host function.
pub trait McpTransport: Send + Sync {
    fn call<'a>(
        &'a self,
        server: &'a str,
        tool: &'a str,
        args_json: &'a str,
    ) -> Pin<Box<dyn Future<Output = McpOutcome> + Send + 'a>>;
}

fn mcp_call_func_type(linker: &Linker<StoreData>) -> wasmtime::Result<FuncType> {
    let engine = linker.engine();
    let string_struct = intrinsic_string_type(engine)?;
    let string = ValType::Ref(RefType::new(false, HeapType::ConcreteStruct(string_struct)));
    let object = crate::runtime::intrinsic_types::build_intrinsic_types(engine)?.object;
    let unknown = ValType::Ref(RefType::new(true, HeapType::ConcreteStruct(object)));
    Ok(FuncType::new(
        engine,
        [string.clone(), string.clone(), string.clone()],
        [unknown],
    ))
}

/// Install the async `submilli:mcp.call` dispatch host fn. Reads `(server, tool,
/// argsJson)`, runs the capability check, dispatches through
/// [`StoreData::mcp_transport`], and builds guest values directly from its tree.
pub fn install_mcp_async(linker: &mut Linker<StoreData>) -> wasmtime::Result<()> {
    let ty = mcp_call_func_type(linker)?;
    register_host_fn_async(
        linker,
        MCP_MODULE_NAME,
        crate::mangle::host(MCP_MODULE_NAME, "call"),
        ty,
        /* deterministic = */ false,
        move |caller, params, results| {
            Box::pin(async move {
                let params: &[Val; 3] = params.try_into().map_err(|_| {
                    crate::runtime::host::fatal_host_error("MCP argument count is invalid")
                })?;
                let result_slot = results.first_mut().ok_or_else(|| {
                    crate::runtime::host::fatal_host_error("MCP result slot is missing")
                })?;
                let server = read_string_arg(&mut *caller, &params[0], "mcp.call (server)")?;
                let tool = read_string_arg(&mut *caller, &params[1], "mcp.call (tool)")?;
                let args_json = read_string_arg(&mut *caller, &params[2], "mcp.call (args)")?;

                // Deny before any network bytes leave. One capability per server
                // (`mcp.<server>`, known when the blueprint is written); the tool is
                // in the filter context so a policy can constrain by tool.
                let ticket = check_security_call(
                    &mut *caller,
                    &format!("mcp.{server}"),
                    serde_json::json!({ "tool": tool, "transport": "streamable_http" }),
                )?;
                record_payload(&*caller, ticket, Side::Request, || {
                    request_payload(&server, &tool, &args_json)
                });

                let transport = caller.data().mcp_transport.clone().ok_or_else(|| {
                    wasmtime::Error::msg(format!(
                        "@mcp/{server}.{tool}: MCP transport not configured"
                    ))
                })?;

                fuel::charge(&mut *caller, fuel::IO, args_json.len() as u64)?;
                let outcome = transport.call(&server, &tool, &args_json).await;
                record_payload(&*caller, ticket, Side::Response, || {
                    let payload = match &outcome.result {
                        Ok(response) => Payload::meta(serde_json::Value::Null).with_owned_body(
                            serde_json::to_vec(response.value()).unwrap_or_default(),
                        ),
                        Err(error) => Payload::meta(error.record()),
                    };
                    payload.with_size(outcome.received_bytes)
                });
                fuel::settle(&mut *caller, fuel::IO, outcome.received_bytes)?;
                fuel::settle(&mut *caller, fuel::PARSE, outcome.parsed_bytes)?;
                fuel::settle_result(caller, |caller| {
                    let result = (|| {
                        let response = outcome
                            .result
                            .map_err(|error| mcp_error_to_throw(&server, &tool, error))?;
                        let value = allocate_response(caller, &response)?;
                        *result_slot = value;
                        Ok(())
                    })();
                    result.map_err(|error| crate::runtime::host::throw_host_error(caller, error))
                })
            })
        },
    )
}

/// A call's request as the recorder keeps it.
fn request_payload<'a>(server: &str, tool: &str, args_json: &'a str) -> Payload<'a> {
    Payload::meta(serde_json::json!({ "server": server, "tool": tool }))
        .with_body(args_json.as_bytes())
}

/// The digest the call log records for a call's request: what a transport sees of it, so a
/// transport can recognize a call it saw recorded.
pub fn request_digest(server: &str, tool: &str, args_json: &str) -> String {
    request_payload(server, tool, args_json).digest()
}

fn allocate_response(
    caller: &mut wasmtime::Caller<'_, StoreData>,
    response: &McpResponse,
) -> wasmtime::Result<Val> {
    fuel::settle_result(caller, |caller| {
        fuel::settle(&mut *caller, fuel::ELEM, response.visited_nodes())?;
        let allocator = crate::runtime::json::JsonUnknownAllocator::new(caller)?;
        let value = crate::runtime::json::JsonValue::from(response.value());
        allocator.allocate(caller, &value)
    })
}

/// Codegen-facing definitions for the internal `submilli:mcp.call` host fn:
/// `call(server: string, tool: string, args: string): unknown`. Lowered as
/// `(ref $string) x3 -> (ref null $Object)` by the generic host-import path.
pub fn mcp_call_package_declaration() -> PackageDeclaration {
    let mut defs = PackageDeclaration::with_package(MCP_MODULE_NAME);
    defs.values.insert(
        "call".to_string(),
        ValueSymbol {
            name: "call".to_string(),
            mangled_name: crate::mangle::host(MCP_MODULE_NAME, "call"),
            declaration_span: Span::at(crate::FileId::MCP),
            kind: ValueKind::Function {
                generics: Vec::new(),
                params: vec![
                    Param::new("server", Type::String),
                    Param::new("tool", Type::String),
                    Param::new("args", Type::String),
                ],
                ret: Type::Unknown,
                type_predicate: None,
                doc: None,
            },
        },
    );
    defs
}

/// Map a transport failure to a catchable script `Error`. `McpAuthExpiredError` is
/// a message-prefixed `Error` (typed error classes aren't landed yet — same shape
/// as the permission-denied throw).
fn mcp_error_to_throw(server: &str, tool: &str, err: McpCallError) -> wasmtime::Error {
    let message = match err {
        McpCallError::Internal { message, .. } => {
            return crate::runtime::host::fatal_host_error(message);
        }
        McpCallError::ResponseTooLarge => {
            return crate::runtime::host::range_error(format!(
                "@mcp/{server}.{tool}: response exceeds the size, node or depth limit; request a smaller result"
            ));
        }
        McpCallError::AuthExpired => format!(
            "McpAuthExpiredError: @mcp/{server}.{tool}: OAuth authentication failed; \
             retry later or authenticate the MCP server again"
        ),
        McpCallError::Mcp { message } => {
            format!("@mcp/{server}.{tool}: {message}")
        }
        McpCallError::Upstream { status, body } => {
            format!("@mcp/{server}.{tool}: server returned HTTP {status}: {body}")
        }
        McpCallError::Transport(detail) | McpCallError::Local(detail) => {
            format!("@mcp/{server}.{tool}: transport error: {detail}")
        }
    };
    wasmtime::Error::msg(message)
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

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

    use wasmtime::{Linker, Module};

    use super::{McpOutcome, McpResponse, McpTransport};
    use crate::runtime::security::{CheckOutcome, SecurityCheck};
    use crate::runtime::{
        RuntimeConfig, StoreData, Vfs, dispatch_main_async, install_runtime_async,
        install_tenant_limits,
    };
    use crate::{
        ObjectField, PackageDeclaration, Param, Shape, Span, Type, ValueKind, ValueSymbol,
        compile_script, mangle,
    };

    /// One `@mcp/test` package: `echo(args: { msg: string }): { reply: string }`
    /// (a typed round-trip) plus `raw(): unknown` (the untyped fallback). Mirrors
    /// what the catalog mapper produces — structural object types, with their
    /// shapes registered so codegen emits the object machinery.
    fn test_mcp_package() -> PackageDeclaration {
        let pkg = "@mcp/test";
        let mut defs = PackageDeclaration::with_package(pkg);
        defs.mcp_server = Some("test".to_string());

        let mut arg_fields = BTreeMap::new();
        arg_fields.insert("msg".to_string(), ObjectField::required(Type::String));
        let arg_obj = Type::Object {
            index: None,
            fields: arg_fields,
        };
        let mut ret_fields = BTreeMap::new();
        ret_fields.insert("reply".to_string(), ObjectField::required(Type::String));
        let ret_obj = Type::Object {
            index: None,
            fields: ret_fields,
        };
        let mut opt_fields = BTreeMap::new();
        opt_fields.insert("user".to_string(), ObjectField::optional(Type::String));
        let opt_arg_obj = Type::Object {
            index: None,
            fields: opt_fields,
        };
        defs.shapes.push(Shape::from_type(&arg_obj).unwrap());
        defs.shapes.push(Shape::from_type(&ret_obj).unwrap());
        defs.shapes.push(Shape::from_type(&opt_arg_obj).unwrap());

        let mut insert_fn = |name: &str, params: Vec<Param>, ret: Type| {
            defs.values.insert(
                name.to_string(),
                ValueSymbol {
                    name: name.to_string(),
                    mangled_name: mangle::host(pkg, name),
                    declaration_span: Span::at(crate::FileId::MCP),
                    kind: ValueKind::Function {
                        generics: Vec::new(),
                        params,
                        ret,
                        type_predicate: None,
                        doc: None,
                    },
                },
            );
        };
        insert_fn("echo", vec![Param::new("args", arg_obj)], ret_obj);
        insert_fn("raw", Vec::new(), Type::Unknown);
        // All-optional args: callable with zero args, the default fills in `{}`.
        insert_fn(
            "getTeams",
            vec![Param::with_default(
                "args",
                opt_arg_obj,
                crate::DefaultValue::EmptyObject,
            )],
            Type::Unknown,
        );
        defs
    }

    /// Records every call and returns a canned result.
    struct MockTransport {
        result: serde_json::Value,
        calls: Arc<Mutex<Vec<(String, String, String)>>>,
    }

    impl McpTransport for MockTransport {
        fn call<'a>(
            &'a self,
            server: &'a str,
            tool: &'a str,
            args_json: &'a str,
        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = McpOutcome> + Send + 'a>> {
            self.calls.lock().unwrap().push((
                server.to_string(),
                tool.to_string(),
                args_json.to_string(),
            ));
            let mut result = self.result.clone();
            Box::pin(async move {
                McpOutcome {
                    result: McpResponse::take(&mut result),
                    received_bytes: 0,
                    parsed_bytes: 0,
                }
            })
        }
    }

    struct DenyAll;
    impl SecurityCheck for DenyAll {
        fn check(&self, _: &str, capability: &str, _: &serde_json::Value) -> CheckOutcome {
            CheckOutcome::Deny {
                rule: None,
                reason: format!("denied {capability}"),
            }
        }
    }

    async fn run(
        source: &str,
        result: serde_json::Value,
        security: Option<Arc<dyn SecurityCheck>>,
    ) -> (
        wasmtime::Result<Option<String>>,
        Arc<Mutex<Vec<(String, String, String)>>>,
    ) {
        let pkg = test_mcp_package();
        let compiled =
            compile_script(source, "<test>", crate::FileId(0), &[], &[&pkg]).expect("compiles");

        let cfg = RuntimeConfig::default();
        let engine = cfg.engine().unwrap();
        let calls = Arc::new(Mutex::new(Vec::new()));
        let mut data = StoreData::with_vfs(Vfs::tempdir().unwrap());
        data.install_type_info(compiled.type_info.clone());
        data.mcp_transport = Some(Arc::new(MockTransport {
            result,
            calls: calls.clone(),
        }));
        if let Some(check) = security {
            data.security_check = check;
        }
        let mut store = cfg.store_async(&engine, data).unwrap();
        install_tenant_limits(&mut store);
        let module = Module::new(&engine, &compiled.wasm).unwrap();
        let mut linker = Linker::<StoreData>::new(&engine);
        install_runtime_async(&mut linker, &mut store)
            .await
            .unwrap();
        let inst = linker.instantiate_async(&mut store, &module).await.unwrap();
        let out = dispatch_main_async(&mut store, &inst).await;
        (out, calls)
    }

    struct FailingTransport {
        internal: bool,
    }

    impl McpTransport for FailingTransport {
        fn call<'a>(
            &'a self,
            _: &'a str,
            _: &'a str,
            _: &'a str,
        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = McpOutcome> + Send + 'a>> {
            Box::pin(async move {
                McpOutcome {
                    result: Err(if self.internal {
                        super::McpCallError::Internal {
                            message: "MCP HTTP client initialization failed",
                        }
                    } else {
                        super::McpCallError::Transport("remote connection failed".into())
                    }),
                    received_bytes: 0,
                    parsed_bytes: 0,
                }
            })
        }
    }

    #[tokio::test]
    async fn internal_client_failure_bypasses_guest_catch_and_keeps_runtime_usable() {
        let package = test_mcp_package();
        let source = r#"import test from "@mcp/test"; function main(): number { try { test.raw(); } catch (e) { return 7; } return 9; }"#;
        let compiled =
            compile_script(source, "<test>", crate::FileId(0), &[], &[&package]).unwrap();
        let healthy = compile_script(
            "function main(): number { return 42; }",
            "<healthy>",
            crate::FileId(0),
            &[],
            &[],
        )
        .unwrap();
        let cfg = RuntimeConfig::default();
        let engine = cfg.engine().unwrap();
        let mut data = StoreData::with_vfs(Vfs::none());
        data.install_type_info(compiled.type_info.clone());
        data.mcp_transport = Some(Arc::new(FailingTransport { internal: true }));
        let mut store = cfg.store_async(&engine, data).unwrap();
        install_tenant_limits(&mut store);
        let mut linker = Linker::<StoreData>::new(&engine);
        install_runtime_async(&mut linker, &mut store)
            .await
            .unwrap();
        let module = Module::new(&engine, &compiled.wasm).unwrap();
        let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
        let error = dispatch_main_async(&mut store, &instance)
            .await
            .unwrap_err();
        assert!(
            error.is::<crate::runtime::host::FatalHostError>(),
            "{error:?}"
        );
        store.data_mut().blocking_work.finish().await.unwrap();
        store.data_mut().mcp_transport = Some(Arc::new(FailingTransport { internal: false }));
        assert_eq!(
            dispatch_main_async(&mut store, &instance).await.unwrap(),
            Some("7".into())
        );
        store
            .data_mut()
            .install_type_info(healthy.type_info.clone());
        let module = Module::new(&engine, &healthy.wasm).unwrap();
        let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
        assert_eq!(
            dispatch_main_async(&mut store, &instance).await.unwrap(),
            Some("42".into())
        );
    }

    #[tokio::test]
    async fn completed_response_is_fully_allocated_before_exhausted_fuel_stops_execution() {
        use wasmtime::{Func, FuncType, Val, ValType};
        let config = RuntimeConfig::default();
        let engine = config.engine().unwrap();
        let mut store = config
            .store_async(&engine, StoreData::with_vfs(Vfs::none()))
            .unwrap();
        let mut linker = Linker::new(&engine);
        install_runtime_async(&mut linker, &mut store)
            .await
            .unwrap();
        let complete = Func::new(
            &mut store,
            FuncType::new(&engine, [ValType::I32], [ValType::I32, ValType::F64]),
            |mut caller, params, results| {
                let count = params[0].i32().unwrap();
                let mut value = serde_json::json!((0..count).collect::<Vec<_>>());
                let response = McpResponse::take(&mut value).unwrap();
                let value = super::allocate_response(&mut caller, &response)?;
                let storage =
                    crate::runtime::array_storage::ArrayStorage::read(&mut caller, &value)?;
                results[0] = Val::I32(storage.len as i32);
                let last = storage.get(&mut caller, (count - 1) as u32)?;
                let Val::AnyRef(Some(last)) = last else {
                    panic!("missing last value")
                };
                results[1] = last.unwrap_struct(&mut caller)?.field(&mut caller, 1)?;
                Ok(())
            },
        );
        let mut result = [Val::I32(0), Val::F64(0)];
        for count in [128, 256] {
            store.set_fuel(100_000).unwrap();
            let before = store.data().host_fuel;
            complete
                .call_async(&mut store, &[Val::I32(count)], &mut result)
                .await
                .unwrap();
            let actual = store.data().host_fuel - before;
            assert_eq!(
                actual,
                2 * crate::runtime::fuel::ELEM.cost(count as u64 + 1)
            );
            // The former path additionally serialized and parsed all result bytes.
            assert_eq!(result[0].i32(), Some(count));
        }
        store.set_fuel(1).unwrap();
        complete
            .call_async(&mut store, &[Val::I32(256)], &mut result)
            .await
            .unwrap();
        assert_eq!(result[0].i32(), Some(256));
        assert_eq!(result[1].f64(), Some(255.0));
        assert_eq!(store.get_fuel().unwrap(), 0);
        assert!(!store.data().settling_host_result);
        assert!(crate::runtime::fuel::charge_call(&mut store).is_err());
    }

    #[tokio::test]
    async fn typed_tool_round_trips_and_serializes_args() {
        let (out, calls) = run(
            r#"import test from "@mcp/test";
               function main(): void {
                 const r = test.echo({ msg: "hi" });
                 assert(r.reply === "yo", "reply");
               }"#,
            serde_json::json!({ "reply": "yo" }),
            None,
        )
        .await;
        out.expect("typed MCP call must succeed");
        let calls = calls.lock().unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].0, "test");
        assert_eq!(calls[0].1, "echo");
        assert!(
            calls[0].2.contains("\"msg\":\"hi\""),
            "args serialized to JSON: {}",
            calls[0].2
        );
    }

    #[tokio::test]
    async fn all_optional_tool_callable_with_zero_args_serializes_empty_object() {
        // The `args` param has a `{}` default, so the zero-arg call type-checks; the
        // synthesized empty object serializes to `{}` (absent optional field skipped).
        let (out, calls) = run(
            r#"import test from "@mcp/test";
               function main(): string {
                 return test.getTeams() as string;
               }"#,
            serde_json::json!("[]"),
            None,
        )
        .await;
        out.expect("zero-arg MCP call must succeed");
        let calls = calls.lock().unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].1, "getTeams");
        assert_eq!(calls[0].2, "{}", "empty optional args serialize to {{}}");
    }

    #[tokio::test]
    async fn named_import_also_routes_to_mcp_call() {
        // `import { echo } from "@mcp/test"; echo(...)` resolves through the
        // bare-identifier call path, not namespace member access — it must still
        // become an McpCall (else codegen panics on an unknown top-level fn).
        let (out, calls) = run(
            r#"import { echo } from "@mcp/test";
               function main(): void {
                 const r = echo({ msg: "hi" });
                 assert(r.reply === "yo", "reply");
               }"#,
            serde_json::json!({ "reply": "yo" }),
            None,
        )
        .await;
        out.expect("named-import MCP call must succeed");
        assert_eq!(calls.lock().unwrap()[0].1, "echo");
    }

    #[tokio::test]
    async fn untyped_tool_returns_unknown_json_value() {
        let (out, _) = run(
            r#"import test from "@mcp/test";
               interface RawResult { k: number; }
               function main(): void {
                 const r = test.raw() as RawResult;
                 assert(r.k === 1, "raw");
               }"#,
            serde_json::json!({ "k": 1 }),
            None,
        )
        .await;
        out.expect("untyped MCP call must succeed");
    }

    #[tokio::test]
    async fn zero_arg_tool_sends_empty_object() {
        let (_out, calls) = run(
            r#"import test from "@mcp/test";
               function main(): void { test.raw(); }"#,
            serde_json::json!({ "k": 1 }),
            None,
        )
        .await;
        assert_eq!(calls.lock().unwrap()[0].2, "{}");
    }

    #[tokio::test]
    async fn denied_call_throws_before_dispatch() {
        let (out, calls) = run(
            r#"import test from "@mcp/test";
               function main(): void { test.echo({ msg: "hi" }); }"#,
            serde_json::json!({ "reply": "yo" }),
            Some(Arc::new(DenyAll)),
        )
        .await;
        let err = out.expect_err("a denied capability must throw");
        assert!(
            err.to_string().contains("permission denied"),
            "unexpected error: {err}"
        );
        assert!(
            calls.lock().unwrap().is_empty(),
            "transport must not be reached when denied"
        );
    }
}