feat(ai-mcp): pass the handle context to createMCPServer setup

.changeset/ai-mcp-setup-context.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai-mcp': minor
+---
+
+`createMCPServer({ setup })` gets the `handle` context of the request as its second argument, so an extension handler or an instrumentation hook can read the caller. A spec 2025 session with `sessions: 'memory'` gets the context of the request that opened it.

docs/config.json

@@ -588,7 +588,7 @@
           "label": "Serve Tools over HTTP",
           "to": "mcp/server",
           "addedAt": "2026-09-22",
-          "updatedAt": "2026-10-21"
+          "updatedAt": "2026-10-22"
         },
         {
           "label": "Serve Resources and Prompts",

docs/mcp/server.md

@@ -183,22 +183,28 @@ import { z } from 'zod'
 const server = createMCPServer({
   name: 'notes',
   version: '1.0.0',
-  setup: (mcp) => {
+  setup: (mcp, context) => {
+    const userId = context?.userId
     mcp.server.registerCapabilities({
       extensions: { 'com.example/notes': {} },
     })
     mcp.server.setRequestHandler(
       'notes/count',
       { params: z.object({}) },
-      () => ({ count: 3 }),
+      () => ({ count: typeof userId === 'string' ? userId.length : 0 }),
     )
   },
 })
+
+export function handleMcp(request: Request, userId: string) {
+  return server.handle(request, { context: { userId } })
+}
 ```
 
 Rules for `setup`:
 
 - It runs each time the server builds an `McpServer`. That is each request, or once for each spec 2025 session when `sessions` is `'memory'`. Keep it fast.
+- The second argument is the `context` object from `server.handle`. On a spec 2025 memory session, that value is from the request that opened the session. A later request on that session does not run `setup` again. A tool on the later request sees the new context.
 - It must be synchronous. The server connects as soon as `setup` returns, so an `async` function fails the request.
 - Add new methods only. Do not replace a method the server already handles, such as `tools/call`.
 

packages/ai-mcp/skills/ai-mcp/SKILL.md

@@ -236,7 +236,7 @@ Use the MCP names: `readOnlyHint`, `destructiveHint`, `idempotentHint`, `openWor
 A host skips its confirmation for a tool with `readOnlyHint: true`.
 Set `metadata._meta` to send the MCP tool `_meta`, for example `{ ui: { resourceUri: 'ui://view' } }` for an MCP Apps view.
 Pass `onerror` to `createMCPServer` to log transport and protocol errors from the SDK. `serveMCPStdio` also sends its transport errors there.
-Pass `setup: (server) => {}` to `createMCPServer` to change each SDK `McpServer` before it connects. It must be synchronous. Serve an MCP extension there: `server.server.registerCapabilities({ extensions: { [id]: {} } })`, then `server.server.setRequestHandler(method, { params }, handler)` for each method.
+Pass `setup: (server, context) => {}` to `createMCPServer` to change each SDK `McpServer` before it connects. It must be synchronous. `context` is the `handle` context of the request that builds the server. A spec 2025 session with `sessions: 'memory'` keeps the context of the request that opened the session. Serve an MCP extension there: `server.server.registerCapabilities({ extensions: { [id]: {} } })`, then `server.server.setRequestHandler(method, { params }, handler)` for each method.
 
 A tool with no `outputSchema` can return an MCP `CallToolResult`.
 The server sends it as is: its content blocks, its `structuredContent`, and its `isError`.

packages/ai-mcp/src/server/create-server.ts

@@ -144,10 +144,18 @@ export type MCPServerOptions = {
    * to serve an MCP extension. Leave the methods this server already
    * handles alone.
    *
+   * `context` is the `handleOptions.context` of the request that builds the
+   * server. A spec 2025 session with `sessions: 'memory'` builds its server
+   * once, so there it is the context of the request that opened the session.
+   * A later request on that session does not run `setup` again.
+   *
    * It must be synchronous. The server connects as soon as it returns, so
    * an async `setup` throws a `TypeError`.
    */
-  setup?: (server: McpServer) => void
+  setup?: (
+    server: McpServer,
+    context: Record<string, unknown> | undefined,
+  ) => void
 }
 
 /**
@@ -214,7 +222,9 @@ type LegacySessions = Map<string, LegacySession>
  * `options.sample` is the model adapter for `ctx.context.sample` on spec 2026.
  * `options.waitUntil` receives the task promise so a worker can stay alive.
  * `options.setup` gets each SDK `McpServer` before it connects, so it can add
- * capabilities and request handlers, such as an MCP extension.
+ * capabilities and request handlers, such as an MCP extension. It also gets
+ * the `handle` context of the request that builds the server. A spec 2025
+ * memory session passes the context of the request that opened it.
  *
  * The result has `fetch(request)`, `handle(request, options)`, `tools`,
  * `resources`, and `prompts`. `handle` takes a token your own middleware
@@ -496,7 +506,7 @@ function buildMcpServer(input: {
   if (input.era === '2025') {
     registerLegacyTaskMethods(server, input.taskStore, input.owner)
   }
-  const setupResult: unknown = input.options.setup?.(server)
+  const setupResult: unknown = input.options.setup?.(server, input.appContext())
   if (setupResult instanceof Promise) {
     throw new TypeError('createMCPServer: setup must be synchronous')
   }

packages/ai-mcp/tests/server/create-server.test.ts

@@ -3,7 +3,10 @@ import {
   StreamableHTTPClientTransport,
 } from '@modelcontextprotocol/client'
 import { OAuthError, OAuthErrorCode } from '@modelcontextprotocol/server'
-import type { OAuthTokenVerifier } from '@modelcontextprotocol/server'
+import type {
+  McpServer,
+  OAuthTokenVerifier,
+} from '@modelcontextprotocol/server'
 import { toolDefinition } from '@tanstack/ai'
 import { describe, expect, it, vi } from 'vitest'
 import { z } from 'zod'
@@ -1380,6 +1383,86 @@ describe('createMCPServer', () => {
     }
   })
 
+  it('gives setup the context of the request that builds the server', async () => {
+    const server = createMCPServer({
+      name: 'tenants',
+      version: '1.0.0',
+      sessions: 'stateless',
+      setup: (built, context) => {
+        built.server.registerCapabilities({
+          extensions: { 'com.example/tenant': {} },
+        })
+        built.server.setRequestHandler(
+          'tenant/get',
+          { params: z.object({}) },
+          () => ({ tenant: context?.tenant }),
+        )
+      },
+    })
+    const tenantResult = z.object({ tenant: z.string() })
+    for (const tenant of ['acme', 'globex']) {
+      await withClient(
+        server,
+        { era: '2026', handleOptions: { context: { tenant } } },
+        async (client) => {
+          const answer = await client.request(
+            { method: 'tenant/get', params: {} },
+            tenantResult,
+          )
+          expect(answer).toEqual({ tenant })
+        },
+      )
+    }
+  })
+
+  it('keeps the opening context on a spec 2025 memory session', async () => {
+    const setup = vi.fn(
+      (built: McpServer, context: Record<string, unknown> | undefined) => {
+        built.server.registerCapabilities({
+          extensions: { 'com.example/tenant': {} },
+        })
+        built.server.setRequestHandler(
+          'tenant/get',
+          { params: z.object({}) },
+          () => ({ tenant: context?.tenant }),
+        )
+      },
+    )
+    const server = createMCPServer({
+      name: 'tenants',
+      version: '1.0.0',
+      sessions: 'memory',
+      tools: [
+        toolDefinition({
+          name: 'tenant',
+          description: 'Reports the tenant',
+          inputSchema: z.object({}),
+        }).server<MCPToolContext<{ tenant: string }>>(
+          async (_args, ctx) => ctx.context.tenant,
+        ),
+      ],
+      setup,
+    })
+    const handleOptions: MCPHandleOptions = { context: { tenant: 'acme' } }
+    await withClient(server, { era: '2025', handleOptions }, async (client) => {
+      const opened = await client.request(
+        { method: 'tenant/get', params: {} },
+        z.object({ tenant: z.string() }),
+      )
+      expect(opened).toEqual({ tenant: 'acme' })
+      handleOptions.context = { tenant: 'globex' }
+      const resumed = await client.request(
+        { method: 'tenant/get', params: {} },
+        z.object({ tenant: z.string() }),
+      )
+      expect(resumed).toEqual({ tenant: 'acme' })
+      const called = await client.callTool({ name: 'tenant', arguments: {} })
+      expect(called.content).toEqual([{ type: 'text', text: 'globex' }])
+    })
+    expect(setup).toHaveBeenCalledTimes(1)
+    expect(setup).toHaveBeenCalledWith(expect.anything(), { tenant: 'acme' })
+  })
+
   it('reports SDK errors to onerror', async () => {
     const onerror = vi.fn()
     const server = createMCPServer({

testing/e2e/src/lib/mcp-typed-server.ts

@@ -73,7 +73,7 @@ export const typedServer = createMCPServer({
  *   the `handle` context, and lists one concrete note.
  * - `show_note` links an MCP Apps view through `metadata._meta`.
  * - `setup` serves the `com.example/notes` extension with a `notes/count`
- *   method.
+ *   method that reads the `tenant` from the `handle` context.
  */
 const note = resourceDefinition({
   name: 'note',
@@ -101,14 +101,14 @@ export const statelessServer = createMCPServer({
   sessions: 'stateless',
   tools: [showNote],
   resources: [note],
-  setup: (server) => {
+  setup: (server, context) => {
     server.server.registerCapabilities({
       extensions: { 'com.example/notes': {} },
     })
     server.server.setRequestHandler(
       'notes/count',
       { params: z.object({}) },
-      () => ({ count: 1 }),
+      () => ({ count: 1, tenant: String(context?.tenant) }),
     )
   },
 })

testing/e2e/src/routes/api.mcp-typed-test.ts

@@ -63,7 +63,7 @@ export const Route = createFileRoute('/api/mcp-typed-test')({
           try {
             const counted = await notes.request(
               { method: 'notes/count', params: {} },
-              z.object({ count: z.number() }),
+              z.object({ count: z.number(), tenant: z.string() }),
             )
             return Response.json({
               extensions: notes.capabilities.extensions ?? null,

testing/e2e/tests/mcp-typed.spec.ts

@@ -232,7 +232,7 @@ test.describe('mcp: createMCPServer with auth, output schema, and a task', () =>
     expect(res.ok(), body).toBe(true)
     expect(JSON.parse(body)).toEqual({
       extensions: { 'com.example/notes': {} },
-      counted: { count: 1 },
+      counted: { count: 1, tenant: 'acme' },
     })
   })