feat(core): export Protocol class + ProtocolSpec generic for typed custom vocabularies

.changeset/export-protocol-spec.md

@@ -0,0 +1,6 @@
+---
+'@modelcontextprotocol/client': minor
+'@modelcontextprotocol/server': minor
+---
+
+Export the abstract `Protocol` class (was reachable in v1 via deep imports) and add `Protocol<ContextT, SpecT extends ProtocolSpec = McpSpec>` for typed custom-method vocabularies. Subclasses supplying a concrete `ProtocolSpec` get method-name autocomplete and result-type correlation on the typed `setRequestHandler`/`setNotificationHandler` overloads (handler param types come from the `paramsSchema` argument; `ProtocolSpec['params']` is informational).

CLAUDE.md

@@ -67,7 +67,7 @@ The SDK is organized into three main layers:
 The SDK has a two-layer export structure to separate internal code from the public API:
 
 - **`@modelcontextprotocol/core`** (main entry, `packages/core/src/index.ts`) — Internal barrel. Exports everything (including Zod schemas, Protocol class, stdio utils). Only consumed by sibling packages within the monorepo (`private: true`).
-- **`@modelcontextprotocol/core/public`** (`packages/core/src/exports/public/index.ts`) — Curated public API. Exports only TypeScript types, error classes, constants, and guards. Re-exported by client and server packages.
+- **`@modelcontextprotocol/core/public`** (`packages/core/src/exports/public/index.ts`) — Curated public API. Exports TypeScript types, error classes, constants, guards, and the `Protocol` class. Re-exported by client and server packages.
 - **`@modelcontextprotocol/client`** and **`@modelcontextprotocol/server`** (`packages/*/src/index.ts`) — Final public surface. Package-specific exports (named explicitly) plus re-exports from `core/public`.
 
 When modifying exports:

packages/client/src/client/client.ts

@@ -344,8 +344,13 @@ export class Client extends Protocol<ClientContext> {
         method: M,
         handler: (request: RequestTypeMap[M], ctx: ClientContext) => ResultTypeMap[M] | Promise<ResultTypeMap[M]>
     ): void;
-    public override setRequestHandler<P extends StandardSchemaV1>(
-        method: string,
+    public override setRequestHandler<M extends RequestMethod, P extends StandardSchemaV1>(
+        method: M,
+        paramsSchema: P,
+        handler: (params: StandardSchemaV1.InferOutput<P>, ctx: ClientContext) => ResultTypeMap[M] | Promise<ResultTypeMap[M]>
+    ): void;
+    public override setRequestHandler<M extends string, P extends StandardSchemaV1>(
+        method: M extends RequestMethod ? never : M,
         paramsSchema: P,
         handler: (params: StandardSchemaV1.InferOutput<P>, ctx: ClientContext) => Result | Promise<Result>
     ): void;

packages/core/src/exports/public/index.ts

@@ -4,9 +4,9 @@
  * This module defines the stable, public-facing API surface. Client and server
  * packages re-export from here so that end users only see supported symbols.
  *
- * Internal utilities (Protocol class, stdio parsing, schema helpers, etc.)
- * remain available via the internal barrel (@modelcontextprotocol/core) for
- * use by client/server packages.
+ * Internal utilities (stdio parsing, schema helpers, etc.) remain available via
+ * the internal barrel (@modelcontextprotocol/core) for use by client/server
+ * packages.
  */
 
 // Auth error classes
@@ -38,17 +38,21 @@ export { checkResourceAllowed, resourceUrlFromServerUrl } from '../../shared/aut
 // Metadata utilities
 export { getDisplayName } from '../../shared/metadataUtils.js';
 
-// Protocol types (NOT the Protocol class itself or mergeCapabilities)
+// Protocol class (abstract; subclass for custom vocabularies via SpecT) + types. NOT mergeCapabilities.
 export type {
     BaseContext,
     ClientContext,
+    McpSpec,
     NotificationOptions,
     ProgressCallback,
     ProtocolOptions,
+    ProtocolSpec,
     RequestOptions,
-    ServerContext
+    ServerContext,
+    SpecNotifications,
+    SpecRequests
 } from '../../shared/protocol.js';
-export { DEFAULT_REQUEST_TIMEOUT_MSEC } from '../../shared/protocol.js';
+export { DEFAULT_REQUEST_TIMEOUT_MSEC, Protocol } from '../../shared/protocol.js';
 export type { ZodLikeRequestSchema } from '../../util/compatSchema.js';
 
 // Task manager types (NOT TaskManager class itself — internal)

packages/core/src/shared/protocol.ts

@@ -305,11 +305,60 @@
     onTimeout: () => void;
 };
 
+/**
+ * Declares the request and notification vocabulary a `Protocol` subclass speaks.
+ *
+ * Supplying a concrete `ProtocolSpec` as `Protocol`'s second type argument gives method-name
+ * autocomplete and result-type correlation on the typed overloads of `setRequestHandler`
+ * and `setNotificationHandler`. `Protocol` defaults to {@linkcode McpSpec}; using the bare
+ * `ProtocolSpec` type leaves methods string-keyed and untyped.
+ *
+ * Only `requests[K].result` is enforced by the type system; `params` shapes are informational
+ * (handler param types come from the `paramsSchema` you pass at the call site).
+ */
+export type ProtocolSpec = {
+    requests?: Record<string, { params?: unknown; result: unknown }>;
+    notifications?: Record<string, { params?: unknown }>;
+};
+
+/**
+ * The {@linkcode ProtocolSpec} that describes the standard MCP method vocabulary, derived from
+ * {@linkcode RequestTypeMap}, {@linkcode ResultTypeMap} and {@linkcode NotificationTypeMap}.
+ */
+export type McpSpec = {
+    requests: { [M in RequestMethod]: { params: RequestTypeMap[M]['params']; result: ResultTypeMap[M] } };
+    notifications: { [M in NotificationMethod]: { params: NotificationTypeMap[M]['params'] } };
+};
+
+type _Requests<SpecT extends ProtocolSpec> = NonNullable<SpecT['requests']>;
+type _Notifications<SpecT extends ProtocolSpec> = NonNullable<SpecT['notifications']>;
+
+/**
+ * Method-name keys from a {@linkcode ProtocolSpec}'s `requests` map, or `never` for the
+ * unconstrained default `ProtocolSpec`. Making the keys `never` for the default disables the
+ * spec-typed overloads on `setRequestHandler` until the caller supplies a concrete `SpecT`.
+ */
+export type SpecRequests<SpecT extends ProtocolSpec> = string extends keyof _Requests<SpecT> ? never : keyof _Requests<SpecT> & string;
+
+/** See {@linkcode SpecRequests}. */
+export type SpecNotifications<SpecT extends ProtocolSpec> = string extends keyof _Notifications<SpecT>
+    ? never
+    : keyof _Notifications<SpecT> & string;
+
 /**
  * Implements MCP protocol framing on top of a pluggable transport, including
  * features like request/response linking, notifications, and progress.
+ *
+ * `Protocol` is abstract; `Client` and `Server` are the concrete role-specific implementations.
+ * Subclasses (such as MCP-dialect protocols like MCP Apps) can supply a {@linkcode ProtocolSpec}
+ * as the second type argument to get method-name autocomplete on their own vocabulary.
+ *
+ * @remarks
+ * Subclassing `Protocol` directly is supported for MCP-dialect frameworks. The protected
+ * surface (`buildContext`, `assertCapability*`, `_setRequestHandlerByMethod`) may evolve in
+ * minor versions; prefer `Client`/`Server` unless you need a custom method vocabulary.
  */
-export abstract class Protocol<ContextT extends BaseContext> {
+export abstract class Protocol<ContextT extends BaseContext = BaseContext, SpecT extends ProtocolSpec = McpSpec> {
     private _transport?: Transport;
     private _requestMessageId = 0;
     private _requestHandlers: Map<string, (request: JSONRPCRequest, ctx: ContextT) => Promise<Result>> = new Map();
@@ -1042,16 +1091,26 @@
      *   Any method string; the supplied schema validates incoming `params`. Absent or undefined
      *   `params` are normalized to `{}` (after stripping `_meta`) before validation, so for
      *   no-params methods use `z.object({})`. `paramsSchema` may be any Standard Schema (Zod,
-     *   Valibot, ArkType, etc.).
+     *   Valibot, ArkType, etc.). The handler's `params` type is inferred from the passed
+     *   `paramsSchema`; when `method` is listed in this instance's {@linkcode ProtocolSpec},
+     *   the handler's result type is constrained to `SpecT`'s declared result.
      * - **Zod schema** — `setRequestHandler(RequestZodSchema, (request, ctx) => …)`. The method
      *   name is read from the schema's `method` literal; the handler receives the parsed request.
      */
+    setRequestHandler<K extends SpecRequests<SpecT>, P extends StandardSchemaV1>(
+        method: K,
+        paramsSchema: P,
+        handler: (
+            params: StandardSchemaV1.InferOutput<P>,
+            ctx: ContextT
+        ) => _Requests<SpecT>[K]['result'] | Promise<_Requests<SpecT>[K]['result']>
+    ): void;
     setRequestHandler<M extends RequestMethod>(
         method: M,
         handler: (request: RequestTypeMap[M], ctx: ContextT) => Result | Promise<Result>
     ): void;
-    setRequestHandler<P extends StandardSchemaV1>(
-        method: string,
+    setRequestHandler<M extends string, P extends StandardSchemaV1>(
+        method: M extends SpecRequests<SpecT> ? never : M,
         paramsSchema: P,
         handler: (params: StandardSchemaV1.InferOutput<P>, ctx: ContextT) => Result | Promise<Result>
     ): void;
@@ -1143,14 +1202,20 @@
      *
      * Mirrors {@linkcode setRequestHandler}: a two-arg spec-method form (handler receives the full
      * notification object), a three-arg form with a `paramsSchema` (handler receives validated
-     * `params`), and a Zod-schema form (method read from the schema's `method` literal).
+     * `params`), and a Zod-schema form (method read from the schema's `method` literal). The
+     * handler's `params` type is always inferred from the passed schema.
      */
+    setNotificationHandler<K extends SpecNotifications<SpecT>, P extends StandardSchemaV1>(
+        method: K,
+        paramsSchema: P,
+        handler: (params: StandardSchemaV1.InferOutput<P>) => void | Promise<void>
+    ): void;
     setNotificationHandler<M extends NotificationMethod>(
         method: M,
         handler: (notification: NotificationTypeMap[M]) => void | Promise<void>
     ): void;
-    setNotificationHandler<P extends StandardSchemaV1>(
-        method: string,
+    setNotificationHandler<M extends string, P extends StandardSchemaV1>(
+        method: M extends SpecNotifications<SpecT> ? never : M,
         paramsSchema: P,
         handler: (params: StandardSchemaV1.InferOutput<P>) => void | Promise<void>
     ): void;

packages/core/test/shared/protocolSpec.test.ts

@@ -0,0 +1,81 @@
+import { describe, expect, it } from 'vitest';
+import { z } from 'zod';
+
+import type { BaseContext, ProtocolSpec, SpecRequests } from '../../src/shared/protocol.js';
+import { Protocol } from '../../src/shared/protocol.js';
+import { InMemoryTransport } from '../../src/util/inMemory.js';
+
+class TestProtocol<SpecT extends ProtocolSpec = ProtocolSpec> extends Protocol<BaseContext, SpecT> {
+    protected assertCapabilityForMethod(): void {}
+    protected assertNotificationCapability(): void {}
+    protected assertRequestHandlerCapability(): void {}
+    protected assertTaskCapability(): void {}
+    protected assertTaskHandlerCapability(): void {}
+    protected buildContext(ctx: BaseContext): BaseContext {
+        return ctx;
+    }
+}
+
+describe('ProtocolSpec typing', () => {
+    type AppSpec = {
+        requests: {
+            'ui/open-link': { params: { url: string }; result: { opened: boolean } };
+        };
+        notifications: {
+            'ui/size-changed': { params: { width: number; height: number } };
+        };
+    };
+
+    type _Assert<T extends true> = T;
+    type _Eq<A, B> = [A] extends [B] ? ([B] extends [A] ? true : false) : false;
+    type _t1 = _Assert<_Eq<SpecRequests<AppSpec>, 'ui/open-link'>>;
+    type _t2 = _Assert<_Eq<SpecRequests<ProtocolSpec>, never>>;
+    void (undefined as unknown as [_t1, _t2]);
+
+    it('typed-SpecT overload infers params/result; string fallback still works', async () => {
+        const [t1, t2] = InMemoryTransport.createLinkedPair();
+        const app = new TestProtocol<AppSpec>();
+        const host = new TestProtocol<AppSpec>();
+        await app.connect(t1);
+        await host.connect(t2);
+
+        host.setRequestHandler('ui/open-link', z.object({ url: z.string() }), p => {
+            const _typed: string = p.url;
+            void _typed;
+            return { opened: true };
+        });
+        const r = await app.request({ method: 'ui/open-link', params: { url: 'https://x' } }, z.object({ opened: z.boolean() }));
+        expect(r.opened).toBe(true);
+
+        host.setRequestHandler('not/in-spec', z.object({ n: z.number() }), p => ({ doubled: p.n * 2 }));
+        const r2 = await app.request({ method: 'not/in-spec', params: { n: 3 } }, z.object({ doubled: z.number() }));
+        expect(r2.doubled).toBe(6);
+    });
+
+    it('typed-SpecT overload types handler from passed schema, not SpecT (regression)', () => {
+        type Spec = { requests: { 'x/y': { params: { a: string; b: string }; result: { ok: boolean } } } };
+        const p = new TestProtocol<Spec>();
+        const Narrow = z.object({ a: z.string() });
+        p.setRequestHandler('x/y', Narrow, params => {
+            const _a: string = params.a;
+            // @ts-expect-error -- params is InferOutput<Narrow>, has no 'b' even though Spec does
+            const _b: string = params.b;
+            void _a;
+            void _b;
+            return { ok: true };
+        });
+    });
+
+    it('typed-SpecT setRequestHandler enforces result type (no fallthrough to loose string overload)', () => {
+        const p = new TestProtocol<AppSpec>();
+        // @ts-expect-error -- result must be { opened: boolean }; string overload is `never`-guarded for spec methods
+        p.setRequestHandler('ui/open-link', z.object({ url: z.string() }), () => ({ ok: 'wrong-type' }));
+        // @ts-expect-error -- empty object doesn't satisfy { opened: boolean }
+        p.setRequestHandler('ui/open-link', z.object({ url: z.string() }), () => ({}));
+        // non-spec methods still allow loose Result
+        p.setRequestHandler('not/in-spec', z.object({}), () => ({ anything: 1 }));
+        // notifications: spec and non-spec both allow any schema and return void
+        p.setNotificationHandler('ui/size-changed', z.object({ width: z.number(), height: z.number() }), () => {});
+        p.setNotificationHandler('not/in-spec', z.object({ x: z.number() }), () => {});
+    });
+});

packages/server/src/server/server.ts

@@ -232,8 +232,13 @@ export class Server extends Protocol<ServerContext> {
         method: M,
         handler: (request: RequestTypeMap[M], ctx: ServerContext) => ResultTypeMap[M] | Promise<ResultTypeMap[M]>
     ): void;
-    public override setRequestHandler<P extends StandardSchemaV1>(
-        method: string,
+    public override setRequestHandler<M extends RequestMethod, P extends StandardSchemaV1>(
+        method: M,
+        paramsSchema: P,
+        handler: (params: StandardSchemaV1.InferOutput<P>, ctx: ServerContext) => ResultTypeMap[M] | Promise<ResultTypeMap[M]>
+    ): void;
+    public override setRequestHandler<M extends string, P extends StandardSchemaV1>(
+        method: M extends RequestMethod ? never : M,
         paramsSchema: P,
         handler: (params: StandardSchemaV1.InferOutput<P>, ctx: ServerContext) => Result | Promise<Result>
     ): void;

packages/server/test/server/setRequestHandlerSchemaParity.test.ts

@@ -63,4 +63,12 @@ describe('Server.setRequestHandler — Zod-schema form parity', () => {
         });
         expect(res.result).toEqual({ reply: 'hi' });
     });
+
+    it('three-arg form on Server enforces spec-method result type (no fallthrough to loose overload)', () => {
+        const s = new Server({ name: 't', version: '1.0' }, { capabilities: { tools: {} } });
+        // @ts-expect-error -- result for 'ping' must be EmptyResult-compatible; loose overload is never-guarded for spec methods
+        s.setRequestHandler('ping', z.object({}), () => ({ ok: 'wrong-type' }) as { ok: string });
+        // non-spec methods still allow loose Result
+        s.setRequestHandler('acme/custom', z.object({}), () => ({ anything: 1 }));
+    });
 });