fix(codemode): preserve MCP request context

.changeset/codemode-mcp-request-context.md

@@ -0,0 +1,5 @@
+---
+"@cloudflare/codemode": patch
+---
+
+Pass the outer MCP tool-call context to `openApiMcpServer` request callbacks so server-to-client requests and notifications can be associated with the originating response stream.

examples/codemode-mcp-openapi/README.md

@@ -9,7 +9,7 @@ Demonstrates how to turn any OpenAPI spec into a pair of MCP tools (`search` + `
 - **`search`** — the LLM queries the spec as a JavaScript object to find endpoints, parameters, and schemas
 - **`execute`** — the LLM calls the API via a host-side `request()` function you provide
 
-Auth tokens and base URLs live in your `request()` function on the host. The sandbox that runs LLM-generated code has no outbound network access and never sees secrets.
+Auth tokens and base URLs live in your `request()` function on the host. The sandbox has no outbound network access and never sees secrets. The callback's second argument is the MCP context for the outer `execute` tool call. Use it when an elicitation, sampling request, roots request, or notification belongs to that call.
 
 This example connects to the live [Cloudflare API](https://api.cloudflare.com/) using the official OpenAPI spec. Pass a Cloudflare API token via the `Authorization` header.
 
@@ -36,9 +36,9 @@ import { openApiMcpServer } from "@cloudflare/codemode/mcp";
 const server = openApiMcpServer({
   spec,
   executor,
-  request: async (opts) => {
-    // Runs on the host — put auth, base URL, and headers here.
-    // The sandbox never sees the token.
+  request: async (opts, context) => {
+    // Runs on the host. Put auth, base URL, and headers here.
+    // The sandbox sees neither the token nor the MCP context.
     const url = new URL(`https://api.example.com${opts.path}`);
     const res = await fetch(url, {
       method: opts.method,
@@ -50,6 +50,57 @@ const server = openApiMcpServer({
 });
 ```
 
+The second argument is the MCP SDK's request-scoped context. For example, a host callback can elicit confirmation through the same response stream as the outer tool call:
+
+```ts
+request: async (opts, context) => {
+  const result = await server.server.elicitInput(
+    {
+      message: `Allow ${opts.method} ${opts.path}?`,
+      requestedSchema: {
+        type: "object",
+        properties: { approved: { type: "boolean" } },
+        required: ["approved"]
+      }
+    },
+    {
+      relatedRequestId: context.requestId,
+      signal: context.signal
+    }
+  );
+
+  if (result.action !== "accept" || !result.content?.approved) {
+    throw new Error("Request declined");
+  }
+
+  return callApi(opts);
+};
+```
+
+The context stays in trusted host code and should only be used while the outer tool call is active. Existing callbacks that only accept `opts` continue to work.
+
+### Timeouts: keep the sandbox budget at or above the elicitation timeout
+
+The `request()` callback runs while the sandbox is still suspended on its
+`codemode.request()` call — that call is a blocking RPC round-trip, so the
+executor's timeout covers the whole wait, including the time a human spends
+answering the elicitation. The default `DynamicWorkerExecutor` timeout is
+**60s**, which matches the MCP elicitation timeout (both the SDK's
+`DEFAULT_REQUEST_TIMEOUT_MSEC` and `McpAgent.elicitInput` default to 60s), so
+the default config already lets an elicitation run to completion.
+
+If you lower the executor timeout, a tool that elicits user input (or samples,
+or lists roots) will abort with `Execution timed out` once the sandbox budget
+expires, before the user can respond. Keep the executor timeout at least as
+long as the elicitation timeout:
+
+```ts
+const executor = new DynamicWorkerExecutor({
+  loader: env.LOADER,
+  timeout: 60_000 // do not drop below the 60s elicitation timeout
+});
+```
+
 The LLM first searches the spec:
 
 ```js

packages/agents/src/tests/agents/index.ts

@@ -1,4 +1,5 @@
 export {
+  TestCodemodeMcpAgent,
   TestMcpAgent,
   TestMcpJurisdiction,
   TestAddMcpServerAgent,

packages/agents/src/tests/agents/mcp.ts

@@ -6,6 +6,8 @@ import type {
   ServerNotification,
   ServerRequest
 } from "@modelcontextprotocol/sdk/types.js";
+import { DynamicWorkerExecutor } from "@cloudflare/codemode";
+import { openApiMcpServer } from "@cloudflare/codemode/mcp";
 import { z } from "zod";
 import { McpAgent } from "../../mcp/index.ts";
 import {
@@ -37,6 +39,36 @@ type Props = {
   testValue: string;
 };
 
+export class TestCodemodeMcpAgent extends McpAgent<
+  Cloudflare.Env & { LOADER: WorkerLoader }
+> {
+  server!: McpServer;
+
+  async init() {
+    this.server = openApiMcpServer({
+      spec: {
+        openapi: "3.1.0",
+        info: { title: "Codemode MCP test", version: "1.0.0" },
+        paths: { "/protected": { delete: { summary: "Protected action" } } }
+      },
+      executor: new DynamicWorkerExecutor({ loader: this.env.LOADER }),
+      request: async (options, context) => {
+        return this.elicitInput(
+          {
+            message: `Allow ${options.method} ${options.path}?`,
+            requestedSchema: {
+              type: "object",
+              properties: { approved: { type: "boolean" } },
+              required: ["approved"]
+            }
+          },
+          { relatedRequestId: context.requestId }
+        );
+      }
+    });
+  }
+}
+
 export class TestMcpAgent extends McpAgent<Cloudflare.Env, unknown, Props> {
   private tempToolHandle?: { remove: () => void };
   private collisionBarrierResolvers: Array<() => void> = [];

packages/agents/src/tests/mcp/codemode-context.test.ts

@@ -0,0 +1,119 @@
+import { env } from "cloudflare:workers";
+import { createExecutionContext, runInDurableObject } from "cloudflare:test";
+import type {
+  CallToolResult,
+  JSONRPCMessage,
+  JSONRPCRequest,
+  JSONRPCResultResponse
+} from "@modelcontextprotocol/sdk/types.js";
+import { describe, expect, it, vi } from "vitest";
+import type { TestCodemodeMcpAgent } from "../agents/mcp";
+import {
+  initializeStreamableHTTPServer,
+  parseSSEData,
+  sendPostRequest
+} from "../shared/test-utils";
+
+async function readOneFrame(
+  reader: ReadableStreamDefaultReader<Uint8Array>
+): Promise<string> {
+  const { value } = await reader.read();
+  if (!value) throw new Error("SSE stream ended before a frame arrived");
+  return new TextDecoder().decode(value);
+}
+
+async function getKeepAliveRefs(sessionId: string): Promise<number> {
+  const id = env.TestCodemodeMcpAgent.idFromName(
+    `streamable-http:${sessionId}`
+  );
+  const stub = env.TestCodemodeMcpAgent.get(id);
+  return runInDurableObject(stub, (instance: TestCodemodeMcpAgent) => {
+    return (
+      instance as unknown as {
+        _keepAliveRefs: number;
+      }
+    )._keepAliveRefs;
+  });
+}
+
+describe("Codemode MCP request context", () => {
+  const baseUrl = "http://example.com/codemode-mcp";
+
+  it("completes an elicitation over the originating POST after Worker Loader RPC", async () => {
+    const ctx = createExecutionContext();
+    const sessionId = await initializeStreamableHTTPServer(ctx, baseUrl);
+
+    const toolCall: JSONRPCMessage = {
+      jsonrpc: "2.0",
+      id: "codemode-call-1",
+      method: "tools/call",
+      params: {
+        name: "execute",
+        arguments: {
+          code: `async () => codemode.request({
+            method: "DELETE",
+            path: "/protected"
+          })`
+        }
+      }
+    };
+
+    const toolResponse = await sendPostRequest(
+      ctx,
+      baseUrl,
+      toolCall,
+      sessionId
+    );
+    expect(toolResponse.status).toBe(200);
+
+    const reader = toolResponse.body?.getReader();
+    if (!reader) throw new Error("No reader available for POST stream");
+
+    const elicitRequest = parseSSEData(
+      await readOneFrame(reader)
+    ) as JSONRPCRequest;
+    expect(elicitRequest).toMatchObject({
+      method: "elicitation/create",
+      params: { message: "Allow DELETE /protected?" }
+    });
+
+    // McpAgent.elicitInput holds a keepAlive lease while its in-memory
+    // response resolver is pending. Without it, an unresolved Promise alone
+    // would not prevent the Durable Object from hibernating.
+    expect(await getKeepAliveRefs(sessionId)).toBe(1);
+
+    const elicitResponse: JSONRPCMessage = {
+      jsonrpc: "2.0",
+      id: elicitRequest.id,
+      result: { action: "accept", content: { approved: true } }
+    } as JSONRPCMessage;
+    const response = await sendPostRequest(
+      ctx,
+      baseUrl,
+      elicitResponse,
+      sessionId
+    );
+    expect(response.status).toBe(202);
+
+    const toolResult = parseSSEData(
+      await readOneFrame(reader)
+    ) as JSONRPCResultResponse;
+    expect(toolResult.id).toBe("codemode-call-1");
+    expect(toolResult.result as CallToolResult).toMatchObject({
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(
+            { action: "accept", content: { approved: true } },
+            null,
+            2
+          )
+        }
+      ]
+    });
+
+    await vi.waitFor(async () => {
+      expect(await getKeepAliveRefs(sessionId)).toBe(0);
+    });
+  });
+});

packages/agents/src/tests/worker.ts

@@ -8,6 +8,7 @@ import {
 // Re-export all test agents so existing imports (e.g. `import { type Env } from "./worker"`)
 // and wrangler bindings continue to work.
 export {
+  TestCodemodeMcpAgent,
   TestMcpAgent,
   TestMcpJurisdiction,
   TestAddMcpServerAgent,
@@ -95,6 +96,7 @@ export {
 // circular dependencies.
 
 import type {
+  TestCodemodeMcpAgent,
   TestRpcMcpClientAgent,
   TestEmailAgent,
   TestCaseSensitiveAgent,
@@ -146,6 +148,7 @@ import type {
 export type Env = {
   LOADER: WorkerLoader;
   MCP_OBJECT: DurableObjectNamespace<McpAgent>;
+  TestCodemodeMcpAgent: DurableObjectNamespace<TestCodemodeMcpAgent>;
   EmailAgent: DurableObjectNamespace<TestEmailAgent>;
   CaseSensitiveAgent: DurableObjectNamespace<TestCaseSensitiveAgent>;
   UserNotificationAgent: DurableObjectNamespace<TestUserNotificationAgent>;
@@ -209,7 +212,10 @@ export type Env = {
 
 // ── Fetch handler ────────────────────────────────────────────────────
 
-import { TestMcpAgent as McpAgentImpl } from "./agents";
+import {
+  TestCodemodeMcpAgent as CodemodeMcpAgentImpl,
+  TestMcpAgent as McpAgentImpl
+} from "./agents";
 
 export default {
   async fetch(request: Request, env: Env, ctx: ExecutionContext) {
@@ -229,6 +235,12 @@ export default {
       return McpAgentImpl.serve("/mcp").fetch(request, env, ctx);
     }
 
+    if (url.pathname === "/codemode-mcp") {
+      return CodemodeMcpAgentImpl.serve("/codemode-mcp", {
+        binding: "TestCodemodeMcpAgent"
+      }).fetch(request, env, ctx);
+    }
+
     if (url.pathname === "/auto" || url.pathname === "/auto/message") {
       return McpAgentImpl.serve("/auto", { transport: "auto" }).fetch(
         request,

packages/agents/src/tests/wrangler.jsonc

@@ -16,6 +16,10 @@
         "class_name": "TestMcpAgent",
         "name": "MCP_OBJECT"
       },
+      {
+        "class_name": "TestCodemodeMcpAgent",
+        "name": "TestCodemodeMcpAgent"
+      },
       {
         "class_name": "TestEmailAgent",
         "name": "EmailAgent"
@@ -317,6 +321,7 @@
     {
       "new_sqlite_classes": [
         "TestMcpAgent",
+        "TestCodemodeMcpAgent",
         "TestEmailAgent",
         "TestCaseSensitiveAgent",
         "TestUserNotificationAgent",

packages/codemode/src/mcp.ts

@@ -2,6 +2,11 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
+import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
+import type {
+  ServerNotification,
+  ServerRequest
+} from "@modelcontextprotocol/sdk/types.js";
 import { z } from "zod";
 import {
   generateTypesFromJsonSchema,
@@ -254,10 +259,38 @@ export interface RequestOptions {
   rawBody?: boolean;
 }
 
+/**
+ * MCP context for the outer `execute` tool call.
+ *
+ * This remains in trusted host code and is never exposed to the sandbox.
+ * Use it only while the outer tool call is active.
+ */
+export type OpenApiMcpRequestContext = RequestHandlerExtra<
+  ServerRequest,
+  ServerNotification
+>;
+
 export interface OpenApiMcpServerOptions {
   spec: Record<string, unknown>;
   executor: Executor;
-  request: (options: RequestOptions) => Promise<unknown>;
+  /**
+   * Handle an API request from sandbox code on the trusted host.
+   *
+   * The context belongs to the outer MCP `execute` tool call. It can be used
+   * to associate elicitation, sampling, roots, and notifications with that
+   * call's response stream.
+   *
+   * This callback runs while the sandbox is still suspended on its
+   * `codemode.request()` call, so the executor timeout covers the whole wait.
+   * The default executor timeout (60s) matches the elicitation timeout, so the
+   * default config already lets an elicitation finish. If you lower the
+   * executor timeout below the elicitation timeout, the sandbox aborts the
+   * call before the user can respond.
+   */
+  request: (
+    options: RequestOptions,
+    context: OpenApiMcpRequestContext
+  ) => Promise<unknown>;
   name?: string;
   version?: string;
   description?: string;
@@ -510,15 +543,16 @@ async () => {
         code: z.string().describe("JavaScript async arrow function to execute")
       }
     },
-    async ({ code }) => {
+    async ({ code }, context) => {
       try {
         const result = await executor.execute(
           createOpenApiSandboxCode(code, spec, true),
           [
             {
               name: "__openapiHost",
               fns: {
-                request: (args: unknown) => requestFn(args as RequestOptions)
+                request: (args: unknown) =>
+                  requestFn(args as RequestOptions, context)
               }
             }
           ]