diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 0171a1e5..f4ec07a4 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -8,10 +8,10 @@ import { McpError, METHOD_NOT_FOUND, Notification, - PING_REQUEST_METHOD, + PingRequestSchema, Progress, - PROGRESS_NOTIFICATION_METHOD, ProgressNotification, + ProgressNotificationSchema, Request, Result, } from "../types.js"; @@ -78,14 +78,14 @@ export class Protocol< constructor() { this.setNotificationHandler( - PROGRESS_NOTIFICATION_METHOD, + ProgressNotificationSchema.shape.method.value, (notification) => { this._onprogress(notification as unknown as ProgressNotification); }, ); this.setRequestHandler( - PING_REQUEST_METHOD, + PingRequestSchema.shape.method.value, // Automatic pong by default. (_request) => ({}) as SendResultT, ); diff --git a/src/types.ts b/src/types.ts index eb001cff..4d09cc32 100644 --- a/src/types.ts +++ b/src/types.ts @@ -263,24 +263,27 @@ export const PingRequestSchema = RequestSchema.extend({ }); /* Progress notifications */ +export const ProgressSchema = z.object({ + /** + * The progress thus far. This should increase every time progress is made, even if the total is unknown. + */ + progress: z.number(), + /** + * Total number of items to process (or total progress required), if known. + */ + total: z.optional(z.number()), +}); + /** * An out-of-band notification used to inform the receiver of a progress update for a long-running request. */ export const ProgressNotificationSchema = NotificationSchema.extend({ method: z.literal("notifications/progress"), - params: z.object({ + params: ProgressSchema.extend({ /** * The progress token which was given in the initial request, used to associate this notification with the request that is proceeding. */ progressToken: ProgressTokenSchema, - /** - * The progress thus far. This should increase every time progress is made, even if the total is unknown. - */ - progress: z.number(), - /** - * Total number of items to process (or total progress required), if known. - */ - total: z.optional(z.number()), }), }); @@ -810,14 +813,3 @@ export class McpError extends Error { super(`MCP error ${code}: ${message}`); } } - -export type Progress = Pick< - z.infer["params"], - "progress" | "total" ->; - -export const PROGRESS_NOTIFICATION_METHOD: z.infer< - typeof ProgressNotificationSchema ->["method"] = "notifications/progress"; -export const PING_REQUEST_METHOD: z.infer["method"] = - "ping";