refactor: remove fixable casts (widen get*Schema/assert*, isStandardSchema guard, union impl sigs)

felixweinberger · felixweinberger · commit 5c25abde892c · 2026-04-28T16:50:18.000Z
diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts
@@ -24,7 +24,6 @@ import type {
     NotificationMethod,
     ProtocolOptions,
     ReadResourceRequest,
-    RequestMethod,
     RequestOptions,
     Result,
     ServerCapabilities,
@@ -565,7 +564,7 @@ export class Client extends Protocol<ClientContext> {
         return this._instructions;
     }
 
-    protected assertCapabilityForMethod(method: RequestMethod): void {
+    protected assertCapabilityForMethod(method: string): void {
         switch (method as ClientRequest['method']) {
             case 'logging/setLevel': {
                 if (!this._serverCapabilities?.logging) {
@@ -628,7 +627,7 @@ export class Client extends Protocol<ClientContext> {
         }
     }
 
-    protected assertNotificationCapability(method: NotificationMethod): void {
+    protected assertNotificationCapability(method: string): void {
         switch (method as ClientNotification['method']) {
             case 'notifications/roots/list_changed': {
                 if (!this._capabilities.roots?.listChanged) {
diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts
@@ -45,7 +45,7 @@ import {
     SUPPORTED_PROTOCOL_VERSIONS
 } from '../types/index.js';
 import type { StandardSchemaV1 } from '../util/standardSchema.js';
-import { validateStandardSchema } from '../util/standardSchema.js';
+import { isStandardSchema, validateStandardSchema } from '../util/standardSchema.js';
 import type { TaskContext, TaskManagerHost, TaskManagerOptions, TaskRequestOptions } from './taskManager.js';
 import { NullTaskManager, TaskManager } from './taskManager.js';
 import type { Transport, TransportSendOptions } from './transport.js';
@@ -609,17 +609,18 @@ export abstract class Protocol<ContextT extends BaseContext> {
                 method: request.method,
                 _meta: request.params?._meta,
                 signal: abortController.signal,
+                // Arrow literals can't carry overload signatures; cast asserts this impl matches the overloaded property type.
                 send: ((r: Request, schemaOrOptions?: StandardSchemaV1 | TaskRequestOptions, maybeOptions?: TaskRequestOptions) => {
-                    if (schemaOrOptions != null && typeof schemaOrOptions === 'object' && '~standard' in schemaOrOptions) {
+                    if (isStandardSchema(schemaOrOptions)) {
                         return sendRequest(r, schemaOrOptions, maybeOptions);
                     }
-                    const resultSchema = getResultSchema(r.method as RequestMethod);
+                    const resultSchema = getResultSchema(r.method);
                     if (!resultSchema) {
                         throw new TypeError(
                             `'${r.method}' is not a spec method; pass a result schema as the second argument to ctx.mcpReq.send().`
                         );
                     }
-                    return sendRequest(r, resultSchema, schemaOrOptions as TaskRequestOptions | undefined);
+                    return sendRequest(r, resultSchema, schemaOrOptions);
                 }) as BaseContext['mcpReq']['send'],
                 notify: sendNotification
             },
@@ -761,14 +762,14 @@ export abstract class Protocol<ContextT extends BaseContext> {
      *
      * This should be implemented by subclasses.
      */
-    protected abstract assertCapabilityForMethod(method: RequestMethod): void;
+    protected abstract assertCapabilityForMethod(method: string): void;
 
     /**
      * A method to check if a notification is supported by the local side, for the given method to be sent.
      *
      * This should be implemented by subclasses.
      */
-    protected abstract assertNotificationCapability(method: NotificationMethod): void;
+    protected abstract assertNotificationCapability(method: string): void;
 
     /**
      * A method to check if a request handler is supported by the local side, for the given method to be handled.
@@ -813,10 +814,10 @@ export abstract class Protocol<ContextT extends BaseContext> {
         options?: RequestOptions
     ): Promise<StandardSchemaV1.InferOutput<T>>;
     request(request: Request, schemaOrOptions?: StandardSchemaV1 | RequestOptions, maybeOptions?: RequestOptions): Promise<unknown> {
-        if (schemaOrOptions != null && typeof schemaOrOptions === 'object' && '~standard' in schemaOrOptions) {
+        if (isStandardSchema(schemaOrOptions)) {
             return this._requestWithSchema(request, schemaOrOptions, maybeOptions);
         }
-        const resultSchema = getResultSchema(request.method as RequestMethod);
+        const resultSchema = getResultSchema(request.method);
         if (!resultSchema) {
             throw new TypeError(`'${request.method}' is not a spec method; pass a result schema as the second argument to request().`);
         }
@@ -852,7 +853,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
 
             if (this._options?.enforceStrictCapabilities === true) {
                 try {
-                    this.assertCapabilityForMethod(request.method as RequestMethod);
+                    this.assertCapabilityForMethod(request.method);
                 } catch (error) {
                     earlyReject(error);
                     return;
@@ -986,7 +987,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
             throw new SdkError(SdkErrorCode.NotConnected, 'Not connected');
         }
 
-        this.assertNotificationCapability(notification.method as NotificationMethod);
+        this.assertNotificationCapability(notification.method);
 
         // Delegate task-related notification routing and JSONRPC building to TaskManager
         const taskResult = await this._taskManager.processOutboundNotification(notification, options);
@@ -1069,35 +1070,35 @@ export abstract class Protocol<ContextT extends BaseContext> {
             ctx: ContextT
         ) => RequestHandlerSchemas.InferResult<R> | Promise<RequestHandlerSchemas.InferResult<R>>
     ): void;
-    setRequestHandler(method: string, schemasOrHandler: unknown, maybeHandler?: unknown): void {
+    setRequestHandler(
+        method: string,
+        schemasOrHandler: RequestHandlerSchemas | ((request: unknown, ctx: ContextT) => Result | Promise<Result>),
+        maybeHandler?: (params: unknown, ctx: ContextT) => Result | Promise<Result>
+    ): void {
         this.assertRequestHandlerCapability(method);
 
         let stored: (request: JSONRPCRequest, ctx: ContextT) => Promise<Result>;
 
-        if (maybeHandler === undefined) {
-            const handler = schemasOrHandler as (request: unknown, ctx: ContextT) => Result | Promise<Result>;
-            const schema = getRequestSchema(method as RequestMethod);
+        if (typeof schemasOrHandler === 'function') {
+            const schema = getRequestSchema(method);
             if (!schema) {
                 throw new TypeError(
                     `'${method}' is not a spec request method; pass schemas as the second argument to setRequestHandler().`
                 );
             }
-            stored = (request, ctx) => {
-                const parsed = schema.parse(request);
-                return Promise.resolve(handler(parsed, ctx));
-            };
-        } else {
-            const schemas = schemasOrHandler as RequestHandlerSchemas;
-            const handler = maybeHandler as (params: unknown, ctx: ContextT) => Result | Promise<Result>;
+            stored = (request, ctx) => Promise.resolve(schemasOrHandler(schema.parse(request), ctx));
+        } else if (maybeHandler) {
             stored = async (request, ctx) => {
-                const userParams = { ...((request.params ?? {}) as Record<string, unknown>) };
+                const userParams = { ...request.params };
                 delete userParams._meta;
-                const parsed = await validateStandardSchema(schemas.params, userParams);
+                const parsed = await validateStandardSchema(schemasOrHandler.params, userParams);
                 if (!parsed.success) {
                     throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Invalid params for ${method}: ${parsed.error}`);
                 }
-                return handler(parsed.data, ctx);
+                return maybeHandler(parsed.data, ctx);
             };
+        } else {
+            throw new TypeError('setRequestHandler: handler is required');
         }
 
         this._requestHandlers.set(method, this._wrapHandler(method, stored));
@@ -1153,32 +1154,33 @@ export abstract class Protocol<ContextT extends BaseContext> {
         schemas: { params: P },
         handler: (params: StandardSchemaV1.InferOutput<P>) => void | Promise<void>
     ): void;
-    setNotificationHandler(method: string, schemasOrHandler: unknown, maybeHandler?: unknown): void {
-        if (maybeHandler !== undefined) {
-            const schemas = schemasOrHandler as { params: StandardSchemaV1 };
-            const handler = maybeHandler as (params: unknown) => void | Promise<void>;
-            this._notificationHandlers.set(method, async notification => {
-                const userParams = { ...((notification.params ?? {}) as Record<string, unknown>) };
-                delete userParams._meta;
-                const parsed = await validateStandardSchema(schemas.params, userParams);
-                if (!parsed.success) {
-                    throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Invalid params for notification ${method}: ${parsed.error}`);
-                }
-                await handler(parsed.data);
-            });
+    setNotificationHandler(
+        method: string,
+        schemasOrHandler: { params: StandardSchemaV1 } | ((notification: unknown) => void | Promise<void>),
+        maybeHandler?: (params: unknown) => void | Promise<void>
+    ): void {
+        if (typeof schemasOrHandler === 'function') {
+            const schema = getNotificationSchema(method);
+            if (!schema) {
+                throw new TypeError(
+                    `'${method}' is not a spec notification method; pass schemas as the second argument to setNotificationHandler().`
+                );
+            }
+            this._notificationHandlers.set(method, notification => Promise.resolve(schemasOrHandler(schema.parse(notification))));
             return;
         }
 
-        const handler = schemasOrHandler as (notification: unknown) => void | Promise<void>;
-        const schema = getNotificationSchema(method as NotificationMethod);
-        if (!schema) {
-            throw new TypeError(
-                `'${method}' is not a spec notification method; pass schemas as the second argument to setNotificationHandler().`
-            );
+        if (!maybeHandler) {
+            throw new TypeError('setNotificationHandler: handler is required');
         }
-        this._notificationHandlers.set(method, notification => {
-            const parsed = schema.parse(notification);
-            return Promise.resolve(handler(parsed));
+        this._notificationHandlers.set(method, async notification => {
+            const userParams = { ...notification.params };
+            delete userParams._meta;
+            const parsed = await validateStandardSchema(schemasOrHandler.params, userParams);
+            if (!parsed.success) {
+                throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Invalid params for notification ${method}: ${parsed.error}`);
+            }
+            await maybeHandler(parsed.data);
         });
     }
 
diff --git a/packages/core/src/shared/taskManager.ts b/packages/core/src/shared/taskManager.ts
@@ -283,6 +283,9 @@ export class TaskManager {
 
         if (!task) {
             try {
+                // TODO: SchemaOutput<T> (Zod) and StandardSchemaV1.InferOutput<T> (host.request's return)
+                // resolve to the same type for Zod schemas, but TS can't unify them generically.
+                // Removing this cast requires aligning ResponseMessage<T extends Result> with StandardSchema.
                 const result = (await host.request(request, resultSchema, options)) as SchemaOutput<T>;
                 yield { type: 'result', result };
             } catch (error) {
@@ -355,6 +358,7 @@ export class TaskManager {
         resultSchema: T,
         options?: RequestOptions
     ): Promise<SchemaOutput<T>> {
+        // TODO: same SchemaOutput<T> vs StandardSchemaV1.InferOutput<T> mismatch as requestStream above.
         return this._requireHost.request({ method: 'tasks/result', params }, resultSchema, options) as Promise<SchemaOutput<T>>;
     }
 
diff --git a/packages/core/src/types/schemas.ts b/packages/core/src/types/schemas.ts
@@ -2181,10 +2181,13 @@ const resultSchemas: Record<string, z.core.$ZodType> = {
 
 /**
  * Gets the Zod schema for validating results of a given request method.
+ * Returns `undefined` for non-spec methods.
  * @see getRequestSchema for explanation of the internal type assertion.
  */
-export function getResultSchema<M extends RequestMethod>(method: M): z.ZodType<ResultTypeMap[M]> {
-    return resultSchemas[method] as unknown as z.ZodType<ResultTypeMap[M]>;
+export function getResultSchema<M extends RequestMethod>(method: M): z.ZodType<ResultTypeMap[M]>;
+export function getResultSchema(method: string): z.ZodType | undefined;
+export function getResultSchema(method: string): z.ZodType | undefined {
+    return resultSchemas[method as RequestMethod] as unknown as z.ZodType | undefined;
 }
 
 /* Runtime schema lookup — request schemas by method */
@@ -2219,14 +2222,19 @@ const notificationSchemas = buildSchemaMap([...ClientNotificationSchema.options,
  * when M is a generic type parameter. Both compute to the same type at
  * instantiation, but TypeScript can't prove this statically.
  */
-export function getRequestSchema<M extends RequestMethod>(method: M): z.ZodType<RequestTypeMap[M]> {
-    return requestSchemas[method] as unknown as z.ZodType<RequestTypeMap[M]>;
+export function getRequestSchema<M extends RequestMethod>(method: M): z.ZodType<RequestTypeMap[M]>;
+export function getRequestSchema(method: string): z.ZodType | undefined;
+export function getRequestSchema(method: string): z.ZodType | undefined {
+    return requestSchemas[method as RequestMethod] as unknown as z.ZodType | undefined;
 }
 
 /**
  * Gets the Zod schema for a given notification method.
+ * Returns `undefined` for non-spec methods.
  * @see getRequestSchema for explanation of the internal type assertion.
  */
-export function getNotificationSchema<M extends NotificationMethod>(method: M): z.ZodType<NotificationTypeMap[M]> {
-    return notificationSchemas[method] as unknown as z.ZodType<NotificationTypeMap[M]>;
+export function getNotificationSchema<M extends NotificationMethod>(method: M): z.ZodType<NotificationTypeMap[M]>;
+export function getNotificationSchema(method: string): z.ZodType | undefined;
+export function getNotificationSchema(method: string): z.ZodType | undefined {
+    return notificationSchemas[method as NotificationMethod] as unknown as z.ZodType | undefined;
 }
diff --git a/packages/server/src/server/server.ts b/packages/server/src/server/server.ts
@@ -19,10 +19,8 @@ import type {
     LoggingLevel,
     LoggingMessageNotification,
     MessageExtraInfo,
-    NotificationMethod,
     NotificationOptions,
     ProtocolOptions,
-    RequestMethod,
     RequestOptions,
     ResourceUpdatedNotification,
     Result,
@@ -266,7 +264,7 @@ export class Server extends Protocol<ServerContext> {
         };
     }
 
-    protected assertCapabilityForMethod(method: RequestMethod): void {
+    protected assertCapabilityForMethod(method: string): void {
         switch (method) {
             case 'sampling/createMessage': {
                 if (!this._clientCapabilities?.sampling) {
@@ -299,7 +297,7 @@ export class Server extends Protocol<ServerContext> {
         }
     }
 
-    protected assertNotificationCapability(method: NotificationMethod): void {
+    protected assertNotificationCapability(method: string): void {
         switch (method) {
             case 'notifications/message': {
                 if (!this._capabilities.logging) {
