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

submilli_engine/runtime/
tools.rs
