feat(mcpx): WASM tool-source spike — wazero host, capability test (#1970)

pkgs/mcpx/internal/wasmtools/testdata/guest/echo.zig

@@ -0,0 +1,52 @@
+// echo.zig — a well-behaved core-WASM tool source for the #1970 spike.
+//
+// Built with: zig build-exe echo.zig -target wasm32-freestanding -O ReleaseSmall -rdynamic -fno-entry -femit-bin=../echo.wasm
+// No WASI, no ambient imports: the only import is the granted host
+// capability `mcpx.log`. If the host does not grant it, this module
+// cannot even be instantiated — that is the capability test's point.
+
+extern "mcpx" fn log(ptr: [*]const u8, len: u32) void;
+
+var input_buf: [8192]u8 align(8) = undefined;
+var output_buf: [8192]u8 align(8) = undefined;
+
+export fn mcpx_input_ptr() [*]u8 {
+    return &input_buf;
+}
+
+export fn mcpx_output_ptr() [*]const u8 {
+    return &output_buf;
+}
+
+fn writeOutput(s: []const u8) u32 {
+    @memcpy(output_buf[0..s.len], s);
+    return @intCast(s.len);
+}
+
+const tools_json = "[{\"name\":\"echo\",\"description\":\"echoes its arguments back\",\"inputSchema\":{\"type\":\"object\"}}]";
+
+export fn mcpx_list_tools() u32 {
+    return writeOutput(tools_json);
+}
+
+export fn mcpx_call_tool(name_len: u32, args_len: u32) u32 {
+    const name = input_buf[0..name_len];
+    const args = input_buf[name_len .. name_len + args_len];
+    // Use the granted capability: log the call through the host.
+    log(name.ptr, name_len);
+    // Result JSON: {"tool":"<name>","arguments":<args raw>}
+    var i: usize = 0;
+    const pre = "{\"tool\":\"";
+    @memcpy(output_buf[i .. i + pre.len], pre);
+    i += pre.len;
+    @memcpy(output_buf[i .. i + name.len], name);
+    i += name.len;
+    const mid = "\",\"arguments\":";
+    @memcpy(output_buf[i .. i + mid.len], mid);
+    i += mid.len;
+    @memcpy(output_buf[i .. i + args.len], args);
+    i += args.len;
+    output_buf[i] = '}';
+    i += 1;
+    return @intCast(i);
+}

pkgs/mcpx/internal/wasmtools/testdata/guest/greedy.zig

@@ -0,0 +1,19 @@
+// greedy.zig — a tool source that reaches for a capability it was not
+// granted: it imports `mcpx.read_file`. The spike host never offers that
+// import, so instantiation must fail. Deny-by-default, proven by the
+// loader, not by a permission check inside the guest.
+//
+// Built with: zig build-exe greedy.zig -target wasm32-freestanding -O ReleaseSmall -rdynamic -fno-entry -femit-bin=../greedy.wasm
+
+extern "mcpx" fn read_file(ptr: [*]const u8, len: u32) u32;
+
+var output_buf: [256]u8 align(8) = undefined;
+
+export fn mcpx_output_ptr() [*]const u8 {
+    return &output_buf;
+}
+
+export fn mcpx_list_tools() u32 {
+    const path = "/etc/passwd";
+    return read_file(path.ptr, path.len);
+}

pkgs/mcpx/internal/wasmtools/wasmtools.go

@@ -0,0 +1,166 @@
+// Package wasmtools is the #1970 spike: a WebAssembly module hosted by
+// wazero as an mcpx tool source (Approach B — a core-WASM .wasm module
+// as a tool source).
+//
+// Capability model: the host grants imports, the guest can reach
+// nothing else. This package instantiates NO WASI at all, so a guest
+// gets no filesystem, no network, no clock and no environment by
+// default; the only imports that can exist are the ones named in the
+// grant set passed to Load (today: "log"). A guest that imports
+// anything ungranted fails to instantiate — deny-by-default is a
+// property of the loader, not a check inside the guest.
+//
+// Guest ABI (core WASM, shaped so a component-model ABI can replace it
+// underneath later, per #1970 Approach C):
+//
+//	exports: memory
+//	         mcpx_input_ptr() i32   — host writes [name][argsJSON] here
+//	         mcpx_output_ptr() i32  — guest writes result JSON here
+//	         mcpx_list_tools() i32  — returns byte length of the JSON
+//	                                  tool list ([mcpclient.Tool])
+//	         mcpx_call_tool(nameLen, argsLen i32) i32
+//	                                — returns byte length of the JSON
+//	                                  result for that call
+//	imports: "mcpx"."log"(ptr, len) — only when "log" is granted
+//
+// The method shapes deliberately mirror mcpclient.Client
+// (ListTools(ctx) ([]Tool, error); CallTool(ctx, name, args)
+// (json.RawMessage, error), internal/mcpclient/client.go) so a WASM
+// source can sit behind the same surface as the stdio/HTTP transports
+// once the tool-source interface tracked by #1885 (child #83) lands.
+// Today there is no such registration point — upstreams are concrete
+// transports in internal/mcpclient and internal/adapter — so this
+// package is a standalone spike, not a fifth composition path.
+package wasmtools
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+
+	"github.com/dezren39/mcpx/internal/mcpclient"
+)
+
+// GrantLog is the host capability that lets a guest write one log line
+// per call through the host (collected on Source.Logs).
+const GrantLog = "log"
+
+const (
+	inputPtrFn  = "mcpx_input_ptr"
+	outputPtrFn = "mcpx_output_ptr"
+	listFn      = "mcpx_list_tools"
+	callFn      = "mcpx_call_tool"
+	maxMsg      = 1 << 20
+)
+
+// Source is one instantiated WASM tool source.
+type Source struct {
+	rt   wazero.Runtime
+	mod  api.Module
+	logs []string
+}
+
+// Load instantiates src (a core-WASM binary) granting exactly the
+// named host capabilities. Any import the guest declares beyond the
+// grant set — WASI included — makes Load fail.
+func Load(ctx context.Context, src []byte, grants ...string) (*Source, error) {
+	s := &Source{}
+	rt := wazero.NewRuntime(ctx)
+	s.rt = rt
+
+	granted := map[string]bool{}
+	for _, g := range grants {
+		granted[g] = true
+	}
+	host := rt.NewHostModuleBuilder("mcpx")
+	if granted[GrantLog] {
+		host.NewFunctionBuilder().WithFunc(func(_ context.Context, m api.Module, ptr, length uint32) {
+			b, ok := m.Memory().Read(ptr, length)
+			if ok {
+				s.logs = append(s.logs, string(b))
+			}
+		}).Export("log")
+	}
+	if _, err := host.Instantiate(ctx); err != nil {
+		rt.Close(ctx)
+		return nil, fmt.Errorf("wasmtools: instantiate host: %w", err)
+	}
+
+	mod, err := rt.Instantiate(ctx, src)
+	if err != nil {
+		rt.Close(ctx)
+		return nil, fmt.Errorf("wasmtools: instantiate guest (ungranted import?): %w", err)
+	}
+	s.mod = mod
+	for _, fn := range []string{inputPtrFn, outputPtrFn, listFn, callFn} {
+		if mod.ExportedFunction(fn) == nil {
+			rt.Close(ctx)
+			return nil, fmt.Errorf("wasmtools: guest does not export %s", fn)
+		}
+	}
+	return s, nil
+}
+
+// Close releases the runtime and every module in it.
+func (s *Source) Close(ctx context.Context) error { return s.rt.Close(ctx) }
+
+// Logs returns the lines the guest wrote through the granted "log"
+// capability, in order.
+func (s *Source) Logs() []string { return append([]string(nil), s.logs...) }
+
+func (s *Source) readOutput(ctx context.Context, length uint64) ([]byte, error) {
+	if length > maxMsg {
+		return nil, fmt.Errorf("wasmtools: guest returned %d bytes, over the %d cap", length, maxMsg)
+	}
+	res, err := s.mod.ExportedFunction(outputPtrFn).Call(ctx)
+	if err != nil {
+		return nil, err
+	}
+	b, ok := s.mod.Memory().Read(uint32(res[0]), uint32(length))
+	if !ok {
+		return nil, fmt.Errorf("wasmtools: guest output out of memory bounds")
+	}
+	return append([]byte(nil), b...), nil
+}
+
+// ListTools returns the guest's tool list.
+func (s *Source) ListTools(ctx context.Context) ([]mcpclient.Tool, error) {
+	res, err := s.mod.ExportedFunction(listFn).Call(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("wasmtools: list_tools: %w", err)
+	}
+	raw, err := s.readOutput(ctx, res[0])
+	if err != nil {
+		return nil, err
+	}
+	var tools []mcpclient.Tool
+	if err := json.Unmarshal(raw, &tools); err != nil {
+		return nil, fmt.Errorf("wasmtools: parse tool list: %w", err)
+	}
+	return tools, nil
+}
+
+// CallTool invokes one tool on the guest and returns its raw JSON result.
+func (s *Source) CallTool(ctx context.Context, name string, args any) (json.RawMessage, error) {
+	argsJSON, err := json.Marshal(args)
+	if err != nil {
+		return nil, err
+	}
+	res, err := s.mod.ExportedFunction(inputPtrFn).Call(ctx)
+	if err != nil {
+		return nil, err
+	}
+	ptr := uint32(res[0])
+	mem := s.mod.Memory()
+	if !mem.Write(ptr, []byte(name)) || !mem.Write(ptr+uint32(len(name)), argsJSON) {
+		return nil, fmt.Errorf("wasmtools: call input out of memory bounds")
+	}
+	out, err := s.mod.ExportedFunction(callFn).Call(ctx, uint64(uint32(len(name))), uint64(uint32(len(argsJSON))))
+	if err != nil {
+		return nil, fmt.Errorf("wasmtools: call_tool: %w", err)
+	}
+	return s.readOutput(ctx, out[0])
+}

pkgs/mcpx/internal/wasmtools/wasmtools_test.go

@@ -0,0 +1,74 @@
+package wasmtools_test
+
+import (
+	"context"
+	_ "embed"
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"github.com/dezren39/mcpx/internal/wasmtools"
+)
+
+//go:embed testdata/echo.wasm
+var echoWasm []byte
+
+//go:embed testdata/greedy.wasm
+var greedyWasm []byte
+
+func TestListAndCallWithGrantedLog(t *testing.T) {
+	ctx := context.Background()
+	src, err := wasmtools.Load(ctx, echoWasm, wasmtools.GrantLog)
+	if err != nil {
+		t.Fatalf("load echo with log grant: %v", err)
+	}
+	defer src.Close(ctx)
+
+	tools, err := src.ListTools(ctx)
+	if err != nil {
+		t.Fatalf("list: %v", err)
+	}
+	if len(tools) != 1 || tools[0].Name != "echo" {
+		t.Fatalf("tools = %+v", tools)
+	}
+
+	out, err := src.CallTool(ctx, "echo", map[string]any{"hello": "world"})
+	if err != nil {
+		t.Fatalf("call: %v", err)
+	}
+	var got struct {
+		Tool      string          `json:"tool"`
+		Arguments json.RawMessage `json:"arguments"`
+	}
+	if err := json.Unmarshal(out, &got); err != nil {
+		t.Fatalf("parse result %s: %v", out, err)
+	}
+	if got.Tool != "echo" || !strings.Contains(string(got.Arguments), "world") {
+		t.Fatalf("result = %s", out)
+	}
+	// The granted import worked: the guest logged the call via the host.
+	if logs := src.Logs(); len(logs) != 1 || logs[0] != "echo" {
+		t.Fatalf("logs = %v", logs)
+	}
+}
+
+func TestSameModuleWithoutGrantIsDenied(t *testing.T) {
+	ctx := context.Background()
+	// echo.wasm imports mcpx.log; with no grant the loader must refuse it.
+	if _, err := wasmtools.Load(ctx, echoWasm); err == nil {
+		t.Fatal("echo.wasm loaded without the log grant; want instantiation failure")
+	}
+}
+
+func TestUngrantedCapabilityImportIsDenied(t *testing.T) {
+	ctx := context.Background()
+	// greedy.wasm imports mcpx.read_file, which the host never offers —
+	// not even when every grantable capability is granted.
+	_, err := wasmtools.Load(ctx, greedyWasm, wasmtools.GrantLog)
+	if err == nil {
+		t.Fatal("greedy.wasm loaded; want instantiation failure for ungranted mcpx.read_file")
+	}
+	if !strings.Contains(err.Error(), "read_file") {
+		t.Fatalf("error should name the missing import, got: %v", err)
+	}
+}