diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 97a972f..8ee060e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -145,6 +145,8 @@ jobs: run: pnpm test:hosted-mcp-artifacts - name: Execute emitted hosted MCP artifacts run: pnpm test:hosted-mcp-artifacts-native + - name: Validate trusted gateway and native MCP Container lifecycle + run: pnpm test:hosted-mcp-container - run: pnpm test:web - run: pnpm test:shell - run: pnpm test:golden-path diff --git a/package.json b/package.json index 8d00156..14401f1 100644 --- a/package.json +++ b/package.json @@ -66,6 +66,7 @@ "test:mcp-resources": "node --test tests/mcp-resources.test.mjs", "test:mcp-invocation": "node --test tests/mcp-invocation.test.mjs", "test:mcp-hosting": "node --test --test-concurrency=1 tests/mcp-hosting-harness.test.mjs tests/mcp-hosting.test.mjs", + "test:hosted-mcp-container": "node --test --test-concurrency=1 tests/hosted-mcp-container-runtime.test.mjs tests/hosted-mcp-container-native.test.mjs", "test:hosted-mcp-contract": "node --test tests/hosted-mcp-contract.test.mjs", "build:hosted-mcp-artifacts": "node scripts/build-hosted-mcp-artifacts.mjs", "test:hosted-mcp-artifacts": "node --test tests/hosted-mcp-artifact.test.mjs tests/hosted-mcp-docker-image.test.mjs tests/hosted-mcp-native-cleanup.test.mjs tests/hosted-mcp-output.test.mjs tests/hosted-mcp-publication-store.test.mjs", diff --git a/tests/fixtures/hosted-mcp-container-worker.ts b/tests/fixtures/hosted-mcp-container-worker.ts new file mode 100644 index 0000000..ae0c24e --- /dev/null +++ b/tests/fixtures/hosted-mcp-container-worker.ts @@ -0,0 +1,247 @@ +import { DurableObject } from "cloudflare:workers"; +import { getContainer } from "@cloudflare/containers"; +import { getAgentByName } from "agents"; +import { + McpHostingClient as FixtureMcpHostingClient, + McpInvocationConversation, +} from "./mcp-hosting/gateway.ts"; +import { + HostedMcpContainer as NativeHostedMcpContainer, + createHostedMcpContainerGateway, + type HostedMcpAdmissionDecision, + type HostedMcpAdmissionRequest, + type HostedMcpAdmissionRevalidationRequest, + type HostedMcpContainerEnv, +} from "../../worker/hosted-mcp-container.ts"; + +type AdmissionStatus = "active" | "revoked" | "suspended"; + +type FixtureEnv = HostedMcpContainerEnv & { + HOSTED_MCP_CONTAINER: DurableObjectNamespace; + HOSTED_MCP_ADMISSION: DurableObjectNamespace; + MCP_CLIENT: DurableObjectNamespace; + MCP_FIXTURE_ADMIN: string; + MCP_GATEWAY_ORIGIN: string; + MCP_TEST_TOKEN: string; +}; + +function bearerMatches(authorization: string, token: string) { + const actual = new TextEncoder().encode(authorization); + const expected = new TextEncoder().encode(`Bearer ${token}`); + return ( + actual.byteLength === expected.byteLength && + crypto.subtle.timingSafeEqual(actual, expected) + ); +} + +export class HostedMcpAdmission extends DurableObject { + async fetch(request: Request) { + const path = new URL(request.url).pathname; + if (path === "/set") { + const { status } = await request.json<{ status: AdmissionStatus }>(); + const revision = + ((await this.ctx.storage.get("authority-revision")) ?? 1) + 1; + await this.ctx.storage.put({ + "authority-revision": revision, + status, + }); + return Response.json({ status }); + } + if (path !== "/admit") return new Response(null, { status: 404 }); + let decision: HostedMcpAdmissionDecision; + if ( + !bearerMatches( + request.headers.get("authorization") ?? "", + this.env.MCP_TEST_TOKEN, + ) + ) { + decision = { status: "unauthorized" }; + } else { + const status = + (await this.ctx.storage.get("status")) ?? "active"; + const authorityRevision = + (await this.ctx.storage.get("authority-revision")) ?? 1; + decision = + status === "active" + ? { status: "admitted", authorityRevision } + : { status }; + } + return Response.json(decision); + } +} + +export class HostedMcpContainer extends NativeHostedMcpContainer { + // A short test-only value proves stream-aware keepalive without weakening the + // reviewed 60..600 second production policy encoded in HOSTED_MCP_RUNTIME. + override sleepAfter = "1s"; + override envVars = { + FLA63_NATIVE_TOOLS: "1", + MCP_STARTUP_DELAY_MARKER: "/flarebot-fla63-started", + MCP_STARTUP_DELAY_MS: "5500", + MCP_TOKEN: this.env.MCP_TEST_TOKEN, + }; + + async stopForTest() { + await this.stop("SIGTERM"); + return this.getState(); + } + + async disableStartupDelayForTest() { + this.envVars.MCP_STARTUP_DELAY_MS = "0"; + await this.ctx.storage.put("startup-fast", true); + } + + override async fetch(request: Request) { + if (await this.ctx.storage.get("startup-fast")) + this.envVars.MCP_STARTUP_DELAY_MS = "0"; + return super.fetch(request); + } + + statsForTest() { + return this.containerFetch( + new Request("http://localhost/__stats", { + headers: { Authorization: `Bearer ${this.env.MCP_TEST_TOKEN}` }, + }), + ); + } +} + +export class McpHostingClient extends FixtureMcpHostingClient {} + +function fixtureAuthorized(request: Request, env: FixtureEnv) { + return request.headers.get("x-fixture-admin") === env.MCP_FIXTURE_ADMIN; +} + +async function clientResponse(request: Request, env: FixtureEnv) { + const agent = await getAgentByName(env.MCP_CLIENT, "fla-63"); + const path = new URL(request.url).pathname; + if (path === "/__fixture/client/runtime") return agent.fetch(request); + const { + action, + value, + name = "fla-63", + } = await request.json<{ + action: string; + value: unknown; + name?: string; + }>(); + try { + if (action === "definitions") + return Response.json(await agent.mcpToolDefinitions()); + if (action === "native-events") + return Response.json(await agent.diagnosticEvents()); + if (action === "policy") { + const policy = String(value); + if (!(["allow", "ask", "never"] as string[]).includes(policy)) + return new Response(null, { status: 400 }); + return Response.json( + await agent.configurePolicy(policy as "allow" | "ask" | "never"), + ); + } + if (action === "start-hosted") + return Response.json( + await agent.startHostedInvocation( + value as Parameters[0], + ), + ); + if (action === "cancel-hosted") + return Response.json(await agent.cancelHostedInvocation()); + if (action === "direct") + return Response.json( + await agent.directHosted( + value as Parameters[0], + ), + ); + return Response.json(await agent.clientRequest(name, action, value)); + } catch (error) { + return Response.json({ + fixtureError: error instanceof Error ? error.message : String(error), + }); + } +} + +function namedContainer(env: FixtureEnv) { + const configuration = env.HOSTED_MCP_RUNTIME as { serverId: string }; + return getContainer(env.HOSTED_MCP_CONTAINER, configuration.serverId); +} + +function admissionAuthority(env: FixtureEnv) { + const configuration = env.HOSTED_MCP_RUNTIME as { serverId: string }; + return env.HOSTED_MCP_ADMISSION.get( + env.HOSTED_MCP_ADMISSION.idFromName(configuration.serverId), + ); +} + +const gateway = createHostedMcpContainerGateway({ + async admit(request, env) { + return authorityDecision(request, env); + }, + async revalidate(request, env) { + return authorityDecision(request, env); + }, +}); + +async function authorityDecision( + request: HostedMcpAdmissionRequest | HostedMcpAdmissionRevalidationRequest, + env: FixtureEnv, +) { + request.authoritySignal.throwIfAborted(); + const response = await admissionAuthority(env).fetch( + new Request("http://admission/admit", { + headers: { + Authorization: request.authorization, + "X-Hosted-Mcp-Server": request.serverId, + }, + signal: request.authoritySignal, + }), + ); + request.authoritySignal.throwIfAborted(); + return response.json(); +} + +export default { + async fetch(request: Request, env: FixtureEnv) { + const url = new URL(request.url); + if (url.pathname.startsWith("/__fixture/client/")) + return clientResponse(request, env); + if (url.pathname === "/__fixture/admission") { + if (request.method !== "POST" || !fixtureAuthorized(request, env)) + return new Response(null, { status: 404 }); + const { status } = await request.json<{ status: AdmissionStatus }>(); + if (status !== "active" && status !== "revoked" && status !== "suspended") + return new Response(null, { status: 400 }); + return admissionAuthority(env).fetch( + new Request("http://admission/set", { + method: "POST", + body: JSON.stringify({ status }), + }), + ); + } + if (url.pathname === "/__fixture/container/state") { + if (request.method !== "GET" || !fixtureAuthorized(request, env)) + return new Response(null, { status: 404 }); + return Response.json(await namedContainer(env).getState()); + } + if (url.pathname === "/__fixture/container/stop") { + if (request.method !== "POST" || !fixtureAuthorized(request, env)) + return new Response(null, { status: 404 }); + return Response.json(await namedContainer(env).stopForTest()); + } + if (url.pathname === "/__fixture/container/startup-fast") { + if (request.method !== "POST" || !fixtureAuthorized(request, env)) + return new Response(null, { status: 404 }); + await namedContainer(env).disableStartupDelayForTest(); + return Response.json({ ok: true }); + } + if (url.pathname === "/__fixture/container/stats") { + if (request.method !== "GET" || !fixtureAuthorized(request, env)) + return new Response(null, { status: 404 }); + return namedContainer(env).statsForTest(); + } + if (url.pathname.startsWith("/__fixture/") || url.pathname === "/mcp") + return new Response(null, { status: 404 }); + return gateway.fetch(request, env); + }, +} satisfies ExportedHandler; + +export { McpInvocationConversation }; diff --git a/tests/fixtures/mcp-hosting/gateway.ts b/tests/fixtures/mcp-hosting/gateway.ts index 72c2407..2f2de90 100644 --- a/tests/fixtures/mcp-hosting/gateway.ts +++ b/tests/fixtures/mcp-hosting/gateway.ts @@ -63,6 +63,8 @@ export class GoMcpContainer extends Container { } export class McpHostingClient extends McpInvocationFixture { + declare env: GatewayEnv; + protected override mcpNetwork(input: RequestInfo | URL, init?: RequestInit) { const request = new Request(input, init); const url = new URL(request.url); @@ -146,8 +148,14 @@ async function clientResponse(request: Request, env: GatewayEnv) { return Response.json(await agent.mcpToolDefinitions()); if (action === "native-events") return Response.json(await agent.diagnosticEvents()); - if (action === "policy") - return Response.json(await agent.configurePolicy(String(value))); + if (action === "policy") { + const policy = String(value); + if (!(["allow", "ask", "never"] as string[]).includes(policy)) + return new Response(null, { status: 400 }); + return Response.json( + await agent.configurePolicy(policy as "allow" | "ask" | "never"), + ); + } if (action === "start-hosted") return Response.json( await agent.startHostedInvocation( diff --git a/tests/fixtures/mcp-hosting/go/main.go b/tests/fixtures/mcp-hosting/go/main.go index 02537d8..e363233 100644 --- a/tests/fixtures/mcp-hosting/go/main.go +++ b/tests/fixtures/mcp-hosting/go/main.go @@ -30,6 +30,7 @@ type requestObservation struct { ContentType string `json:"contentType"` Accept string `json:"accept"` Origin string `json:"origin,omitempty"` + ParamRegion string `json:"paramRegion,omitempty"` BodyVersion string `json:"bodyVersion,omitempty"` BodyMethod string `json:"bodyMethod,omitempty"` ClientName string `json:"clientName,omitempty"` @@ -154,6 +155,15 @@ type slowInput struct { DelayMS int `json:"delayMs" jsonschema:"required,delay in milliseconds"` } +type headerInput struct { + Region string `json:"region"` +} + +type longStreamInput struct { + Chunks int `json:"chunks"` + DelayMS int `json:"delayMs"` +} + func createServer(state *fixtureState) *mcp.Server { server := mcp.NewServer( &mcp.Implementation{Name: "Flarebot Go hosting fixture", Version: "1.0.0"}, @@ -262,6 +272,65 @@ func createServer(state *fixtureState) *mcp.Server { }, func(_ context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { return textResult("intercepted by fixture middleware"), nil, nil }) + if os.Getenv("FLA63_NATIVE_TOOLS") == "1" { + mcp.AddTool(server, &mcp.Tool{ + Name: "header_echo", + Description: "Validate modern x-mcp-header transport metadata.", + InputSchema: json.RawMessage(`{"type":"object","properties":{"region":{"type":"string","x-mcp-header":"Region"}},"required":["region"],"additionalProperties":false}`), + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: readOnly, OpenWorldHint: &closedWorld}, + }, func(_ context.Context, _ *mcp.CallToolRequest, input headerInput) (*mcp.CallToolResult, any, error) { + state.recordTool("header_echo") + return textResult("region:" + input.Region), nil, nil + }) + mcp.AddTool(server, &mcp.Tool{ + Name: "large_response", + Description: "Return a response beyond the hosted gateway budget.", + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: readOnly, OpenWorldHint: &closedWorld}, + }, func(_ context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + state.recordTool("large_response") + return textResult(string(bytes.Repeat([]byte("x"), 70_000))), nil, nil + }) + mcp.AddTool(server, &mcp.Tool{ + Name: "long_stream", + Description: "Keep an SSE response active beyond the test idle policy.", + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: readOnly, OpenWorldHint: &closedWorld}, + }, func(ctx context.Context, request *mcp.CallToolRequest, input longStreamInput) (*mcp.CallToolResult, any, error) { + state.recordTool("long_stream") + if input.Chunks < 2 || input.Chunks > 20 || input.DelayMS < 50 || input.DelayMS > 500 { + return nil, nil, errors.New("invalid long stream") + } + for index := 0; index < input.Chunks; index++ { + if token := request.Params.GetProgressToken(); token != nil { + state.mu.Lock() + state.progress++ + state.mu.Unlock() + if err := request.Session.NotifyProgress(ctx, &mcp.ProgressNotificationParams{ + ProgressToken: token, + Progress: float64(index + 1), + Total: float64(input.Chunks), + Message: "native runtime keepalive", + }); err != nil { + return nil, nil, err + } + } + select { + case <-ctx.Done(): + return nil, nil, ctx.Err() + case <-time.After(time.Duration(input.DelayMS) * time.Millisecond): + } + } + return textResult("long-stream-complete"), nil, nil + }) + mcp.AddTool(server, &mcp.Tool{ + Name: "crash_process", + Description: "Terminate the native fixture after request dispatch.", + Annotations: &mcp.ToolAnnotations{OpenWorldHint: &closedWorld}, + }, func(_ context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + state.recordTool("crash_process") + os.Exit(17) + return nil, nil, errors.New("unreachable") + }) + } return server } @@ -274,6 +343,7 @@ func observeRequest(request *http.Request, payload []byte) requestObservation { ContentType: request.Header.Get("Content-Type"), Accept: request.Header.Get("Accept"), Origin: request.Header.Get("Origin"), + ParamRegion: request.Header.Get("Mcp-Param-Region"), } var envelope struct { Method string `json:"method"` @@ -345,6 +415,18 @@ func mcpMiddleware(state *fixtureState, next http.Handler) http.Handler { } func main() { + if delay, err := strconv.Atoi(os.Getenv("MCP_STARTUP_DELAY_MS")); err == nil && delay > 0 { + marker := os.Getenv("MCP_STARTUP_DELAY_MARKER") + if marker == "" { + marker = "/tmp/flarebot-mcp-started" + } + if _, err := os.Stat(marker); errors.Is(err, os.ErrNotExist) { + if err := os.WriteFile(marker, []byte("ready"), 0o600); err != nil { + log.Fatal(err) + } + time.Sleep(time.Duration(delay) * time.Millisecond) + } + } state := newFixtureState() server := createServer(state) mcpHandler := mcp.NewStreamableHTTPHandler( diff --git a/tests/hosted-mcp-container-native.test.mjs b/tests/hosted-mcp-container-native.test.mjs new file mode 100644 index 0000000..0f886d7 --- /dev/null +++ b/tests/hosted-mcp-container-native.test.mjs @@ -0,0 +1,551 @@ +import assert from "node:assert/strict"; +import { randomUUID } from "node:crypto"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { test } from "node:test"; +import { unstable_dev } from "wrangler"; +import { buildHostingArtifacts } from "./fixtures/mcp-hosting/artifact-build.mjs"; +import { + ResourceScope, + boundedFetch, + eventually, + execute, + freePort, + withDeadline, +} from "./fixtures/mcp-hosting/harness.mjs"; + +const token = "fla-63-local-credential-claim"; +const admin = "fla-63-local-admin"; +const protocolRevision = "2026-07-28"; +const serverId = `hosted-mcp-${"6".repeat(32)}`; +const phase = (name) => console.log(`FLA63_NATIVE_PHASE=${name}`); + +async function removeOwnedImage(executeDocker, imageTag) { + try { + await executeDocker("docker", ["image", "rm", imageTag]); + } catch (error) { + const message = error instanceof Error ? error.message : ""; + const missingReference = message.match(/No such image: ([^\s]+)/)?.[1]; + if ( + missingReference === imageTag || + missingReference === `${imageTag}:latest` || + missingReference === `docker.io/library/${imageTag}:latest` + ) + return; + throw error; + } +} + +function textOf(result) { + return result.content?.find((part) => part.type === "text")?.text; +} + +function modernEnvelope(method, params, id = randomUUID()) { + return { + jsonrpc: "2.0", + id, + method, + params: { + ...params, + _meta: { + ...params?._meta, + "io.modelcontextprotocol/protocolVersion": protocolRevision, + "io.modelcontextprotocol/clientCapabilities": {}, + "io.modelcontextprotocol/clientInfo": { + name: "Flarebot FLA-63 native runtime test", + version: "1.0.0", + }, + }, + }, + }; +} + +test("owned image cleanup ignores only its exact absent tag", async () => { + const imageTag = "flarebot-fla63-owned-absent"; + await removeOwnedImage(async () => { + throw new Error( + `Error response from daemon: No such image: ${imageTag}:latest`, + ); + }, imageTag); +}); + +test("owned image cleanup surfaces daemon and timeout failures", async () => { + for (const error of [ + new Error("Cannot connect to the Docker daemon"), + new Error("docker image cleanup exceeded its deadline"), + new Error("conflict: image is being used"), + new Error("No such image: a-different-tag"), + ]) { + await assert.rejects( + removeOwnedImage(async () => { + throw error; + }, "flarebot-fla63-owned"), + (actual) => actual === error, + ); + } +}); + +test( + "trusted gateway drives the native Container runtime through lifecycle failures", + { timeout: 180_000 }, + async (context) => { + const suiteSignal = AbortSignal.timeout(165_000); + const scope = new ResourceScope(); + context.after(() => scope.dispose()); + const temporary = await mkdtemp(join(tmpdir(), "flarebot-fla63-")); + scope.defer(() => rm(temporary, { recursive: true, force: true })); + const port = await freePort(suiteSignal); + const origin = `http://127.0.0.1:${port}`; + const imageTag = `flarebot-fla63-${process.pid}-${randomUUID().slice(0, 8)}`; + const executeBounded = (command, args, options = {}, timeout) => + execute(command, args, options, suiteSignal, timeout); + const cleanupExecute = (command, args, options = {}) => + execute(command, args, options, AbortSignal.timeout(30_000), 30_000); + const fetchBounded = (path, init = {}, timeout) => + boundedFetch(`${origin}${path}`, init, suiteSignal, timeout); + const adminFetch = (path, init = {}) => + fetchBounded(path, { + ...init, + headers: { "X-Fixture-Admin": admin, ...init.headers }, + }); + + const manifest = JSON.parse( + await readFile("tests/fixtures/hosted-mcp/container.json", "utf8"), + ); + let tools; + + const request = async (path, body) => { + const response = await fetchBounded(path, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }); + const result = await response.json(); + assert.equal(response.ok, true, JSON.stringify(result)); + assert.ok(!result.fixtureError, JSON.stringify(result)); + return result; + }; + const runtime = async (action, args = {}) => { + const result = await request("/__fixture/client/runtime", { + action, + ...args, + }); + assert.equal(result.ok, true, JSON.stringify(result)); + return result.value; + }; + const control = (action, value, name) => + request("/__fixture/client/control", { action, value, name }); + const connections = async () => (await runtime("list")).connections; + const settled = (id) => + eventually( + connections, + (items) => { + const item = items.find((candidate) => candidate.id === id); + return item?.state !== "connecting" ? item : undefined; + }, + `native MCP connection ${id} settlement`, + suiteSignal, + ); + const containerState = async () => { + const response = await adminFetch("/__fixture/container/state"); + assert.equal(response.status, 200); + return response.json(); + }; + const containerStats = async () => { + const response = await adminFetch("/__fixture/container/stats"); + assert.equal(response.status, 200); + return response.json(); + }; + const setAdmission = async (status) => { + const response = await adminFetch("/__fixture/admission", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ status }), + }); + assert.equal(response.status, 200); + assert.deepEqual(await response.json(), { status }); + }; + const rawMcp = (method, params, options = {}) => { + const body = modernEnvelope(method, params, options.id); + const headers = { + Accept: "application/json, text/event-stream", + Authorization: `Bearer ${token}`, + "Content-Type": "application/json; charset=utf-8", + "MCP-Protocol-Version": protocolRevision, + "Mcp-Method": method, + Origin: origin, + ...options.headers, + }; + if (method === "tools/call") headers["Mcp-Name"] = params.name; + return fetchBounded( + "/container/mcp", + { + method: "POST", + headers, + body: JSON.stringify(body), + signal: options.signal, + }, + options.timeout, + ); + }; + const tool = (name) => { + const value = tools.find((candidate) => candidate.metadata.name === name); + assert.ok(value, `Missing native fixture tool ${name}`); + return value; + }; + const invoke = async (name, input) => { + const definition = tool(name); + const response = await control( + "direct", + { + reference: { + source: definition.metadata.source, + id: definition.id, + fingerprint: definition.fingerprint, + }, + input, + approvedOnce: true, + }, + randomUUID(), + ); + return response.result; + }; + + try { + phase("fla63-artifact-build"); + scope.defer(() => removeOwnedImage(cleanupExecute, imageTag)); + const { imageDirectory } = await buildHostingArtifacts({ + execute: executeBounded, + imageTag, + temporary, + }); + const configPath = join(temporary, "wrangler.json"); + await writeFile( + configPath, + JSON.stringify({ + name: `flarebot-fla63-native-${process.pid}`, + main: resolve("tests/fixtures/hosted-mcp-container-worker.ts"), + compatibility_date: "2026-09-04", + compatibility_flags: ["nodejs_compat"], + vars: { + HOSTED_MCP_RUNTIME: { + manifest, + publicOrigin: origin, + serverId, + }, + MCP_FIXTURE_ADMIN: admin, + MCP_GATEWAY_ORIGIN: origin, + MCP_TEST_TOKEN: token, + }, + containers: [ + { + class_name: "HostedMcpContainer", + image: join(imageDirectory, "Dockerfile"), + image_build_context: imageDirectory, + instance_type: "lite", + max_instances: 1, + }, + ], + durable_objects: { + bindings: [ + { + name: "HOSTED_MCP_CONTAINER", + class_name: "HostedMcpContainer", + }, + { + name: "HOSTED_MCP_ADMISSION", + class_name: "HostedMcpAdmission", + }, + { name: "MCP_CLIENT", class_name: "McpHostingClient" }, + ], + }, + migrations: [ + { + tag: "v1", + new_sqlite_classes: [ + "HostedMcpContainer", + "HostedMcpAdmission", + "McpHostingClient", + ], + }, + ], + }), + ); + await scope.acquire( + "FLA-63 native workerd startup", + () => + unstable_dev("tests/fixtures/hosted-mcp-container-worker.ts", { + config: configPath, + local: true, + ip: "127.0.0.1", + port, + inspectorPort: 0, + logLevel: "error", + experimental: { + disableExperimentalWarning: true, + enableContainers: true, + watch: false, + }, + }), + (worker) => + withDeadline( + worker.stop(), + "FLA-63 native workerd teardown", + AbortSignal.timeout(15_000), + 10_000, + ), + suiteSignal, + 120_000, + ); + + phase("fla63-admission-before-wake"); + assert.equal((await containerState()).status, "stopped"); + for (const [status, expectedStatus] of [ + ["suspended", 503], + ["revoked", 401], + ]) { + await setAdmission(status); + const rejected = await rawMcp("server/discover", {}); + const rejectedBody = await rejected.clone().json(); + assert.equal( + rejectedBody.error.code, + status === "suspended" ? "hosted_mcp_suspended" : "invalid_token", + JSON.stringify(rejectedBody), + ); + assert.equal( + rejected.status, + expectedStatus, + await rejected.clone().text(), + ); + assert.ok( + ["stopped", "stopped_with_code"].includes( + (await containerState()).status, + ), + ); + } + await setAdmission("active"); + const wrongCredential = await rawMcp( + "server/discover", + {}, + { + headers: { Authorization: "Bearer wrong" }, + }, + ); + assert.equal(wrongCredential.status, 401); + assert.ok( + ["stopped", "stopped_with_code"].includes( + (await containerState()).status, + ), + ); + assert.equal((await fetchBounded("/mcp")).status, 404); + + phase("fla63-stalled-readiness"); + const readinessStarted = performance.now(); + const notReady = await rawMcp("server/discover", {}, { timeout: 10_000 }); + const readinessElapsed = performance.now() - readinessStarted; + assert.equal(notReady.status, 503); + assert.equal((await notReady.json()).error.code, "hosted_mcp_not_ready"); + assert.ok(readinessElapsed >= 4_500 && readinessElapsed < 8_000); + await new Promise((resolveDelay) => setTimeout(resolveDelay, 700)); + phase("fla63-readiness-retry"); + const explicitReadinessRetry = await rawMcp("server/discover", {}); + assert.equal( + explicitReadinessRetry.status, + 200, + await explicitReadinessRetry.clone().text(), + ); + await explicitReadinessRetry.arrayBuffer(); + assert.equal((await containerState()).status, "healthy"); + const startupFast = await adminFetch( + "/__fixture/container/startup-fast", + { method: "POST" }, + ); + assert.equal(startupFast.status, 200); + await startupFast.arrayBuffer(); + + phase("fla63-client-discovery"); + let connection = await runtime("add", { + value: { + name: "FLA-63 native Container", + endpoint: "https://gateway.fla57.test/container/mcp", + authMode: "headers", + }, + }); + connection = await settled(connection.id); + assert.equal(connection.state, "authenticating"); + connection = await runtime("headers", { + ...connection, + value: [{ name: "Authorization", value: `Bearer ${token}` }], + }); + connection = await settled(connection.id); + assert.equal(connection.state, "ready", JSON.stringify(connection)); + const enabled = connection.capabilities + .filter((capability) => capability.kind === "tool" && !capability.issue) + .map(({ id, fingerprint }) => ({ id, fingerprint })); + assert.equal(enabled.length, 10, JSON.stringify(connection.capabilities)); + await runtime("tools", { + id: connection.id, + value: { + settingsRevision: connection.settingsRevision, + enabled: true, + tools: enabled, + }, + }); + tools = (await control("definitions")).filter( + (definition) => definition.metadata.source.id === connection.id, + ); + + phase("fla63-client-calls"); + assert.equal( + textOf(await invoke("read_status", { key: "status" })), + "container-ready", + ); + assert.equal( + textOf(await invoke("header_echo", { region: "eu-central-1" })), + "region:eu-central-1", + ); + const afterHeader = await containerStats(); + assert.equal(afterHeader.requests.at(-1).paramRegion, "eu-central-1"); + + assert.equal( + textOf(await invoke("header_echo", { region: "日本語" })), + "region:日本語", + ); + const afterEncodedHeader = await containerStats(); + assert.equal( + afterEncodedHeader.requests.at(-1).paramRegion, + "=?base64?5pel5pys6Kqe?=", + ); + + const mismatch = await rawMcp( + "tools/call", + { + name: "header_echo", + arguments: { region: "eu-central-1" }, + }, + { + headers: { "Mcp-Param-Region": "us-west-2" }, + }, + ); + assert.equal(mismatch.status, 400); + assert.equal((await mismatch.json()).error.code, -32020); + + const activeResponse = await rawMcp("tools/call", { + name: "long_stream", + arguments: { chunks: 10, delayMs: 250 }, + _meta: { progressToken: "keepalive" }, + }); + assert.equal(activeResponse.status, 200); + await new Promise((resolveDelay) => setTimeout(resolveDelay, 1_300)); + assert.equal((await containerState()).status, "healthy"); + assert.match(await activeResponse.text(), /long-stream-complete/); + + phase("fla63-concurrency"); + const concurrent = await Promise.all( + Array.from({ length: 5 }, (_, index) => + rawMcp("tools/call", { + name: "long_stream", + arguments: { chunks: 8, delayMs: 250 }, + _meta: { progressToken: `concurrency-${index}` }, + }), + ), + ); + assert.deepEqual( + concurrent.map(({ status }) => status).sort((a, b) => a - b), + [200, 200, 200, 200, 429], + ); + await Promise.all(concurrent.map((response) => response.arrayBuffer())); + + phase("fla63-cancellation"); + const slowTool = tool("slow_mutation"); + const mutationId = "fla63-cancel-once"; + const beforeCancel = await containerStats(); + const active = control("start-hosted", { + reference: { + source: slowTool.metadata.source, + id: slowTool.id, + fingerprint: slowTool.fingerprint, + }, + input: { mutationId, delayMs: 3_000 }, + toolCallId: randomUUID(), + }); + await eventually( + containerStats, + (snapshot) => snapshot.mutations.includes(mutationId), + "native mutation dispatch before cancellation", + suiteSignal, + ); + assert.deepEqual(await control("cancel-hosted"), { cancelled: true }); + assert.deepEqual(await active, { + error: { + code: "outcome_uncertain", + message: + "The MCP action was dispatched but no reliable terminal outcome was received", + }, + outcome: "uncertain", + replay: "forbidden", + }); + const afterCancel = await containerStats(); + assert.equal( + afterCancel.toolCalls.slow_mutation, + (beforeCancel.toolCalls.slow_mutation ?? 0) + 1, + ); + + assert.deepEqual(await invoke("large_response", {}), { + error: { + code: "outcome_uncertain", + message: + "The MCP action was dispatched but no reliable terminal outcome was received", + }, + outcome: "uncertain", + replay: "forbidden", + }); + + phase("fla63-crash-restart"); + const wakeAfterBudget = await invoke("read_status", { key: "status" }); + assert.equal( + textOf(wakeAfterBudget), + "container-ready", + JSON.stringify(wakeAfterBudget), + ); + const beforeCrash = await containerStats(); + assert.deepEqual(await invoke("crash_process", {}), { + error: { + code: "outcome_uncertain", + message: + "The MCP action was dispatched but no reliable terminal outcome was received", + }, + outcome: "uncertain", + replay: "forbidden", + }); + await eventually( + containerState, + (state) => ["stopped", "stopped_with_code"].includes(state.status), + "native process crash", + suiteSignal, + ); + assert.equal( + textOf(await invoke("read_status", { key: "status" })), + "container-ready", + ); + assert.notEqual((await containerStats()).bootId, beforeCrash.bootId); + + phase("fla63-sleep-wake"); + await eventually( + containerState, + (state) => ["stopped", "stopped_with_code"].includes(state.status), + "idle native Container sleep", + suiteSignal, + 5_000, + ); + assert.equal( + textOf(await invoke("read_status", { key: "status" })), + "container-ready", + ); + assert.equal((await containerState()).status, "healthy"); + } finally { + await scope.dispose(); + } + }, +); diff --git a/tests/hosted-mcp-container-runtime.test.mjs b/tests/hosted-mcp-container-runtime.test.mjs new file mode 100644 index 0000000..072c3bf --- /dev/null +++ b/tests/hosted-mcp-container-runtime.test.mjs @@ -0,0 +1,705 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import { test } from "node:test"; +import { build } from "esbuild"; + +const protocolRevision = "2026-07-28"; +const serverId = `hosted-mcp-${"a".repeat(32)}`; +const manifest = JSON.parse( + await readFile( + new URL("fixtures/hosted-mcp/container.json", import.meta.url), + "utf8", + ), +); + +const bundle = await build({ + entryPoints: [ + new URL("../worker/hosted-mcp-container.ts", import.meta.url).pathname, + ], + bundle: true, + format: "esm", + platform: "neutral", + write: false, + plugins: [ + { + name: "container-sdk-fixture", + setup(builder) { + builder.onResolve({ filter: /^@cloudflare\/containers$/ }, () => ({ + path: "container-sdk", + namespace: "fixture", + })); + builder.onLoad({ filter: /.*/, namespace: "fixture" }, () => ({ + contents: ` + export class Container { + constructor(ctx, env) { + this.ctx = ctx; + this.env = env; + this.container = ctx.container; + } + startAndWaitForPorts(options) { + return this.ctx.startAndWaitForPorts(options); + } + containerFetch(request) { + return this.ctx.containerFetch(request); + } + renewActivityTimeout() { + this.ctx.renewActivityTimeout(); + } + async onActivityExpired() { + this.ctx.activityExpired(); + } + } + `, + loader: "js", + })); + }, + }, + ], +}); +const runtime = await import( + `data:text/javascript;base64,${Buffer.from(bundle.outputFiles[0].contents).toString("base64")}` +); + +function configuration(overrides = {}) { + return { + manifest: structuredClone(manifest), + publicOrigin: "https://mcp.test", + serverId, + ...overrides, + }; +} + +function envelope({ + method = "tools/call", + name = "read_status", + arguments: toolArguments = { key: "status" }, + id = "request-1", +} = {}) { + const params = { + _meta: { + "io.modelcontextprotocol/protocolVersion": protocolRevision, + "io.modelcontextprotocol/clientCapabilities": {}, + "io.modelcontextprotocol/clientInfo": { + name: "Flarebot runtime contract test", + version: "1.0.0", + }, + }, + }; + if (method === "tools/call") { + params.name = name; + params.arguments = toolArguments; + } + return { jsonrpc: "2.0", id, method, params }; +} + +function mcpRequest(options = {}) { + const body = options.body ?? envelope(options); + const method = body.method ?? "tools/call"; + const headers = new Headers({ + Accept: "application/json, text/event-stream", + Authorization: "Bearer fixture-token", + "Content-Type": "application/json; charset=utf-8", + Host: "mcp.test", + "MCP-Protocol-Version": protocolRevision, + "Mcp-Method": method, + Origin: "https://mcp.test", + ...options.headers, + }); + if (method === "tools/call" && !headers.has("Mcp-Name")) + headers.set("Mcp-Name", body.params?.name ?? "read_status"); + return new Request(options.url ?? "https://mcp.test/container/mcp", { + method: options.requestMethod ?? "POST", + headers, + body: + options.rawBody === undefined ? JSON.stringify(body) : options.rawBody, + signal: options.signal, + ...(options.rawBody instanceof ReadableStream ? { duplex: "half" } : {}), + }); +} + +function fixture(options = {}) { + const calls = []; + const names = []; + const stub = { + async fetch(request) { + calls.push(request); + return ( + options.response?.(request) ?? + Response.json({ jsonrpc: "2.0", id: "request-1", result: {} }) + ); + }, + }; + const env = { + HOSTED_MCP_CONTAINER: { + idFromName(name) { + names.push(name); + return name; + }, + get() { + return stub; + }, + }, + HOSTED_MCP_RUNTIME: options.configuration ?? configuration(), + }; + let admissionCalls = 0; + let revalidationCalls = 0; + const admissionRequests = []; + const revalidationRequests = []; + const gateway = runtime.createHostedMcpContainerGateway({ + async admit(request) { + admissionCalls++; + admissionRequests.push(request); + return ( + options.admission ?? { + status: "admitted", + authorityRevision: 1, + } + ); + }, + async revalidate(request) { + revalidationCalls++; + revalidationRequests.push(request); + return ( + options.revalidation ?? + options.admission ?? { + status: "admitted", + authorityRevision: 1, + } + ); + }, + }); + return { + calls, + env, + gateway, + names, + admissionCalls: () => admissionCalls, + admissionRequests, + revalidationCalls: () => revalidationCalls, + revalidationRequests, + }; +} + +async function errorCode(response) { + return (await response.json()).error.code; +} + +test("parses only the reviewed Container manifest and stable server identity", () => { + const parsed = runtime.parseHostedMcpContainerConfiguration(configuration()); + assert.equal(parsed.serverId, serverId); + assert.equal(parsed.manifest.runtime.type, "container"); + assert.equal(parsed.manifest.protocol.protocolRevision, protocolRevision); + + for (const invalid of [ + { ...configuration(), extra: true }, + configuration({ serverId: "shared-sandbox" }), + configuration({ publicOrigin: "https://MCP.test" }), + configuration({ publicOrigin: "https://mcp.test/path" }), + configuration({ publicOrigin: "http://mcp.test" }), + configuration({ manifest: { ...manifest, runtime: { type: "sandbox" } } }), + ]) { + assert.throws( + () => runtime.parseHostedMcpContainerConfiguration(invalid), + runtime.HostedMcpContainerConfigurationError, + ); + } +}); + +test("rejects authorization, revocation, suspension, host and origin before wake", async () => { + for (const [admission, expectedStatus] of [ + ["unauthorized", 401], + ["revoked", 401], + ["suspended", 503], + ["unavailable", 503], + ]) { + const instance = fixture({ admission: { status: admission } }); + const response = await instance.gateway.fetch(mcpRequest(), instance.env); + assert.equal(response.status, expectedStatus); + assert.equal(instance.calls.length, 0); + assert.equal(instance.names.length, 0); + } + + for (const request of [ + mcpRequest({ headers: { Host: "attacker.test" } }), + mcpRequest({ headers: { Origin: "https://attacker.test" } }), + mcpRequest({ + url: "http://mcp.test/container/mcp", + headers: { Host: "mcp.test", Origin: "http://mcp.test" }, + }), + mcpRequest({ + url: "https://mcp.test:8443/container/mcp", + headers: { Host: "mcp.test:8443", Origin: "https://mcp.test:8443" }, + }), + mcpRequest({ url: "https://mcp.test/container/mcp?wake=true" }), + ]) { + const instance = fixture(); + const response = await instance.gateway.fetch(request, instance.env); + assert.ok(response.status >= 400); + assert.equal(instance.admissionCalls(), 0); + assert.equal(instance.revalidationCalls(), 0); + assert.equal(instance.calls.length, 0); + } +}); + +test("does not consume an unauthorized or pre-aborted streaming body", async () => { + let pulls = 0; + const unreadBody = () => + new ReadableStream( + { + pull() { + pulls++; + throw new Error("body must remain unread"); + }, + }, + { highWaterMark: 0 }, + ); + + const noCredential = fixture(); + const unauthorizedRequest = mcpRequest({ + headers: { Authorization: "" }, + rawBody: unreadBody(), + }); + assert.equal( + (await noCredential.gateway.fetch(unauthorizedRequest, noCredential.env)) + .status, + 401, + ); + assert.equal(pulls, 0); + assert.equal(noCredential.calls.length, 0); + + const controller = new AbortController(); + controller.abort(new Error("cancel before admission")); + const aborted = fixture(); + await assert.rejects( + aborted.gateway.fetch( + mcpRequest({ rawBody: unreadBody(), signal: controller.signal }), + aborted.env, + ), + /cancel before admission/, + ); + assert.equal(pulls, 0); + assert.equal(aborted.admissionCalls(), 0); + assert.equal(aborted.calls.length, 0); +}); + +test("cancels a stalled admitted upload at the total gateway deadline before wake", async () => { + const timeout = AbortSignal.timeout; + const deadline = new AbortController(); + let cancelled = false; + AbortSignal.timeout = (milliseconds) => { + assert.equal(milliseconds, 25_000); + return deadline.signal; + }; + try { + const body = new ReadableStream( + { + pull() { + return new Promise(() => {}); + }, + cancel() { + cancelled = true; + }, + }, + { highWaterMark: 0 }, + ); + const instance = fixture(); + const pending = instance.gateway.fetch( + mcpRequest({ rawBody: body }), + instance.env, + ); + await new Promise((resolve) => setImmediate(resolve)); + deadline.abort(new DOMException("request deadline", "TimeoutError")); + const response = await pending; + assert.equal(response.status, 504); + assert.equal(await errorCode(response), "hosted_mcp_deadline_exceeded"); + assert.equal(cancelled, true); + assert.equal(instance.admissionCalls(), 1); + assert.equal(instance.revalidationCalls(), 0); + assert.equal(instance.names.length, 0); + assert.equal(instance.calls.length, 0); + } finally { + AbortSignal.timeout = timeout; + } +}); + +test("revalidates the admitted authority revision after reading the body", async () => { + for (const status of ["revoked", "suspended"]) { + let current = { status: "admitted", authorityRevision: 7 }; + const options = { + admission: current, + get revalidation() { + return current; + }, + }; + const body = new ReadableStream( + { + pull(controller) { + current = { status }; + controller.enqueue( + new TextEncoder().encode(JSON.stringify(envelope())), + ); + controller.close(); + }, + }, + { highWaterMark: 0 }, + ); + const instance = fixture(options); + const response = await instance.gateway.fetch( + mcpRequest({ rawBody: body }), + instance.env, + ); + assert.equal(response.status, status === "revoked" ? 401 : 503); + assert.equal(instance.admissionCalls(), 1); + assert.equal(instance.revalidationCalls(), 1); + assert.equal(instance.revalidationRequests[0].authorityRevision, 7); + assert.equal(instance.names.length, 0); + assert.equal(instance.calls.length, 0); + } + + const changedRevision = fixture({ + admission: { status: "admitted", authorityRevision: 4 }, + revalidation: { status: "admitted", authorityRevision: 5 }, + }); + const changed = await changedRevision.gateway.fetch( + mcpRequest(), + changedRevision.env, + ); + assert.equal(await errorCode(changed), "hosted_mcp_admission_unavailable"); + assert.equal(changedRevision.calls.length, 0); +}); + +test("enforces body, revision, session and routing headers before dispatch", async () => { + const cases = [ + { + request: mcpRequest({ rawBody: "x".repeat(65_537) }), + status: 413, + code: "hosted_mcp_request_too_large", + }, + { + request: mcpRequest({ + headers: { "MCP-Protocol-Version": "2025-11-25" }, + }), + status: 400, + code: "invalid_mcp_request", + }, + { + request: mcpRequest({ headers: { "Mcp-Session-Id": "legacy-session" } }), + status: 400, + code: "invalid_mcp_request", + }, + { + request: mcpRequest({ headers: { "Mcp-Method": "initialize" } }), + status: 400, + code: "invalid_mcp_request", + }, + { + request: mcpRequest({ headers: { "Mcp-Name": "record_mutation" } }), + status: 400, + code: "invalid_mcp_request", + }, + { + request: mcpRequest({ headers: { Accept: "application/json" } }), + status: 400, + code: "invalid_mcp_request", + }, + { + request: mcpRequest({ + headers: { + Accept: "application/json;q=0, text/event-stream", + }, + }), + status: 400, + code: "invalid_mcp_request", + }, + ]; + for (const entry of cases) { + const instance = fixture(); + const response = await instance.gateway.fetch(entry.request, instance.env); + assert.equal(response.status, entry.status); + assert.equal(await errorCode(response), entry.code); + assert.equal(instance.calls.length, 0); + assert.equal(instance.names.length, 0); + } + + let oversizedCancelled = false; + const oversizedStream = new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(65_537)); + }, + cancel() { + oversizedCancelled = true; + }, + }); + const streamed = fixture(); + const streamedResponse = await streamed.gateway.fetch( + mcpRequest({ rawBody: oversizedStream }), + streamed.env, + ); + assert.equal(streamedResponse.status, 413); + assert.equal(oversizedCancelled, true); + assert.equal(streamed.calls.length, 0); +}); + +test("uses one private hop and an exclusive reviewed forwarding allowlist", async () => { + const instance = fixture(); + const request = mcpRequest({ + headers: { + Cookie: "must-not-cross", + "Mcp-Param-Region": "eu-central-1", + Traceparent: "00-11111111111111111111111111111111-2222222222222222-01", + "X-Unreviewed": "must-not-cross", + }, + }); + const source = await request.clone().text(); + const response = await instance.gateway.fetch(request, instance.env); + assert.equal(response.status, 200); + assert.equal( + Object.keys(instance.admissionRequests[0]).sort().join(","), + "authoritySignal,authorization,serverId", + ); + assert.equal( + Object.keys(instance.revalidationRequests[0]).sort().join(","), + "authorityRevision,authoritySignal,authorization,serverId", + ); + assert.equal( + instance.admissionRequests[0].authorization, + "Bearer fixture-token", + ); + assert.equal(instance.admissionRequests[0].serverId, serverId); + assert.deepEqual(instance.names, [serverId]); + assert.equal(instance.calls.length, 1); + const forwarded = instance.calls[0]; + assert.equal(forwarded.url, "http://localhost/mcp"); + assert.equal(await forwarded.text(), source); + assert.equal(forwarded.headers.get("authorization"), "Bearer fixture-token"); + assert.equal(forwarded.headers.get("mcp-param-region"), "eu-central-1"); + assert.ok(forwarded.headers.has("traceparent")); + assert.equal(forwarded.headers.get("origin"), null); + assert.equal(forwarded.headers.get("cookie"), null); + assert.equal(forwarded.headers.get("x-unreviewed"), null); + assert.equal(forwarded.headers.get("x-flarebot-hosted-mcp-server"), serverId); +}); + +test("streams responses and cancels stalled or oversized upstream bodies", async () => { + let oversizedCancelled = false; + const oversized = fixture({ + response() { + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(65_537)); + }, + cancel() { + oversizedCancelled = true; + }, + }), + { headers: { "Content-Type": "text/event-stream" } }, + ); + }, + }); + const oversizedResponse = await oversized.gateway.fetch( + mcpRequest(), + oversized.env, + ); + await assert.rejects(oversizedResponse.arrayBuffer(), /byte limit/); + assert.equal(oversizedCancelled, true); + + let stalledCancelled = false; + const abortController = new AbortController(); + const stalled = fixture({ + response() { + return new Response( + new ReadableStream({ + pull() { + return new Promise(() => {}); + }, + cancel() { + stalledCancelled = true; + }, + }), + { headers: { "Content-Type": "text/event-stream" } }, + ); + }, + }); + const stalledResponse = await stalled.gateway.fetch( + mcpRequest({ signal: abortController.signal }), + stalled.env, + ); + const read = stalledResponse.body.getReader().read(); + abortController.abort(new Error("client disconnected")); + await assert.rejects(read, /client disconnected/); + assert.equal(stalledCancelled, true); +}); + +test("never follows redirects or retries an uncertain dispatch", async () => { + let attempts = 0; + const uncertain = fixture({ + response() { + attempts++; + throw new Error("container RPC disappeared after admission"); + }, + }); + const response = await uncertain.gateway.fetch(mcpRequest(), uncertain.env); + assert.equal(response.status, 502); + assert.deepEqual(await response.json(), { + error: { + code: "outcome_uncertain", + message: + "The MCP action was dispatched but no reliable terminal outcome was received", + }, + outcome: "uncertain", + replay: "forbidden", + }); + assert.equal(attempts, 1); + + const redirect = fixture({ + response: () => + new Response(null, { + status: 307, + headers: { Location: "https://attacker.test/mcp" }, + }), + }); + assert.equal( + await errorCode(await redirect.gateway.fetch(mcpRequest(), redirect.env)), + "hosted_mcp_invalid_response", + ); + assert.equal(redirect.calls.length, 1); +}); + +function internalContainerRequest(signal) { + return new Request("http://localhost/mcp", { + method: "POST", + headers: { "X-Flarebot-Hosted-Mcp-Server": serverId }, + body: "{}", + signal, + }); +} + +function containerFixture(response) { + let starts = 0; + let implicitFetches = 0; + let dispatches = 0; + let expirations = 0; + let renewals = 0; + const id = { equals: (candidate) => candidate === id }; + const ctx = { + id, + container: { + getTcpPort(port) { + assert.equal(port, 8080); + return { + fetch(request, init) { + dispatches++; + return Promise.resolve().then(() => response(request, init)); + }, + }; + }, + }, + startAndWaitForPorts() { + starts++; + }, + containerFetch() { + implicitFetches++; + throw new Error("the implicit SDK readiness path must not be used"); + }, + renewActivityTimeout() { + renewals++; + }, + activityExpired() { + expirations++; + }, + }; + const env = { + HOSTED_MCP_CONTAINER: { + get() { + throw new Error("gateway binding is not used inside the DO"); + }, + idFromName() { + return id; + }, + }, + HOSTED_MCP_RUNTIME: configuration(), + }; + return { + container: new runtime.HostedMcpContainer(ctx, env), + counts: () => ({ + dispatches, + expirations, + implicitFetches, + renewals, + starts, + }), + }; +} + +test("uses one readiness path and treats a crash before the direct hop as uncertain", async () => { + const fixture = containerFixture(() => { + throw new Error("process exited between readiness and dispatch"); + }); + const response = await fixture.container.fetch(internalContainerRequest()); + assert.equal(response.status, 502); + assert.equal(await errorCode(response), "outcome_uncertain"); + assert.deepEqual(fixture.counts(), { + dispatches: 1, + expirations: 0, + implicitFetches: 0, + renewals: 2, + starts: 1, + }); +}); + +test("renews the idle window on warm dispatch and final response settlement", async () => { + const fixture = containerFixture(() => + Response.json({ jsonrpc: "2.0", id: "ok", result: {} }), + ); + const response = await fixture.container.fetch(internalContainerRequest()); + assert.equal(response.status, 200); + assert.equal(fixture.counts().renewals, 1); + await response.arrayBuffer(); + assert.equal(fixture.counts().renewals, 2); +}); + +test("releases unread response lifetimes on abort without awaiting upstream cancel", async () => { + let completeImmediately = false; + let cancellations = 0; + const fixture = containerFixture(() => { + if (completeImmediately) + return Response.json({ jsonrpc: "2.0", id: "ok", result: {} }); + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array([1])); + }, + cancel() { + cancellations++; + return new Promise(() => {}); + }, + }), + { headers: { "Content-Type": "text/event-stream" } }, + ); + }); + const controllers = Array.from({ length: 4 }, () => new AbortController()); + const unread = await Promise.all( + controllers.map((controller) => + fixture.container.fetch(internalContainerRequest(controller.signal)), + ), + ); + assert.equal( + (await fixture.container.fetch(internalContainerRequest())).status, + 429, + ); + controllers.forEach((controller) => + controller.abort(new Error("downstream stopped reading")), + ); + await new Promise((resolve) => setImmediate(resolve)); + assert.equal(cancellations, 4); + + completeImmediately = true; + const admitted = await fixture.container.fetch(internalContainerRequest()); + assert.equal(admitted.status, 200); + await admitted.arrayBuffer(); + await fixture.container.onActivityExpired(); + assert.equal(fixture.counts().expirations, 1); + assert.equal(fixture.counts().implicitFetches, 0); + assert.equal(unread.length, 4); +}); diff --git a/tsconfig.worker.json b/tsconfig.worker.json index 73d6639..91af849 100644 --- a/tsconfig.worker.json +++ b/tsconfig.worker.json @@ -31,6 +31,7 @@ "tests/fixtures/bridge-customer-worker.ts", "tests/fixtures/upgrade-customer-worker.ts", "tests/fixtures/golden-customer-worker.ts", - "tests/fixtures/golden-control-worker.ts" + "tests/fixtures/golden-control-worker.ts", + "tests/fixtures/hosted-mcp-container-worker.ts" ] } diff --git a/worker/hosted-mcp-container.ts b/worker/hosted-mcp-container.ts new file mode 100644 index 0000000..ef39d77 --- /dev/null +++ b/worker/hosted-mcp-container.ts @@ -0,0 +1,156 @@ +import { Container } from "@cloudflare/containers"; +import { + awaitWithSignal, + cancelHostedMcpBody, + hostedMcpError, + hostedMcpUncertainOutcome, + limitHostedMcpResponse, +} from "./hosted-mcp-http.ts"; +import { + HOSTED_MCP_CONTAINER_MAX_CONCURRENCY, + HOSTED_MCP_CONTAINER_MAX_RESPONSE_BYTES, + HOSTED_MCP_CONTAINER_READINESS_TIMEOUT_MS, + HOSTED_MCP_CONTAINER_REQUEST_TIMEOUT_MS, + parseHostedMcpContainerConfiguration, + type HostedMcpContainerEnv, +} from "./hosted-mcp-gateway.ts"; + +export * from "./hosted-mcp-gateway.ts"; + +export class HostedMcpContainer< + Env extends HostedMcpContainerEnv = HostedMcpContainerEnv, +> extends Container { + defaultPort = 8080; + requiredPorts = [8080]; + enableInternet = false; + pingEndpoint = "localhost/healthz"; + private activeRequests = 0; + + constructor(ctx: DurableObjectState<{}>, env: Env) { + const configuration = parseHostedMcpContainerConfiguration( + env.HOSTED_MCP_RUNTIME, + ); + super(ctx, env, { + defaultPort: configuration.manifest.runtime.instance.defaultPort, + enableInternet: false, + sleepAfter: configuration.manifest.runtime.instance.idleTimeoutSeconds, + }); + this.defaultPort = configuration.manifest.runtime.instance.defaultPort; + this.requiredPorts = [ + ...configuration.manifest.runtime.instance.requiredPorts, + ]; + this.sleepAfter = + configuration.manifest.runtime.instance.idleTimeoutSeconds; + } + + override async onActivityExpired() { + if (this.activeRequests > 0) { + this.renewActivityTimeout(); + return; + } + await super.onActivityExpired(); + } + + protected dispatchToReadyContainer(request: Request) { + const container = this.ctx.container; + if (!container) throw new Error("Container context is unavailable"); + return container + .getTcpPort(this.defaultPort) + .fetch(request.url, new Request(request, { signal: request.signal })); + } + + override async fetch(request: Request): Promise { + request.signal.throwIfAborted(); + const requestDeadline = AbortSignal.any([ + request.signal, + AbortSignal.timeout(HOSTED_MCP_CONTAINER_REQUEST_TIMEOUT_MS), + ]); + const configuration = parseHostedMcpContainerConfiguration( + this.env.HOSTED_MCP_RUNTIME, + ); + const url = new URL(request.url); + if ( + !this.ctx.id.equals( + this.env.HOSTED_MCP_CONTAINER.idFromName(configuration.serverId), + ) || + request.method !== "POST" || + url.origin !== "http://localhost" || + url.pathname !== configuration.manifest.runtime.containerMcpPath || + request.headers.get("x-flarebot-hosted-mcp-server") !== + configuration.serverId || + request.headers.has("origin") || + request.headers.has("mcp-session-id") + ) + return new Response(null, { status: 404 }); + if (this.activeRequests >= HOSTED_MCP_CONTAINER_MAX_CONCURRENCY) + return hostedMcpError( + "hosted_mcp_busy", + "The hosted MCP Container is at its concurrency limit", + 429, + true, + { "Retry-After": "1" }, + ); + this.activeRequests++; + let released = false; + let acceptedActivity = false; + const release = () => { + if (released) return; + released = true; + this.activeRequests--; + if (acceptedActivity && this.activeRequests === 0) + this.renewActivityTimeout(); + }; + try { + const readinessSignal = AbortSignal.any([ + requestDeadline, + AbortSignal.timeout(HOSTED_MCP_CONTAINER_READINESS_TIMEOUT_MS), + ]); + await this.startAndWaitForPorts({ + ports: this.requiredPorts, + cancellationOptions: { + abort: readinessSignal, + instanceGetTimeoutMS: HOSTED_MCP_CONTAINER_READINESS_TIMEOUT_MS, + portReadyTimeoutMS: HOSTED_MCP_CONTAINER_READINESS_TIMEOUT_MS, + waitInterval: 100, + }, + }); + } catch { + release(); + if (request.signal.aborted) throw request.signal.reason; + return hostedMcpError( + "hosted_mcp_not_ready", + "The hosted MCP Container did not become ready in time", + 503, + true, + { "Retry-After": "1" }, + ); + } + acceptedActivity = true; + this.renewActivityTimeout(); + let response: Response; + try { + response = await awaitWithSignal( + this.dispatchToReadyContainer( + new Request(request, { signal: requestDeadline }), + ), + requestDeadline, + ); + } catch { + release(); + if (request.signal.aborted) throw request.signal.reason; + return hostedMcpUncertainOutcome(); + } + if (response.status === 500) { + release(); + if (response.body) + cancelHostedMcpBody(response.body, "Container transport failed"); + return hostedMcpUncertainOutcome(); + } + return limitHostedMcpResponse( + response, + requestDeadline, + HOSTED_MCP_CONTAINER_MAX_RESPONSE_BYTES, + release, + ); + } +} diff --git a/worker/hosted-mcp-gateway.ts b/worker/hosted-mcp-gateway.ts new file mode 100644 index 0000000..c9e7bbe --- /dev/null +++ b/worker/hosted-mcp-gateway.ts @@ -0,0 +1,428 @@ +import { + HOSTED_MCP_PROTOCOL_REVISION, + parseHostedMcpTemplateManifest, + type HostedMcpTemplateManifest, +} from "../shared/hosted-mcp.ts"; +import { safeJsonSnapshot } from "../shared/safe-json.ts"; +import { + awaitWithSignal, + cancelHostedMcpBody, + hostedMcpError, + hostedMcpResponseAccepted, + hostedMcpUncertainOutcome, + limitHostedMcpResponse, + readHostedMcpHttpRequest, +} from "./hosted-mcp-http.ts"; + +export const HOSTED_MCP_CONTAINER_MAX_REQUEST_BYTES = 65_536; +export const HOSTED_MCP_CONTAINER_MAX_RESPONSE_BYTES = 65_536; +export const HOSTED_MCP_CONTAINER_MAX_CONCURRENCY = 4; +export const HOSTED_MCP_CONTAINER_READINESS_TIMEOUT_MS = 5_000; +export const HOSTED_MCP_CONTAINER_REQUEST_TIMEOUT_MS = 25_000; + +const serverIdPattern = /^hosted-mcp-[a-f0-9]{32}$/; + +type ContainerManifest = HostedMcpTemplateManifest & { + runtime: Extract; +}; + +export interface HostedMcpContainerConfiguration { + manifest: ContainerManifest; + publicOrigin: string; + serverId: string; +} + +export type HostedMcpAdmissionDecision = + | { status: "admitted"; authorityRevision: number } + | { status: "unauthorized" | "revoked" | "suspended" | "unavailable" }; + +export interface HostedMcpAdmissionRequest { + authorization: string; + authoritySignal: AbortSignal; + serverId: string; +} + +export interface HostedMcpAdmissionRevalidationRequest extends HostedMcpAdmissionRequest { + authorityRevision: number; +} + +export interface HostedMcpContainerEnv { + HOSTED_MCP_CONTAINER: { + get(id: DurableObjectId): Fetcher; + idFromName(name: string): DurableObjectId; + }; + HOSTED_MCP_RUNTIME: unknown; +} + +export interface HostedMcpContainerGatewayOptions< + Env extends HostedMcpContainerEnv, +> { + admit( + request: HostedMcpAdmissionRequest, + env: Env, + ): HostedMcpAdmissionDecision | Promise; + revalidate( + request: HostedMcpAdmissionRevalidationRequest, + env: Env, + ): HostedMcpAdmissionDecision | Promise; + configuration?(env: Env): unknown; +} + +export class HostedMcpContainerConfigurationError extends Error { + constructor() { + super("The hosted MCP Container runtime configuration is invalid"); + this.name = "HostedMcpContainerConfigurationError"; + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function hasOnlyKeys(value: Record, expected: string[]) { + const keys = Object.keys(value); + return ( + keys.length === expected.length && + expected.every((key) => keys.includes(key)) + ); +} + +export function parseHostedMcpContainerConfiguration( + value: unknown, +): HostedMcpContainerConfiguration { + try { + const snapshot = safeJsonSnapshot(value); + if ( + !isRecord(snapshot) || + !hasOnlyKeys(snapshot, ["manifest", "publicOrigin", "serverId"]) || + typeof snapshot.publicOrigin !== "string" || + typeof snapshot.serverId !== "string" || + !serverIdPattern.test(snapshot.serverId) + ) + throw new Error("invalid runtime configuration"); + const publicUrl = new URL(snapshot.publicOrigin); + const localHttp = + publicUrl.protocol === "http:" && + (publicUrl.hostname === "localhost" || + publicUrl.hostname === "127.0.0.1" || + publicUrl.hostname === "[::1]"); + if ( + publicUrl.origin !== snapshot.publicOrigin || + (publicUrl.protocol !== "https:" && !localHttp) + ) + throw new Error("invalid public origin"); + const manifest = parseHostedMcpTemplateManifest(snapshot.manifest); + if (manifest.runtime.type !== "container") + throw new Error("invalid Container manifest"); + const containerManifest: ContainerManifest = { + ...manifest, + runtime: manifest.runtime, + }; + if ( + containerManifest.protocol.protocolRevision !== + HOSTED_MCP_PROTOCOL_REVISION || + containerManifest.protocol.mcpPath !== + containerManifest.runtime.publicMcpPath || + containerManifest.runtime.instance.maxInstances !== 1 + ) + throw new Error("invalid Container manifest"); + return { + manifest: containerManifest, + publicOrigin: snapshot.publicOrigin, + serverId: snapshot.serverId, + }; + } catch (error) { + if (error instanceof HostedMcpContainerConfigurationError) throw error; + throw new HostedMcpContainerConfigurationError(); + } +} + +function unauthorized() { + return hostedMcpError( + "invalid_token", + "The hosted MCP credential was not accepted", + 401, + false, + { "WWW-Authenticate": 'Bearer realm="hosted-mcp"' }, + ); +} + +function admissionResponse(decision: HostedMcpAdmissionDecision) { + if (decision.status === "admitted") return undefined; + if (decision.status === "unauthorized" || decision.status === "revoked") + return unauthorized(); + if (decision.status === "suspended") + return hostedMcpError( + "hosted_mcp_suspended", + "The hosted MCP server is suspended", + 503, + false, + ); + return hostedMcpError( + "hosted_mcp_admission_unavailable", + "Hosted MCP admission is temporarily unavailable", + 503, + true, + { "Retry-After": "1" }, + ); +} + +function deadlineResponse() { + return hostedMcpError( + "hosted_mcp_deadline_exceeded", + "The hosted MCP request exceeded its pre-dispatch deadline", + 504, + false, + ); +} + +function parseAdmissionDecision(value: unknown): HostedMcpAdmissionDecision { + const snapshot = safeJsonSnapshot(value); + if (!isRecord(snapshot) || typeof snapshot.status !== "string") + throw new Error("invalid admission decision"); + if (snapshot.status === "admitted") { + if ( + !hasOnlyKeys(snapshot, ["authorityRevision", "status"]) || + typeof snapshot.authorityRevision !== "number" || + !Number.isSafeInteger(snapshot.authorityRevision) || + snapshot.authorityRevision < 1 + ) + throw new Error("invalid admission decision"); + return { + status: "admitted", + authorityRevision: snapshot.authorityRevision, + }; + } + if ( + !hasOnlyKeys(snapshot, ["status"]) || + !["unauthorized", "revoked", "suspended", "unavailable"].includes( + snapshot.status, + ) + ) + throw new Error("invalid admission decision"); + switch (snapshot.status) { + case "unauthorized": + case "revoked": + case "suspended": + case "unavailable": + return { status: snapshot.status }; + default: + throw new Error("invalid admission decision"); + } +} + +function requestOriginResponse( + request: Request, + configuration: HostedMcpContainerConfiguration, +) { + const url = new URL(request.url); + const expected = new URL(configuration.publicOrigin); + if ( + url.origin !== configuration.publicOrigin || + request.headers.get("host") !== expected.host + ) + return hostedMcpError( + "invalid_public_host", + "The hosted MCP request host is not accepted", + 403, + false, + ); + const origin = request.headers.get("origin"); + if (origin !== null && origin !== configuration.publicOrigin) + return hostedMcpError( + "invalid_origin", + "The hosted MCP request origin is not accepted", + 403, + false, + ); + return undefined; +} + +async function admissionDecision( + operation: () => + HostedMcpAdmissionDecision | Promise, + deadline: AbortSignal, +) { + try { + return parseAdmissionDecision( + await awaitWithSignal(Promise.resolve(operation()), deadline), + ); + } catch { + if (deadline.aborted) throw deadline.reason; + return { status: "unavailable" } satisfies HostedMcpAdmissionDecision; + } +} + +async function gatewayFetch( + request: Request, + env: Env, + options: HostedMcpContainerGatewayOptions, +) { + request.signal.throwIfAborted(); + let configuration: HostedMcpContainerConfiguration; + try { + configuration = parseHostedMcpContainerConfiguration( + options.configuration?.(env) ?? env.HOSTED_MCP_RUNTIME, + ); + } catch { + return hostedMcpError( + "hosted_mcp_configuration_invalid", + "The hosted MCP runtime is unavailable", + 503, + false, + ); + } + const url = new URL(request.url); + const originFailure = requestOriginResponse(request, configuration); + if (originFailure) return originFailure; + if ( + url.pathname === configuration.manifest.runtime.healthPath && + url.search === "" + ) { + if (request.method !== "GET" && request.method !== "HEAD") + return new Response(null, { + status: 405, + headers: { Allow: "GET, HEAD" }, + }); + return new Response( + request.method === "HEAD" + ? null + : JSON.stringify({ + status: "ok", + transport: "streamable-http", + protocol: HOSTED_MCP_PROTOCOL_REVISION, + }), + { headers: { "Content-Type": "application/json" } }, + ); + } + if ( + url.pathname !== configuration.manifest.runtime.publicMcpPath || + url.search !== "" + ) + return new Response(null, { status: 404 }); + if (request.method !== "POST") + return new Response(null, { status: 405, headers: { Allow: "POST" } }); + const authorization = request.headers.get("authorization"); + if (!authorization || authorization.length > 4_096) return unauthorized(); + const deadline = AbortSignal.any([ + request.signal, + AbortSignal.timeout(HOSTED_MCP_CONTAINER_REQUEST_TIMEOUT_MS), + ]); + let decision: HostedMcpAdmissionDecision; + try { + decision = await admissionDecision( + () => + options.admit( + { + authorization, + authoritySignal: deadline, + serverId: configuration.serverId, + }, + env, + ), + deadline, + ); + } catch { + if (request.signal.aborted) throw request.signal.reason; + return deadlineResponse(); + } + const rejected = admissionResponse(decision); + if (rejected) return rejected; + if (decision.status !== "admitted") + throw new Error("unreachable admission decision"); + let validated: Awaited>; + try { + validated = await readHostedMcpHttpRequest( + request, + deadline, + HOSTED_MCP_CONTAINER_MAX_REQUEST_BYTES, + ); + } catch { + if (request.signal.aborted) throw request.signal.reason; + return deadlineResponse(); + } + if (validated instanceof Response) return validated; + + let current: HostedMcpAdmissionDecision; + try { + current = await admissionDecision( + () => + options.revalidate( + { + authorization, + authorityRevision: decision.authorityRevision, + authoritySignal: deadline, + serverId: configuration.serverId, + }, + env, + ), + deadline, + ); + } catch { + if (request.signal.aborted) throw request.signal.reason; + return deadlineResponse(); + } + const revalidationRejected = admissionResponse(current); + if (revalidationRejected) return revalidationRejected; + if (current.status !== "admitted") + throw new Error("unreachable revalidation decision"); + if (current.authorityRevision !== decision.authorityRevision) + return admissionResponse({ status: "unavailable" })!; + + validated.privateHeaders.set( + "x-flarebot-hosted-mcp-server", + configuration.serverId, + ); + const internalRequest = new Request( + `http://localhost${configuration.manifest.runtime.containerMcpPath}`, + { + method: "POST", + headers: validated.privateHeaders, + body: validated.body, + signal: deadline, + redirect: "manual", + }, + ); + const container = env.HOSTED_MCP_CONTAINER.get( + env.HOSTED_MCP_CONTAINER.idFromName(configuration.serverId), + ); + let response: Response; + try { + response = await awaitWithSignal( + container.fetch(internalRequest), + deadline, + ); + } catch (error) { + if (request.signal.aborted) throw request.signal.reason; + return hostedMcpUncertainOutcome(); + } + if ( + (response.status >= 300 && response.status < 400) || + response.headers.has("mcp-session-id") || + !hostedMcpResponseAccepted(response) + ) { + if (response.body) + cancelHostedMcpBody(response.body, "Invalid hosted MCP response"); + return hostedMcpError( + "hosted_mcp_invalid_response", + "The hosted MCP Container returned an invalid response", + 502, + false, + ); + } + return limitHostedMcpResponse( + response, + deadline, + HOSTED_MCP_CONTAINER_MAX_RESPONSE_BYTES, + ); +} + +export function createHostedMcpContainerGateway< + Env extends HostedMcpContainerEnv, +>(options: HostedMcpContainerGatewayOptions) { + return { + fetch(request: Request, env: Env) { + return gatewayFetch(request, env, options); + }, + } satisfies ExportedHandler; +} diff --git a/worker/hosted-mcp-http.ts b/worker/hosted-mcp-http.ts new file mode 100644 index 0000000..edabfbc --- /dev/null +++ b/worker/hosted-mcp-http.ts @@ -0,0 +1,399 @@ +import { + HOSTED_MCP_OUTCOME_UNCERTAIN, + HOSTED_MCP_PROTOCOL_REVISION, +} from "../shared/hosted-mcp.ts"; +import { safeJsonSnapshot } from "../shared/safe-json.ts"; + +const forwardedHeaders = new Set([ + "accept", + "authorization", + "baggage", + "content-type", + "mcp-method", + "mcp-name", + "mcp-protocol-version", + "traceparent", + "tracestate", +]); + +const namedParameters: Readonly> = { + "prompts/get": "name", + "resources/read": "uri", + "tools/call": "name", +}; + +export type HostedMcpHttpRequest = { + body: Uint8Array; + privateHeaders: Headers; +}; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +export function hostedMcpError( + code: string, + message: string, + status: number, + retryable: boolean, + headers?: HeadersInit, +) { + return Response.json( + { + error: { code, message }, + retryable, + replay: "forbidden", + }, + { status, headers }, + ); +} + +export function hostedMcpUncertainOutcome() { + return Response.json(HOSTED_MCP_OUTCOME_UNCERTAIN, { status: 502 }); +} + +function invalidRequest(message: string) { + return hostedMcpError("invalid_mcp_request", message, 400, false); +} + +function isJsonContentType(value: string | null) { + if (value === null) return false; + const [mediaType, ...parameters] = value.split(";"); + if (mediaType.trim().toLowerCase() !== "application/json") return false; + return parameters.every((parameter) => { + const [name, candidate, ...rest] = parameter.split("="); + return ( + rest.length === 0 && + name.trim().toLowerCase() === "charset" && + candidate?.trim().toLowerCase() === "utf-8" + ); + }); +} + +function isEventStreamContentType(value: string) { + const [mediaType, ...parameters] = value.split(";"); + if (mediaType.trim().toLowerCase() !== "text/event-stream") return false; + return parameters.every((parameter) => { + const [name, candidate, ...rest] = parameter.split("="); + return ( + rest.length === 0 && + name.trim().toLowerCase() === "charset" && + candidate?.trim().toLowerCase() === "utf-8" + ); + }); +} + +function acceptsRequiredMediaTypes(value: string | null) { + if (value === null) return false; + const mediaTypes = new Set(); + for (const entry of value.split(",")) { + const [rawMediaType, ...rawParameters] = entry.split(";"); + const mediaType = rawMediaType?.trim().toLowerCase(); + if (!mediaType) return false; + let quality = 1; + for (const parameter of rawParameters) { + const [rawName, rawValue, ...rest] = parameter.split("="); + if ( + rest.length > 0 || + rawName?.trim().toLowerCase() !== "q" || + !/^(?:0(?:\.\d{0,3})?|1(?:\.0{0,3})?)$/.test(rawValue?.trim() ?? "") + ) + return false; + quality = Number(rawValue); + } + if (quality > 0) mediaTypes.add(mediaType); + } + return ( + mediaTypes.has("application/json") && mediaTypes.has("text/event-stream") + ); +} + +export function awaitWithSignal(value: Promise, signal: AbortSignal) { + signal.throwIfAborted(); + return new Promise((resolve, reject) => { + const aborted = () => reject(signal.reason); + signal.addEventListener("abort", aborted, { once: true }); + value.then( + (result) => { + signal.removeEventListener("abort", aborted); + resolve(result); + }, + (error) => { + signal.removeEventListener("abort", aborted); + reject(error); + }, + ); + }); +} + +function releaseReader(reader: ReadableStreamDefaultReader) { + try { + reader.releaseLock(); + } catch {} +} + +function cancelReader( + reader: ReadableStreamDefaultReader, + reason: unknown, +) { + let cancellation: Promise; + try { + cancellation = reader.cancel(reason); + } catch { + releaseReader(reader); + return; + } + void cancellation.catch(() => undefined).finally(() => releaseReader(reader)); +} + +async function readCappedRequest( + request: Request, + signal: AbortSignal, + limit: number, +) { + const declared = request.headers.get("content-length"); + if (declared !== null) { + if (!/^\d+$/.test(declared) || Number(declared) > limit) return undefined; + } + if (!request.body) return new Uint8Array(); + const reader = request.body.getReader(); + const chunks: Uint8Array[] = []; + let length = 0; + try { + for (;;) { + signal.throwIfAborted(); + const { done, value } = await awaitWithSignal(reader.read(), signal); + if (done) break; + length += value.byteLength; + if (length > limit) { + cancelReader(reader, "Hosted MCP request exceeded its byte limit"); + return undefined; + } + chunks.push(value); + } + } catch (error) { + cancelReader(reader, error); + throw error; + } + releaseReader(reader); + const body = new Uint8Array(length); + let offset = 0; + for (const chunk of chunks) { + body.set(chunk, offset); + offset += chunk.byteLength; + } + return body; +} + +function decodeRequestBody(body: Uint8Array) { + try { + const text = new TextDecoder("utf-8", { + fatal: true, + ignoreBOM: false, + }).decode(body); + return safeJsonSnapshot(JSON.parse(text)); + } catch { + return undefined; + } +} + +function validateProtocolRequest(request: Request, body: Uint8Array) { + if (!isJsonContentType(request.headers.get("content-type"))) + return invalidRequest("Content-Type must be application/json with UTF-8"); + if (!acceptsRequiredMediaTypes(request.headers.get("accept"))) + return invalidRequest( + "Accept must include application/json and text/event-stream", + ); + if ( + request.headers.get("mcp-protocol-version") !== HOSTED_MCP_PROTOCOL_REVISION + ) + return invalidRequest("MCP-Protocol-Version is not supported"); + if (request.headers.has("mcp-session-id")) + return invalidRequest("Mcp-Session-Id is not supported"); + const envelope = decodeRequestBody(body); + if ( + !isRecord(envelope) || + envelope.jsonrpc !== "2.0" || + !( + typeof envelope.id === "string" || + (typeof envelope.id === "number" && Number.isFinite(envelope.id)) + ) || + typeof envelope.method !== "string" || + envelope.method.length < 1 || + envelope.method === "initialize" || + !isRecord(envelope.params) || + !isRecord(envelope.params._meta) + ) + return invalidRequest("The body is not one modern JSON-RPC request"); + const meta = envelope.params._meta; + if ( + meta["io.modelcontextprotocol/protocolVersion"] !== + HOSTED_MCP_PROTOCOL_REVISION || + !isRecord(meta["io.modelcontextprotocol/clientCapabilities"]) || + !isRecord(meta["io.modelcontextprotocol/clientInfo"]) + ) + return invalidRequest("Required per-request MCP metadata is missing"); + const clientInfo = meta["io.modelcontextprotocol/clientInfo"]; + if ( + typeof clientInfo.name !== "string" || + clientInfo.name.length < 1 || + typeof clientInfo.version !== "string" || + clientInfo.version.length < 1 + ) + return invalidRequest("Required client identity is missing"); + if (request.headers.get("mcp-method") !== envelope.method) + return invalidRequest("Mcp-Method does not match the request body"); + const namedParameter = namedParameters[envelope.method]; + const headerName = request.headers.get("mcp-name"); + if (namedParameter) { + const bodyName = envelope.params[namedParameter]; + if ( + typeof bodyName !== "string" || + bodyName.length < 1 || + headerName !== bodyName + ) + return invalidRequest("Mcp-Name does not match the request body"); + } else if (headerName !== null) { + return invalidRequest("Mcp-Name is not valid for this method"); + } + return undefined; +} + +function privateHeaders(request: Request) { + const result = new Headers(); + let parameterCount = 0; + let totalBytes = 0; + for (const [name, value] of request.headers) { + const lower = name.toLowerCase(); + const parameter = lower.startsWith("mcp-param-"); + if (!forwardedHeaders.has(lower) && !parameter) continue; + if (parameter) { + parameterCount++; + if ( + parameterCount > 32 || + !/^mcp-param-[!#$%&'*+.^_`|~0-9a-z-]+$/.test(lower) + ) + return undefined; + } + totalBytes += name.length + value.length; + if (value.length > 4_096 || totalBytes > 16_384) return undefined; + result.set(name, value); + } + return result; +} + +export async function readHostedMcpHttpRequest( + request: Request, + signal: AbortSignal, + limit: number, +): Promise { + const body = await readCappedRequest(request, signal, limit); + if (!body) + return hostedMcpError( + "hosted_mcp_request_too_large", + "The hosted MCP request exceeds its byte limit", + 413, + false, + ); + const protocolFailure = validateProtocolRequest(request, body); + if (protocolFailure) return protocolFailure; + const headers = privateHeaders(request); + if (!headers) + return invalidRequest("The forwarded MCP headers exceed their limits"); + return { body, privateHeaders: headers }; +} + +function safeResponseHeaders(headers: Headers) { + const safe = new Headers(); + for (const name of ["cache-control", "content-type", "retry-after"]) { + const value = headers.get(name); + if (value !== null) safe.set(name, value); + } + return safe; +} + +export function hostedMcpResponseAccepted(response: Response) { + if ( + response.body === null || + response.status === 202 || + response.status === 204 + ) + return true; + const value = response.headers.get("content-type"); + if (value === null) return false; + return isJsonContentType(value) || isEventStreamContentType(value); +} + +export function cancelHostedMcpBody(body: ReadableStream, reason: unknown) { + void body.cancel(reason).catch(() => undefined); +} + +export function limitHostedMcpResponse( + response: Response, + signal: AbortSignal, + limit: number, + onSettled?: () => void, +) { + const headers = safeResponseHeaders(response.headers); + if (!response.body) { + onSettled?.(); + return new Response(null, { status: response.status, headers }); + } + + const reader = response.body.getReader(); + let controller: ReadableStreamDefaultController | undefined; + let settled = false; + let total = 0; + + const settle = () => { + if (settled) return false; + settled = true; + signal.removeEventListener("abort", aborted); + onSettled?.(); + return true; + }; + const fail = (reason: unknown) => { + if (!settle()) return; + cancelReader(reader, reason); + try { + controller?.error(reason); + } catch {} + }; + const aborted = () => fail(signal.reason); + + const body = new ReadableStream({ + start(value) { + controller = value; + signal.addEventListener("abort", aborted, { once: true }); + if (signal.aborted) aborted(); + }, + async pull(value) { + if (settled) return; + try { + const result = await awaitWithSignal(reader.read(), signal); + if (settled) return; + if (result.done) { + if (settle()) { + releaseReader(reader); + value.close(); + } + return; + } + total += result.value.byteLength; + if (total > limit) { + fail(new Error("Hosted MCP response exceeded its byte limit")); + return; + } + value.enqueue(result.value); + } catch (error) { + fail(error); + } + }, + cancel(reason) { + if (!settle()) return; + cancelReader(reader, reason); + }, + }); + + return new Response(body, { status: response.status, headers }); +}