feat: Add RequestIdGenerator for custom request ID generation

mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

@@ -23,6 +23,7 @@
 import io.modelcontextprotocol.spec.McpClientSession.RequestHandler;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.spec.RequestIdGenerator;
 import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
 import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
 import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
@@ -182,6 +183,24 @@ public class McpAsyncClient {
 	 */
 	McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout,
 			JsonSchemaValidator jsonSchemaValidator, McpClientFeatures.Async features) {
+		this(transport, requestTimeout, initializationTimeout, jsonSchemaValidator, features, null);
+	}
+
+	/**
+	 * Create a new McpAsyncClient with the given transport and session request-response
+	 * timeout.
+	 * @param transport the transport to use.
+	 * @param requestTimeout the session request-response timeout.
+	 * @param initializationTimeout the max timeout to await for the client-server
+	 * @param jsonSchemaValidator the JSON schema validator to use for validating tool
+	 * @param features the MCP Client supported features. responses against output
+	 * schemas.
+	 * @param requestIdGenerator the generator for creating unique request IDs. If null, a
+	 * default generator will be used.
+	 */
+	McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout,
+			JsonSchemaValidator jsonSchemaValidator, McpClientFeatures.Async features,
+			RequestIdGenerator requestIdGenerator) {
 
 		Assert.notNull(transport, "Transport must not be null");
 		Assert.notNull(requestTimeout, "Request timeout must not be null");
@@ -315,9 +334,11 @@ public class McpAsyncClient {
 			}).then();
 		};
 
+		RequestIdGenerator effectiveIdGenerator = requestIdGenerator != null ? requestIdGenerator
+				: RequestIdGenerator.ofDefault();
 		this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo, transport.protocolVersions(),
 				initializationTimeout, ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
-						notificationHandlers, con -> con.contextWrite(ctx)),
+						notificationHandlers, con -> con.contextWrite(ctx), effectiveIdGenerator),
 				postInitializationHook);
 
 		this.transport.setExceptionHandler(this.initializer::handleException);

mcp-core/src/main/java/io/modelcontextprotocol/client/McpClient.java

@@ -17,6 +17,7 @@
 import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.spec.RequestIdGenerator;
 import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
 import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
 import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
@@ -193,6 +194,8 @@ class SyncSpec {
 
 		private boolean enableCallToolSchemaCaching = false; // Default to false
 
+		private RequestIdGenerator requestIdGenerator;
+
 		private SyncSpec(McpClientTransport transport) {
 			Assert.notNull(transport, "Transport must not be null");
 			this.transport = transport;
@@ -461,6 +464,31 @@ public SyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching)
 			return this;
 		}
 
+		/**
+		 * Sets a custom request ID generator for creating unique request IDs. This is
+		 * useful for MCP servers that require specific ID formats, such as numeric-only
+		 * IDs.
+		 *
+		 * <p>
+		 * Example usage with a numeric ID generator:
+		 *
+		 * <pre>{@code
+		 * AtomicLong counter = new AtomicLong(0);
+		 * McpClient.sync(transport)
+		 *     .requestIdGenerator(() -> String.valueOf(counter.incrementAndGet()))
+		 *     .build();
+		 * }</pre>
+		 * @param requestIdGenerator The generator for creating unique request IDs. If
+		 * null, a default UUID-prefixed generator will be used.
+		 * @return This builder instance for method chaining
+		 * @see RequestIdGenerator#ofIncremental()
+		 * @see RequestIdGenerator#ofDefault()
+		 */
+		public SyncSpec requestIdGenerator(RequestIdGenerator requestIdGenerator) {
+			this.requestIdGenerator = requestIdGenerator;
+			return this;
+		}
+
 		/**
 		 * Create an instance of {@link McpSyncClient} with the provided configurations or
 		 * sensible defaults.
@@ -475,8 +503,8 @@ public McpSyncClient build() {
 			McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
 
 			return new McpSyncClient(new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout,
-					jsonSchemaValidator != null ? jsonSchemaValidator : JsonSchemaValidator.getDefault(),
-					asyncFeatures), this.contextProvider);
+					jsonSchemaValidator != null ? jsonSchemaValidator : JsonSchemaValidator.getDefault(), asyncFeatures,
+					this.requestIdGenerator), this.contextProvider);
 		}
 
 	}
@@ -531,6 +559,8 @@ class AsyncSpec {
 
 		private boolean enableCallToolSchemaCaching = false; // Default to false
 
+		private RequestIdGenerator requestIdGenerator;
+
 		private AsyncSpec(McpClientTransport transport) {
 			Assert.notNull(transport, "Transport must not be null");
 			this.transport = transport;
@@ -802,6 +832,31 @@ public AsyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching
 			return this;
 		}
 
+		/**
+		 * Sets a custom request ID generator for creating unique request IDs. This is
+		 * useful for MCP servers that require specific ID formats, such as numeric-only
+		 * IDs.
+		 *
+		 * <p>
+		 * Example usage with a numeric ID generator:
+		 *
+		 * <pre>{@code
+		 * AtomicLong counter = new AtomicLong(0);
+		 * McpClient.async(transport)
+		 *     .requestIdGenerator(() -> String.valueOf(counter.incrementAndGet()))
+		 *     .build();
+		 * }</pre>
+		 * @param requestIdGenerator The generator for creating unique request IDs. If
+		 * null, a default UUID-prefixed generator will be used.
+		 * @return This builder instance for method chaining
+		 * @see RequestIdGenerator#ofIncremental()
+		 * @see RequestIdGenerator#ofDefault()
+		 */
+		public AsyncSpec requestIdGenerator(RequestIdGenerator requestIdGenerator) {
+			this.requestIdGenerator = requestIdGenerator;
+			return this;
+		}
+
 		/**
 		 * Create an instance of {@link McpAsyncClient} with the provided configurations
 		 * or sensible defaults.
@@ -815,7 +870,8 @@ public McpAsyncClient build() {
 					new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots,
 							this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
 							this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers,
-							this.samplingHandler, this.elicitationHandler, this.enableCallToolSchemaCaching));
+							this.samplingHandler, this.elicitationHandler, this.enableCallToolSchemaCaching),
+					this.requestIdGenerator);
 		}
 
 	}

mcp-core/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java

@@ -14,9 +14,7 @@
 
 import java.time.Duration;
 import java.util.Map;
-import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
 
 /**
@@ -56,11 +54,8 @@ public class McpClientSession implements McpSession {
 	/** Map of notification handlers keyed by method name */
 	private final ConcurrentHashMap<String, NotificationHandler> notificationHandlers = new ConcurrentHashMap<>();
 
-	/** Session-specific prefix for request IDs */
-	private final String sessionPrefix = UUID.randomUUID().toString().substring(0, 8);
-
-	/** Atomic counter for generating unique request IDs */
-	private final AtomicLong requestCounter = new AtomicLong(0);
+	/** Generator for creating unique request IDs */
+	private final RequestIdGenerator requestIdGenerator;
 
 	/**
 	 * Functional interface for handling incoming JSON-RPC requests. Implementations
@@ -123,6 +118,26 @@ public McpClientSession(Duration requestTimeout, McpClientTransport transport,
 	public McpClientSession(Duration requestTimeout, McpClientTransport transport,
 			Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers,
 			Function<? super Mono<Void>, ? extends Publisher<Void>> connectHook) {
+		this(requestTimeout, transport, requestHandlers, notificationHandlers, connectHook,
+				RequestIdGenerator.ofDefault());
+	}
+
+	/**
+	 * Creates a new McpClientSession with the specified configuration, handlers, and
+	 * custom request ID generator.
+	 * @param requestTimeout Duration to wait for responses
+	 * @param transport Transport implementation for message exchange
+	 * @param requestHandlers Map of method names to request handlers
+	 * @param notificationHandlers Map of method names to notification handlers
+	 * @param connectHook Hook that allows transforming the connection Publisher prior to
+	 * subscribing
+	 * @param requestIdGenerator Generator for creating unique request IDs. If null, a
+	 * default generator will be used.
+	 */
+	public McpClientSession(Duration requestTimeout, McpClientTransport transport,
+			Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers,
+			Function<? super Mono<Void>, ? extends Publisher<Void>> connectHook,
+			RequestIdGenerator requestIdGenerator) {
 
 		Assert.notNull(requestTimeout, "The requestTimeout can not be null");
 		Assert.notNull(transport, "The transport can not be null");
@@ -133,6 +148,7 @@ public McpClientSession(Duration requestTimeout, McpClientTransport transport,
 		this.transport = transport;
 		this.requestHandlers.putAll(requestHandlers);
 		this.notificationHandlers.putAll(notificationHandlers);
+		this.requestIdGenerator = requestIdGenerator != null ? requestIdGenerator : RequestIdGenerator.ofDefault();
 
 		this.transport.connect(mono -> mono.doOnNext(this::handle)).transform(connectHook).subscribe();
 	}
@@ -243,12 +259,11 @@ private Mono<Void> handleIncomingNotification(McpSchema.JSONRPCNotification noti
 	}
 
 	/**
-	 * Generates a unique request ID in a non-blocking way. Combines a session-specific
-	 * prefix with an atomic counter to ensure uniqueness.
+	 * Generates a unique request ID using the configured request ID generator.
 	 * @return A unique request ID string
 	 */
 	private String generateRequestId() {
-		return this.sessionPrefix + "-" + this.requestCounter.getAndIncrement();
+		return this.requestIdGenerator.generate();
 	}
 
 	/**

mcp-core/src/main/java/io/modelcontextprotocol/spec/RequestIdGenerator.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2024-2025 the original author or authors.
+ */
+
+package io.modelcontextprotocol.spec;
+
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * A generator for creating unique request IDs for JSON-RPC messages.
+ *
+ * <p>
+ * Implementations of this interface are responsible for generating unique IDs that are
+ * used to correlate requests with their corresponding responses in JSON-RPC
+ * communication.
+ *
+ * <p>
+ * The MCP specification requires that:
+ * <ul>
+ * <li>Request IDs MUST be a string or integer</li>
+ * <li>Request IDs MUST NOT be null</li>
+ * <li>Request IDs MUST NOT have been previously used within the same session</li>
+ * </ul>
+ *
+ * <p>
+ * Example usage with a simple numeric ID generator:
+ *
+ * <pre>{@code
+ * AtomicLong counter = new AtomicLong(0);
+ * RequestIdGenerator generator = () -> String.valueOf(counter.incrementAndGet());
+ * }</pre>
+ *
+ * @author Christian Tzolov
+ * @see McpClientSession
+ */
+@FunctionalInterface
+public interface RequestIdGenerator {
+
+	/**
+	 * Generates a unique request ID.
+	 *
+	 * <p>
+	 * The generated ID must be unique within the session and must not be null.
+	 * Implementations should ensure thread-safety if the generator may be called from
+	 * multiple threads.
+	 * @return a unique request ID as a String
+	 */
+	String generate();
+
+	/**
+	 * Creates a default request ID generator that produces UUID-prefixed incrementing
+	 * IDs.
+	 *
+	 * <p>
+	 * The generated IDs follow the format: {@code <8-char-uuid>-<counter>}, for example:
+	 * {@code "a1b2c3d4-0"}, {@code "a1b2c3d4-1"}, etc.
+	 * @return a new default request ID generator
+	 */
+	static RequestIdGenerator ofDefault() {
+		String sessionPrefix = UUID.randomUUID().toString().substring(0, 8);
+		AtomicLong counter = new AtomicLong(0);
+		return () -> sessionPrefix + "-" + counter.getAndIncrement();
+	}
+
+	/**
+	 * Creates a request ID generator that produces simple incrementing numeric IDs.
+	 *
+	 * <p>
+	 * This generator is useful for MCP servers that require strictly numeric request IDs
+	 * (such as the Snowflake MCP server).
+	 *
+	 * <p>
+	 * The generated IDs are: {@code "1"}, {@code "2"}, {@code "3"}, etc.
+	 * @return a new numeric request ID generator
+	 */
+	static RequestIdGenerator ofIncremental() {
+		AtomicLong counter = new AtomicLong(0);
+		return () -> String.valueOf(counter.incrementAndGet());
+	}
+
+}