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

submilli_engine/runtime/
tools.rs
