feat(mcp): add configurable request timeout to MCP server (#134)
Adds the ability to configure request timeouts for MCP server operations. This enhancement allows setting a custom duration to wait for server responses before timing out requests, which applies to all requests made through the client including tool calls, resource access, and prompt operations. - Add requestTimeout parameter to McpServerSession constructor - Add requestTimeout field and builder method to server classes - Pass timeout configuration through to session creation - Add tests for both success and failure scenarios across different transport implementations - Default timeout is set to 10 seconds if not explicitly configured.
This commit is contained in:
parent
263e3741b2
commit
472f07ec6d
3 changed files with 491 additions and 12 deletions
|
|
@ -4,6 +4,7 @@
|
|||
|
||||
package io.modelcontextprotocol.server;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
|
@ -90,8 +91,8 @@ public class McpAsyncServer {
|
|||
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
|
||||
*/
|
||||
McpAsyncServer(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
|
||||
McpServerFeatures.Async features) {
|
||||
this.delegate = new AsyncServerImpl(mcpTransportProvider, objectMapper, features);
|
||||
McpServerFeatures.Async features, Duration requestTimeout) {
|
||||
this.delegate = new AsyncServerImpl(mcpTransportProvider, objectMapper, requestTimeout, features);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -271,7 +272,7 @@ public class McpAsyncServer {
|
|||
private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);
|
||||
|
||||
AsyncServerImpl(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
|
||||
McpServerFeatures.Async features) {
|
||||
Duration requestTimeout, McpServerFeatures.Async features) {
|
||||
this.mcpTransportProvider = mcpTransportProvider;
|
||||
this.objectMapper = objectMapper;
|
||||
this.serverInfo = features.serverInfo();
|
||||
|
|
@ -330,9 +331,9 @@ public class McpAsyncServer {
|
|||
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED,
|
||||
asyncRootsListChangedNotificationHandler(rootsChangeConsumers));
|
||||
|
||||
mcpTransportProvider
|
||||
.setSessionFactory(transport -> new McpServerSession(UUID.randomUUID().toString(), transport,
|
||||
this::asyncInitializeRequestHandler, Mono::empty, requestHandlers, notificationHandlers));
|
||||
mcpTransportProvider.setSessionFactory(
|
||||
transport -> new McpServerSession(UUID.randomUUID().toString(), requestTimeout, transport,
|
||||
this::asyncInitializeRequestHandler, Mono::empty, requestHandlers, notificationHandlers));
|
||||
}
|
||||
|
||||
// ---------------------------------------
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
|
||||
package io.modelcontextprotocol.server;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.HashMap;
|
||||
|
|
@ -193,11 +194,28 @@ public interface McpServer {
|
|||
|
||||
private final List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeHandlers = new ArrayList<>();
|
||||
|
||||
private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout
|
||||
|
||||
private AsyncSpecification(McpServerTransportProvider transportProvider) {
|
||||
Assert.notNull(transportProvider, "Transport provider must not be null");
|
||||
this.transportProvider = transportProvider;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the duration to wait for server responses before timing out requests. This
|
||||
* timeout applies to all requests made through the client, including tool calls,
|
||||
* resource access, and prompt operations.
|
||||
* @param requestTimeout The duration to wait before timing out requests. Must not
|
||||
* be null.
|
||||
* @return This builder instance for method chaining
|
||||
* @throws IllegalArgumentException if requestTimeout is null
|
||||
*/
|
||||
public AsyncSpecification requestTimeout(Duration requestTimeout) {
|
||||
Assert.notNull(requestTimeout, "Request timeout must not be null");
|
||||
this.requestTimeout = requestTimeout;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the server implementation information that will be shared with clients
|
||||
* during connection initialization. This helps with version compatibility,
|
||||
|
|
@ -565,7 +583,7 @@ public interface McpServer {
|
|||
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools,
|
||||
this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers, this.instructions);
|
||||
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
|
||||
return new McpAsyncServer(this.transportProvider, mapper, features);
|
||||
return new McpAsyncServer(this.transportProvider, mapper, features, this.requestTimeout);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -619,11 +637,28 @@ public interface McpServer {
|
|||
|
||||
private final List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeHandlers = new ArrayList<>();
|
||||
|
||||
private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout
|
||||
|
||||
private SyncSpecification(McpServerTransportProvider transportProvider) {
|
||||
Assert.notNull(transportProvider, "Transport provider must not be null");
|
||||
this.transportProvider = transportProvider;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the duration to wait for server responses before timing out requests. This
|
||||
* timeout applies to all requests made through the client, including tool calls,
|
||||
* resource access, and prompt operations.
|
||||
* @param requestTimeout The duration to wait before timing out requests. Must not
|
||||
* be null.
|
||||
* @return This builder instance for method chaining
|
||||
* @throws IllegalArgumentException if requestTimeout is null
|
||||
*/
|
||||
public SyncSpecification requestTimeout(Duration requestTimeout) {
|
||||
Assert.notNull(requestTimeout, "Request timeout must not be null");
|
||||
this.requestTimeout = requestTimeout;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the server implementation information that will be shared with clients
|
||||
* during connection initialization. This helps with version compatibility,
|
||||
|
|
@ -992,7 +1027,7 @@ public interface McpServer {
|
|||
this.instructions);
|
||||
McpServerFeatures.Async asyncFeatures = McpServerFeatures.Async.fromSync(syncFeatures);
|
||||
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
|
||||
var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures);
|
||||
var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures, this.requestTimeout);
|
||||
|
||||
return new McpSyncServer(asyncServer);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -27,6 +27,9 @@ public class McpServerSession implements McpSession {
|
|||
|
||||
private final String id;
|
||||
|
||||
/** Duration to wait for request responses before timing out */
|
||||
private final Duration requestTimeout;
|
||||
|
||||
private final AtomicLong requestCounter = new AtomicLong(0);
|
||||
|
||||
private final InitRequestHandler initRequestHandler;
|
||||
|
|
@ -65,10 +68,11 @@ public class McpServerSession implements McpSession {
|
|||
* @param requestHandlers map of request handlers to use
|
||||
* @param notificationHandlers map of notification handlers to use
|
||||
*/
|
||||
public McpServerSession(String id, McpServerTransport transport, InitRequestHandler initHandler,
|
||||
InitNotificationHandler initNotificationHandler, Map<String, RequestHandler<?>> requestHandlers,
|
||||
Map<String, NotificationHandler> notificationHandlers) {
|
||||
public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
|
||||
InitRequestHandler initHandler, InitNotificationHandler initNotificationHandler,
|
||||
Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers) {
|
||||
this.id = id;
|
||||
this.requestTimeout = requestTimeout;
|
||||
this.transport = transport;
|
||||
this.initRequestHandler = initHandler;
|
||||
this.initNotificationHandler = initNotificationHandler;
|
||||
|
|
@ -116,7 +120,7 @@ public class McpServerSession implements McpSession {
|
|||
this.pendingResponses.remove(requestId);
|
||||
sink.error(error);
|
||||
});
|
||||
}).timeout(Duration.ofSeconds(10)).handle((jsonRpcResponse, sink) -> {
|
||||
}).timeout(requestTimeout).handle((jsonRpcResponse, sink) -> {
|
||||
if (jsonRpcResponse.error() != null) {
|
||||
sink.error(new McpError(jsonRpcResponse.error()));
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue