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

Exports the abstract Protocol class (was reachable in v1 via deep imports) and adds
Protocol<ContextT, SpecT extends ProtocolSpec = McpSpec>. Subclasses supplying a concrete
ProtocolSpec get method-name autocomplete and params/result correlation on the typed
setRequestHandler/setNotificationHandler overloads.

Author: felixweinberger

.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 params/result correlation on the typed `setRequestHandler`/`setNotificationHandler` overloads.

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

@@ -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) + 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,51 @@ type TimeoutInfo = {
     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 params/result correlation on the typed overloads of `setRequestHandler`
+ * and `setNotificationHandler`. The default leaves them string-keyed and untyped.
+ */
+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.
  */
-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,10 +1082,19 @@ export abstract class Protocol<ContextT extends BaseContext> {
      *   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.). When `method` is listed in this instance's
+     *   {@linkcode ProtocolSpec}, params and result types are inferred from `SpecT`.
      * - **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<_Requests<SpecT>[K]['params']>>(
+        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>
@@ -1143,8 +1192,15 @@ export abstract class Protocol<ContextT extends BaseContext> {
      *
      * 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). When the
+     * three-arg form's `method` is listed in this instance's {@linkcode ProtocolSpec}, the params
+     * type is inferred from `SpecT`.
      */
+    setNotificationHandler<K extends SpecNotifications<SpecT>, P extends StandardSchemaV1<_Notifications<SpecT>[K]['params']>>(
+        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>

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

@@ -0,0 +1,68 @@
+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 };
+        });
+    });
+});