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
2 changed files with 491 additions and 12 deletions

View file

@ -8,6 +8,7 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.stream.Collectors;
@ -48,6 +49,7 @@ import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.server.RouterFunctions;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
import static org.awaitility.Awaitility.await;
import static org.mockito.Mockito.mock;
@ -196,6 +198,153 @@ class WebFluxSseIntegrationTests {
mcpServer.close();
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "httpclient", "webflux" })
void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws InterruptedException {
// Client
var clientBuilder = clientBuilders.get(clientType);
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
assertThat(request.messages()).hasSize(1);
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
try {
TimeUnit.SECONDS.sleep(2);
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
CreateMessageResult.StopReason.STOP_SEQUENCE);
};
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.build();
// Server
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
null);
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
new McpSchema.TextContent("Test message"))))
.modelPreferences(ModelPreferences.builder()
.hints(List.of())
.costPriority(1.0)
.speedPriority(1.0)
.intelligencePriority(1.0)
.build())
.build();
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
assertThat(result).isNotNull();
assertThat(result.role()).isEqualTo(Role.USER);
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
assertThat(result.model()).isEqualTo("MockModelName");
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
}).verifyComplete();
return Mono.just(callResponse);
});
var mcpServer = McpServer.async(mcpServerTransportProvider)
.requestTimeout(Duration.ofSeconds(4))
.serverInfo("test-server", "1.0.0")
.tools(tool)
.build();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
assertThat(response).isNotNull();
assertThat(response).isEqualTo(callResponse);
mcpClient.close();
mcpServer.close();
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "httpclient", "webflux" })
void testCreateMessageWithRequestTimeoutFail(String clientType) throws InterruptedException {
// Client
var clientBuilder = clientBuilders.get(clientType);
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
assertThat(request.messages()).hasSize(1);
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
try {
TimeUnit.SECONDS.sleep(3);
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
CreateMessageResult.StopReason.STOP_SEQUENCE);
};
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.build();
// Server
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
null);
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
new McpSchema.TextContent("Test message"))))
.modelPreferences(ModelPreferences.builder()
.hints(List.of())
.costPriority(1.0)
.speedPriority(1.0)
.intelligencePriority(1.0)
.build())
.build();
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
assertThat(result).isNotNull();
assertThat(result.role()).isEqualTo(Role.USER);
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
assertThat(result.model()).isEqualTo("MockModelName");
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
}).verifyComplete();
return Mono.just(callResponse);
});
var mcpServer = McpServer.async(mcpServerTransportProvider)
.requestTimeout(Duration.ofSeconds(1))
.serverInfo("test-server", "1.0.0")
.tools(tool)
.build();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
}).withMessageContaining("Timeout");
mcpClient.close();
mcpServer.close();
}
// ---------------------------------------
// Roots Tests
// ---------------------------------------

View file

@ -6,6 +6,7 @@ package io.modelcontextprotocol.server;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
@ -41,6 +42,7 @@ import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerResponse;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
import static org.awaitility.Awaitility.await;
import static org.mockito.Mockito.mock;
@ -212,6 +214,149 @@ class WebMvcSseIntegrationTests {
mcpServer.close();
}
@Test
void testCreateMessageWithRequestTimeoutSuccess() throws InterruptedException {
// Client
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
assertThat(request.messages()).hasSize(1);
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
try {
TimeUnit.SECONDS.sleep(2);
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
CreateMessageResult.StopReason.STOP_SEQUENCE);
};
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.build();
// Server
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
null);
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
new McpSchema.TextContent("Test message"))))
.modelPreferences(ModelPreferences.builder()
.hints(List.of())
.costPriority(1.0)
.speedPriority(1.0)
.intelligencePriority(1.0)
.build())
.build();
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
assertThat(result).isNotNull();
assertThat(result.role()).isEqualTo(Role.USER);
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
assertThat(result.model()).isEqualTo("MockModelName");
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
}).verifyComplete();
return Mono.just(callResponse);
});
var mcpServer = McpServer.async(mcpServerTransportProvider)
.serverInfo("test-server", "1.0.0")
.requestTimeout(Duration.ofSeconds(4))
.tools(tool)
.build();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
assertThat(response).isNotNull();
assertThat(response).isEqualTo(callResponse);
mcpClient.close();
mcpServer.close();
}
@Test
void testCreateMessageWithRequestTimeoutFail() throws InterruptedException {
// Client
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
assertThat(request.messages()).hasSize(1);
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
try {
TimeUnit.SECONDS.sleep(2);
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
CreateMessageResult.StopReason.STOP_SEQUENCE);
};
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.build();
// Server
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
null);
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
new McpSchema.TextContent("Test message"))))
.modelPreferences(ModelPreferences.builder()
.hints(List.of())
.costPriority(1.0)
.speedPriority(1.0)
.intelligencePriority(1.0)
.build())
.build();
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
assertThat(result).isNotNull();
assertThat(result.role()).isEqualTo(Role.USER);
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
assertThat(result.model()).isEqualTo("MockModelName");
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
}).verifyComplete();
return Mono.just(callResponse);
});
var mcpServer = McpServer.async(mcpServerTransportProvider)
.serverInfo("test-server", "1.0.0")
.requestTimeout(Duration.ofSeconds(1))
.tools(tool)
.build();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
}).withMessageContaining("Timeout");
mcpClient.close();
mcpServer.close();
}
// ---------------------------------------
// Roots Tests
// ---------------------------------------