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