Skip to content

Commit 472f07e

Browse files
mipengcheng3tzolov
mipengcheng3
authored and
tzolov
committed
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.
1 parent 263e374 commit 472f07e

File tree

6 files changed

+491
-12
lines changed

6 files changed

+491
-12
lines changed

‎mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java

+149
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import java.util.List;
99
import java.util.Map;
1010
import java.util.concurrent.ConcurrentHashMap;
11+
import java.util.concurrent.TimeUnit;
1112
import java.util.concurrent.atomic.AtomicReference;
1213
import java.util.function.Function;
1314
import java.util.stream.Collectors;
@@ -48,6 +49,7 @@
4849
import org.springframework.web.reactive.function.server.RouterFunctions;
4950

5051
import static org.assertj.core.api.Assertions.assertThat;
52+
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
5153
import static org.awaitility.Awaitility.await;
5254
import static org.mockito.Mockito.mock;
5355

@@ -196,6 +198,153 @@ void testCreateMessageSuccess(String clientType) {
196198
mcpServer.close();
197199
}
198200

201+
@ParameterizedTest(name = "{0} : {displayName} ")
202+
@ValueSource(strings = { "httpclient", "webflux" })
203+
void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws InterruptedException {
204+
205+
// Client
206+
var clientBuilder = clientBuilders.get(clientType);
207+
208+
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
209+
assertThat(request.messages()).hasSize(1);
210+
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
211+
try {
212+
TimeUnit.SECONDS.sleep(2);
213+
}
214+
catch (InterruptedException e) {
215+
throw new RuntimeException(e);
216+
}
217+
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
218+
CreateMessageResult.StopReason.STOP_SEQUENCE);
219+
};
220+
221+
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
222+
.capabilities(ClientCapabilities.builder().sampling().build())
223+
.sampling(samplingHandler)
224+
.build();
225+
226+
// Server
227+
228+
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
229+
null);
230+
231+
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
232+
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
233+
234+
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
235+
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
236+
new McpSchema.TextContent("Test message"))))
237+
.modelPreferences(ModelPreferences.builder()
238+
.hints(List.of())
239+
.costPriority(1.0)
240+
.speedPriority(1.0)
241+
.intelligencePriority(1.0)
242+
.build())
243+
.build();
244+
245+
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
246+
assertThat(result).isNotNull();
247+
assertThat(result.role()).isEqualTo(Role.USER);
248+
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
249+
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
250+
assertThat(result.model()).isEqualTo("MockModelName");
251+
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
252+
}).verifyComplete();
253+
254+
return Mono.just(callResponse);
255+
});
256+
257+
var mcpServer = McpServer.async(mcpServerTransportProvider)
258+
.requestTimeout(Duration.ofSeconds(4))
259+
.serverInfo("test-server", "1.0.0")
260+
.tools(tool)
261+
.build();
262+
263+
InitializeResult initResult = mcpClient.initialize();
264+
assertThat(initResult).isNotNull();
265+
266+
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
267+
268+
assertThat(response).isNotNull();
269+
assertThat(response).isEqualTo(callResponse);
270+
271+
mcpClient.close();
272+
mcpServer.close();
273+
}
274+
275+
@ParameterizedTest(name = "{0} : {displayName} ")
276+
@ValueSource(strings = { "httpclient", "webflux" })
277+
void testCreateMessageWithRequestTimeoutFail(String clientType) throws InterruptedException {
278+
279+
// Client
280+
var clientBuilder = clientBuilders.get(clientType);
281+
282+
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
283+
assertThat(request.messages()).hasSize(1);
284+
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
285+
try {
286+
TimeUnit.SECONDS.sleep(3);
287+
}
288+
catch (InterruptedException e) {
289+
throw new RuntimeException(e);
290+
}
291+
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
292+
CreateMessageResult.StopReason.STOP_SEQUENCE);
293+
};
294+
295+
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
296+
.capabilities(ClientCapabilities.builder().sampling().build())
297+
.sampling(samplingHandler)
298+
.build();
299+
300+
// Server
301+
302+
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
303+
null);
304+
305+
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
306+
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
307+
308+
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
309+
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
310+
new McpSchema.TextContent("Test message"))))
311+
.modelPreferences(ModelPreferences.builder()
312+
.hints(List.of())
313+
.costPriority(1.0)
314+
.speedPriority(1.0)
315+
.intelligencePriority(1.0)
316+
.build())
317+
.build();
318+
319+
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
320+
assertThat(result).isNotNull();
321+
assertThat(result.role()).isEqualTo(Role.USER);
322+
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
323+
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
324+
assertThat(result.model()).isEqualTo("MockModelName");
325+
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
326+
}).verifyComplete();
327+
328+
return Mono.just(callResponse);
329+
});
330+
331+
var mcpServer = McpServer.async(mcpServerTransportProvider)
332+
.requestTimeout(Duration.ofSeconds(1))
333+
.serverInfo("test-server", "1.0.0")
334+
.tools(tool)
335+
.build();
336+
337+
InitializeResult initResult = mcpClient.initialize();
338+
assertThat(initResult).isNotNull();
339+
340+
assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
341+
mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
342+
}).withMessageContaining("Timeout");
343+
344+
mcpClient.close();
345+
mcpServer.close();
346+
}
347+
199348
// ---------------------------------------
200349
// Roots Tests
201350
// ---------------------------------------

‎mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java

+145
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import java.time.Duration;
77
import java.util.List;
88
import java.util.Map;
9+
import java.util.concurrent.TimeUnit;
910
import java.util.concurrent.atomic.AtomicReference;
1011
import java.util.function.Function;
1112

@@ -41,6 +42,7 @@
4142
import org.springframework.web.servlet.function.ServerResponse;
4243

4344
import static org.assertj.core.api.Assertions.assertThat;
45+
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
4446
import static org.awaitility.Awaitility.await;
4547
import static org.mockito.Mockito.mock;
4648

@@ -212,6 +214,149 @@ void testCreateMessageSuccess() {
212214
mcpServer.close();
213215
}
214216

217+
@Test
218+
void testCreateMessageWithRequestTimeoutSuccess() throws InterruptedException {
219+
220+
// Client
221+
222+
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
223+
assertThat(request.messages()).hasSize(1);
224+
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
225+
try {
226+
TimeUnit.SECONDS.sleep(2);
227+
}
228+
catch (InterruptedException e) {
229+
throw new RuntimeException(e);
230+
}
231+
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
232+
CreateMessageResult.StopReason.STOP_SEQUENCE);
233+
};
234+
235+
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
236+
.capabilities(ClientCapabilities.builder().sampling().build())
237+
.sampling(samplingHandler)
238+
.build();
239+
240+
// Server
241+
242+
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
243+
null);
244+
245+
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
246+
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
247+
248+
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
249+
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
250+
new McpSchema.TextContent("Test message"))))
251+
.modelPreferences(ModelPreferences.builder()
252+
.hints(List.of())
253+
.costPriority(1.0)
254+
.speedPriority(1.0)
255+
.intelligencePriority(1.0)
256+
.build())
257+
.build();
258+
259+
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
260+
assertThat(result).isNotNull();
261+
assertThat(result.role()).isEqualTo(Role.USER);
262+
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
263+
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
264+
assertThat(result.model()).isEqualTo("MockModelName");
265+
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
266+
}).verifyComplete();
267+
268+
return Mono.just(callResponse);
269+
});
270+
271+
var mcpServer = McpServer.async(mcpServerTransportProvider)
272+
.serverInfo("test-server", "1.0.0")
273+
.requestTimeout(Duration.ofSeconds(4))
274+
.tools(tool)
275+
.build();
276+
277+
InitializeResult initResult = mcpClient.initialize();
278+
assertThat(initResult).isNotNull();
279+
280+
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
281+
282+
assertThat(response).isNotNull();
283+
assertThat(response).isEqualTo(callResponse);
284+
285+
mcpClient.close();
286+
mcpServer.close();
287+
}
288+
289+
@Test
290+
void testCreateMessageWithRequestTimeoutFail() throws InterruptedException {
291+
292+
// Client
293+
294+
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
295+
assertThat(request.messages()).hasSize(1);
296+
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
297+
try {
298+
TimeUnit.SECONDS.sleep(2);
299+
}
300+
catch (InterruptedException e) {
301+
throw new RuntimeException(e);
302+
}
303+
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
304+
CreateMessageResult.StopReason.STOP_SEQUENCE);
305+
};
306+
307+
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
308+
.capabilities(ClientCapabilities.builder().sampling().build())
309+
.sampling(samplingHandler)
310+
.build();
311+
312+
// Server
313+
314+
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
315+
null);
316+
317+
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
318+
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
319+
320+
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
321+
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
322+
new McpSchema.TextContent("Test message"))))
323+
.modelPreferences(ModelPreferences.builder()
324+
.hints(List.of())
325+
.costPriority(1.0)
326+
.speedPriority(1.0)
327+
.intelligencePriority(1.0)
328+
.build())
329+
.build();
330+
331+
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
332+
assertThat(result).isNotNull();
333+
assertThat(result.role()).isEqualTo(Role.USER);
334+
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
335+
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
336+
assertThat(result.model()).isEqualTo("MockModelName");
337+
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
338+
}).verifyComplete();
339+
340+
return Mono.just(callResponse);
341+
});
342+
343+
var mcpServer = McpServer.async(mcpServerTransportProvider)
344+
.serverInfo("test-server", "1.0.0")
345+
.requestTimeout(Duration.ofSeconds(1))
346+
.tools(tool)
347+
.build();
348+
349+
InitializeResult initResult = mcpClient.initialize();
350+
assertThat(initResult).isNotNull();
351+
352+
assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
353+
mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
354+
}).withMessageContaining("Timeout");
355+
356+
mcpClient.close();
357+
mcpServer.close();
358+
}
359+
215360
// ---------------------------------------
216361
// Roots Tests
217362
// ---------------------------------------

‎mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java

+7-6
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
package io.modelcontextprotocol.server;
66

7+
import java.time.Duration;
78
import java.util.HashMap;
89
import java.util.List;
910
import java.util.Map;
@@ -90,8 +91,8 @@ public class McpAsyncServer {
9091
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
9192
*/
9293
McpAsyncServer(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
93-
McpServerFeatures.Async features) {
94-
this.delegate = new AsyncServerImpl(mcpTransportProvider, objectMapper, features);
94+
McpServerFeatures.Async features, Duration requestTimeout) {
95+
this.delegate = new AsyncServerImpl(mcpTransportProvider, objectMapper, requestTimeout, features);
9596
}
9697

9798
/**
@@ -271,7 +272,7 @@ private static class AsyncServerImpl extends McpAsyncServer {
271272
private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);
272273

273274
AsyncServerImpl(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
274-
McpServerFeatures.Async features) {
275+
Duration requestTimeout, McpServerFeatures.Async features) {
275276
this.mcpTransportProvider = mcpTransportProvider;
276277
this.objectMapper = objectMapper;
277278
this.serverInfo = features.serverInfo();
@@ -330,9 +331,9 @@ private static class AsyncServerImpl extends McpAsyncServer {
330331
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED,
331332
asyncRootsListChangedNotificationHandler(rootsChangeConsumers));
332333

333-
mcpTransportProvider
334-
.setSessionFactory(transport -> new McpServerSession(UUID.randomUUID().toString(), transport,
335-
this::asyncInitializeRequestHandler, Mono::empty, requestHandlers, notificationHandlers));
334+
mcpTransportProvider.setSessionFactory(
335+
transport -> new McpServerSession(UUID.randomUUID().toString(), requestTimeout, transport,
336+
this::asyncInitializeRequestHandler, Mono::empty, requestHandlers, notificationHandlers));
336337
}
337338

338339
// ---------------------------------------

0 commit comments

Comments
 (0)