diff --git a/.changeset/extension-registrar.md b/.changeset/extension-registrar.md new file mode 100644 index 000000000..48379ccc2 --- /dev/null +++ b/.changeset/extension-registrar.md @@ -0,0 +1,6 @@ +--- +'@modelcontextprotocol/client': minor +'@modelcontextprotocol/server': minor +--- + +Add `Client.extension()` / `Server.extension()` registrar for SEP-2133 capability-aware custom methods. Declares an extension in `capabilities.extensions[id]` and returns an `ExtensionHandle` whose `setRequestHandler`/`sendRequest`/`setNotificationHandler`/`sendNotification` calls are tied to that declared capability. `getPeerSettings()` returns the peer's extension settings, optionally validated against a `peerSchema`. diff --git a/docs/migration.md b/docs/migration.md index 8a63a1162..6a56cdabe 100644 --- a/docs/migration.md +++ b/docs/migration.md @@ -434,6 +434,36 @@ before sending and gives typed `params`; passing a bare result schema sends para For larger sub-protocols where neither side is semantically an MCP client or server, prefer composition: hold a `Client` (or `Server`) instance, register custom handlers on it, and expose typed facade methods. See `examples/server/src/customMethodExample.ts` and `examples/client/src/customMethodExample.ts` for runnable examples. +#### Declaring extension capabilities (SEP-2133) + +When your custom methods constitute a formal extension with an SEP-2133 identifier (e.g. +`io.modelcontextprotocol/ui`), use `Client.extension()` / `Server.extension()` instead of the flat +`*Custom*` methods. This declares the extension in `capabilities.extensions[id]` so it is +negotiated during `initialize`, and returns a scoped `ExtensionHandle` whose `setRequestHandler` / +`sendRequest` calls are tied to that declared capability: + +```typescript +import { Client } from '@modelcontextprotocol/client'; + +const client = new Client({ name: 'app', version: '1.0.0' }); +const ui = client.extension( + 'io.modelcontextprotocol/ui', + { availableDisplayModes: ['inline'] }, + { peerSchema: HostCapabilitiesSchema } +); + +ui.setRequestHandler('ui/resource-teardown', TeardownParams, p => onTeardown(p)); + +await client.connect(transport); +ui.getPeerSettings(); // server's capabilities.extensions['io.modelcontextprotocol/ui'], typed via peerSchema +await ui.sendRequest('ui/open-link', { url }, OpenLinkResult); +``` + +`handle.sendRequest`/`sendNotification` respect `enforceStrictCapabilities`: when strict, sending +throws if the peer did not advertise the same extension ID. The flat `setCustomRequestHandler` / +`sendCustomRequest` methods remain available as the ungated escape hatch for one-off vendor +methods that do not warrant a SEP-2133 entry. + ### `Protocol.request()`, `ctx.mcpReq.send()`, and `Client.callTool()` no longer take a schema parameter The public `Protocol.request()`, `BaseContext.mcpReq.send()`, and `Client.callTool()` methods no longer accept a Zod result schema argument. The SDK now resolves the correct result schema internally based on the method name. This means you no longer need to import result schemas diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts index 21a43bd15..2f52aebb2 100644 --- a/packages/client/src/client/client.ts +++ b/packages/client/src/client/client.ts @@ -1,5 +1,6 @@ import { DefaultJsonSchemaValidator } from '@modelcontextprotocol/client/_shims'; import type { + AnySchema, BaseContext, CallToolRequest, ClientCapabilities, @@ -8,8 +9,10 @@ import type { ClientRequest, ClientResult, CompleteRequest, + ExtensionOptions, GetPromptRequest, Implementation, + JSONObject, JsonSchemaType, JsonSchemaValidator, jsonSchemaValidator, @@ -28,6 +31,7 @@ import type { RequestOptions, RequestTypeMap, ResultTypeMap, + SchemaOutput, ServerCapabilities, SubscribeRequest, TaskManagerOptions, @@ -47,6 +51,7 @@ import { ElicitRequestSchema, ElicitResultSchema, EmptyResultSchema, + ExtensionHandle, extractTaskManagerOptions, GetPromptResultSchema, InitializeResultSchema, @@ -307,6 +312,43 @@ export class Client extends Protocol { this._capabilities = mergeCapabilities(this._capabilities, capabilities); } + /** + * Declares an SEP-2133 extension and returns a scoped {@linkcode ExtensionHandle} for + * registering and sending its custom JSON-RPC methods. + * + * Merges `settings` into `capabilities.extensions[id]`, which is advertised to the server + * during `initialize`. Must be called before {@linkcode connect}. After connecting, + * {@linkcode ExtensionHandle.getPeerSettings | handle.getPeerSettings()} returns the server's + * `capabilities.extensions[id]` blob (validated against `peerSchema` if provided). + */ + public extension(id: string, settings: L): ExtensionHandle; + public extension( + id: string, + settings: L, + opts: ExtensionOptions

+ ): ExtensionHandle, ClientContext>; + public extension( + id: string, + settings: L, + opts?: ExtensionOptions

+ ): ExtensionHandle | JSONObject, ClientContext> { + if (this.transport) { + throw new SdkError(SdkErrorCode.AlreadyConnected, 'Cannot register extension after connecting to transport'); + } + if (this._capabilities.extensions && Object.hasOwn(this._capabilities.extensions, id)) { + throw new Error(`Extension "${id}" is already registered`); + } + this._capabilities.extensions = { ...this._capabilities.extensions, [id]: settings }; + return new ExtensionHandle( + this, + id, + settings, + () => this._serverCapabilities?.extensions?.[id], + this._enforceStrictCapabilities, + opts?.peerSchema + ); + } + /** * Registers a handler for server-initiated requests (sampling, elicitation, roots). * The client must declare the corresponding capability for the handler to be accepted. diff --git a/packages/client/test/client/extension.test.ts b/packages/client/test/client/extension.test.ts new file mode 100644 index 000000000..8cbe05438 --- /dev/null +++ b/packages/client/test/client/extension.test.ts @@ -0,0 +1,111 @@ +import { InMemoryTransport, type JSONRPCMessage, SdkError, SdkErrorCode } from '@modelcontextprotocol/core'; +import { describe, expect, test } from 'vitest'; +import * as z from 'zod/v4'; + +import { Client } from '../../src/client/client.js'; + +/** + * These tests exercise the `Client.extension()` factory and the client side of the + * `capabilities.extensions` round-trip via `initialize`. The `ExtensionHandle` class itself is + * unit-tested in `@modelcontextprotocol/core/test/shared/extensionHandle.test.ts`. + */ + +interface RawServerHarness { + serverSide: InMemoryTransport; + capturedInitParams: Promise>; +} + +function rawServer(serverCapabilities: Record = {}): RawServerHarness { + const [clientSide, serverSide] = InMemoryTransport.createLinkedPair(); + let resolveInit: (p: Record) => void; + const capturedInitParams = new Promise>(r => { + resolveInit = r; + }); + serverSide.onmessage = (msg: JSONRPCMessage) => { + if ('method' in msg && msg.method === 'initialize' && 'id' in msg) { + resolveInit((msg.params ?? {}) as Record); + void serverSide.send({ + jsonrpc: '2.0', + id: msg.id, + result: { + protocolVersion: '2025-11-25', + capabilities: serverCapabilities, + serverInfo: { name: 'raw-server', version: '0.0.0' } + } + }); + } + }; + void serverSide.start(); + // Expose clientSide via the harness's serverSide.peer for the test to connect to. + return { serverSide: clientSide, capturedInitParams }; +} + +describe('Client.extension()', () => { + test('merges settings into capabilities.extensions and advertises them in initialize request', async () => { + const client = new Client({ name: 'c', version: '1.0.0' }, { capabilities: {} }); + client.extension('io.example/ui', { contentTypes: ['text/html'] }); + client.extension('com.acme/widgets', { v: 2 }); + + const harness = rawServer(); + await client.connect(harness.serverSide); + const initParams = await harness.capturedInitParams; + + const caps = initParams.capabilities as Record; + expect(caps.extensions).toEqual({ + 'io.example/ui': { contentTypes: ['text/html'] }, + 'com.acme/widgets': { v: 2 } + }); + }); + + test('throws AlreadyConnected after connect()', async () => { + const client = new Client({ name: 'c', version: '1.0.0' }); + const harness = rawServer(); + await client.connect(harness.serverSide); + + expect(() => client.extension('io.example/ui', {})).toThrow(SdkError); + try { + client.extension('io.example/ui', {}); + expect.fail('should have thrown'); + } catch (e) { + expect(e).toBeInstanceOf(SdkError); + expect((e as SdkError).code).toBe(SdkErrorCode.AlreadyConnected); + } + }); + + test('throws on duplicate extension id', () => { + const client = new Client({ name: 'c', version: '1.0.0' }); + client.extension('io.example/ui', { v: 1 }); + expect(() => client.extension('io.example/ui', { v: 2 })).toThrow(/already registered/); + expect(() => client.extension('com.other/thing', {})).not.toThrow(); + }); + + test("getPeerSettings() reads the server's capabilities.extensions[id] from initialize result", async () => { + const PeerSchema = z.object({ availableDisplayModes: z.array(z.string()) }); + const client = new Client({ name: 'c', version: '1.0.0' }); + const handle = client.extension('io.example/ui', { clientSide: true }, { peerSchema: PeerSchema }); + + expect(handle.getPeerSettings()).toBeUndefined(); + + const harness = rawServer({ + extensions: { 'io.example/ui': { availableDisplayModes: ['inline', 'fullscreen'] } } + }); + await client.connect(harness.serverSide); + + expect(handle.getPeerSettings()).toEqual({ availableDisplayModes: ['inline', 'fullscreen'] }); + }); + + test('getPeerSettings() reflects reconnect to a different server', async () => { + const client = new Client({ name: 'c', version: '1.0.0' }); + const handle = client.extension('io.example/ui', {}); + + const harnessA = rawServer({ extensions: { 'io.example/ui': { v: 1 } } }); + await client.connect(harnessA.serverSide); + expect(handle.getPeerSettings()).toEqual({ v: 1 }); + + await client.close(); + + const harnessB = rawServer({ extensions: { 'io.example/ui': { v: 2 } } }); + await client.connect(harnessB.serverSide); + expect(handle.getPeerSettings()).toEqual({ v: 2 }); + }); +}); diff --git a/packages/core/src/exports/public/index.ts b/packages/core/src/exports/public/index.ts index 353567ba5..05a409b49 100644 --- a/packages/core/src/exports/public/index.ts +++ b/packages/core/src/exports/public/index.ts @@ -35,6 +35,10 @@ export type { // Auth utilities export { checkResourceAllowed, resourceUrlFromServerUrl } from '../../shared/authUtils.js'; +// Extension registrar (SEP-2133 capability-aware custom methods) +export type { ExtensionOptions } from '../../shared/extensionHandle.js'; +export { ExtensionHandle } from '../../shared/extensionHandle.js'; + // Metadata utilities export { getDisplayName } from '../../shared/metadataUtils.js'; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index e707d9939..277d28f4b 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -2,6 +2,7 @@ export * from './auth/errors.js'; export * from './errors/sdkErrors.js'; export * from './shared/auth.js'; export * from './shared/authUtils.js'; +export * from './shared/extensionHandle.js'; export * from './shared/metadataUtils.js'; export * from './shared/protocol.js'; export * from './shared/responseMessage.js'; diff --git a/packages/core/src/shared/extensionHandle.ts b/packages/core/src/shared/extensionHandle.ts new file mode 100644 index 000000000..8324778f4 --- /dev/null +++ b/packages/core/src/shared/extensionHandle.ts @@ -0,0 +1,160 @@ +import { SdkError, SdkErrorCode } from '../errors/sdkErrors.js'; +import type { JSONObject, Result } from '../types/types.js'; +import type { AnySchema, SchemaOutput } from '../util/schema.js'; +import { parseSchema } from '../util/schema.js'; +import type { BaseContext, NotificationOptions, RequestOptions } from './protocol.js'; + +/** + * The subset of `Client`/`Server` that {@linkcode ExtensionHandle} delegates to. + * + * @internal + */ +export interface ExtensionHost { + setCustomRequestHandler

( + method: string, + paramsSchema: P, + handler: (params: SchemaOutput

, ctx: ContextT) => Result | Promise + ): void; + setCustomNotificationHandler

( + method: string, + paramsSchema: P, + handler: (params: SchemaOutput

) => void | Promise + ): void; + sendCustomRequest( + method: string, + params: Record | undefined, + resultSchema: R, + options?: RequestOptions + ): Promise>; + sendCustomNotification(method: string, params?: Record, options?: NotificationOptions): Promise; +} + +/** + * Options for {@linkcode @modelcontextprotocol/client!client/client.Client#extension | Client.extension} / + * {@linkcode @modelcontextprotocol/server!server/server.Server#extension | Server.extension}. + */ +export interface ExtensionOptions

{ + /** + * Schema to validate the peer's `capabilities.extensions[id]` blob against. When provided, + * {@linkcode ExtensionHandle.getPeerSettings | getPeerSettings()} returns the parsed value + * (typed as `SchemaOutput

`) or `undefined` if the peer's blob does not match. + */ + peerSchema: P; +} + +/** + * A scoped handle for registering and sending custom JSON-RPC methods belonging to a single + * SEP-2133 extension. + * + * Obtained via {@linkcode @modelcontextprotocol/client!client/client.Client#extension | Client.extension} or + * {@linkcode @modelcontextprotocol/server!server/server.Server#extension | Server.extension}. Creating a handle + * declares the extension in `capabilities.extensions[id]` so it is advertised during `initialize`. + * Handlers registered through the handle are thus structurally guaranteed to belong to a declared + * extension. + * + * Send-side methods respect `enforceStrictCapabilities`: when strict, sending throws if the peer + * did not advertise the same extension ID; when lax (the default), sends proceed regardless and + * {@linkcode getPeerSettings} returns `undefined`. + */ +export class ExtensionHandle { + /** + * @internal Use `Client.extension()` or `Server.extension()` to construct. + */ + constructor( + private readonly _host: ExtensionHost, + /** The SEP-2133 extension identifier (e.g. `io.modelcontextprotocol/ui`). */ + public readonly id: string, + /** The local settings object advertised in `capabilities.extensions[id]`. */ + public readonly settings: Local, + private readonly _getPeerExtensionSettings: () => JSONObject | undefined, + private readonly _enforceStrictCapabilities: boolean, + private readonly _peerSchema?: AnySchema + ) {} + + /** + * Returns the peer's `capabilities.extensions[id]` settings, or `undefined` if the peer did not + * advertise this extension or (when `peerSchema` was provided) if the peer's blob fails + * validation. Reads the current peer capabilities on each call (no caching), so it reflects + * reconnects. + */ + getPeerSettings(): Peer | undefined { + const raw = this._getPeerExtensionSettings(); + if (raw === undefined) { + return undefined; + } + if (this._peerSchema === undefined) { + return raw as Peer; + } + const parsed = parseSchema(this._peerSchema, raw); + if (!parsed.success) { + console.warn( + `[ExtensionHandle] Peer's capabilities.extensions["${this.id}"] failed schema validation: ${parsed.error.message}` + ); + return undefined; + } + return parsed.data as Peer; + } + + /** + * Registers a request handler for a custom method belonging to this extension. Delegates to + * the underlying `setCustomRequestHandler`; the collision guard + * against standard MCP methods applies. + */ + setRequestHandler

( + method: string, + paramsSchema: P, + handler: (params: SchemaOutput

, ctx: ContextT) => Result | Promise + ): void { + this._host.setCustomRequestHandler(method, paramsSchema, handler); + } + + /** + * Registers a notification handler for a custom method belonging to this extension. Delegates + * to the underlying `setCustomNotificationHandler`. + */ + setNotificationHandler

( + method: string, + paramsSchema: P, + handler: (params: SchemaOutput

) => void | Promise + ): void { + this._host.setCustomNotificationHandler(method, paramsSchema, handler); + } + + /** + * Sends a custom request belonging to this extension and waits for a response. + * + * When `enforceStrictCapabilities` is enabled and the peer did not advertise + * `capabilities.extensions[id]`, throws {@linkcode SdkError} with + * {@linkcode SdkErrorCode.CapabilityNotSupported}. + */ + async sendRequest( + method: string, + params: Record | undefined, + resultSchema: R, + options?: RequestOptions + ): Promise> { + this._assertPeerCapability(method); + return this._host.sendCustomRequest(method, params, resultSchema, options); + } + + /** + * Sends a custom notification belonging to this extension. + * + * When `enforceStrictCapabilities` is enabled and the peer did not advertise + * `capabilities.extensions[id]`, throws {@linkcode SdkError} with + * {@linkcode SdkErrorCode.CapabilityNotSupported}. + */ + async sendNotification(method: string, params?: Record, options?: NotificationOptions): Promise { + this._assertPeerCapability(method); + return this._host.sendCustomNotification(method, params, options); + } + + private _assertPeerCapability(method: string): void { + if (this._enforceStrictCapabilities && this._getPeerExtensionSettings() === undefined) { + throw new SdkError( + SdkErrorCode.CapabilityNotSupported, + `Peer does not support extension "${this.id}" (required for ${method})` + ); + } + } +} diff --git a/packages/core/test/shared/extensionHandle.test.ts b/packages/core/test/shared/extensionHandle.test.ts new file mode 100644 index 000000000..990463f80 --- /dev/null +++ b/packages/core/test/shared/extensionHandle.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, test, vi } from 'vitest'; +import * as z from 'zod/v4'; + +import { SdkError, SdkErrorCode } from '../../src/errors/sdkErrors.js'; +import type { ExtensionHost } from '../../src/shared/extensionHandle.js'; +import { ExtensionHandle } from '../../src/shared/extensionHandle.js'; +import type { BaseContext } from '../../src/shared/protocol.js'; +import type { JSONObject } from '../../src/types/types.js'; + +type MockHost = { + setCustomRequestHandler: ReturnType; + setCustomNotificationHandler: ReturnType; + sendCustomRequest: ReturnType; + sendCustomNotification: ReturnType; +}; + +function makeMockHost(): MockHost { + return { + setCustomRequestHandler: vi.fn(), + setCustomNotificationHandler: vi.fn(), + sendCustomRequest: vi.fn().mockResolvedValue({ ok: true }), + sendCustomNotification: vi.fn().mockResolvedValue(undefined) + }; +} + +function makeHandle(opts: { peer?: JSONObject | undefined; strict?: boolean; peerSchema?: z.core.$ZodType }): { + host: MockHost; + handle: ExtensionHandle; +} { + const host = makeMockHost(); + const handle = new ExtensionHandle( + host as unknown as ExtensionHost, + 'io.example/ui', + { local: true }, + () => opts.peer, + opts.strict ?? false, + opts.peerSchema + ); + return { host, handle }; +} + +describe('ExtensionHandle.getPeerSettings', () => { + test('returns raw blob when no peerSchema given', () => { + const { handle } = makeHandle({ peer: { feature: 'x' } }); + expect(handle.getPeerSettings()).toEqual({ feature: 'x' }); + }); + + test('returns undefined when peer did not advertise', () => { + const { handle } = makeHandle({ peer: undefined }); + expect(handle.getPeerSettings()).toBeUndefined(); + }); + + test('parses and returns typed value when peerSchema matches', () => { + const PeerSchema = z.object({ openLinks: z.boolean(), maxSize: z.number() }); + const { handle } = makeHandle({ peer: { openLinks: true, maxSize: 5 }, peerSchema: PeerSchema }); + expect(handle.getPeerSettings()).toEqual({ openLinks: true, maxSize: 5 }); + }); + + test('returns undefined and warns when peerSchema does not match', () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const PeerSchema = z.object({ openLinks: z.boolean() }); + const { handle } = makeHandle({ peer: { openLinks: 'yes' }, peerSchema: PeerSchema }); + expect(handle.getPeerSettings()).toBeUndefined(); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0]?.[0]).toMatch(/io\.example\/ui.*failed schema validation/); + warn.mockRestore(); + }); + + test('reflects current peer settings on each call (no caching across reconnects)', () => { + let peer: JSONObject | undefined; + const getter = vi.fn(() => peer); + const host = makeMockHost() as unknown as ExtensionHost; + const handle = new ExtensionHandle(host, 'io.example/ui', {}, getter, false); + + expect(handle.getPeerSettings()).toBeUndefined(); + peer = { v: 1 }; + expect(handle.getPeerSettings()).toEqual({ v: 1 }); + peer = { v: 2 }; + expect(handle.getPeerSettings()).toEqual({ v: 2 }); + peer = undefined; + expect(handle.getPeerSettings()).toBeUndefined(); + expect(getter).toHaveBeenCalledTimes(4); + }); +}); + +describe('ExtensionHandle.setRequestHandler / setNotificationHandler', () => { + test('delegates to host setCustom* (anytime)', () => { + const { host, handle } = makeHandle({ peer: undefined }); + const params = z.object({ q: z.string() }); + const reqHandler = vi.fn(); + const notifHandler = vi.fn(); + + handle.setRequestHandler('ui/search', params, reqHandler); + expect(host.setCustomRequestHandler).toHaveBeenCalledWith('ui/search', params, reqHandler); + + handle.setNotificationHandler('ui/ping', params, notifHandler); + expect(host.setCustomNotificationHandler).toHaveBeenCalledWith('ui/ping', params, notifHandler); + }); +}); + +describe('ExtensionHandle.sendRequest / sendNotification — peer gating', () => { + const Result = z.object({ ok: z.boolean() }); + + test('lax mode (default): sends even when peer did not advertise', async () => { + const { host, handle } = makeHandle({ peer: undefined, strict: false }); + await handle.sendRequest('ui/do', { x: 1 }, Result); + expect(host.sendCustomRequest).toHaveBeenCalledWith('ui/do', { x: 1 }, Result, undefined); + await handle.sendNotification('ui/ping', {}); + expect(host.sendCustomNotification).toHaveBeenCalledWith('ui/ping', {}, undefined); + }); + + test('strict mode: rejects with CapabilityNotSupported when peer did not advertise', async () => { + const { host, handle } = makeHandle({ peer: undefined, strict: true }); + await expect(handle.sendRequest('ui/do', {}, Result)).rejects.toSatisfy( + (e: unknown) => + e instanceof SdkError && e.code === SdkErrorCode.CapabilityNotSupported && /io\.example\/ui.*ui\/do/.test(e.message) + ); + expect(host.sendCustomRequest).not.toHaveBeenCalled(); + await expect(handle.sendNotification('ui/ping')).rejects.toSatisfy( + (e: unknown) => e instanceof SdkError && e.code === SdkErrorCode.CapabilityNotSupported + ); + expect(host.sendCustomNotification).not.toHaveBeenCalled(); + }); + + test('strict mode: sends when peer did advertise', async () => { + const { host, handle } = makeHandle({ peer: { ok: true }, strict: true }); + await handle.sendRequest('ui/do', {}, Result); + expect(host.sendCustomRequest).toHaveBeenCalledTimes(1); + await handle.sendNotification('ui/ping'); + expect(host.sendCustomNotification).toHaveBeenCalledTimes(1); + }); +}); + +describe('ExtensionHandle — id and settings', () => { + test('exposes id and local settings as readonly fields', () => { + const { handle } = makeHandle({ peer: undefined }); + expect(handle.id).toBe('io.example/ui'); + expect(handle.settings).toEqual({ local: true }); + }); +}); diff --git a/packages/server/src/server/server.ts b/packages/server/src/server/server.ts index 4361f3e1e..b5b936bb2 100644 --- a/packages/server/src/server/server.ts +++ b/packages/server/src/server/server.ts @@ -1,4 +1,5 @@ import type { + AnySchema, BaseContext, ClientCapabilities, CreateMessageRequest, @@ -9,9 +10,11 @@ import type { ElicitRequestFormParams, ElicitRequestURLParams, ElicitResult, + ExtensionOptions, Implementation, InitializeRequest, InitializeResult, + JSONObject, JsonSchemaType, jsonSchemaValidator, ListRootsRequest, @@ -26,6 +29,7 @@ import type { RequestTypeMap, ResourceUpdatedNotification, ResultTypeMap, + SchemaOutput, ServerCapabilities, ServerContext, ServerResult, @@ -43,6 +47,7 @@ import { CreateTaskResultSchema, ElicitResultSchema, EmptyResultSchema, + ExtensionHandle, extractTaskManagerOptions, LATEST_PROTOCOL_VERSION, ListRootsResultSchema, @@ -102,6 +107,7 @@ export class Server extends Protocol { private _capabilities: ServerCapabilities; private _instructions?: string; private _jsonSchemaValidator: jsonSchemaValidator; + private _enforceStrictCapabilities: boolean; private _experimental?: { tasks: ExperimentalServerTasks }; /** @@ -123,6 +129,7 @@ export class Server extends Protocol { this._capabilities = options?.capabilities ? { ...options.capabilities } : {}; this._instructions = options?.instructions; this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new DefaultJsonSchemaValidator(); + this._enforceStrictCapabilities = options?.enforceStrictCapabilities ?? false; // Strip runtime-only fields from advertised capabilities if (options?.capabilities?.tasks) { @@ -219,6 +226,43 @@ export class Server extends Protocol { } } + /** + * Declares an SEP-2133 extension and returns a scoped {@linkcode ExtensionHandle} for + * registering and sending its custom JSON-RPC methods. + * + * Merges `settings` into `capabilities.extensions[id]`, which is advertised to the client + * in the `initialize` result. Must be called before {@linkcode connect}. After connecting, + * {@linkcode ExtensionHandle.getPeerSettings | handle.getPeerSettings()} returns the client's + * `capabilities.extensions[id]` blob (validated against `peerSchema` if provided). + */ + public extension(id: string, settings: L): ExtensionHandle; + public extension( + id: string, + settings: L, + opts: ExtensionOptions

+ ): ExtensionHandle, ServerContext>; + public extension( + id: string, + settings: L, + opts?: ExtensionOptions

+ ): ExtensionHandle | JSONObject, ServerContext> { + if (this.transport) { + throw new SdkError(SdkErrorCode.AlreadyConnected, 'Cannot register extension after connecting to transport'); + } + if (this._capabilities.extensions && Object.hasOwn(this._capabilities.extensions, id)) { + throw new Error(`Extension "${id}" is already registered`); + } + this._capabilities.extensions = { ...this._capabilities.extensions, [id]: settings }; + return new ExtensionHandle( + this, + id, + settings, + () => this._clientCapabilities?.extensions?.[id], + this._enforceStrictCapabilities, + opts?.peerSchema + ); + } + /** * Override request handler registration to enforce server-side validation for `tools/call`. */ diff --git a/packages/server/test/server/extension.test.ts b/packages/server/test/server/extension.test.ts new file mode 100644 index 000000000..e87ea4b63 --- /dev/null +++ b/packages/server/test/server/extension.test.ts @@ -0,0 +1,102 @@ +import { InMemoryTransport, type JSONRPCMessage, SdkError, SdkErrorCode } from '@modelcontextprotocol/core'; +import { describe, expect, test } from 'vitest'; +import * as z from 'zod/v4'; + +import { Server } from '../../src/server/server.js'; + +/** + * These tests exercise the `Server.extension()` factory and the server side of the + * `capabilities.extensions` round-trip via `initialize`. The `ExtensionHandle` class itself is + * unit-tested in `@modelcontextprotocol/core/test/shared/extensionHandle.test.ts`. + */ + +async function rawInitialize( + clientSide: InMemoryTransport, + clientCapabilities: Record = {} +): Promise> { + const result = new Promise>((resolve, reject) => { + clientSide.onmessage = (msg: JSONRPCMessage) => { + if ('id' in msg && msg.id === 1) { + if ('result' in msg) resolve(msg.result as Record); + else if ('error' in msg) reject(new Error(JSON.stringify(msg.error))); + } + }; + }); + await clientSide.start(); + await clientSide.send({ + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { + protocolVersion: '2025-11-25', + capabilities: clientCapabilities, + clientInfo: { name: 'raw-client', version: '0.0.0' } + } + }); + return result; +} + +describe('Server.extension()', () => { + test('merges settings into capabilities.extensions and advertises them in initialize result', async () => { + const server = new Server({ name: 's', version: '1.0.0' }, { capabilities: {} }); + server.extension('io.example/ui', { contentTypes: ['text/html'] }); + server.extension('com.acme/widgets', { v: 2 }); + + const [clientSide, serverSide] = InMemoryTransport.createLinkedPair(); + await server.connect(serverSide); + const result = await rawInitialize(clientSide); + + const caps = result.capabilities as Record; + expect(caps.extensions).toEqual({ + 'io.example/ui': { contentTypes: ['text/html'] }, + 'com.acme/widgets': { v: 2 } + }); + }); + + test('throws AlreadyConnected after connect()', async () => { + const server = new Server({ name: 's', version: '1.0.0' }); + const [, serverSide] = InMemoryTransport.createLinkedPair(); + await server.connect(serverSide); + + expect(() => server.extension('io.example/ui', {})).toThrow(SdkError); + try { + server.extension('io.example/ui', {}); + expect.fail('should have thrown'); + } catch (e) { + expect(e).toBeInstanceOf(SdkError); + expect((e as SdkError).code).toBe(SdkErrorCode.AlreadyConnected); + } + }); + + test('throws on duplicate extension id', () => { + const server = new Server({ name: 's', version: '1.0.0' }); + server.extension('io.example/ui', { v: 1 }); + expect(() => server.extension('io.example/ui', { v: 2 })).toThrow(/already registered/); + expect(() => server.extension('com.other/thing', {})).not.toThrow(); + }); + + test("getPeerSettings() reads the client's capabilities.extensions[id] after initialize", async () => { + const PeerSchema = z.object({ availableDisplayModes: z.array(z.string()) }); + const server = new Server({ name: 's', version: '1.0.0' }); + const handle = server.extension('io.example/ui', { hostSide: true }, { peerSchema: PeerSchema }); + + expect(handle.getPeerSettings()).toBeUndefined(); + + const [clientSide, serverSide] = InMemoryTransport.createLinkedPair(); + await server.connect(serverSide); + await rawInitialize(clientSide, { + extensions: { 'io.example/ui': { availableDisplayModes: ['inline', 'fullscreen'] } } + }); + + expect(handle.getPeerSettings()).toEqual({ availableDisplayModes: ['inline', 'fullscreen'] }); + }); + + test('handle.setRequestHandler can be called after connect()', async () => { + const server = new Server({ name: 's', version: '1.0.0' }); + const handle = server.extension('io.example/ui', {}); + const [, serverSide] = InMemoryTransport.createLinkedPair(); + await server.connect(serverSide); + + expect(() => handle.setRequestHandler('ui/late', z.object({}), () => ({}))).not.toThrow(); + }); +});