diff --git a/.changeset/server-extensions.md b/.changeset/server-extensions.md new file mode 100644 index 0000000000..a5c7c939db --- /dev/null +++ b/.changeset/server-extensions.md @@ -0,0 +1,9 @@ +--- +'@modelcontextprotocol/core-internal': minor +'@modelcontextprotocol/client': minor +'@modelcontextprotocol/server': minor +--- + +Server and client extensions. `ServerOptions.extensions` takes `ServerExtension` objects (`{ id, install(server) }`): each is advertised under `capabilities.extensions[id]` and installed at construction; everything else an extension does, settings included, happens in `install`. Extensions register custom methods with `setRequestHandler(method, { params, result }, handler)` and intercept spec methods with the new `Protocol.use(method, (request, ctx, next) => …)` — middleware around the registered handler, composed at dispatch time in registration order (so middleware on `tools/call` applies even though `McpServer` registers that handler lazily); it returns a remover. `ClientOptions.extensions` takes the symmetric `ClientExtension` objects, advertised under the client's `capabilities.extensions[id]` (in `initialize` on a legacy connection, in every request's client-capabilities envelope on 2026-07-28) and installed with the `Client`. `Protocol.acceptResultType(method, resultType)` declares an extension result kind for a method: a raw response carrying that `resultType` bypasses the era codec's closed vocabulary and is validated against the caller's explicit result schema as-is, which is how a client extension receives shapes such as the Tasks extension's `resultType: "task"` on `tools/call`. + +`McpServer` tool dispatch now re-throws `MissingRequiredClientCapabilityError` (`-32021`) as a JSON-RPC error instead of converting it into an `isError` tool result, matching the existing `UrlElicitationRequiredError` passthrough. diff --git a/docs/.vitepress/nav.ts b/docs/.vitepress/nav.ts index 02ab4d2ff6..f1d5f96cdf 100644 --- a/docs/.vitepress/nav.ts +++ b/docs/.vitepress/nav.ts @@ -66,6 +66,7 @@ export const guideSidebar: DefaultTheme.SidebarItem[] = [ items: [ { text: 'Low-level server', link: '/advanced/low-level-server' }, { text: 'Custom methods', link: '/advanced/custom-methods' }, + { text: 'Server extensions', link: '/advanced/extensions' }, { text: 'Schema libraries', link: '/advanced/schema-libraries' }, { text: 'Custom transports', link: '/advanced/custom-transports' }, { text: 'Wire schemas', link: '/advanced/wire-schemas' }, diff --git a/docs/advanced/extensions.md b/docs/advanced/extensions.md new file mode 100644 index 0000000000..14c4ceb039 --- /dev/null +++ b/docs/advanced/extensions.md @@ -0,0 +1,85 @@ +--- +shape: how-to +--- + +# Server extensions + +A **server extension** packages protocol behaviour outside the core specification — an MCP extension such as `io.modelcontextprotocol/tasks`, or a vendor feature — as one object you pass to the server. The SDK advertises it, installs it, and gives it two hooks: custom methods and middleware on spec methods. What the extension does behind those hooks is its own business. + +## Write an extension + +An extension is an `id` and an `install` function that receives the low-level `Server`. Everything it does to the protocol happens in `install`. + +```ts +import type { ServerExtension } from '@modelcontextprotocol/server'; +import { MissingRequiredClientCapabilityError, CLIENT_CAPABILITIES_META_KEY } from '@modelcontextprotocol/server'; +import * as z from 'zod/v4'; + +const GATE = 'com.example/gate'; + +export const gate: ServerExtension = { + id: GATE, + install(server) { + // Settings for the advertised capability, if the extension has any. + server.registerCapabilities({ extensions: { [GATE]: { exampleData: true } } }); + + // A custom method, exactly as in Custom methods. + server.setRequestHandler('gate/status', { params: z.looseObject({}) }, () => ({ armed: true })); + + // Middleware on a spec method: runs around the registered handler, + // may answer, transform, or refuse. + server.use('tools/call', (request, ctx, next) => { + const declared = ctx.mcpReq.envelope?.[CLIENT_CAPABILITIES_META_KEY]?.extensions ?? {}; + if (!(GATE in declared)) { + throw new MissingRequiredClientCapabilityError({ requiredCapabilities: { extensions: { [GATE]: {} } } }, 'declare the gate'); + } + return next(request, ctx); + }); + } +}; +``` + +## Install it + +Pass extensions at construction. Each is advertised under `capabilities.extensions[id]` as `{}` — legacy connections see it in the `initialize` result, 2026-07-28 connections in `server/discover` — and then installed in order, after the built-in handlers exist. An extension with settings registers them in `install`, as above; the maps merge. + +```ts +const server = new McpServer({ name: 'gated', version: '1.0.0' }, { extensions: [gate] }); +``` + +The same option exists on the low-level `Server`. + +## Client extensions + +The client half is symmetric: `ClientExtension` is `{ id, install(client) }`, passed in `ClientOptions.extensions`. The client advertises it under its own `capabilities.extensions[id]` — in `initialize` on a legacy connection, and in every request's `_meta` client-capabilities envelope on a 2026-07-28 connection, which is where a server extension reads it — and `install` receives the `Client` to register handlers for server-to-client requests and notifications, or wrap the ones the SDK installs with middleware. + +```ts +import type { ClientExtension } from '@modelcontextprotocol/client'; + +const gateClient: ClientExtension = { + id: GATE, + install(client) { + client.setRequestHandler('gate/ping', { params: z.looseObject({}) }, () => ({ pong: true })); + // Results of tools/call may carry the extension's own kind; the + // explicit-schema request() path then hands them to the caller's + // schema as-is instead of rejecting the unknown resultType. + client.acceptResultType('tools/call', 'gate'); + } +}; + +const client = new Client({ name: 'gated-client', version: '1.0.0' }, { extensions: [gateClient] }); +``` + +## How middleware composes + +`setRequestHandler` is the route handler; `use(method, middleware)` is the middleware around it, Koa-shaped: `await next(request, ctx)` yields the result and the middleware returns what goes on the wire. It wraps whatever handler serves `method` at dispatch time. That matters for `tools/call`, which `McpServer` registers on the first tool registration: middleware installed at construction still applies. With no underlying handler, `next` throws `MethodNotFound`. Several middleware nest in registration order, the first installed outermost. The returned function removes the middleware. Method names are exact; there is no wildcard. Something every peer must know about the extension belongs in its advertised capability (`registerCapabilities`), not in per-request `_meta`. Where a result does carry extension `_meta`, the middleware adds its own namespaced key and the encoder still stamps the SDK's reserved keys and `resultType` on top. + +A thrown `ProtocolError` becomes the JSON-RPC error response. Inside a tool handler, `McpServer` converts most throws into an `isError` tool result; the exceptions are protocol-level errors the client must see as errors — `UrlElicitationRequiredError` and `MissingRequiredClientCapabilityError` (`-32021`). + +## Recap + +- `ServerExtension` is `{ id, install(server) }`; pass it in `ServerOptions.extensions`. `ClientExtension` mirrors it on `ClientOptions.extensions`. +- `install` gets the low-level `Server`: `setRequestHandler` for custom methods, `use` for middleware on spec methods, `registerCapabilities` for the extension's settings. +- Middleware composes at dispatch time, in registration order, and applies to handlers registered later. +- `acceptResultType(method, resultType)` lets a client extension receive a result kind outside `complete` / `input_required` through the explicit-schema `request()` path. +- The SDK owns the hooks and the capability advertisement, not the extension's state or execution. diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts index 0b386a63e8..bb1b825878 100644 --- a/packages/client/src/client/client.ts +++ b/packages/client/src/client/client.ts @@ -80,6 +80,7 @@ import { SUPPORTED_MODERN_PROTOCOL_VERSIONS } from '@modelcontextprotocol/core-internal'; +import type { ClientExtension } from './extension'; import type { PriorDiscovery } from './probeClassifier'; import type { CacheMode, CacheScope, ResponseCacheStore } from './responseCache'; import { ClientResponseCache, InMemoryResponseCacheStore, MAX_CACHE_TTL_MS } from './responseCache'; @@ -181,6 +182,13 @@ export function getSupportedElicitationModes(capabilities: ClientCapabilities['e } export type ClientOptions = ProtocolOptions & { + /** + * Extensions to install at construction. Each is advertised under + * `capabilities.extensions[extension.id]` and then installed, in order — + * see {@linkcode ClientExtension}. + */ + extensions?: ClientExtension[]; + /** * Capabilities to advertise as being supported by this client. */ @@ -658,6 +666,11 @@ export class Client extends Protocol { if (options?.listChanged) { this._listChangedConfig = options.listChanged; } + + for (const extension of options?.extensions ?? []) { + this.registerCapabilities({ extensions: { [extension.id]: {} } }); + extension.install(this); + } } protected override buildContext(ctx: BaseContext, _transportInfo?: MessageExtraInfo): ClientContext { diff --git a/packages/client/src/client/extension.ts b/packages/client/src/client/extension.ts new file mode 100644 index 0000000000..7e0f4c59b5 --- /dev/null +++ b/packages/client/src/client/extension.ts @@ -0,0 +1,28 @@ +import type { Client } from './client'; + +/** + * A client extension: the client half of protocol behaviour outside the + * core specification (an MCP extension such as `io.modelcontextprotocol/tasks`, + * or a vendor feature) that installs itself onto a {@linkcode Client}. + * + * Pass extensions at construction — `new Client(info, { extensions: [ext] })`. + * The client advertises each extension under `capabilities.extensions[id]` + * (as `{}`: supported, no settings) + * (in `initialize` on a legacy connection, in every request's + * `_meta` client-capabilities envelope on a 2026-07-28 connection) and then + * calls `install`, which is where the extension registers handlers for + * server-to-client requests and notifications, or wraps the ones the SDK + * installs with middleware (`client.use('elicitation/create', …)`). + * The SDK provides the hooks; what an extension does behind them is its own. + */ +export interface ClientExtension { + /** + * The extension identifier, prefix-qualified (`io.modelcontextprotocol/tasks`, + * `com.example/feature-flags`). Advertised as the key under + * `capabilities.extensions`. Settings for that key, if any, are the + * extension's to register in `install` via `registerCapabilities`. + */ + readonly id: string; + /** Installs the extension's handlers and middleware onto the client. Called once, at construction. */ + install(client: Client): void; +} diff --git a/packages/client/src/index.ts b/packages/client/src/index.ts index 0b5b6e86ea..12c2667b7f 100644 --- a/packages/client/src/index.ts +++ b/packages/client/src/index.ts @@ -68,6 +68,7 @@ export { StaticPrivateKeyJwtProvider } from './client/authExtensions'; export type { CacheableRequestOptions, CallToolRequestOptions, ClientOptions, ConnectOptions, McpSubscription } from './client/client'; +export type { ClientExtension } from './client/extension'; export { Client } from './client/client'; export { getSupportedElicitationModes } from './client/client'; export type { DiscoverAndRequestJwtAuthGrantOptions, JwtAuthGrantResult, RequestJwtAuthGrantOptions } from './client/crossAppAccess'; diff --git a/packages/client/test/client/extensions.test.ts b/packages/client/test/client/extensions.test.ts new file mode 100644 index 0000000000..065af3d88c --- /dev/null +++ b/packages/client/test/client/extensions.test.ts @@ -0,0 +1,145 @@ +/** + * `ClientOptions.extensions`: an extension is advertised under the client's + * `capabilities.extensions[id]` — in `initialize` on a legacy connection, in + * every request's `_meta` client-capabilities envelope on a 2026-07-28 + * connection — and installed at construction, where it can register + * handlers for server-to-client requests. + */ +import type { JSONRPCMessage } from '@modelcontextprotocol/core-internal'; +import { CLIENT_CAPABILITIES_META_KEY, InMemoryTransport, LATEST_PROTOCOL_VERSION } from '@modelcontextprotocol/core-internal'; +import { describe, expect, it } from 'vitest'; +import * as z from 'zod/v4'; + +import { Client } from '../../src/client/client'; +import type { ClientExtension } from '../../src/client/extension'; + +const MODERN = '2026-07-28'; +const EXT_ID = 'com.example/gate'; + +const flush = () => new Promise(resolve => setTimeout(resolve, 20)); + +function gateExtension(log: string[]): ClientExtension { + return { + id: EXT_ID, + install(client) { + log.push('installed'); + client.registerCapabilities({ extensions: { [EXT_ID]: { exampleData: true } } }); + client.setRequestHandler('gate/ping', { params: z.looseObject({}) }, () => ({ pong: true })); + client.acceptResultType('tools/call', 'task'); + } + }; +} + +/** A scripted server side: answers the handshake for the requested era and records everything the client writes. */ +async function scriptedServer(era: 'modern' | 'legacy') { + const [clientTx, serverTx] = InMemoryTransport.createLinkedPair(); + const written: JSONRPCMessage[] = []; + serverTx.onmessage = message => { + written.push(message); + const request = message as { id?: number | string; method?: string }; + if (request.id === undefined) return; + if (request.method === 'server/discover') { + void serverTx.send( + era === 'modern' + ? { + jsonrpc: '2.0', + id: request.id, + result: { + resultType: 'complete', + supportedVersions: [MODERN], + capabilities: { tools: {} }, + _meta: { 'io.modelcontextprotocol/serverInfo': { name: 'scripted', version: '1.0.0' } } + } + } + : { jsonrpc: '2.0', id: request.id, error: { code: -32_601, message: 'Method not found' } } + ); + } else if (request.method === 'initialize') { + void serverTx.send({ + jsonrpc: '2.0', + id: request.id, + result: { + protocolVersion: LATEST_PROTOCOL_VERSION, + capabilities: { tools: {} }, + serverInfo: { name: 'scripted', version: '1.0.0' } + } + }); + } else if (request.method === 'tools/call') { + void serverTx.send({ + jsonrpc: '2.0', + id: request.id, + result: { resultType: 'task', taskId: 't-1', status: 'working', createdAt: 'now', lastUpdatedAt: 'now', ttlMs: null } + }); + } else if (request.method === 'tools/list') { + void serverTx.send({ + jsonrpc: '2.0', + id: request.id, + result: era === 'modern' ? { resultType: 'complete', tools: [], ttlMs: 0, cacheScope: 'public' } : { tools: [] } + }); + } + }; + await serverTx.start(); + return { clientTx, serverTx, written }; +} + +const paramsOf = (message: JSONRPCMessage): Record => (message as { params?: Record }).params ?? {}; + +describe('ClientOptions.extensions', () => { + it('installs the extension at construction', () => { + const log: string[] = []; + new Client({ name: 'c', version: '1' }, { extensions: [gateExtension(log)] }); + expect(log).toEqual(['installed']); + }); + + it('advertises {} when install sets no settings', async () => { + const { clientTx, written } = await scriptedServer('legacy'); + const client = new Client({ name: 'c', version: '1' }, { extensions: [{ id: 'com.example/plain', install: () => {} }] }); + await client.connect(clientTx); + const initialize = written.find(message => (message as { method?: string }).method === 'initialize'); + expect(paramsOf(initialize as JSONRPCMessage)['capabilities']).toMatchObject({ extensions: { 'com.example/plain': {} } }); + await client.close(); + }); + + it('sends the extension in initialize on a legacy connection', async () => { + const { clientTx, written } = await scriptedServer('legacy'); + const client = new Client({ name: 'c', version: '1' }, { extensions: [gateExtension([])] }); + await client.connect(clientTx); + const initialize = written.find(message => (message as { method?: string }).method === 'initialize'); + expect(paramsOf(initialize as JSONRPCMessage)['capabilities']).toMatchObject({ extensions: { [EXT_ID]: { exampleData: true } } }); + await client.close(); + }); + + it('stamps the extension into every request envelope on a 2026-07-28 connection', async () => { + const { clientTx, written } = await scriptedServer('modern'); + const client = new Client({ name: 'c', version: '1' }, { versionNegotiation: { mode: 'auto' }, extensions: [gateExtension([])] }); + await client.connect(clientTx); + await client.listTools(); + await flush(); + const toolsList = written.find(message => (message as { method?: string }).method === 'tools/list'); + const meta = paramsOf(toolsList as JSONRPCMessage)['_meta'] as Record; + expect(meta[CLIENT_CAPABILITIES_META_KEY]).toMatchObject({ extensions: { [EXT_ID]: { exampleData: true } } }); + await client.close(); + }); + + it('receives an extension result kind the extension accepted, discriminator included', async () => { + const { clientTx } = await scriptedServer('modern'); + const client = new Client({ name: 'c', version: '1' }, { versionNegotiation: { mode: 'auto' }, extensions: [gateExtension([])] }); + await client.connect(clientTx); + const taskSchema = z.looseObject({ resultType: z.literal('task'), taskId: z.string(), status: z.string() }); + const result = await client.request({ method: 'tools/call', params: { name: 'slow', arguments: {} } }, taskSchema); + expect(result).toMatchObject({ resultType: 'task', taskId: 't-1', status: 'working' }); + // callTool validates against CallToolResultSchema, which a task handle does not satisfy. + await expect(client.callTool({ name: 'slow', arguments: {} })).rejects.toThrow(/Invalid result for tools\/call/); + await client.close(); + }); + + it('serves the handler the extension installed for a server-to-client request', async () => { + const { clientTx, serverTx, written } = await scriptedServer('legacy'); + const client = new Client({ name: 'c', version: '1' }, { extensions: [gateExtension([])] }); + await client.connect(clientTx); + await serverTx.send({ jsonrpc: '2.0', id: 'srv-1', method: 'gate/ping', params: {} }); + await flush(); + const response = written.find(message => 'id' in message && message.id === 'srv-1'); + expect((response as { result?: unknown }).result).toEqual({ pong: true }); + await client.close(); + }); +}); diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 637be389aa..59d7e7f95d 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -198,6 +198,17 @@ const RESERVED_ENVELOPE_META_KEYS: readonly string[] = [ */ const RETRY_PARAMS_KEYS = ['inputResponses', 'requestState'] as const; +/** + * Request middleware (see `Protocol.use`): receives the parsed request, the + * context, and `next` — the handler it wraps — and returns the result that + * goes on the wire. + */ +export type RequestMiddleware = ( + request: JSONRPCRequest, + ctx: ContextT, + next: (request: JSONRPCRequest, ctx: ContextT) => Promise +) => Result | Promise; + /** * Lift wire-only material out of an inbound message so handlers see exactly * the 2025-era shape, and surface it for the protocol layer (requests: via @@ -559,6 +570,10 @@ export abstract class Protocol { private _transport?: Transport; private _requestMessageId = 0; private _requestHandlers: Map Promise> = new Map(); + /** Middleware installed by `use`, in registration order across every method (see `_resolveRequestHandler`). */ + private _requestMiddleware: Array<{ method: string; middleware: RequestMiddleware }> = []; + /** Extension result kinds accepted per method (see `acceptResultType`). */ + private _acceptedResultTypes: Map> = new Map(); private _requestHandlerAbortControllers: Map = new Map(); private _notificationHandlers: Map Promise> = new Map(); private _responseHandlers: Map void> = new Map(); @@ -1005,7 +1020,7 @@ export abstract class Protocol { return; } - const handler = this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler; + const handler = this._resolveRequestHandler(request.method); if (handler === undefined) { sendErrorResponse(ProtocolErrorCode.MethodNotFound, 'Method not found'); @@ -1514,6 +1529,21 @@ export abstract class Protocol { // `_onresponse`, so a throw out of the decode hop would // otherwise propagate into the transport's onmessage instead // of failing this request. + // Extension result kinds (see `acceptResultType`): a raw + // `resultType` the caller declared for this method bypasses + // the codec's closed vocabulary and reaches the caller's + // schema as-is, discriminator included. + if (this._isAcceptedResultType(request.method, response.result)) { + validateStandardSchema(resultSchema, response.result).then(parseResult => { + if (parseResult.success) { + resolve(parseResult.data); + } else { + reject(new SdkError(SdkErrorCode.InvalidResult, `Invalid result for ${request.method}: ${parseResult.error}`)); + } + }, reject); + return; + } + let decoded: ReturnType; try { decoded = codec.decodeResult(request.method, response.result); @@ -1754,6 +1784,94 @@ export abstract class Protocol { this._requestHandlers.set(method, this._wrapHandler(method, stored)); } + /** + * Installs middleware on the request handler for `method`. The middleware + * receives the parsed request, the context, and `next` — the handler it + * wraps — and may answer itself, transform what `next` returns, or throw + * a `ProtocolError` that becomes the JSON-RPC error response. + * `setRequestHandler` is the route handler; this is the middleware + * around it, Koa-shaped: `await next(request, ctx)` yields the result and + * the middleware returns what goes on the wire. + * + * Middleware composes at dispatch time, so middleware installed before + * the underlying handler exists (e.g. `tools/call`, which `McpServer` + * registers on the first tool registration) still applies; with no + * underlying handler and no fallback, `next` throws `MethodNotFound`. + * Middleware runs in registration order: the first installed is the + * outermost. Method names are exact — there is no wildcard; what an + * extension needs every peer to know goes through `registerCapabilities`, + * not per-request middleware. This is the hook extensions use to + * intercept spec methods. + * + * @returns A function that removes the middleware. + */ + use(method: RequestMethod | string, middleware: RequestMiddleware): () => void { + if (typeof middleware !== 'function') { + throw new TypeError('use: middleware is required'); + } + const entry = { method, middleware }; + this._requestMiddleware.push(entry); + return () => { + const index = this._requestMiddleware.indexOf(entry); + if (index !== -1) this._requestMiddleware.splice(index, 1); + }; + } + + /** + * Declares that results of `method` may carry `resultType` — a kind an + * extension defines beyond the spec's `complete` / `input_required` + * vocabulary (the Tasks extension answers `tools/call` with + * `resultType: "task"`). A raw response with that discriminator skips the + * era codec's decode and is validated against the caller's explicit + * result schema as-is, discriminator included. Only the explicit-schema + * `request(request, resultSchema)` path consults this; typed spec calls + * keep the closed vocabulary. This is the hook client extensions use to + * receive extension result shapes. + * + * @returns A function that withdraws the declaration. + */ + acceptResultType(method: string, resultType: string): () => void { + const accepted = this._acceptedResultTypes.get(method) ?? new Set(); + accepted.add(resultType); + this._acceptedResultTypes.set(method, accepted); + return () => { + const current = this._acceptedResultTypes.get(method); + current?.delete(resultType); + if (current?.size === 0) this._acceptedResultTypes.delete(method); + }; + } + + private _isAcceptedResultType(method: string, raw: unknown): boolean { + const accepted = this._acceptedResultTypes.get(method); + if (accepted === undefined || !isPlainObject(raw)) return false; + const resultType = raw['resultType']; + return typeof resultType === 'string' && accepted.has(resultType); + } + + /** + * The handler `_onrequest` dispatches to for `method`: the registered + * handler (or the fallback), with any middleware for the method composed + * around it in registration order (first registered outermost). + * `undefined` when nothing is registered and no middleware applies. + */ + private _resolveRequestHandler(method: string): ((request: JSONRPCRequest, ctx: ContextT) => Promise) | undefined { + const base = this._requestHandlers.get(method) ?? this.fallbackRequestHandler; + const stack = this._requestMiddleware.filter(entry => entry.method === method); + if (stack.length === 0) return base; + let composed: (request: JSONRPCRequest, ctx: ContextT) => Promise = + base ?? + (async () => { + throw new ProtocolError(ProtocolErrorCode.MethodNotFound, 'Method not found'); + }); + // Wrap from the innermost (last registered) outwards so the first + // registered middleware ends up outermost. + for (const { middleware } of stack.toReversed()) { + const next = composed; + composed = async (request, ctx) => middleware(request, ctx, next); + } + return composed; + } + /** * Hook for subclasses to wrap a registered request handler with role-specific * validation or behavior (e.g. `Server` validates `tools/call` results, `Client` diff --git a/packages/core-internal/test/shared/requestMiddleware.test.ts b/packages/core-internal/test/shared/requestMiddleware.test.ts new file mode 100644 index 0000000000..f3d803fc35 --- /dev/null +++ b/packages/core-internal/test/shared/requestMiddleware.test.ts @@ -0,0 +1,153 @@ +/** + * `Protocol.use`: middleware composes around the registered handler at + * dispatch time, may answer or throw itself, applies to handlers registered + * later, runs in registration order, and can be removed. + */ +import { describe, expect, it } from 'vitest'; +import * as z from 'zod/v4'; + +import type { BaseContext } from '../../src/shared/protocol'; +import { Protocol } from '../../src/shared/protocol'; +import type { JSONRPCErrorResponse, JSONRPCMessage, JSONRPCResultResponse, Result } from '../../src/types/index'; +import { ProtocolError, ProtocolErrorCode } from '../../src/types/index'; +import { InMemoryTransport } from '../../src/util/inMemory'; + +class TestProtocol extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + protected buildContext(ctx: BaseContext): BaseContext { + return ctx; + } +} + +const flush = () => new Promise(resolve => setTimeout(resolve, 5)); + +async function harness() { + const [peerTx, protocolTx] = InMemoryTransport.createLinkedPair(); + const sent: JSONRPCMessage[] = []; + peerTx.onmessage = message => void sent.push(message); + await peerTx.start(); + const protocol = new TestProtocol(); + await protocol.connect(protocolTx); + let nextId = 0; + const call = async (method: string, params: Record = {}) => { + const id = ++nextId; + await peerTx.send({ jsonrpc: '2.0', id, method, params }); + await flush(); + const response = sent.find(message => 'id' in message && message.id === id); + return response as JSONRPCResultResponse | JSONRPCErrorResponse; + }; + return { protocol, call }; +} + +describe('Protocol.use (request middleware)', () => { + it('wraps the registered handler; next reaches it and the middleware may transform the result', async () => { + const { protocol, call } = await harness(); + protocol.setRequestHandler('acme/op', { params: z.looseObject({}) }, () => ({ value: 1 }) as Result); + protocol.use('acme/op', async (request, ctx, next) => { + const result = (await next(request, ctx)) as { value: number }; + return { value: result.value + 1, wrapped: true } as Result; + }); + const response = await call('acme/op'); + expect((response as JSONRPCResultResponse).result).toEqual({ value: 2, wrapped: true }); + }); + + it('may answer without calling next, and a thrown ProtocolError becomes the JSON-RPC error', async () => { + const { protocol, call } = await harness(); + let handlerRan = false; + protocol.setRequestHandler('acme/op', { params: z.looseObject({ deny: z.boolean().optional() }) }, () => { + handlerRan = true; + return {} as Result; + }); + protocol.use('acme/op', (request, ctx, next) => { + if ((request.params as { deny?: boolean }).deny) { + throw new ProtocolError(ProtocolErrorCode.InvalidParams, 'denied by middleware'); + } + return next(request, ctx); + }); + const denied = await call('acme/op', { deny: true }); + expect((denied as JSONRPCErrorResponse).error).toMatchObject({ + code: ProtocolErrorCode.InvalidParams, + message: 'denied by middleware' + }); + expect(handlerRan).toBe(false); + await call('acme/op'); + expect(handlerRan).toBe(true); + }); + + it('applies to a handler registered after the middleware was installed', async () => { + const { protocol, call } = await harness(); + const order: string[] = []; + protocol.use('acme/late', async (request, ctx, next) => { + order.push('middleware'); + return next(request, ctx); + }); + protocol.setRequestHandler('acme/late', { params: z.looseObject({}) }, () => { + order.push('handler'); + return {} as Result; + }); + await call('acme/late'); + expect(order).toEqual(['middleware', 'handler']); + }); + + it('keeps one registration order across methods and stamps _meta through next', async () => { + const { protocol, call } = await harness(); + const order: string[] = []; + for (const method of ['acme/a', 'acme/b']) { + protocol.setRequestHandler(method, { params: z.looseObject({}) }, () => { + order.push(`handler:${method}`); + return { method } as Result; + }); + } + protocol.use('acme/a', async (request, ctx, next) => { + order.push('a-1'); + const result = await next(request, ctx); + return { ...result, _meta: { 'acme/stamped': true } } as Result; + }); + protocol.use('acme/b', async (request, ctx, next) => { + order.push('b-1'); + return next(request, ctx); + }); + protocol.use('acme/a', async (request, ctx, next) => { + order.push('a-2'); + return next(request, ctx); + }); + const a = await call('acme/a'); + expect((a as JSONRPCResultResponse).result).toEqual({ method: 'acme/a', _meta: { 'acme/stamped': true } }); + expect(order).toEqual(['a-1', 'a-2', 'handler:acme/a']); + order.length = 0; + await call('acme/b'); + expect(order).toEqual(['b-1', 'handler:acme/b']); + }); + + it('next throws MethodNotFound when nothing underlies the middleware', async () => { + const { protocol, call } = await harness(); + protocol.use('acme/missing', (request, ctx, next) => next(request, ctx)); + const response = await call('acme/missing'); + expect((response as JSONRPCErrorResponse).error).toMatchObject({ code: ProtocolErrorCode.MethodNotFound }); + }); + + it('runs in registration order (first registered outermost), and removal restores the handler', async () => { + const { protocol, call } = await harness(); + const order: string[] = []; + protocol.setRequestHandler('acme/op', { params: z.looseObject({}) }, () => { + order.push('handler'); + return {} as Result; + }); + const removeFirst = protocol.use('acme/op', async (request, ctx, next) => { + order.push('first'); + return next(request, ctx); + }); + protocol.use('acme/op', async (request, ctx, next) => { + order.push('second'); + return next(request, ctx); + }); + await call('acme/op'); + expect(order).toEqual(['first', 'second', 'handler']); + order.length = 0; + removeFirst(); + await call('acme/op'); + expect(order).toEqual(['second', 'handler']); + }); +}); diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts index 9b3196a80b..eca691891d 100644 --- a/packages/server/src/index.ts +++ b/packages/server/src/index.ts @@ -55,6 +55,7 @@ export type { PerRequestHTTPServerTransportOptions, PerRequestMessageExtra, PerR export { PerRequestHTTPServerTransport } from './server/perRequestTransport'; // Opt-in HMAC sealing for the multi-round-trip requestState (SEP-2322): the // convenience codec consumers drop into ServerOptions.requestState.verify. +export type { ServerExtension } from './server/extension'; export type { RequestStateCodec, RequestStateCodecOptions } from './server/requestStateCodec'; export { createRequestStateCodec } from './server/requestStateCodec'; export type { ServerOptions } from './server/server'; diff --git a/packages/server/src/server/extension.ts b/packages/server/src/server/extension.ts new file mode 100644 index 0000000000..c190bbb4ad --- /dev/null +++ b/packages/server/src/server/extension.ts @@ -0,0 +1,28 @@ +import type { Server } from './server'; + +/** + * A server extension: a unit of protocol behaviour outside the core + * specification (an MCP extension such as `io.modelcontextprotocol/tasks`, + * or a vendor feature) that installs itself onto a {@linkcode Server}. + * + * Pass extensions at construction — `new McpServer(info, { extensions: [ext] })` + * or `new Server(info, { extensions: [ext] })`. The server advertises each + * extension under `capabilities.extensions[id]` and then calls `install`, + * which is where the extension registers its custom methods + * (`server.setRequestHandler(method, { params, result }, handler)`), + * installs middleware on spec methods it needs to intercept + * (`server.use('tools/call', …)`), and adds notification + * handlers. The SDK provides the hooks; what an extension does behind them + * — how it stores state, where its work runs — is the extension's own. + */ +export interface ServerExtension { + /** + * The extension identifier, prefix-qualified (`io.modelcontextprotocol/tasks`, + * `com.example/feature-flags`). Advertised as the key under + * `capabilities.extensions`. Settings for that key, if any, are the + * extension's to register in `install` via `registerCapabilities`. + */ + readonly id: string; + /** Installs the extension's handlers and middleware onto the server. Called once, at construction. */ + install(server: Server): void; +} diff --git a/packages/server/src/server/mcp.ts b/packages/server/src/server/mcp.ts index 70a5539bfb..113f17f30a 100644 --- a/packages/server/src/server/mcp.ts +++ b/packages/server/src/server/mcp.ts @@ -280,8 +280,12 @@ export class McpServer { // wrapped) so the listing and the call cannot diverge. return this.server.projectCallToolResult(result, tool.outputSchemaJson); } catch (error) { - if (error instanceof ProtocolError && error.code === ProtocolErrorCode.UrlElicitationRequired) { - throw error; // Return the error to the caller without wrapping in CallToolResult + if ( + error instanceof ProtocolError && + (error.code === ProtocolErrorCode.UrlElicitationRequired || + error.code === ProtocolErrorCode.MissingRequiredClientCapability) + ) { + throw error; // Protocol-level: return the error to the caller without wrapping in CallToolResult } return this.createToolError(error instanceof Error ? error.message : String(error)); } diff --git a/packages/server/src/server/server.ts b/packages/server/src/server/server.ts index 5de0d8919c..49d5dbfc1a 100644 --- a/packages/server/src/server/server.ts +++ b/packages/server/src/server/server.ts @@ -72,7 +72,16 @@ import { coerceEmbeddedInputRequest, LegacyInputRequiredShim, resolveLegacyShimO */ const INPUT_REQUIRED_CAPABLE_METHODS: ReadonlySet = new Set(['tools/call', 'prompts/get', 'resources/read']); +import type { ServerExtension } from './extension'; + export type ServerOptions = ProtocolOptions & { + /** + * Extensions to install at construction. Each is advertised under + * `capabilities.extensions[extension.id]` and then installed, in order, + * after the built-in handlers exist — see {@linkcode ServerExtension}. + */ + extensions?: ServerExtension[]; + /** * Capabilities to advertise as being supported by this server. * @@ -349,6 +358,11 @@ export class Server extends Protocol { if (this._capabilities.logging) { this._registerLoggingHandler(); } + + for (const extension of options?.extensions ?? []) { + this.registerCapabilities({ extensions: { [extension.id]: {} } }); + extension.install(this); + } } /** diff --git a/packages/server/test/server/extensions.test.ts b/packages/server/test/server/extensions.test.ts new file mode 100644 index 0000000000..9a388d6e73 --- /dev/null +++ b/packages/server/test/server/extensions.test.ts @@ -0,0 +1,118 @@ +/** + * `ServerOptions.extensions`: an extension is advertised under + * `capabilities.extensions[id]` and installed at construction, where it can + * register custom methods and install middleware on spec methods such as `tools/call`. + */ +import type { JSONRPCRequest, MessageClassification } from '@modelcontextprotocol/core-internal'; +import { + CLIENT_CAPABILITIES_META_KEY, + CLIENT_INFO_META_KEY, + MissingRequiredClientCapabilityError, + PROTOCOL_VERSION_META_KEY, + setNegotiatedProtocolVersion +} from '@modelcontextprotocol/core-internal'; +import { describe, expect, it } from 'vitest'; +import * as z from 'zod/v4'; + +import type { ServerExtension } from '../../src/server/extension'; +import { invoke } from '../../src/server/invoke'; +import { McpServer } from '../../src/server/mcp'; +import { Server } from '../../src/server/server'; + +const MODERN_REVISION = '2026-07-28'; +const MODERN: MessageClassification = { era: 'modern', revision: MODERN_REVISION }; +const EXT_ID = 'com.example/gate'; + +const modernRequest = ( + method: string, + params: Record = {}, + clientCapabilities: Record = {} +): JSONRPCRequest => + ({ + jsonrpc: '2.0', + id: 1, + method, + params: { + ...params, + _meta: { + [PROTOCOL_VERSION_META_KEY]: MODERN_REVISION, + [CLIENT_INFO_META_KEY]: { name: 'ext-client', version: '1.0.0' }, + [CLIENT_CAPABILITIES_META_KEY]: clientCapabilities + } + } + }) as JSONRPCRequest; + +async function exchange(server: Server, request: JSONRPCRequest): Promise> { + setNegotiatedProtocolVersion(server, MODERN_REVISION); + const response = await invoke(server, request, { classification: MODERN }); + return (await response.json()) as Record; +} + +/** An extension that gates `tools/call` on a client capability and adds one custom method. */ +function gateExtension(log: string[]): ServerExtension { + return { + id: EXT_ID, + install(server) { + log.push('installed'); + server.registerCapabilities({ extensions: { [EXT_ID]: { exampleData: true } } }); + server.setRequestHandler('gate/status', { params: z.looseObject({}) }, () => ({ armed: true })); + server.use('tools/call', (request, ctx, next) => { + const envelope = ctx.mcpReq.envelope as Record> | undefined; + const extensions = envelope?.[CLIENT_CAPABILITIES_META_KEY]?.['extensions'] as Record | undefined; + if (extensions === undefined || !(EXT_ID in extensions)) { + throw new MissingRequiredClientCapabilityError( + { requiredCapabilities: { extensions: { [EXT_ID]: {} } } }, + 'declare the gate' + ); + } + return next(request, ctx); + }); + } + }; +} + +describe('ServerOptions.extensions', () => { + it('advertises the extension and lets install set its settings', () => { + const log: string[] = []; + const server = new Server({ name: 's', version: '1' }, { extensions: [gateExtension(log)] }); + expect(log).toEqual(['installed']); + expect(server.getCapabilities().extensions).toEqual({ [EXT_ID]: { exampleData: true } }); + }); + + it('advertises {} when install sets no settings, and passes through McpServer', () => { + const ext: ServerExtension = { id: 'com.example/plain', install: () => {} }; + const mcp = new McpServer({ name: 's', version: '1' }, { extensions: [ext] }); + expect(mcp.server.getCapabilities().extensions).toEqual({ 'com.example/plain': {} }); + }); + + it('serves the extension custom method', async () => { + const server = new Server({ name: 's', version: '1' }, { extensions: [gateExtension([])] }); + const body = await exchange(server, modernRequest('gate/status')); + expect(body['result']).toMatchObject({ armed: true }); + }); + + it('middleware on tools/call registered later by McpServer: refuses without the capability, passes through with it', async () => { + const mcp = new McpServer({ name: 's', version: '1' }, { extensions: [gateExtension([])] }); + mcp.registerTool('echo', { inputSchema: z.object({ text: z.string() }) }, async ({ text }) => ({ + content: [{ type: 'text', text }] + })); + + const refused = await exchange(mcp.server, modernRequest('tools/call', { name: 'echo', arguments: { text: 'hi' } })); + expect(refused['error']).toMatchObject({ code: -32_021, message: 'declare the gate' }); + + const allowed = await exchange( + mcp.server, + modernRequest('tools/call', { name: 'echo', arguments: { text: 'hi' } }, { extensions: { [EXT_ID]: {} } }) + ); + expect(allowed['result']).toMatchObject({ content: [{ type: 'text', text: 'hi' }] }); + }); + + it('a MissingRequiredClientCapabilityError thrown by a tool handler is a JSON-RPC error, not an isError result', async () => { + const mcp = new McpServer({ name: 's', version: '1' }); + mcp.registerTool('needs-ext', {}, async () => { + throw new MissingRequiredClientCapabilityError({ requiredCapabilities: { extensions: { [EXT_ID]: {} } } }, 'declare the gate'); + }); + const body = await exchange(mcp.server, modernRequest('tools/call', { name: 'needs-ext', arguments: {} })); + expect(body['error']).toMatchObject({ code: -32_021, data: { requiredCapabilities: { extensions: { [EXT_ID]: {} } } } }); + }); +});