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
2 changed files with 491 additions and 12 deletions
|
|
@ -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
|
||||
// ---------------------------------------
|
||||
|
|
|
|||
|
|
@ -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
|
||||
// ---------------------------------------
|
||||
|
|
|
|||
Loading…
Reference in a new issue