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        move |caller, params, results| {
107            Box::pin(async move {
108                let params: &[Val; 2] = params.try_into().map_err(|_| {
109                    crate::runtime::host::fatal_host_error("tools argument count is invalid")
110                })?;
111                let result_slot = results.first_mut().ok_or_else(|| {
112                    crate::runtime::host::fatal_host_error("tools result slot is missing")
113                })?;
114                let tool = read_string_arg(&mut *caller, &params[0], "tools.call (tool)")?;
115                let args_json = read_string_arg(&mut *caller, &params[1], "tools.call (args)")?;
116
117                // Deny before the harness hears of the call. The tool is in the
118                // filter context so a policy can choose which tools a program reaches.
119                let ticket = check_security_call(
120                    &mut *caller,
121                    TOOLS_CAPABILITY,
122                    serde_json::json!({ "tool": tool }),
123                )?;
124                record_payload(&*caller, ticket, Side::Request, || {
125                    Payload::meta(serde_json::json!({ "tool": tool }))
126                        .with_body(args_json.as_bytes())
127                });
128                let request = ToolRequest { tool, args_json };
129                let response = super::effects::execute(
130                    caller,
131                    "tools.call",
132                    request,
133                    move |caller, _, request| {
134                        Box::pin(async move {
135                            fuel::charge(&mut *caller, fuel::IO, request.args_json.len() as u64)?;
136                            let ToolRequest { tool, args_json } = request;
137                            // Looked up after the gate, so a denied caller cannot learn
138                            // whether the session offers tools.
139                            let provider = caller.data().tool_provider.clone().ok_or_else(|| {
140                    wasmtime::Error::msg(format!(
141                        "{TOOLS_PACKAGE}.{tool}: this session offers no application tools"
142                    ))
143                })?;
144                            let principal = running_package(&*caller)
145                                .map_err(|error| error.into_denial(TOOLS_CAPABILITY))?;
146
147                            let waited = Instant::now();
148                            let outcome = provider.call(&tool, &args_json).await;
149                            let data = caller.data_mut();
150                            data.clock_paused = data.clock_paused.saturating_add(waited.elapsed());
151
152                            // A refusal by the harness is a denial in the audit trail, inside the
153                            // call the policy allowed.
154                            if let Err(ToolCallError::Denied { .. }) = &outcome {
155                                audit_denial_in(
156                                    &*caller,
157                                    ticket,
158                                    &principal,
159                                    TOOLS_CAPABILITY,
160                                    &serde_json::json!({ "tool": tool }),
161                                    crate::runtime::host::DenialSource::Harness.as_str(),
162                                    "the application refused the call",
163                                );
164                            }
165                            record_payload(&*caller, ticket, Side::Response, || match &outcome {
166                                Ok(response) => Payload::meta(serde_json::Value::Null)
167                                    .with_owned_body(
168                                        serde_json::to_vec(response.value()).unwrap_or_default(),
169                                    ),
170                                Err(error) => Payload::meta(error.record()),
171                            });
172                            outcome.map_err(|error| tool_error_to_throw(&principal, &tool, error))
173                        })
174                    },
175                )
176                .await?;
177                *result_slot = super::mcp::allocate_response(caller, &response)?;
178                Ok(())
179            })
180        },
181    )
182}
183
184#[derive(serde::Serialize)]
185struct ToolRequest {
186    tool: String,
187    args_json: String,
188}
189
190fn tool_error_to_throw(principal: &str, tool: &str, error: ToolCallError) -> wasmtime::Error {
191    match error {
192        ToolCallError::Failed { message } => {
193            wasmtime::Error::msg(format!("{TOOLS_PACKAGE}.{tool}: {message}"))
194        }
195        ToolCallError::Denied { message } => crate::runtime::host::permission_denied_by_harness(
196            principal,
197            TOOLS_CAPABILITY,
198            format!("{TOOLS_PACKAGE}.{tool}: {message}"),
199        ),
200        ToolCallError::TimedOut => crate::runtime::host::timeout_error(format!(
201            "{TOOLS_PACKAGE}.{tool}: the application did not answer in time"
202        )),
203        ToolCallError::QuotaExceeded { message } => {
204            crate::runtime::host::quota_exceeded_error(format!("{TOOLS_PACKAGE}.{tool}: {message}"))
205        }
206        ToolCallError::Cancelled => crate::runtime::host::fatal_host_error(format!(
207            "{TOOLS_PACKAGE}.{tool}: the call was cancelled"
208        )),
209    }
210}
211
212/// Codegen-facing definitions for the internal `submilli:tools.call` host fn:
213/// `call(tool: string, args: string): unknown`.
214pub fn tools_call_package_declaration() -> PackageDeclaration {
215    let mut defs = PackageDeclaration::with_package(TOOLS_MODULE_NAME);
216    defs.values.insert(
217        "call".to_string(),
218        ValueSymbol {
219            name: "call".to_string(),
220            mangled_name: crate::mangle::host(TOOLS_MODULE_NAME, "call"),
221            declaration_span: Span::at(crate::FileId::MCP),
222            kind: ValueKind::Function {
223                generics: Vec::new(),
224                params: vec![
225                    Param::new("tool", Type::String),
226                    Param::new("args", Type::String),
227                ],
228                ret: Type::Unknown,
229                type_predicate: None,
230                doc: None,
231            },
232        },
233    );
234    defs
235}
236
237#[cfg(test)]
238mod tests {
239    use std::collections::BTreeMap;
240    use std::sync::{Arc, Mutex};
241    use std::time::Duration;
242
243    use wasmtime::{Linker, Module};
244
245    use super::{ToolCallError, ToolProvider};
246    use crate::runtime::McpResponse;
247    use crate::runtime::host::DenialSource;
248    use crate::runtime::security::{CheckOutcome, SecurityCheck};
249    use crate::runtime::{
250        RuntimeConfig, StoreData, Vfs, dispatch_main_async, install_runtime_async,
251        install_tenant_limits,
252    };
253    use crate::{
254        ObjectField, PackageDeclaration, Param, Shape, Span, Type, ValueKind, ValueSymbol,
255        VirtualDispatch, compile_script, mangle,
256    };
257
258    /// `@tools` with `echo(args: { msg: string }): { reply: string }` and
259    /// `raw(): unknown`, shaped as the catalog mapper builds them.
260    fn tools_package() -> PackageDeclaration {
261        let pkg = super::TOOLS_PACKAGE;
262        let mut defs = PackageDeclaration::with_package(pkg);
263        defs.virtual_dispatch = Some(VirtualDispatch::Tools);
264        let object = |field: &str| Type::Object {
265            index: None,
266            fields: BTreeMap::from([(field.to_string(), ObjectField::required(Type::String))]),
267        };
268        let (args, ret) = (object("msg"), object("reply"));
269        defs.shapes.push(Shape::from_type(&args).unwrap());
270        defs.shapes.push(Shape::from_type(&ret).unwrap());
271        let mut insert = |name: &str, params: Vec<Param>, ret: Type| {
272            defs.values.insert(
273                name.to_string(),
274                ValueSymbol {
275                    name: name.to_string(),
276                    mangled_name: mangle::host(pkg, name),
277                    declaration_span: Span::at(crate::FileId::MCP),
278                    kind: ValueKind::Function {
279                        generics: Vec::new(),
280                        params,
281                        ret,
282                        type_predicate: None,
283                        doc: None,
284                    },
285                },
286            );
287        };
288        insert("echo", vec![Param::new("args", args)], ret);
289        insert("raw", Vec::new(), Type::Unknown);
290        defs
291    }
292
293    type Calls = Arc<Mutex<Vec<(String, String)>>>;
294
295    /// Records every call, waits `delay`, then answers `answer`.
296    struct Answer {
297        answer: Result<serde_json::Value, ToolCallError>,
298        delay: Duration,
299        calls: Calls,
300    }
301
302    impl ToolProvider for Answer {
303        fn call<'a>(
304            &'a self,
305            tool: &'a str,
306            args_json: &'a str,
307        ) -> std::pin::Pin<
308            Box<dyn std::future::Future<Output = Result<McpResponse, ToolCallError>> + Send + 'a>,
309        > {
310            self.calls
311                .lock()
312                .unwrap()
313                .push((tool.to_string(), args_json.to_string()));
314            Box::pin(async move {
315                tokio::time::sleep(self.delay).await;
316                let mut value = self.answer.clone()?;
317                Ok(McpResponse::take(&mut value).expect("bounded"))
318            })
319        }
320    }
321
322    struct DenyAll;
323    impl SecurityCheck for DenyAll {
324        fn check(&self, _: &str, capability: &str, _: &serde_json::Value) -> CheckOutcome {
325            CheckOutcome::Deny {
326                rule: None,
327                reason: format!("denied {capability}"),
328            }
329        }
330    }
331
332    struct Run {
333        result: wasmtime::Result<Option<String>>,
334        calls: Calls,
335        clock_paused: Duration,
336    }
337
338    async fn run_with(
339        source: &str,
340        answer: Result<serde_json::Value, ToolCallError>,
341        delay: Duration,
342        security: Option<Arc<dyn SecurityCheck>>,
343        provided: bool,
344    ) -> Run {
345        let pkg = tools_package();
346        let compiled =
347            compile_script(source, "<test>", crate::FileId(0), &[], &[&pkg]).expect("compiles");
348        let cfg = RuntimeConfig::default();
349        let engine = cfg.engine().unwrap();
350        let calls = Calls::default();
351        let mut data = StoreData::with_vfs("test-run".into(), Vfs::none());
352        data.install_type_info(compiled.type_info.clone());
353        if provided {
354            data.tool_provider = Some(Arc::new(Answer {
355                answer,
356                delay,
357                calls: calls.clone(),
358            }));
359        }
360        if let Some(check) = security {
361            data.security_check = check;
362        }
363        let mut store = cfg.store_async(&engine, data).unwrap();
364        install_tenant_limits(&mut store);
365        let module = Module::new(&engine, &compiled.wasm).unwrap();
366        let mut linker = Linker::<StoreData>::new(&engine);
367        install_runtime_async(&mut linker, &mut store)
368            .await
369            .unwrap();
370        let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
371        let result = dispatch_main_async(&mut store, &instance).await;
372        Run {
373            result,
374            calls,
375            clock_paused: store.data().clock_paused,
376        }
377    }
378
379    async fn run(source: &str, answer: Result<serde_json::Value, ToolCallError>) -> Run {
380        run_with(source, answer, Duration::ZERO, None, true).await
381    }
382
383    fn escaped_source(error: &wasmtime::Error) -> Option<DenialSource> {
384        error
385            .downcast_ref::<crate::backtrace::ThrownError>()
386            .and_then(|thrown| thrown.denial.as_ref())
387            .map(|denial| denial.source)
388    }
389
390    #[tokio::test]
391    async fn typed_tool_round_trips_through_the_provider() {
392        let run = run(
393            r#"import tools from "@tools";
394               function main(): string { return tools.echo({ msg: "hi" }).reply; }"#,
395            Ok(serde_json::json!({ "reply": "hello" })),
396        )
397        .await;
398        assert_eq!(run.result.unwrap().as_deref(), Some("hello"));
399        assert_eq!(
400            *run.calls.lock().unwrap(),
401            [("echo".to_string(), r#"{"msg":"hi"}"#.to_string())]
402        );
403    }
404
405    #[tokio::test]
406    async fn a_zero_argument_tool_sends_an_empty_object_and_returns_unknown() {
407        let run = run(
408            r#"import { raw } from "@tools";
409               function main(): string { return JSON.stringify(raw()) ?? ""; }"#,
410            Ok(serde_json::json!([1, "two"])),
411        )
412        .await;
413        assert_eq!(run.result.unwrap().as_deref(), Some(r#"[1,"two"]"#));
414        assert_eq!(run.calls.lock().unwrap()[0].1, "{}");
415    }
416
417    #[tokio::test]
418    async fn a_value_that_does_not_match_the_declared_type_throws_type_error() {
419        let run = run(
420            r#"import tools from "@tools";
421               function main(): string {
422                   try { return tools.echo({ msg: "hi" }).reply; }
423                   catch (e) { return e instanceof TypeError ? "type" : "other"; }
424               }"#,
425            Ok(serde_json::json!({ "reply": 5 })),
426        )
427        .await;
428        assert_eq!(run.result.unwrap().as_deref(), Some("type"));
429    }
430
431    #[tokio::test]
432    async fn a_policy_denial_throws_before_the_harness_hears_of_the_call() {
433        let run = run_with(
434            r#"import tools from "@tools";
435               function main(): void { tools.raw(); }"#,
436            Ok(serde_json::Value::Null),
437            Duration::ZERO,
438            Some(Arc::new(DenyAll)),
439            true,
440        )
441        .await;
442        let error = run.result.unwrap_err();
443        assert_eq!(escaped_source(&error), Some(DenialSource::Policy));
444        assert!(run.calls.lock().unwrap().is_empty());
445    }
446
447    #[tokio::test]
448    async fn each_harness_answer_throws_its_own_error_class() {
449        let cases = [
450            (
451                ToolCallError::Failed {
452                    message: "no such order".into(),
453                },
454                "e instanceof Error && e.message == \"@tools.raw: no such order\"",
455            ),
456            (
457                ToolCallError::Denied {
458                    message: "refused".into(),
459                },
460                "e instanceof PermissionDeniedError",
461            ),
462            (ToolCallError::TimedOut, "e instanceof TimeoutError"),
463            (
464                ToolCallError::QuotaExceeded {
465                    message: "full".into(),
466                },
467                "e instanceof QuotaExceededError",
468            ),
469        ];
470        for (answer, check) in cases {
471            let source = format!(
472                r#"import tools from "@tools";
473                   function main(): string {{
474                       try {{ tools.raw(); return "returned"; }}
475                       catch (e) {{ return {check} ? "caught" : "wrong"; }}
476                   }}"#
477            );
478            let run = run(&source, Err(answer.clone())).await;
479            assert_eq!(run.result.unwrap().as_deref(), Some("caught"), "{answer:?}");
480        }
481    }
482
483    #[tokio::test]
484    async fn an_escaped_harness_refusal_reports_source_harness() {
485        let run = run(
486            r#"import tools from "@tools";
487               function main(): void { tools.raw(); }"#,
488            Err(ToolCallError::Denied {
489                message: "a person said no".into(),
490            }),
491        )
492        .await;
493        let error = run.result.unwrap_err();
494        assert_eq!(escaped_source(&error), Some(DenialSource::Harness));
495    }
496
497    #[tokio::test]
498    async fn a_cancelled_call_ends_the_run_uncaught() {
499        let run = run(
500            r#"import tools from "@tools";
501               function main(): string { try { tools.raw(); } catch (e) { return "caught"; } return "returned"; }"#,
502            Err(ToolCallError::Cancelled),
503        )
504        .await;
505        let error = run.result.unwrap_err();
506        assert!(
507            error.is::<crate::runtime::host::FatalHostError>(),
508            "{error:?}"
509        );
510    }
511
512    #[tokio::test]
513    async fn without_a_provider_the_call_throws_a_catchable_error() {
514        let run = run_with(
515            r#"import tools from "@tools";
516               function main(): string { try { tools.raw(); } catch (e) { return e.message; } return ""; }"#,
517            Ok(serde_json::Value::Null),
518            Duration::ZERO,
519            None,
520            false,
521        )
522        .await;
523        assert_eq!(
524            run.result.unwrap().as_deref(),
525            Some("@tools.raw: this session offers no application tools")
526        );
527    }
528
529    #[tokio::test]
530    async fn time_spent_waiting_for_the_harness_is_recorded_as_paused() {
531        let run = run_with(
532            r#"import tools from "@tools";
533               function main(): void { tools.raw(); tools.raw(); }"#,
534            Ok(serde_json::Value::Null),
535            Duration::from_millis(30),
536            None,
537            true,
538        )
539        .await;
540        run.result.unwrap();
541        assert!(
542            run.clock_paused >= Duration::from_millis(60),
543            "{:?}",
544            run.clock_paused
545        );
546    }
547
548    #[test]
549    fn importing_tools_without_a_catalog_is_a_compile_error() {
550        let error = compile_script(
551            r#"import tools from "@tools"; function main(): void {}"#,
552            "<test>",
553            crate::FileId(0),
554            &[],
555            &[],
556        )
557        .unwrap_err();
558        assert!(
559            format!("{error:?}").contains("`@tools` is unavailable"),
560            "{error:?}"
561        );
562    }
563
564    #[test]
565    fn the_import_is_reported_as_tools() {
566        let parsed = crate::compile::parse_script(
567            r#"import tools from "@tools"; import x from "@mcp/x"; function main(): void {}"#,
568            crate::FileId(0),
569        );
570        let imports = parsed.external_imports().unwrap();
571        assert!(imports.tools);
572        assert!(imports.registry_packages.is_empty());
573        assert_eq!(imports.mcp_servers.len(), 1);
574    }
575}
submilli_engine/runtime/tools.rs

submilli_engine/runtime/
tools.rs
