1use std::future::Future;
16use std::pin::Pin;
17
18use wasmtime::{FuncType, HeapType, Linker, RefType, Val, ValType};
19
20use crate::runtime::StoreData;
21use crate::runtime::call_log::{Payload, Side, record_payload};
22use crate::runtime::fuel;
23use crate::runtime::host::{intrinsic_string_type, read_string_arg, register_host_fn_async};
24use crate::stdlib::shared::check_security_call;
25use crate::{PackageDeclaration, Param, Span, Type, ValueKind, ValueSymbol};
26
27pub const MCP_MODULE_NAME: &str = "submilli:mcp";
29pub const MCP_MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024;
30
31mod response;
32pub use response::McpResponse;
33
34#[derive(Debug)]
37pub enum McpCallError {
38 ResponseTooLarge,
40 Internal { message: &'static str },
42 AuthExpired,
45 Mcp { message: String },
47 Upstream { status: u16, body: String },
49 Transport(String),
51 Local(String),
54}
55
56impl McpCallError {
57 fn record(&self) -> serde_json::Value {
60 let (kind, message, status) = match self {
61 McpCallError::ResponseTooLarge => ("response-too-large", String::new(), None),
62 McpCallError::Internal { message } => ("internal", (*message).to_owned(), None),
63 McpCallError::AuthExpired => ("auth-expired", String::new(), None),
64 McpCallError::Mcp { message } => ("mcp", message.clone(), None),
65 McpCallError::Upstream { status, body } => ("upstream", body.clone(), Some(*status)),
66 McpCallError::Transport(detail) => ("transport", detail.clone(), None),
67 McpCallError::Local(detail) => ("local", detail.clone(), None),
68 };
69 serde_json::json!({ "kind": kind, "message": message, "status": status })
70 }
71}
72
73pub struct McpOutcome {
75 pub result: Result<McpResponse, McpCallError>,
76 pub received_bytes: u64,
77 pub parsed_bytes: u64,
78}
79
80pub trait McpTransport: Send + Sync {
89 fn call<'a>(
90 &'a self,
91 server: &'a str,
92 tool: &'a str,
93 args_json: &'a str,
94 ) -> Pin<Box<dyn Future<Output = McpOutcome> + Send + 'a>>;
95}
96
97fn mcp_call_func_type(linker: &Linker<StoreData>) -> wasmtime::Result<FuncType> {
98 let engine = linker.engine();
99 let string_struct = intrinsic_string_type(engine)?;
100 let string = ValType::Ref(RefType::new(false, HeapType::ConcreteStruct(string_struct)));
101 let object = crate::runtime::intrinsic_types::build_intrinsic_types(engine)?.object;
102 let unknown = ValType::Ref(RefType::new(true, HeapType::ConcreteStruct(object)));
103 Ok(FuncType::new(
104 engine,
105 [string.clone(), string.clone(), string.clone()],
106 [unknown],
107 ))
108}
109
110pub fn install_mcp_async(linker: &mut Linker<StoreData>) -> wasmtime::Result<()> {
114 let ty = mcp_call_func_type(linker)?;
115 register_host_fn_async(
116 linker,
117 MCP_MODULE_NAME,
118 crate::mangle::host(MCP_MODULE_NAME, "call"),
119 ty,
120 move |caller, params, results| {
121 Box::pin(async move {
122 let params: &[Val; 3] = params.try_into().map_err(|_| {
123 crate::runtime::host::fatal_host_error("MCP argument count is invalid")
124 })?;
125 let result_slot = results.first_mut().ok_or_else(|| {
126 crate::runtime::host::fatal_host_error("MCP result slot is missing")
127 })?;
128 let server = read_string_arg(&mut *caller, ¶ms[0], "mcp.call (server)")?;
129 let tool = read_string_arg(&mut *caller, ¶ms[1], "mcp.call (tool)")?;
130 let args_json = read_string_arg(&mut *caller, ¶ms[2], "mcp.call (args)")?;
131
132 let ticket = check_security_call(
136 &mut *caller,
137 &format!("mcp.{server}"),
138 serde_json::json!({ "tool": tool, "transport": "streamable_http" }),
139 )?;
140 record_payload(&*caller, ticket, Side::Request, || {
141 request_payload(&server, &tool, &args_json)
142 });
143
144 let request = ToolRequest {
145 server,
146 tool,
147 args_json,
148 };
149 let response = super::effects::execute(
150 caller,
151 "mcp.call",
152 request,
153 move |caller, _, request| {
154 Box::pin(async move {
155 fuel::charge(&mut *caller, fuel::IO, request.args_json.len() as u64)?;
156 let ToolRequest {
157 server,
158 tool,
159 args_json,
160 } = request;
161 let transport =
162 caller.data().mcp_transport.clone().ok_or_else(|| {
163 wasmtime::Error::msg(format!(
164 "@mcp/{server}.{tool}: MCP transport not configured"
165 ))
166 })?;
167
168 let outcome = transport.call(&server, &tool, &args_json).await;
169 record_payload(&*caller, ticket, Side::Response, || {
170 let payload = match &outcome.result {
171 Ok(response) => Payload::meta(serde_json::Value::Null)
172 .with_owned_body(
173 serde_json::to_vec(response.value())
174 .unwrap_or_default(),
175 ),
176 Err(error) => Payload::meta(error.record()),
177 };
178 payload.with_size(outcome.received_bytes)
179 });
180 fuel::settle(&mut *caller, fuel::IO, outcome.received_bytes)?;
181 fuel::settle(&mut *caller, fuel::PARSE, outcome.parsed_bytes)?;
182 outcome
183 .result
184 .map_err(|error| mcp_error_to_throw(&server, &tool, error))
185 })
186 },
187 )
188 .await?;
189 *result_slot = allocate_response(caller, &response)?;
190 Ok(())
191 })
192 },
193 )
194}
195
196#[derive(serde::Serialize)]
197struct ToolRequest {
198 server: String,
199 tool: String,
200 args_json: String,
201}
202
203fn request_payload<'a>(server: &str, tool: &str, args_json: &'a str) -> Payload<'a> {
205 Payload::meta(serde_json::json!({ "server": server, "tool": tool }))
206 .with_body(args_json.as_bytes())
207}
208
209pub fn request_digest(server: &str, tool: &str, args_json: &str) -> String {
212 request_payload(server, tool, args_json).digest()
213}
214
215pub(crate) fn allocate_response(
216 caller: &mut wasmtime::Caller<'_, StoreData>,
217 response: &McpResponse,
218) -> wasmtime::Result<Val> {
219 fuel::charge(&mut *caller, fuel::ELEM, response.visited_nodes())?;
220 let allocator = crate::runtime::json::JsonUnknownAllocator::new(caller)?;
221 let value = crate::runtime::json::JsonValue::from(response.value());
222 allocator.allocate(caller, &value)
223}
224
225pub fn mcp_call_package_declaration() -> PackageDeclaration {
229 let mut defs = PackageDeclaration::with_package(MCP_MODULE_NAME);
230 defs.values.insert(
231 "call".to_string(),
232 ValueSymbol {
233 name: "call".to_string(),
234 mangled_name: crate::mangle::host(MCP_MODULE_NAME, "call"),
235 declaration_span: Span::at(crate::FileId::MCP),
236 kind: ValueKind::Function {
237 generics: Vec::new(),
238 params: vec![
239 Param::new("server", Type::String),
240 Param::new("tool", Type::String),
241 Param::new("args", Type::String),
242 ],
243 ret: Type::Unknown,
244 type_predicate: None,
245 doc: None,
246 },
247 },
248 );
249 defs
250}
251
252fn mcp_error_to_throw(server: &str, tool: &str, err: McpCallError) -> wasmtime::Error {
256 let message = match err {
257 McpCallError::Internal { message, .. } => {
258 return crate::runtime::host::fatal_host_error(message);
259 }
260 McpCallError::ResponseTooLarge => {
261 return crate::runtime::host::range_error(format!(
262 "@mcp/{server}.{tool}: response exceeds the size, node or depth limit; request a smaller result"
263 ));
264 }
265 McpCallError::AuthExpired => format!(
266 "McpAuthExpiredError: @mcp/{server}.{tool}: OAuth authentication failed; \
267 retry later or authenticate the MCP server again"
268 ),
269 McpCallError::Mcp { message } => {
270 format!("@mcp/{server}.{tool}: {message}")
271 }
272 McpCallError::Upstream { status, body } => {
273 format!("@mcp/{server}.{tool}: server returned HTTP {status}: {body}")
274 }
275 McpCallError::Transport(detail) | McpCallError::Local(detail) => {
276 format!("@mcp/{server}.{tool}: transport error: {detail}")
277 }
278 };
279 wasmtime::Error::msg(message)
280}
281
282#[cfg(test)]
283mod tests {
284 use std::collections::BTreeMap;
285
286 use std::sync::{Arc, Mutex};
287
288 use wasmtime::{Linker, Module};
289
290 use super::{McpOutcome, McpResponse, McpTransport};
291 use crate::runtime::security::{CheckOutcome, SecurityCheck};
292 use crate::runtime::{
293 RuntimeConfig, StoreData, Vfs, dispatch_main_async, install_runtime_async,
294 install_tenant_limits,
295 };
296 use crate::{
297 ObjectField, PackageDeclaration, Param, Shape, Span, Type, ValueKind, ValueSymbol,
298 compile_script, mangle,
299 };
300
301 fn test_mcp_package() -> PackageDeclaration {
306 let pkg = "@mcp/test";
307 let mut defs = PackageDeclaration::with_package(pkg);
308 defs.virtual_dispatch = Some(crate::VirtualDispatch::Mcp {
309 server: "test".to_string(),
310 });
311
312 let mut arg_fields = BTreeMap::new();
313 arg_fields.insert("msg".to_string(), ObjectField::required(Type::String));
314 let arg_obj = Type::Object {
315 index: None,
316 fields: arg_fields,
317 };
318 let mut ret_fields = BTreeMap::new();
319 ret_fields.insert("reply".to_string(), ObjectField::required(Type::String));
320 let ret_obj = Type::Object {
321 index: None,
322 fields: ret_fields,
323 };
324 let mut opt_fields = BTreeMap::new();
325 opt_fields.insert("user".to_string(), ObjectField::optional(Type::String));
326 let opt_arg_obj = Type::Object {
327 index: None,
328 fields: opt_fields,
329 };
330 defs.shapes.push(Shape::from_type(&arg_obj).unwrap());
331 defs.shapes.push(Shape::from_type(&ret_obj).unwrap());
332 defs.shapes.push(Shape::from_type(&opt_arg_obj).unwrap());
333
334 let mut insert_fn = |name: &str, params: Vec<Param>, ret: Type| {
335 defs.values.insert(
336 name.to_string(),
337 ValueSymbol {
338 name: name.to_string(),
339 mangled_name: mangle::host(pkg, name),
340 declaration_span: Span::at(crate::FileId::MCP),
341 kind: ValueKind::Function {
342 generics: Vec::new(),
343 params,
344 ret,
345 type_predicate: None,
346 doc: None,
347 },
348 },
349 );
350 };
351 insert_fn("echo", vec![Param::new("args", arg_obj)], ret_obj);
352 insert_fn("raw", Vec::new(), Type::Unknown);
353 insert_fn(
355 "getTeams",
356 vec![Param::with_default(
357 "args",
358 opt_arg_obj,
359 crate::DefaultValue::EmptyObject,
360 )],
361 Type::Unknown,
362 );
363 defs
364 }
365
366 struct MockTransport {
368 result: serde_json::Value,
369 calls: Arc<Mutex<Vec<(String, String, String)>>>,
370 }
371
372 impl McpTransport for MockTransport {
373 fn call<'a>(
374 &'a self,
375 server: &'a str,
376 tool: &'a str,
377 args_json: &'a str,
378 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = McpOutcome> + Send + 'a>> {
379 self.calls.lock().unwrap().push((
380 server.to_string(),
381 tool.to_string(),
382 args_json.to_string(),
383 ));
384 let mut result = self.result.clone();
385 Box::pin(async move {
386 McpOutcome {
387 result: McpResponse::take(&mut result),
388 received_bytes: 0,
389 parsed_bytes: 0,
390 }
391 })
392 }
393 }
394
395 struct DenyAll;
396 impl SecurityCheck for DenyAll {
397 fn check(&self, _: &str, capability: &str, _: &serde_json::Value) -> CheckOutcome {
398 CheckOutcome::Deny {
399 rule: None,
400 reason: format!("denied {capability}"),
401 }
402 }
403 }
404
405 async fn run(
406 source: &str,
407 result: serde_json::Value,
408 security: Option<Arc<dyn SecurityCheck>>,
409 ) -> (
410 wasmtime::Result<Option<String>>,
411 Arc<Mutex<Vec<(String, String, String)>>>,
412 ) {
413 let pkg = test_mcp_package();
414 let compiled =
415 compile_script(source, "<test>", crate::FileId(0), &[], &[&pkg]).expect("compiles");
416
417 let cfg = RuntimeConfig::default();
418 let engine = cfg.engine().unwrap();
419 let calls = Arc::new(Mutex::new(Vec::new()));
420 let mut data = StoreData::with_vfs("test-run".into(), Vfs::tempdir().unwrap());
421 data.install_type_info(compiled.type_info.clone());
422 data.mcp_transport = Some(Arc::new(MockTransport {
423 result,
424 calls: calls.clone(),
425 }));
426 if let Some(check) = security {
427 data.security_check = check;
428 }
429 let mut store = cfg.store_async(&engine, data).unwrap();
430 install_tenant_limits(&mut store);
431 let module = Module::new(&engine, &compiled.wasm).unwrap();
432 let mut linker = Linker::<StoreData>::new(&engine);
433 install_runtime_async(&mut linker, &mut store)
434 .await
435 .unwrap();
436 let inst = linker.instantiate_async(&mut store, &module).await.unwrap();
437 let out = dispatch_main_async(&mut store, &inst).await;
438 (out, calls)
439 }
440
441 struct FailingTransport {
442 internal: bool,
443 }
444
445 impl McpTransport for FailingTransport {
446 fn call<'a>(
447 &'a self,
448 _: &'a str,
449 _: &'a str,
450 _: &'a str,
451 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = McpOutcome> + Send + 'a>> {
452 Box::pin(async move {
453 McpOutcome {
454 result: Err(if self.internal {
455 super::McpCallError::Internal {
456 message: "MCP HTTP client initialization failed",
457 }
458 } else {
459 super::McpCallError::Transport("remote connection failed".into())
460 }),
461 received_bytes: 0,
462 parsed_bytes: 0,
463 }
464 })
465 }
466 }
467
468 #[tokio::test]
469 async fn internal_client_failure_bypasses_guest_catch_and_keeps_runtime_usable() {
470 let package = test_mcp_package();
471 let source = r#"import test from "@mcp/test"; function main(): number { try { test.raw(); } catch (e) { return 7; } return 9; }"#;
472 let compiled =
473 compile_script(source, "<test>", crate::FileId(0), &[], &[&package]).unwrap();
474 let healthy = compile_script(
475 "function main(): number { return 42; }",
476 "<healthy>",
477 crate::FileId(0),
478 &[],
479 &[],
480 )
481 .unwrap();
482 let cfg = RuntimeConfig::default();
483 let engine = cfg.engine().unwrap();
484 let mut data = StoreData::with_vfs("test-run".into(), Vfs::none());
485 data.install_type_info(compiled.type_info.clone());
486 data.mcp_transport = Some(Arc::new(FailingTransport { internal: true }));
487 let mut store = cfg.store_async(&engine, data).unwrap();
488 install_tenant_limits(&mut store);
489 let mut linker = Linker::<StoreData>::new(&engine);
490 install_runtime_async(&mut linker, &mut store)
491 .await
492 .unwrap();
493 let module = Module::new(&engine, &compiled.wasm).unwrap();
494 let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
495 let error = dispatch_main_async(&mut store, &instance)
496 .await
497 .unwrap_err();
498 assert!(
499 error.is::<crate::runtime::host::FatalHostError>(),
500 "{error:?}"
501 );
502 store.data_mut().blocking_work.finish().await.unwrap();
503 store.data_mut().mcp_transport = Some(Arc::new(FailingTransport { internal: false }));
504 assert_eq!(
505 dispatch_main_async(&mut store, &instance).await.unwrap(),
506 Some("7".into())
507 );
508 store
509 .data_mut()
510 .install_type_info(healthy.type_info.clone());
511 let module = Module::new(&engine, &healthy.wasm).unwrap();
512 let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
513 assert_eq!(
514 dispatch_main_async(&mut store, &instance).await.unwrap(),
515 Some("42".into())
516 );
517 }
518
519 #[tokio::test]
520 async fn response_allocation_obeys_the_fuel_budget() {
521 use wasmtime::{Func, FuncType, Val, ValType};
522 let config = RuntimeConfig::default();
523 let engine = config.engine().unwrap();
524 let mut store = config
525 .store_async(&engine, StoreData::with_vfs("test-run".into(), Vfs::none()))
526 .unwrap();
527 let mut linker = Linker::new(&engine);
528 install_runtime_async(&mut linker, &mut store)
529 .await
530 .unwrap();
531 let complete = Func::new(
532 &mut store,
533 FuncType::new(&engine, [ValType::I32], [ValType::I32, ValType::F64]),
534 |mut caller, params, results| {
535 let count = params[0].i32().unwrap();
536 let mut value = serde_json::json!((0..count).collect::<Vec<_>>());
537 let response = McpResponse::take(&mut value).unwrap();
538 let value = super::allocate_response(&mut caller, &response)?;
539 let storage =
540 crate::runtime::array_storage::ArrayStorage::read(&mut caller, &value)?;
541 results[0] = Val::I32(storage.len as i32);
542 let last = storage.get(&mut caller, (count - 1) as u32)?;
543 let Val::AnyRef(Some(last)) = last else {
544 panic!("missing last value")
545 };
546 results[1] = last.unwrap_struct(&mut caller)?.field(&mut caller, 1)?;
547 Ok(())
548 },
549 );
550 let mut result = [Val::I32(0), Val::F64(0)];
551 for count in [128, 256] {
552 store.set_fuel(100_000).unwrap();
553 let before = store.data().host_fuel;
554 complete
555 .call_async(&mut store, &[Val::I32(count)], &mut result)
556 .await
557 .unwrap();
558 let actual = store.data().host_fuel - before;
559 assert_eq!(
560 actual,
561 2 * crate::runtime::fuel::ELEM.cost(count as u64 + 1)
562 );
563 assert_eq!(result[0].i32(), Some(count));
565 }
566 store.set_fuel(1).unwrap();
567 let mut result = [Val::I32(-1), Val::F64(0)];
568 let error = complete
569 .call_async(&mut store, &[Val::I32(256)], &mut result)
570 .await
571 .unwrap_err();
572 assert_eq!(
573 error.downcast_ref::<wasmtime::Trap>(),
574 Some(&wasmtime::Trap::OutOfFuel)
575 );
576 assert_eq!(result[0].i32(), Some(-1));
577 assert_eq!(store.get_fuel().unwrap(), 0);
578 assert!(!store.data().settling_host_result);
579 store.set_fuel(100_000).unwrap();
580 complete
581 .call_async(&mut store, &[Val::I32(256)], &mut result)
582 .await
583 .unwrap();
584 assert_eq!(result[0].i32(), Some(256));
585 assert_eq!(result[1].f64(), Some(255.0));
586 }
587
588 #[tokio::test]
589 async fn typed_tool_round_trips_and_serializes_args() {
590 let (out, calls) = run(
591 r#"import test from "@mcp/test";
592 function main(): void {
593 const r = test.echo({ msg: "hi" });
594 assert(r.reply === "yo", "reply");
595 }"#,
596 serde_json::json!({ "reply": "yo" }),
597 None,
598 )
599 .await;
600 out.expect("typed MCP call must succeed");
601 let calls = calls.lock().unwrap();
602 assert_eq!(calls.len(), 1);
603 assert_eq!(calls[0].0, "test");
604 assert_eq!(calls[0].1, "echo");
605 assert!(
606 calls[0].2.contains("\"msg\":\"hi\""),
607 "args serialized to JSON: {}",
608 calls[0].2
609 );
610 }
611
612 #[tokio::test]
613 async fn all_optional_tool_callable_with_zero_args_serializes_empty_object() {
614 let (out, calls) = run(
617 r#"import test from "@mcp/test";
618 function main(): string {
619 return test.getTeams() as string;
620 }"#,
621 serde_json::json!("[]"),
622 None,
623 )
624 .await;
625 out.expect("zero-arg MCP call must succeed");
626 let calls = calls.lock().unwrap();
627 assert_eq!(calls.len(), 1);
628 assert_eq!(calls[0].1, "getTeams");
629 assert_eq!(calls[0].2, "{}", "empty optional args serialize to {{}}");
630 }
631
632 #[tokio::test]
633 async fn named_import_also_routes_to_mcp_call() {
634 let (out, calls) = run(
638 r#"import { echo } from "@mcp/test";
639 function main(): void {
640 const r = echo({ msg: "hi" });
641 assert(r.reply === "yo", "reply");
642 }"#,
643 serde_json::json!({ "reply": "yo" }),
644 None,
645 )
646 .await;
647 out.expect("named-import MCP call must succeed");
648 assert_eq!(calls.lock().unwrap()[0].1, "echo");
649 }
650
651 #[tokio::test]
652 async fn untyped_tool_returns_unknown_json_value() {
653 let (out, _) = run(
654 r#"import test from "@mcp/test";
655 interface RawResult { k: number; }
656 function main(): void {
657 const r = test.raw() as RawResult;
658 assert(r.k === 1, "raw");
659 }"#,
660 serde_json::json!({ "k": 1 }),
661 None,
662 )
663 .await;
664 out.expect("untyped MCP call must succeed");
665 }
666
667 #[tokio::test]
668 async fn zero_arg_tool_sends_empty_object() {
669 let (_out, calls) = run(
670 r#"import test from "@mcp/test";
671 function main(): void { test.raw(); }"#,
672 serde_json::json!({ "k": 1 }),
673 None,
674 )
675 .await;
676 assert_eq!(calls.lock().unwrap()[0].2, "{}");
677 }
678
679 #[tokio::test]
680 async fn denied_call_throws_before_dispatch() {
681 let (out, calls) = run(
682 r#"import test from "@mcp/test";
683 function main(): void { test.echo({ msg: "hi" }); }"#,
684 serde_json::json!({ "reply": "yo" }),
685 Some(Arc::new(DenyAll)),
686 )
687 .await;
688 let err = out.expect_err("a denied capability must throw");
689 assert!(
690 err.to_string().contains("permission denied"),
691 "unexpected error: {err}"
692 );
693 assert!(
694 calls.lock().unwrap().is_empty(),
695 "transport must not be reached when denied"
696 );
697 }
698}