submilli_engine/runtime/
tools.rs

1//! The embedder-supplied seam for `@tools`: a program calls one of the
2//! application's own tools, and the harness that declared it answers.
3//!
4//! Every `@tools.<tool>(...)` call lowers (in codegen) to **one** host import,
5//! `submilli:tools.call(tool, argsJson) -> unknown`, exactly as `@mcp/<server>`
6//! calls lower to `submilli:mcp.call`. The host fn runs the `tools.call`
7//! capability check before anything reaches the harness, then awaits the
8//! embedder's [`ToolProvider`]. Normal runtime casts validate declared return
9//! types.
10//!
11//! **Waiting is not running.** The provider may take as long as the harness
12//! needs, for example while a person approves the call. The time spent waiting
13//! is added to [`StoreData::clock_paused`], which an embedder's execution
14//! timeout subtracts, so a program's time limit counts its own work only. The
15//! provider bounds the wait itself. Dropping the future returned by
16//! [`ToolProvider::call`] is how the engine cancels it.
17
18use std::future::Future;
19use std::pin::Pin;
20use std::time::Instant;
21
22use serde::{Deserialize, Serialize};
23use wasmtime::{FuncType, HeapType, Linker, RefType, Val, ValType};
24
25use crate::runtime::StoreData;
26use crate::runtime::call_log::{Payload, Side, record_payload};
27use crate::runtime::fuel;
28use crate::runtime::host::{intrinsic_string_type, read_string_arg, register_host_fn_async};
29use crate::runtime::mcp::McpResponse;
30use crate::stdlib::shared::{audit_denial_in, check_security_call, running_package};
31use crate::{PackageDeclaration, Param, Span, Type, ValueKind, ValueSymbol};
32
33/// The internal host module every `@tools` call dispatches through.
34pub const TOOLS_MODULE_NAME: &str = "submilli:tools";
35/// The source-level name of the virtual package.
36pub const TOOLS_PACKAGE: &str = "@tools";
37/// The capability every call is checked against, with the tool in its `tool` field.
38pub const TOOLS_CAPABILITY: &str = "tools.call";
39
40/// Why the harness did not return a value. Each maps to a catchable guest error,
41/// except [`ToolCallError::Cancelled`], which ends the run.
42#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
43pub enum ToolCallError {
44    /// The tool ran and failed. Thrown as `Error(message)`.
45    Failed { message: String },
46    /// The harness, or a person it asked, refused the call. Thrown as
47    /// `PermissionDeniedError`, reported with source `harness`.
48    Denied { message: String },
49    /// The harness did not answer in time. Thrown as `TimeoutError`.
50    TimedOut,
51    /// The embedder could not hold another waiting call. Thrown as
52    /// `QuotaExceededError`.
53    QuotaExceeded { message: String },
54    /// The wait was abandoned, as when the session ends. The run ends with it.
55    Cancelled,
56}
57
58impl ToolCallError {
59    /// The failure as the call log keeps it: a stable kind and a message.
60    fn record(&self) -> serde_json::Value {
61        let (kind, message) = match self {
62            ToolCallError::Failed { message } => ("failed", message.as_str()),
63            ToolCallError::Denied { message } => ("denied", message.as_str()),
64            ToolCallError::TimedOut => ("timeout", ""),
65            ToolCallError::QuotaExceeded { message } => ("quota", message.as_str()),
66            ToolCallError::Cancelled => ("cancelled", ""),
67        };
68        serde_json::json!({ "kind": kind, "message": message })
69    }
70}
71
72/// The embedder-provided answerer for `@tools` calls. The value is a bounded
73/// JSON tree ([`McpResponse::take`] admits it), so guest conversion stays
74/// bounded whoever produced it.
75///
76/// Object-safe boxed-future trait (rather than `async fn`), like
77/// [`McpTransport`](crate::runtime::McpTransport), so a harness may answer in
78/// process or forward the call over the wire.
79pub trait ToolProvider: Send + Sync {
80    fn call<'a>(
81        &'a self,
82        tool: &'a str,
83        args_json: &'a str,
84    ) -> Pin<Box<dyn Future<Output = Result<McpResponse, ToolCallError>> + Send + 'a>>;
85}
86
87fn tools_call_func_type(linker: &Linker<StoreData>) -> wasmtime::Result<FuncType> {
88    let engine = linker.engine();
89    let string_struct = intrinsic_string_type(engine)?;
90    let string = ValType::Ref(RefType::new(false, HeapType::ConcreteStruct(string_struct)));
91    let object = crate::runtime::intrinsic_types::build_intrinsic_types(engine)?.object;
92    let unknown = ValType::Ref(RefType::new(true, HeapType::ConcreteStruct(object)));
93    Ok(FuncType::new(engine, [string.clone(), string], [unknown]))
94}
95
96/// Install the async `submilli:tools.call` host fn. Reads `(tool, argsJson)`,
97/// runs the capability check, awaits [`StoreData::tool_provider`] with the
98/// clock paused, and builds guest values directly from the returned tree.
99pub fn install_tools_async(linker: &mut Linker<StoreData>) -> wasmtime::Result<()> {
100    let ty = tools_call_func_type(linker)?;
101    register_host_fn_async(
102        linker,
103        TOOLS_MODULE_NAME,
104        crate::mangle::host(TOOLS_MODULE_NAME, "call"),
105        ty,
106        /* deterministic = */ false,
107        move |caller, params, results| {
108            Box::pin(async move {
109                let params: &[Val; 2] = params.try_into().map_err(|_| {
110                    crate::runtime::host::fatal_host_error("tools argument count is invalid")
111                })?;
112                let result_slot = results.first_mut().ok_or_else(|| {
113                    crate::runtime::host::fatal_host_error("tools result slot is missing")
114                })?;
115                let tool = read_string_arg(&mut *caller, &params[0], "tools.call (tool)")?;
116                let args_json = read_string_arg(&mut *caller, &params[1], "tools.call (args)")?;
117
118                // Deny before the harness hears of the call. The tool is in the
119                // filter context so a policy can choose which tools a program reaches.
120                let ticket = check_security_call(
121                    &mut *caller,
122                    TOOLS_CAPABILITY,
123                    serde_json::json!({ "tool": tool }),
124                )?;
125                record_payload(&*caller, ticket, Side::Request, || {
126                    Payload::meta(serde_json::json!({ "tool": tool }))
127                        .with_body(args_json.as_bytes())
128                });
129                // Looked up after the gate, so a denied caller cannot learn
130                // whether the session offers tools.
131                let provider = caller.data().tool_provider.clone().ok_or_else(|| {
132                    wasmtime::Error::msg(format!(
133                        "{TOOLS_PACKAGE}.{tool}: this session offers no application tools"
134                    ))
135                })?;
136                let principal = running_package(&*caller)
137                    .map_err(|error| error.into_denial(TOOLS_CAPABILITY))?;
138
139                fuel::charge(&mut *caller, fuel::IO, args_json.len() as u64)?;
140                let waited = Instant::now();
141                let outcome = provider.call(&tool, &args_json).await;
142                let data = caller.data_mut();
143                data.clock_paused = data.clock_paused.saturating_add(waited.elapsed());
144
145                // A refusal by the harness is a denial in the audit trail, inside the
146                // call the policy allowed.
147                if let Err(ToolCallError::Denied { .. }) = &outcome {
148                    audit_denial_in(
149                        &*caller,
150                        ticket,
151                        &principal,
152                        TOOLS_CAPABILITY,
153                        &serde_json::json!({ "tool": tool }),
154                        crate::runtime::host::DenialSource::Harness.as_str(),
155                        "the application refused the call",
156                    );
157                }
158                record_payload(&*caller, ticket, Side::Response, || match &outcome {
159                    Ok(response) => Payload::meta(serde_json::Value::Null)
160                        .with_owned_body(serde_json::to_vec(response.value()).unwrap_or_default()),
161                    Err(error) => Payload::meta(error.record()),
162                });
163                fuel::settle_result(caller, |caller| {
164                    let result = (|| {
165                        let response = outcome
166                            .map_err(|error| tool_error_to_throw(&principal, &tool, error))?;
167                        *result_slot = super::mcp::allocate_response(caller, &response)?;
168                        Ok(())
169                    })();
170                    result.map_err(|error| crate::runtime::host::throw_host_error(caller, error))
171                })
172            })
173        },
174    )
175}
176
177fn tool_error_to_throw(principal: &str, tool: &str, error: ToolCallError) -> wasmtime::Error {
178    match error {
179        ToolCallError::Failed { message } => {
180            wasmtime::Error::msg(format!("{TOOLS_PACKAGE}.{tool}: {message}"))
181        }
182        ToolCallError::Denied { message } => crate::runtime::host::permission_denied_by_harness(
183            principal,
184            TOOLS_CAPABILITY,
185            format!("{TOOLS_PACKAGE}.{tool}: {message}"),
186        ),
187        ToolCallError::TimedOut => crate::runtime::host::timeout_error(format!(
188            "{TOOLS_PACKAGE}.{tool}: the application did not answer in time"
189        )),
190        ToolCallError::QuotaExceeded { message } => {
191            crate::runtime::host::quota_exceeded_error(format!("{TOOLS_PACKAGE}.{tool}: {message}"))
192        }
193        ToolCallError::Cancelled => crate::runtime::host::fatal_host_error(format!(
194            "{TOOLS_PACKAGE}.{tool}: the call was cancelled"
195        )),
196    }
197}
198
199/// Codegen-facing definitions for the internal `submilli:tools.call` host fn:
200/// `call(tool: string, args: string): unknown`.
201pub fn tools_call_package_declaration() -> PackageDeclaration {
202    let mut defs = PackageDeclaration::with_package(TOOLS_MODULE_NAME);
203    defs.values.insert(
204        "call".to_string(),
205        ValueSymbol {
206            name: "call".to_string(),
207            mangled_name: crate::mangle::host(TOOLS_MODULE_NAME, "call"),
208            declaration_span: Span::at(crate::FileId::MCP),
209            kind: ValueKind::Function {
210                generics: Vec::new(),
211                params: vec![
212                    Param::new("tool", Type::String),
213                    Param::new("args", Type::String),
214                ],
215                ret: Type::Unknown,
216                type_predicate: None,
217                doc: None,
218            },
219        },
220    );
221    defs
222}
223
224#[cfg(test)]
225mod tests {
226    use std::collections::BTreeMap;
227    use std::sync::{Arc, Mutex};
228    use std::time::Duration;
229
230    use wasmtime::{Linker, Module};
231
232    use super::{ToolCallError, ToolProvider};
233    use crate::runtime::McpResponse;
234    use crate::runtime::host::DenialSource;
235    use crate::runtime::security::{CheckOutcome, SecurityCheck};
236    use crate::runtime::{
237        RuntimeConfig, StoreData, Vfs, dispatch_main_async, install_runtime_async,
238        install_tenant_limits,
239    };
240    use crate::{
241        ObjectField, PackageDeclaration, Param, Shape, Span, Type, ValueKind, ValueSymbol,
242        VirtualDispatch, compile_script, mangle,
243    };
244
245    /// `@tools` with `echo(args: { msg: string }): { reply: string }` and
246    /// `raw(): unknown`, shaped as the catalog mapper builds them.
247    fn tools_package() -> PackageDeclaration {
248        let pkg = super::TOOLS_PACKAGE;
249        let mut defs = PackageDeclaration::with_package(pkg);
250        defs.virtual_dispatch = Some(VirtualDispatch::Tools);
251        let object = |field: &str| Type::Object {
252            index: None,
253            fields: BTreeMap::from([(field.to_string(), ObjectField::required(Type::String))]),
254        };
255        let (args, ret) = (object("msg"), object("reply"));
256        defs.shapes.push(Shape::from_type(&args).unwrap());
257        defs.shapes.push(Shape::from_type(&ret).unwrap());
258        let mut insert = |name: &str, params: Vec<Param>, ret: Type| {
259            defs.values.insert(
260                name.to_string(),
261                ValueSymbol {
262                    name: name.to_string(),
263                    mangled_name: mangle::host(pkg, name),
264                    declaration_span: Span::at(crate::FileId::MCP),
265                    kind: ValueKind::Function {
266                        generics: Vec::new(),
267                        params,
268                        ret,
269                        type_predicate: None,
270                        doc: None,
271                    },
272                },
273            );
274        };
275        insert("echo", vec![Param::new("args", args)], ret);
276        insert("raw", Vec::new(), Type::Unknown);
277        defs
278    }
279
280    type Calls = Arc<Mutex<Vec<(String, String)>>>;
281
282    /// Records every call, waits `delay`, then answers `answer`.
283    struct Answer {
284        answer: Result<serde_json::Value, ToolCallError>,
285        delay: Duration,
286        calls: Calls,
287    }
288
289    impl ToolProvider for Answer {
290        fn call<'a>(
291            &'a self,
292            tool: &'a str,
293            args_json: &'a str,
294        ) -> std::pin::Pin<
295            Box<dyn std::future::Future<Output = Result<McpResponse, ToolCallError>> + Send + 'a>,
296        > {
297            self.calls
298                .lock()
299                .unwrap()
300                .push((tool.to_string(), args_json.to_string()));
301            Box::pin(async move {
302                tokio::time::sleep(self.delay).await;
303                let mut value = self.answer.clone()?;
304                Ok(McpResponse::take(&mut value).expect("bounded"))
305            })
306        }
307    }
308
309    struct DenyAll;
310    impl SecurityCheck for DenyAll {
311        fn check(&self, _: &str, capability: &str, _: &serde_json::Value) -> CheckOutcome {
312            CheckOutcome::Deny {
313                rule: None,
314                reason: format!("denied {capability}"),
315            }
316        }
317    }
318
319    struct Run {
320        result: wasmtime::Result<Option<String>>,
321        calls: Calls,
322        clock_paused: Duration,
323    }
324
325    async fn run_with(
326        source: &str,
327        answer: Result<serde_json::Value, ToolCallError>,
328        delay: Duration,
329        security: Option<Arc<dyn SecurityCheck>>,
330        provided: bool,
331    ) -> Run {
332        let pkg = tools_package();
333        let compiled =
334            compile_script(source, "<test>", crate::FileId(0), &[], &[&pkg]).expect("compiles");
335        let cfg = RuntimeConfig::default();
336        let engine = cfg.engine().unwrap();
337        let calls = Calls::default();
338        let mut data = StoreData::with_vfs(Vfs::none());
339        data.install_type_info(compiled.type_info.clone());
340        if provided {
341            data.tool_provider = Some(Arc::new(Answer {
342                answer,
343                delay,
344                calls: calls.clone(),
345            }));
346        }
347        if let Some(check) = security {
348            data.security_check = check;
349        }
350        let mut store = cfg.store_async(&engine, data).unwrap();
351        install_tenant_limits(&mut store);
352        let module = Module::new(&engine, &compiled.wasm).unwrap();
353        let mut linker = Linker::<StoreData>::new(&engine);
354        install_runtime_async(&mut linker, &mut store)
355            .await
356            .unwrap();
357        let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
358        let result = dispatch_main_async(&mut store, &instance).await;
359        Run {
360            result,
361            calls,
362            clock_paused: store.data().clock_paused,
363        }
364    }
365
366    async fn run(source: &str, answer: Result<serde_json::Value, ToolCallError>) -> Run {
367        run_with(source, answer, Duration::ZERO, None, true).await
368    }
369
370    fn escaped_source(error: &wasmtime::Error) -> Option<DenialSource> {
371        error
372            .downcast_ref::<crate::backtrace::ThrownError>()
373            .and_then(|thrown| thrown.denial.as_ref())
374            .map(|denial| denial.source)
375    }
376
377    #[tokio::test]
378    async fn typed_tool_round_trips_through_the_provider() {
379        let run = run(
380            r#"import tools from "@tools";
381               function main(): string { return tools.echo({ msg: "hi" }).reply; }"#,
382            Ok(serde_json::json!({ "reply": "hello" })),
383        )
384        .await;
385        assert_eq!(run.result.unwrap().as_deref(), Some("hello"));
386        assert_eq!(
387            *run.calls.lock().unwrap(),
388            [("echo".to_string(), r#"{"msg":"hi"}"#.to_string())]
389        );
390    }
391
392    #[tokio::test]
393    async fn a_zero_argument_tool_sends_an_empty_object_and_returns_unknown() {
394        let run = run(
395            r#"import { raw } from "@tools";
396               function main(): string { return JSON.stringify(raw()) ?? ""; }"#,
397            Ok(serde_json::json!([1, "two"])),
398        )
399        .await;
400        assert_eq!(run.result.unwrap().as_deref(), Some(r#"[1,"two"]"#));
401        assert_eq!(run.calls.lock().unwrap()[0].1, "{}");
402    }
403
404    #[tokio::test]
405    async fn a_value_that_does_not_match_the_declared_type_throws_type_error() {
406        let run = run(
407            r#"import tools from "@tools";
408               function main(): string {
409                   try { return tools.echo({ msg: "hi" }).reply; }
410                   catch (e) { return e instanceof TypeError ? "type" : "other"; }
411               }"#,
412            Ok(serde_json::json!({ "reply": 5 })),
413        )
414        .await;
415        assert_eq!(run.result.unwrap().as_deref(), Some("type"));
416    }
417
418    #[tokio::test]
419    async fn a_policy_denial_throws_before_the_harness_hears_of_the_call() {
420        let run = run_with(
421            r#"import tools from "@tools";
422               function main(): void { tools.raw(); }"#,
423            Ok(serde_json::Value::Null),
424            Duration::ZERO,
425            Some(Arc::new(DenyAll)),
426            true,
427        )
428        .await;
429        let error = run.result.unwrap_err();
430        assert_eq!(escaped_source(&error), Some(DenialSource::Policy));
431        assert!(run.calls.lock().unwrap().is_empty());
432    }
433
434    #[tokio::test]
435    async fn each_harness_answer_throws_its_own_error_class() {
436        let cases = [
437            (
438                ToolCallError::Failed {
439                    message: "no such order".into(),
440                },
441                "e instanceof Error && e.message == \"@tools.raw: no such order\"",
442            ),
443            (
444                ToolCallError::Denied {
445                    message: "refused".into(),
446                },
447                "e instanceof PermissionDeniedError",
448            ),
449            (ToolCallError::TimedOut, "e instanceof TimeoutError"),
450            (
451                ToolCallError::QuotaExceeded {
452                    message: "full".into(),
453                },
454                "e instanceof QuotaExceededError",
455            ),
456        ];
457        for (answer, check) in cases {
458            let source = format!(
459                r#"import tools from "@tools";
460                   function main(): string {{
461                       try {{ tools.raw(); return "returned"; }}
462                       catch (e) {{ return {check} ? "caught" : "wrong"; }}
463                   }}"#
464            );
465            let run = run(&source, Err(answer.clone())).await;
466            assert_eq!(run.result.unwrap().as_deref(), Some("caught"), "{answer:?}");
467        }
468    }
469
470    #[tokio::test]
471    async fn an_escaped_harness_refusal_reports_source_harness() {
472        let run = run(
473            r#"import tools from "@tools";
474               function main(): void { tools.raw(); }"#,
475            Err(ToolCallError::Denied {
476                message: "a person said no".into(),
477            }),
478        )
479        .await;
480        let error = run.result.unwrap_err();
481        assert_eq!(escaped_source(&error), Some(DenialSource::Harness));
482    }
483
484    #[tokio::test]
485    async fn a_cancelled_call_ends_the_run_uncaught() {
486        let run = run(
487            r#"import tools from "@tools";
488               function main(): string { try { tools.raw(); } catch (e) { return "caught"; } return "returned"; }"#,
489            Err(ToolCallError::Cancelled),
490        )
491        .await;
492        let error = run.result.unwrap_err();
493        assert!(
494            error.is::<crate::runtime::host::FatalHostError>(),
495            "{error:?}"
496        );
497    }
498
499    #[tokio::test]
500    async fn without_a_provider_the_call_throws_a_catchable_error() {
501        let run = run_with(
502            r#"import tools from "@tools";
503               function main(): string { try { tools.raw(); } catch (e) { return e.message; } return ""; }"#,
504            Ok(serde_json::Value::Null),
505            Duration::ZERO,
506            None,
507            false,
508        )
509        .await;
510        assert_eq!(
511            run.result.unwrap().as_deref(),
512            Some("@tools.raw: this session offers no application tools")
513        );
514    }
515
516    #[tokio::test]
517    async fn time_spent_waiting_for_the_harness_is_recorded_as_paused() {
518        let run = run_with(
519            r#"import tools from "@tools";
520               function main(): void { tools.raw(); tools.raw(); }"#,
521            Ok(serde_json::Value::Null),
522            Duration::from_millis(30),
523            None,
524            true,
525        )
526        .await;
527        run.result.unwrap();
528        assert!(
529            run.clock_paused >= Duration::from_millis(60),
530            "{:?}",
531            run.clock_paused
532        );
533    }
534
535    #[test]
536    fn importing_tools_without_a_catalog_is_a_compile_error() {
537        let error = compile_script(
538            r#"import tools from "@tools"; function main(): void {}"#,
539            "<test>",
540            crate::FileId(0),
541            &[],
542            &[],
543        )
544        .unwrap_err();
545        assert!(
546            format!("{error:?}").contains("`@tools` is unavailable"),
547            "{error:?}"
548        );
549    }
550
551    #[test]
552    fn the_import_is_reported_as_tools() {
553        let parsed = crate::compile::parse_script(
554            r#"import tools from "@tools"; import x from "@mcp/x"; function main(): void {}"#,
555            crate::FileId(0),
556        );
557        let imports = parsed.external_imports().unwrap();
558        assert!(imports.tools);
559        assert!(imports.registry_packages.is_empty());
560        assert_eq!(imports.mcp_servers.len(), 1);
561    }
562}
submilli_engine/runtime/tools.rs

submilli_engine/runtime/
tools.rs
