From d740877dd12aeb2f9542cfb6e5883b6ae925de66 Mon Sep 17 00:00:00 2001 From: Argjent Sahiti Date: Wed, 21 Oct 2026 22:29:34 +0200 Subject: [PATCH 1/2] feat(ai-mcp): pass the request context to createMCPServer setup setup(server, context) now gets the handleOptions.context of the request that builds the server. One createMCPServer serves many callers, so an extension handler or an instrumentation hook can read the caller from the request instead of from shared state. A spec 2025 session with sessions: 'memory' gets the context of the request that opened it. --- .changeset/ai-mcp-setup-context.md | 5 ++ docs/config.json | 2 +- docs/mcp/server.md | 10 ++- packages/ai-mcp/skills/ai-mcp/SKILL.md | 2 +- packages/ai-mcp/src/server/create-server.ts | 19 ++++- .../ai-mcp/tests/server/create-server.test.ts | 85 ++++++++++++++++++- testing/e2e/src/lib/mcp-typed-server.ts | 6 +- testing/e2e/src/routes/api.mcp-typed-test.ts | 2 +- testing/e2e/tests/mcp-typed.spec.ts | 2 +- 9 files changed, 120 insertions(+), 13 deletions(-) create mode 100644 .changeset/ai-mcp-setup-context.md 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..2c36d75f8c 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,10 @@ 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' }, }) }) From 8eb7b15da1c053027bfb3c62c283b304a4de7cfa Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 04:00:06 +0000 Subject: [PATCH 2/2] ci: apply automated fixes --- packages/ai-mcp/src/server/create-server.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/packages/ai-mcp/src/server/create-server.ts b/packages/ai-mcp/src/server/create-server.ts index 2c36d75f8c..66c61fa2fa 100644 --- a/packages/ai-mcp/src/server/create-server.ts +++ b/packages/ai-mcp/src/server/create-server.ts @@ -506,10 +506,7 @@ function buildMcpServer(input: { if (input.era === '2025') { registerLegacyTaskMethods(server, input.taskStore, input.owner) } - const setupResult: unknown = input.options.setup?.( - server, - input.appContext(), - ) + const setupResult: unknown = input.options.setup?.(server, input.appContext()) if (setupResult instanceof Promise) { throw new TypeError('createMCPServer: setup must be synchronous') }