diff --git a/.changeset/ai-mcp-setup-context.md b/.changeset/ai-mcp-setup-context.md new file mode 100644 index 0000000000..de7f3d4c86 --- /dev/null +++ b/.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. diff --git a/docs/config.json b/docs/config.json index ae37dc60fb..56ca07e6a8 100644 --- a/docs/config.json +++ b/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", diff --git a/docs/mcp/server.md b/docs/mcp/server.md index bbb05e1d45..59f268b704 100644 --- a/docs/mcp/server.md +++ b/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`. diff --git a/packages/ai-mcp/skills/ai-mcp/SKILL.md b/packages/ai-mcp/skills/ai-mcp/SKILL.md index feeae587a2..3b20362285 100644 --- a/packages/ai-mcp/skills/ai-mcp/SKILL.md +++ b/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`. diff --git a/packages/ai-mcp/src/server/create-server.ts b/packages/ai-mcp/src/server/create-server.ts index c0790023a4..66c61fa2fa 100644 --- a/packages/ai-mcp/src/server/create-server.ts +++ b/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 | undefined, + ) => void } /** @@ -214,7 +222,9 @@ type LegacySessions = Map * `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') } diff --git a/packages/ai-mcp/tests/server/create-server.test.ts b/packages/ai-mcp/tests/server/create-server.test.ts index 38a2df2163..3fc6020c1e 100644 --- a/packages/ai-mcp/tests/server/create-server.test.ts +++ b/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 | 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>( + 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({ diff --git a/testing/e2e/src/lib/mcp-typed-server.ts b/testing/e2e/src/lib/mcp-typed-server.ts index becf3ff785..2de6e4a011 100644 --- a/testing/e2e/src/lib/mcp-typed-server.ts +++ b/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) }), ) }, }) diff --git a/testing/e2e/src/routes/api.mcp-typed-test.ts b/testing/e2e/src/routes/api.mcp-typed-test.ts index 05a37f45a5..0c017c767a 100644 --- a/testing/e2e/src/routes/api.mcp-typed-test.ts +++ b/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, diff --git a/testing/e2e/tests/mcp-typed.spec.ts b/testing/e2e/tests/mcp-typed.spec.ts index 0ab4654763..9fba47615c 100644 --- a/testing/e2e/tests/mcp-typed.spec.ts +++ b/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' }, }) })