Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

feat: Add RequestIdGenerator for custom request ID generation #732

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rhtnr wants to merge 1 commit into modelcontextprotocol:main
base: main
Choose a base branch
Loading
from rhtnr:feature/custom-request-id-generator
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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);
Expand Down
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand All @@ -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);
}

}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand All @@ -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);
}

}
Expand Down
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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");
Expand All @@ -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();
}
Expand Down Expand Up @@ -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();
}

/**
Expand Down
View file Open in desktop
Original file line number Diff line number Diff line change
@@ -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());
}

}
Loading

AltStyle によって変換されたページ (->オリジナル) /