Skip to content

Commit 278cfb1

Browse files
Merge pull request #114 from modelcontextprotocol/justin/improve-type-inference
Improve type inference from Zod schemas
2 parents 3f178b2 + 0c5d6c4 commit 278cfb1

File tree

2 files changed

+96
-101
lines changed

2 files changed

+96
-101
lines changed

package-lock.json

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

src/types.ts

Lines changed: 94 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { z } from "zod";
1+
import { z, ZodTypeAny } from "zod";
22

33
export const LATEST_PROTOCOL_VERSION = "2024-11-05";
44
export const SUPPORTED_PROTOCOL_VERSIONS = [
@@ -106,7 +106,7 @@ export enum ErrorCode {
106106
// SDK error codes
107107
ConnectionClosed = -32000,
108108
RequestTimeout = -32001,
109-
109+
110110
// Standard JSON-RPC error codes
111111
ParseError = -32700,
112112
InvalidRequest = -32600,
@@ -1110,131 +1110,126 @@ export class McpError extends Error {
11101110
}
11111111
}
11121112

1113+
type Primitive = string | number | boolean | bigint | null | undefined;
1114+
type Flatten<T> = T extends Primitive
1115+
? T
1116+
: T extends Array<infer U>
1117+
? Array<Flatten<U>>
1118+
: T extends Set<infer U>
1119+
? Set<Flatten<U>>
1120+
: T extends Map<infer K, infer V>
1121+
? Map<Flatten<K>, Flatten<V>>
1122+
: T extends object
1123+
? { [K in keyof T]: Flatten<T[K]> }
1124+
: T;
1125+
1126+
type Infer<Schema extends ZodTypeAny> = Flatten<z.infer<Schema>>;
1127+
11131128
/* JSON-RPC types */
1114-
export type ProgressToken = z.infer<typeof ProgressTokenSchema>;
1115-
export type Cursor = z.infer<typeof CursorSchema>;
1116-
export type Request = z.infer<typeof RequestSchema>;
1117-
export type Notification = z.infer<typeof NotificationSchema>;
1118-
export type Result = z.infer<typeof ResultSchema>;
1119-
export type RequestId = z.infer<typeof RequestIdSchema>;
1120-
export type JSONRPCRequest = z.infer<typeof JSONRPCRequestSchema>;
1121-
export type JSONRPCNotification = z.infer<typeof JSONRPCNotificationSchema>;
1122-
export type JSONRPCResponse = z.infer<typeof JSONRPCResponseSchema>;
1123-
export type JSONRPCError = z.infer<typeof JSONRPCErrorSchema>;
1124-
export type JSONRPCMessage = z.infer<typeof JSONRPCMessageSchema>;
1129+
export type ProgressToken = Infer<typeof ProgressTokenSchema>;
1130+
export type Cursor = Infer<typeof CursorSchema>;
1131+
export type Request = Infer<typeof RequestSchema>;
1132+
export type Notification = Infer<typeof NotificationSchema>;
1133+
export type Result = Infer<typeof ResultSchema>;
1134+
export type RequestId = Infer<typeof RequestIdSchema>;
1135+
export type JSONRPCRequest = Infer<typeof JSONRPCRequestSchema>;
1136+
export type JSONRPCNotification = Infer<typeof JSONRPCNotificationSchema>;
1137+
export type JSONRPCResponse = Infer<typeof JSONRPCResponseSchema>;
1138+
export type JSONRPCError = Infer<typeof JSONRPCErrorSchema>;
1139+
export type JSONRPCMessage = Infer<typeof JSONRPCMessageSchema>;
11251140

11261141
/* Empty result */
1127-
export type EmptyResult = z.infer<typeof EmptyResultSchema>;
1142+
export type EmptyResult = Infer<typeof EmptyResultSchema>;
11281143

11291144
/* Cancellation */
1130-
export type CancelledNotification = z.infer<typeof CancelledNotificationSchema>;
1145+
export type CancelledNotification = Infer<typeof CancelledNotificationSchema>;
11311146

11321147
/* Initialization */
1133-
export type Implementation = z.infer<typeof ImplementationSchema>;
1134-
export type ClientCapabilities = z.infer<typeof ClientCapabilitiesSchema>;
1135-
export type InitializeRequest = z.infer<typeof InitializeRequestSchema>;
1136-
export type ServerCapabilities = z.infer<typeof ServerCapabilitiesSchema>;
1137-
export type InitializeResult = z.infer<typeof InitializeResultSchema>;
1138-
export type InitializedNotification = z.infer<
1139-
typeof InitializedNotificationSchema
1140-
>;
1148+
export type Implementation = Infer<typeof ImplementationSchema>;
1149+
export type ClientCapabilities = Infer<typeof ClientCapabilitiesSchema>;
1150+
export type InitializeRequest = Infer<typeof InitializeRequestSchema>;
1151+
export type ServerCapabilities = Infer<typeof ServerCapabilitiesSchema>;
1152+
export type InitializeResult = Infer<typeof InitializeResultSchema>;
1153+
export type InitializedNotification = Infer<typeof InitializedNotificationSchema>;
11411154

11421155
/* Ping */
1143-
export type PingRequest = z.infer<typeof PingRequestSchema>;
1156+
export type PingRequest = Infer<typeof PingRequestSchema>;
11441157

11451158
/* Progress notifications */
1146-
export type Progress = z.infer<typeof ProgressSchema>;
1147-
export type ProgressNotification = z.infer<typeof ProgressNotificationSchema>;
1159+
export type Progress = Infer<typeof ProgressSchema>;
1160+
export type ProgressNotification = Infer<typeof ProgressNotificationSchema>;
11481161

11491162
/* Pagination */
1150-
export type PaginatedRequest = z.infer<typeof PaginatedRequestSchema>;
1151-
export type PaginatedResult = z.infer<typeof PaginatedResultSchema>;
1163+
export type PaginatedRequest = Infer<typeof PaginatedRequestSchema>;
1164+
export type PaginatedResult = Infer<typeof PaginatedResultSchema>;
11521165

11531166
/* Resources */
1154-
export type ResourceContents = z.infer<typeof ResourceContentsSchema>;
1155-
export type TextResourceContents = z.infer<typeof TextResourceContentsSchema>;
1156-
export type BlobResourceContents = z.infer<typeof BlobResourceContentsSchema>;
1157-
export type Resource = z.infer<typeof ResourceSchema>;
1158-
export type ResourceTemplate = z.infer<typeof ResourceTemplateSchema>;
1159-
export type ListResourcesRequest = z.infer<typeof ListResourcesRequestSchema>;
1160-
export type ListResourcesResult = z.infer<typeof ListResourcesResultSchema>;
1161-
export type ListResourceTemplatesRequest = z.infer<
1162-
typeof ListResourceTemplatesRequestSchema
1163-
>;
1164-
export type ListResourceTemplatesResult = z.infer<
1165-
typeof ListResourceTemplatesResultSchema
1166-
>;
1167-
export type ReadResourceRequest = z.infer<typeof ReadResourceRequestSchema>;
1168-
export type ReadResourceResult = z.infer<typeof ReadResourceResultSchema>;
1169-
export type ResourceListChangedNotification = z.infer<
1170-
typeof ResourceListChangedNotificationSchema
1171-
>;
1172-
export type SubscribeRequest = z.infer<typeof SubscribeRequestSchema>;
1173-
export type UnsubscribeRequest = z.infer<typeof UnsubscribeRequestSchema>;
1174-
export type ResourceUpdatedNotification = z.infer<
1175-
typeof ResourceUpdatedNotificationSchema
1176-
>;
1167+
export type ResourceContents = Infer<typeof ResourceContentsSchema>;
1168+
export type TextResourceContents = Infer<typeof TextResourceContentsSchema>;
1169+
export type BlobResourceContents = Infer<typeof BlobResourceContentsSchema>;
1170+
export type Resource = Infer<typeof ResourceSchema>;
1171+
export type ResourceTemplate = Infer<typeof ResourceTemplateSchema>;
1172+
export type ListResourcesRequest = Infer<typeof ListResourcesRequestSchema>;
1173+
export type ListResourcesResult = Infer<typeof ListResourcesResultSchema>;
1174+
export type ListResourceTemplatesRequest = Infer<typeof ListResourceTemplatesRequestSchema>;
1175+
export type ListResourceTemplatesResult = Infer<typeof ListResourceTemplatesResultSchema>;
1176+
export type ReadResourceRequest = Infer<typeof ReadResourceRequestSchema>;
1177+
export type ReadResourceResult = Infer<typeof ReadResourceResultSchema>;
1178+
export type ResourceListChangedNotification = Infer<typeof ResourceListChangedNotificationSchema>;
1179+
export type SubscribeRequest = Infer<typeof SubscribeRequestSchema>;
1180+
export type UnsubscribeRequest = Infer<typeof UnsubscribeRequestSchema>;
1181+
export type ResourceUpdatedNotification = Infer<typeof ResourceUpdatedNotificationSchema>;
11771182

11781183
/* Prompts */
1179-
export type PromptArgument = z.infer<typeof PromptArgumentSchema>;
1180-
export type Prompt = z.infer<typeof PromptSchema>;
1181-
export type ListPromptsRequest = z.infer<typeof ListPromptsRequestSchema>;
1182-
export type ListPromptsResult = z.infer<typeof ListPromptsResultSchema>;
1183-
export type GetPromptRequest = z.infer<typeof GetPromptRequestSchema>;
1184-
export type TextContent = z.infer<typeof TextContentSchema>;
1185-
export type ImageContent = z.infer<typeof ImageContentSchema>;
1186-
export type EmbeddedResource = z.infer<typeof EmbeddedResourceSchema>;
1187-
export type PromptMessage = z.infer<typeof PromptMessageSchema>;
1188-
export type GetPromptResult = z.infer<typeof GetPromptResultSchema>;
1189-
export type PromptListChangedNotification = z.infer<
1190-
typeof PromptListChangedNotificationSchema
1191-
>;
1184+
export type PromptArgument = Infer<typeof PromptArgumentSchema>;
1185+
export type Prompt = Infer<typeof PromptSchema>;
1186+
export type ListPromptsRequest = Infer<typeof ListPromptsRequestSchema>;
1187+
export type ListPromptsResult = Infer<typeof ListPromptsResultSchema>;
1188+
export type GetPromptRequest = Infer<typeof GetPromptRequestSchema>;
1189+
export type TextContent = Infer<typeof TextContentSchema>;
1190+
export type ImageContent = Infer<typeof ImageContentSchema>;
1191+
export type EmbeddedResource = Infer<typeof EmbeddedResourceSchema>;
1192+
export type PromptMessage = Infer<typeof PromptMessageSchema>;
1193+
export type GetPromptResult = Infer<typeof GetPromptResultSchema>;
1194+
export type PromptListChangedNotification = Infer<typeof PromptListChangedNotificationSchema>;
11921195

11931196
/* Tools */
1194-
export type Tool = z.infer<typeof ToolSchema>;
1195-
export type ListToolsRequest = z.infer<typeof ListToolsRequestSchema>;
1196-
export type ListToolsResult = z.infer<typeof ListToolsResultSchema>;
1197-
export type CallToolResult = z.infer<typeof CallToolResultSchema>;
1198-
export type CompatibilityCallToolResult = z.infer<
1199-
typeof CompatibilityCallToolResultSchema
1200-
>;
1201-
export type CallToolRequest = z.infer<typeof CallToolRequestSchema>;
1202-
export type ToolListChangedNotification = z.infer<
1203-
typeof ToolListChangedNotificationSchema
1204-
>;
1197+
export type Tool = Infer<typeof ToolSchema>;
1198+
export type ListToolsRequest = Infer<typeof ListToolsRequestSchema>;
1199+
export type ListToolsResult = Infer<typeof ListToolsResultSchema>;
1200+
export type CallToolResult = Infer<typeof CallToolResultSchema>;
1201+
export type CompatibilityCallToolResult = Infer<typeof CompatibilityCallToolResultSchema>;
1202+
export type CallToolRequest = Infer<typeof CallToolRequestSchema>;
1203+
export type ToolListChangedNotification = Infer<typeof ToolListChangedNotificationSchema>;
12051204

12061205
/* Logging */
1207-
export type LoggingLevel = z.infer<typeof LoggingLevelSchema>;
1208-
export type SetLevelRequest = z.infer<typeof SetLevelRequestSchema>;
1209-
export type LoggingMessageNotification = z.infer<
1210-
typeof LoggingMessageNotificationSchema
1211-
>;
1206+
export type LoggingLevel = Infer<typeof LoggingLevelSchema>;
1207+
export type SetLevelRequest = Infer<typeof SetLevelRequestSchema>;
1208+
export type LoggingMessageNotification = Infer<typeof LoggingMessageNotificationSchema>;
12121209

12131210
/* Sampling */
1214-
export type SamplingMessage = z.infer<typeof SamplingMessageSchema>;
1215-
export type CreateMessageRequest = z.infer<typeof CreateMessageRequestSchema>;
1216-
export type CreateMessageResult = z.infer<typeof CreateMessageResultSchema>;
1211+
export type SamplingMessage = Infer<typeof SamplingMessageSchema>;
1212+
export type CreateMessageRequest = Infer<typeof CreateMessageRequestSchema>;
1213+
export type CreateMessageResult = Infer<typeof CreateMessageResultSchema>;
12171214

12181215
/* Autocomplete */
1219-
export type ResourceReference = z.infer<typeof ResourceReferenceSchema>;
1220-
export type PromptReference = z.infer<typeof PromptReferenceSchema>;
1221-
export type CompleteRequest = z.infer<typeof CompleteRequestSchema>;
1222-
export type CompleteResult = z.infer<typeof CompleteResultSchema>;
1216+
export type ResourceReference = Infer<typeof ResourceReferenceSchema>;
1217+
export type PromptReference = Infer<typeof PromptReferenceSchema>;
1218+
export type CompleteRequest = Infer<typeof CompleteRequestSchema>;
1219+
export type CompleteResult = Infer<typeof CompleteResultSchema>;
12231220

12241221
/* Roots */
1225-
export type Root = z.infer<typeof RootSchema>;
1226-
export type ListRootsRequest = z.infer<typeof ListRootsRequestSchema>;
1227-
export type ListRootsResult = z.infer<typeof ListRootsResultSchema>;
1228-
export type RootsListChangedNotification = z.infer<
1229-
typeof RootsListChangedNotificationSchema
1230-
>;
1222+
export type Root = Infer<typeof RootSchema>;
1223+
export type ListRootsRequest = Infer<typeof ListRootsRequestSchema>;
1224+
export type ListRootsResult = Infer<typeof ListRootsResultSchema>;
1225+
export type RootsListChangedNotification = Infer<typeof RootsListChangedNotificationSchema>;
12311226

12321227
/* Client messages */
1233-
export type ClientRequest = z.infer<typeof ClientRequestSchema>;
1234-
export type ClientNotification = z.infer<typeof ClientNotificationSchema>;
1235-
export type ClientResult = z.infer<typeof ClientResultSchema>;
1228+
export type ClientRequest = Infer<typeof ClientRequestSchema>;
1229+
export type ClientNotification = Infer<typeof ClientNotificationSchema>;
1230+
export type ClientResult = Infer<typeof ClientResultSchema>;
12361231

12371232
/* Server messages */
1238-
export type ServerRequest = z.infer<typeof ServerRequestSchema>;
1239-
export type ServerNotification = z.infer<typeof ServerNotificationSchema>;
1240-
export type ServerResult = z.infer<typeof ServerResultSchema>;
1233+
export type ServerRequest = Infer<typeof ServerRequestSchema>;
1234+
export type ServerNotification = Infer<typeof ServerNotificationSchema>;
1235+
export type ServerResult = Infer<typeof ServerResultSchema>;

0 commit comments

Comments
 (0)