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:
mipengcheng3 2025-04-08 17:53:59 +08:00 • committed by Christian Tzolov
commit 472f07ec6d
3 changed files with 491 additions and 12 deletions

View file

@ -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));
}
// ---------------------------------------

View file

@ -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);
}

View file

@ -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()));
}