tmcp 1.4.0 → 1.6.0

package/src/index.js

@@ -1,21 +1,61 @@
 /**
  * @import { StandardSchemaV1 } from "@standard-schema/spec";
+ * @import SqidsType from "sqids";
  * @import { JSONRPCRequest, JSONRPCParams } from "json-rpc-2.0";
  * @import { ExtractURITemplateVariables } from "./internal/uri-template.js";
- * @import { CallToolResult, ReadResourceResult, GetPromptResult, ClientCapabilities as ClientCapabilitiesType } from "./validation/index.js";
- * @import { Tool, Completion, Prompt, Resource, ServerOptions, ServerInfo, SubscriptionsKeys, McpEvents } from "./internal/internal.js";
+ * @import { CallToolResult, ReadResourceResult, GetPromptResult, ClientCapabilities as ClientCapabilitiesType, JSONRPCRequest as JSONRPCRequestType, JSONRPCResponse, CreateMessageRequestParams, CreateMessageResult, Resource, LoggingLevel } from "./validation/index.js";
+ * @import { Tool, Completion, Prompt, StoredResource, ServerOptions, ServerInfo, SubscriptionsKeys, McpEvents } from "./internal/internal.js";
  */
-import { JSONRPCServer, JSONRPCClient } from 'json-rpc-2.0';
+import { JSONRPCClient, JSONRPCServer } from 'json-rpc-2.0';
+import { AsyncLocalStorage } from 'node:async_hooks';
 import { UriTemplateMatcher } from 'uri-template-matcher';
+import * as v from 'valibot';
 import {
 	CallToolResultSchema,
-	ReadResourceResultSchema,
-	GetPromptResultSchema,
 	CompleteResultSchema,
-	InitializeRequestSchema,
+	CreateMessageRequestParamsSchema,
+	CreateMessageResultSchema,
+	GetPromptResultSchema,
+	InitializeRequestParamsSchema,
+	JSONRPCRequestSchema,
+	JSONRPCResponseSchema,
+	McpError,
+	ReadResourceResultSchema,
 } from './validation/index.js';
-import { AsyncLocalStorage } from 'node:async_hooks';
-import * as v from 'valibot';
+import {
+	get_supported_versions,
+	negotiate_protocol_version,
+} from './validation/version.js';
+import { should_version_negotiation_fail } from './validation/version.js';
+
+/**
+ * @type {SqidsType | undefined}
+ */
+let Sqids;
+
+async function get_sqids() {
+	if (!Sqids) {
+		Sqids = new (await import('sqids')).default();
+	}
+	return Sqids;
+}
+
+/**
+ * Encode a cursor for pagination
+ * @param {number} offset
+ */
+async function encode_cursor(offset) {
+	return (await get_sqids()).encode([offset]);
+}
+
+/**
+ * Decode a cursor from pagination
+ * @param {string} cursor
+ */
+async function decode_cursor(cursor) {
+	const [decoded] = (await get_sqids()).decode(cursor);
+	return decoded;
+}
 
 /**
  * @typedef {ClientCapabilitiesType} ClientCapabilities
@@ -40,10 +80,14 @@ export class McpServer {
 	 */
 	#prompts = new Map();
 	/**
-	 * @type {Map<string, Resource>}
+	 * @type {Map<string, StoredResource>}
 	 */
 	#resources = new Map();
 	#templates = new UriTemplateMatcher();
+	/**
+	 * @type {Array<{uri: string, name?: string}>}
+	 */
+	roots = [];
 	/**
 	 * @type {{ [ref: string]: Map<string, Partial<Record<string, Completion>>> }}
 	 */
@@ -74,6 +118,16 @@ export class McpServer {
 	 */
 	#client_capabilities_map = new Map();
 
+	/**
+	 * @type {Map<string|undefined, string>}
+	 */
+	#negotiated_protocol_versions = new Map();
+
+	/**
+	 * @type {Map<string|undefined, LoggingLevel>}
+	 */
+	#session_log_levels = new Map();
+
 	/**
 	 * @param {ServerInfo} server_info
 	 * @param {ServerOptions<StandardSchema>} options
@@ -81,24 +135,90 @@ export class McpServer {
 	constructor(server_info, options) {
 		this.#options = options;
 		this.#server.addMethod('initialize', (initialize_request) => {
-			const validated_initialize = v.parse(
-				InitializeRequestSchema,
-				initialize_request,
-			);
-			this.#client_capabilities_map.set(
-				this.#session_id,
-				validated_initialize.capabilities,
-			);
-			this.#event_target.dispatchEvent(
-				new CustomEvent('initialize', {
-					detail: validated_initialize,
-				}),
-			);
-			return {
-				protocolVersion: validated_initialize.protocolVersion,
-				...options,
-				serverInfo: server_info,
-			};
+			try {
+				// Validate basic request format
+				const validated_initialize = v.parse(
+					InitializeRequestParamsSchema,
+					initialize_request,
+				);
+
+				const session_id = this.#session_id;
+
+				// Validate protocol version format
+				if (
+					should_version_negotiation_fail(
+						validated_initialize.protocolVersion,
+					)
+				) {
+					// Return JSON-RPC error for invalid protocol version format
+					const error = new McpError(
+						-32602,
+						'Invalid protocol version format',
+					);
+					throw error;
+				}
+
+				// Negotiate protocol version
+				const negotiated_version = negotiate_protocol_version(
+					validated_initialize.protocolVersion,
+				);
+
+				// Store negotiated protocol version
+				this.#negotiated_protocol_versions.set(
+					session_id,
+					negotiated_version,
+				);
+
+				// Store client capabilities
+				this.#client_capabilities_map.set(
+					session_id,
+					validated_initialize.capabilities,
+				);
+
+				// Dispatch initialization event
+				this.#event_target.dispatchEvent(
+					new CustomEvent('initialize', {
+						detail: validated_initialize,
+					}),
+				);
+
+				// Trigger initial roots request if client supports it
+				if (validated_initialize.capabilities?.roots) {
+					this.#refresh_roots();
+				}
+
+				// Return server response with negotiated version and capabilities
+				return {
+					protocolVersion: negotiated_version,
+					...options,
+					serverInfo: server_info,
+				};
+			} catch (error) {
+				// Enhanced error handling for initialization failures
+				if (error instanceof McpError) {
+					// Already has JSON-RPC error code, re-throw
+					throw error;
+				}
+
+				if (
+					/** @type {Error} */ (error).message?.includes(
+						'Protocol version',
+					)
+				) {
+					const rpc_error = new McpError(
+						-32602,
+						`Protocol version validation failed: ${/** @type {Error} */ (error).message}. Server supports: ${get_supported_versions().join(', ')}`,
+					);
+					throw rpc_error;
+				}
+
+				// General initialization error
+				const rpc_error = new McpError(
+					-32603,
+					`Initialization failed: ${/** @type {Error} */ (error).message}`,
+				);
+				throw rpc_error;
+			}
 		});
 		this.#server.addMethod('ping', () => {
 			return {};
@@ -109,7 +229,9 @@ export class McpServer {
 		this.#init_tools();
 		this.#init_prompts();
 		this.#init_resources();
+		this.#init_roots();
 		this.#init_completion();
+		this.#init_logging();
 	}
 
 	get #session_id() {
@@ -175,7 +297,7 @@ export class McpServer {
 				[...this.#tools].map(async ([name, tool]) => {
 					return {
 						name,
-						title: tool.description,
+						title: tool.title || tool.description,
 						description: tool.description,
 						inputSchema: tool.schema
 							? await this.#options.adapter.toJsonSchema(
@@ -219,8 +341,8 @@ export class McpServer {
 	 */
 	#init_prompts() {
 		if (!this.#options.capabilities?.prompts) return;
-		this.#server.addMethod('prompts/list', async () => {
-			const available_prompts = await Promise.all(
+		this.#server.addMethod('prompts/list', async ({ cursor } = {}) => {
+			const all_prompts = await Promise.all(
 				[...this.#prompts].map(async ([name, prompt]) => {
 					const arguments_schema = prompt.schema
 						? await this.#options.adapter.toJsonSchema(
@@ -236,7 +358,7 @@ export class McpServer {
 					const required = arguments_schema.required ?? [];
 					return {
 						name,
-						title: prompt.description,
+						title: prompt.title || prompt.description,
 						description: prompt.description,
 						arguments: keys.map((key) => {
 							const property = arguments_schema.properties?.[key];
@@ -253,8 +375,26 @@ export class McpServer {
 					};
 				}),
 			);
+
+			const pagination_options = this.#options.pagination?.prompts;
+			if (!pagination_options || pagination_options.size == null) {
+				return { prompts: all_prompts };
+			}
+
+			const page_length = pagination_options.size;
+			const offset = cursor ? await decode_cursor(cursor) : 0;
+			const start_index = offset;
+			const end_index = start_index + page_length;
+
+			const prompts = all_prompts.slice(start_index, end_index);
+			const has_next = end_index < all_prompts.length;
+			const next_cursor = has_next
+				? await encode_cursor(end_index)
+				: null;
+
 			return {
-				prompts: available_prompts,
+				prompts,
+				...(next_cursor && { nextCursor: next_cursor }),
 			};
 		});
 		this.#server.addMethod(
@@ -303,37 +443,61 @@ export class McpServer {
 			});
 		}
 
-		this.#server.addMethod('resources/list', async () => {
+		this.#server.addMethod('resources/list', async ({ cursor } = {}) => {
+			const all_resources = [];
+
+			// Add static resources
+			for (const [uri, { description, name, title, ...resource }] of this
+				.#resources) {
+				if (!resource.template) {
+					all_resources.push({
+						name,
+						title: title || description,
+						description,
+						uri,
+					});
+				} else if (resource.list_resources) {
+					const template_resources = await resource.list_resources();
+					all_resources.push(...template_resources);
+				}
+			}
+
+			const pagination_options = this.#options.pagination?.resources;
+			if (!pagination_options || pagination_options.size == null) {
+				return { resources: all_resources };
+			}
+
+			const page_length = pagination_options.size;
+			const offset = cursor ? await decode_cursor(cursor) : 0;
+			const start_index = offset;
+			const end_index = start_index + page_length;
+
+			const resources = all_resources.slice(start_index, end_index);
+			const has_next = end_index < all_resources.length;
+			const next_cursor = has_next
+				? await encode_cursor(end_index)
+				: null;
+
 			return {
-				resources: [...this.#resources].reduce(
-					(arr, [uri, { description, name, template }]) => {
-						if (!template) {
-							arr.push({
-								name,
-								description,
-								uri,
-							});
-						}
-						return arr;
-					},
-					/** @type {Array<{name: string, description: string, uri: string}>} */ ([]),
-				),
+				resources,
+				...(next_cursor && { nextCursor: next_cursor }),
 			};
 		});
 		this.#server.addMethod('resources/templates/list', async () => {
 			return {
 				resourceTemplates: [...this.#resources].reduce(
-					(arr, [uri, { description, name, template }]) => {
+					(arr, [uri, { description, name, title, template }]) => {
 						if (template) {
 							arr.push({
 								name,
+								title: title || description,
 								description,
 								uriTemplate: uri,
 							});
 						}
 						return arr;
 					},
-					/** @type {Array<{name: string, description: string, uriTemplate: string}>} */ ([]),
+					/** @type {Array<{name: string, title: string, description: string, uriTemplate: string}>} */ ([]),
 				),
 			};
 		});
@@ -364,6 +528,37 @@ export class McpServer {
 			);
 		});
 	}
+	/**
+	 *
+	 */
+	#init_roots() {
+		this.#server.addMethod('notifications/roots/list_changed', () => {
+			this.#refresh_roots();
+			return null;
+		});
+	}
+
+	/**
+	 * Request roots list from client
+	 */
+	async #refresh_roots() {
+		if (!this.#client_capabilities?.roots) return;
+
+		this.#lazyily_create_client();
+		try {
+			const response = await this.#client?.request(
+				'roots/list',
+				undefined,
+				{
+					sessions: [this.#session_id],
+				},
+			);
+			this.roots = response?.roots || [];
+		} catch {
+			// Client doesn't support roots or request failed
+			this.roots = [];
+		}
+	}
 	/**
 	 *
 	 */
@@ -384,27 +579,39 @@ export class McpServer {
 			},
 		);
 	}
+	/**
+	 *
+	 */
+	#init_logging() {
+		if (!this.#options.capabilities?.logging) return;
+
+		this.#server.addMethod('logging/setLevel', ({ level }) => {
+			this.#session_log_levels.set(this.#session_id, level);
+			return {};
+		});
+	}
 	/**
 	 * @template {StandardSchema | undefined} [TSchema=undefined]
-	 * @param {{ name: string; description: string; schema?: StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema> extends Record<string, unknown> ? TSchema : never }} options
+	 * @param {{ name: string; description: string; title?: string; schema?: StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema> extends Record<string, unknown> ? TSchema : never }} options
 	 * @param {TSchema extends undefined ? (()=>Promise<CallToolResult> | CallToolResult) : ((input: StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema>) => Promise<CallToolResult> | CallToolResult)} execute
 	 */
-	tool({ name, description, schema }, execute) {
+	tool({ name, description, title, schema }, execute) {
 		if (this.#options.capabilities?.tools?.listChanged) {
 			this.#notify('notifications/tools/list_changed', {});
 		}
 		this.#tools.set(name, {
 			description,
+			title,
 			schema,
 			execute,
 		});
 	}
 	/**
 	 * @template {StandardSchema | undefined} [TSchema=undefined]
-	 * @param {{ name: string; description: string; schema?: StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema> extends Record<string, unknown> ? TSchema : never; complete?: NoInfer<TSchema extends undefined ? never : Partial<Record<keyof (StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema>), Completion>>> }} options
+	 * @param {{ name: string; description: string; title?: string; schema?: StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema> extends Record<string, unknown> ? TSchema : never; complete?: NoInfer<TSchema extends undefined ? never : Partial<Record<keyof (StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema>), Completion>>> }} options
 	 * @param {TSchema extends undefined ? (()=>Promise<GetPromptResult> | GetPromptResult) : (input: StandardSchemaV1.InferInput<TSchema extends undefined ? never : TSchema>) => Promise<GetPromptResult> | GetPromptResult} execute
 	 */
-	prompt({ name, description, schema, complete }, execute) {
+	prompt({ name, description, title, schema, complete }, execute) {
 		if (complete) {
 			this.#completions['ref/prompt'].set(name, complete);
 		}
@@ -413,12 +620,13 @@ export class McpServer {
 		}
 		this.#prompts.set(name, {
 			description,
+			title,
 			schema,
 			execute,
 		});
 	}
 	/**
-	 * @type {(resource: Resource & {complete?: Partial<Record<string,Completion>>, uri: string})=> void}
+	 * @type {(resource: StoredResource & {complete?: Partial<Record<string,Completion>>, uri: string})=> void}
 	 */
 	#resource({ uri, complete, ...resource }) {
 		if (complete) {
@@ -433,13 +641,14 @@ export class McpServer {
 		this.#resources.set(uri, resource);
 	}
 	/**
-	 * @param {{ name: string; description: string; uri: string }} options
+	 * @param {{ name: string; description: string; title?: string; uri: string }} options
 	 * @param {(uri: string) => Promise<ReadResourceResult> | ReadResourceResult} execute
 	 */
-	resource({ name, description, uri }, execute) {
+	resource({ name, description, title, uri }, execute) {
 		this.#resource({
 			name,
 			description,
+			title,
 			uri,
 			execute,
 			template: false,
@@ -448,28 +657,58 @@ export class McpServer {
 	/**
 	 * @template {string} TUri
 	 * @template {ExtractURITemplateVariables<TUri>} TVariables
-	 * @param {{ name: string; description: string; uri: TUri; complete?: NoInfer<TVariables extends never ? never : Partial<Record<TVariables, Completion>>> }} options
+	 * @param {{ name: string; description: string; title?: string; uri: TUri; complete?: NoInfer<TVariables extends never ? never : Partial<Record<TVariables, Completion>>>; list?: () => Promise<Array<Resource>> | Array<Resource> }} options
 	 * @param {(uri: string, params: Record<TVariables, string | string[]>) => Promise<ReadResourceResult> | ReadResourceResult} execute
 	 */
-	template({ name, description, uri, complete }, execute) {
+	template(
+		{ name, description, title, uri, complete, list: list_resources },
+		execute,
+	) {
 		this.#resource({
 			name,
 			description,
+			title,
 			uri,
 			execute,
 			complete,
+			list_resources,
 			template: true,
 		});
 	}
 	/**
-	 * @param {JSONRPCRequest} request
+	 * @param {JSONRPCResponse | JSONRPCRequest} message
 	 * @param {string} [session_id]
-	 * @returns {ReturnType<JSONRPCServer['receive']>}
+	 * @returns {ReturnType<JSONRPCServer['receive']> | ReturnType<JSONRPCClient['receive'] | undefined>}
 	 */
-	receive(request, session_id) {
-		return this.#session_storage.run(
-			session_id,
-			async () => await this.#server.receive(request),
+	receive(message, session_id) {
+		// Validate the message first
+		const validated_message = v.safeParse(JSONRPCRequestSchema, message);
+
+		// Check if it's a request or response
+		if (validated_message.success) {
+			// we want to set a default log level for each session
+			// so that if the user calls `log` it can send a notification
+			// to each client that didn't explicitly set a log level too
+			if (
+				!!this.#options.capabilities?.logging &&
+				!this.#session_log_levels.has(session_id)
+			) {
+				this.#session_log_levels.set(
+					session_id,
+					this.#options.logging?.default ?? 'info',
+				);
+			}
+			return this.#session_storage.run(
+				session_id,
+				async () =>
+					await this.#server.receive(validated_message.output),
+			);
+		}
+		// It's a response - handle with client
+		const validated_response = v.parse(JSONRPCResponseSchema, message);
+		this.#lazyily_create_client();
+		return this.#session_storage.run(session_id, async () =>
+			this.#client?.receive(validated_response),
 		);
 	}
 
@@ -496,6 +735,13 @@ export class McpServer {
 		}
 	}
 
+	/**
+	 * Refresh roots list from client
+	 */
+	async refreshRoots() {
+		await this.#refresh_roots();
+	}
+
 	/**
 	 * @template {StandardSchema} TSchema
 	 * @param {TSchema} schema
@@ -526,4 +772,96 @@ export class McpServer {
 		}
 		return validated_result.value;
 	}
+
+	/**
+	 * Request language model sampling from the client
+	 * @param {CreateMessageRequestParams} request
+	 * @returns {Promise<CreateMessageResult>}
+	 */
+	async message(request) {
+		if (!this.#client_capabilities?.sampling)
+			throw new Error("Client doesn't support sampling");
+
+		this.#lazyily_create_client();
+
+		// Validate the request
+		const validated_request = v.parse(
+			CreateMessageRequestParamsSchema,
+			request,
+		);
+
+		// Make the request to the client
+		const response = await this.#client?.request(
+			'sampling/createMessage',
+			validated_request,
+			{
+				sessions: [this.#session_id],
+			},
+		);
+
+		// Validate and return the response
+		return v.parse(CreateMessageResultSchema, response);
+	}
+
+	/**
+	 * Log a message to the client if logging is enabled and the level is appropriate
+	 * @param {LoggingLevel} level
+	 * @param {unknown} data
+	 * @param {string} [logger]
+	 */
+	log(level, data, logger) {
+		if (!this.#options.capabilities?.logging) {
+			throw new Error(
+				"The server doesn't support logging, please enable it in capabilities",
+			);
+		}
+
+		const sessions = [];
+
+		for (let [session_id, session_log_level] of this.#session_log_levels) {
+			// Check if the current log level should be sent based on severity
+			if (this.#should_log(level, session_log_level)) {
+				sessions.push(session_id);
+			}
+		}
+
+		if (sessions.length === 0) return;
+
+		// Send the log notification to the client
+		this.#notify(
+			'notifications/message',
+			{
+				level,
+				data,
+				logger,
+			},
+			{
+				sessions,
+			},
+		);
+	}
+
+	/**
+	 * Check if a log message should be sent based on severity levels
+	 * @param {LoggingLevel} message_level
+	 * @param {LoggingLevel} session_level
+	 * @returns {boolean}
+	 */
+	#should_log(message_level, session_level) {
+		const levels = [
+			'debug',
+			'info',
+			'notice',
+			'warning',
+			'error',
+			'critical',
+			'alert',
+			'emergency',
+		];
+		const message_severity = levels.indexOf(message_level);
+		const session_severity = levels.indexOf(session_level);
+
+		// Send if message severity is equal to or higher than session level
+		return message_severity >= session_severity;
+	}
 }