Add option for immediate execution in McpSyncServer
- The McpSyncServer wraps an async server. By default, reactive operations are scheduled on a bounded-elastic scheduler, to offload blocking work and prevent accidental blocking of non-blocking operations. - With the default behavior, there will be thead ops, even in a blocking context, which means thread-locals from the request thread will be lost. This is inconenvient for frameworks that store state in thread-locals. - This commit adds the ability to avoid offloading, when the user is sure they are executing code in a blocking environment. Work happens in the calling thread, and thread-locals are available throughout the execution.
This commit is contained in:
parent
c711f8315a
commit
b5a1e3b8f4
6 changed files with 172 additions and 34 deletions
|
|
@ -1,5 +1,5 @@
|
||||||
/*
|
/*
|
||||||
* Copyright 2024-2024 the original author or authors.
|
* Copyright 2024-2025 the original author or authors.
|
||||||
*/
|
*/
|
||||||
|
|
||||||
package io.modelcontextprotocol.server;
|
package io.modelcontextprotocol.server;
|
||||||
|
|
@ -695,6 +695,8 @@ public interface McpServer {
|
||||||
|
|
||||||
private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout
|
private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout
|
||||||
|
|
||||||
|
private boolean immediateExecution = false;
|
||||||
|
|
||||||
private SyncSpecification(McpServerTransportProvider transportProvider) {
|
private SyncSpecification(McpServerTransportProvider transportProvider) {
|
||||||
Assert.notNull(transportProvider, "Transport provider must not be null");
|
Assert.notNull(transportProvider, "Transport provider must not be null");
|
||||||
this.transportProvider = transportProvider;
|
this.transportProvider = transportProvider;
|
||||||
|
|
@ -1116,6 +1118,22 @@ public interface McpServer {
|
||||||
return this;
|
return this;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Enable on "immediate execution" of the operations on the underlying
|
||||||
|
* {@link McpAsyncServer}. Defaults to false, which does blocking code offloading
|
||||||
|
* to prevent accidental blocking of the non-blocking transport.
|
||||||
|
* <p>
|
||||||
|
* Do NOT set to true if the underlying transport is a non-blocking
|
||||||
|
* implementation.
|
||||||
|
* @param immediateExecution When true, do not offload work asynchronously.
|
||||||
|
* @return This builder instance for method chaining.
|
||||||
|
*
|
||||||
|
*/
|
||||||
|
public SyncSpecification immediateExecution(boolean immediateExecution) {
|
||||||
|
this.immediateExecution = immediateExecution;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Builds a synchronous MCP server that provides blocking operations.
|
* Builds a synchronous MCP server that provides blocking operations.
|
||||||
* @return A new instance of {@link McpSyncServer} configured with this builder's
|
* @return A new instance of {@link McpSyncServer} configured with this builder's
|
||||||
|
|
@ -1125,12 +1143,13 @@ public interface McpServer {
|
||||||
McpServerFeatures.Sync syncFeatures = new McpServerFeatures.Sync(this.serverInfo, this.serverCapabilities,
|
McpServerFeatures.Sync syncFeatures = new McpServerFeatures.Sync(this.serverInfo, this.serverCapabilities,
|
||||||
this.tools, this.resources, this.resourceTemplates, this.prompts, this.completions,
|
this.tools, this.resources, this.resourceTemplates, this.prompts, this.completions,
|
||||||
this.rootsChangeHandlers, this.instructions);
|
this.rootsChangeHandlers, this.instructions);
|
||||||
McpServerFeatures.Async asyncFeatures = McpServerFeatures.Async.fromSync(syncFeatures);
|
McpServerFeatures.Async asyncFeatures = McpServerFeatures.Async.fromSync(syncFeatures,
|
||||||
|
this.immediateExecution);
|
||||||
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
|
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
|
||||||
var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures, this.requestTimeout,
|
var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures, this.requestTimeout,
|
||||||
this.uriTemplateManagerFactory);
|
this.uriTemplateManagerFactory);
|
||||||
|
|
||||||
return new McpSyncServer(asyncServer);
|
return new McpSyncServer(asyncServer, this.immediateExecution);
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,5 @@
|
||||||
/*
|
/*
|
||||||
* Copyright 2024-2024 the original author or authors.
|
* Copyright 2024-2025 the original author or authors.
|
||||||
*/
|
*/
|
||||||
|
|
||||||
package io.modelcontextprotocol.server;
|
package io.modelcontextprotocol.server;
|
||||||
|
|
@ -95,28 +95,30 @@ public class McpServerFeatures {
|
||||||
* blocking code offloading to prevent accidental blocking of the non-blocking
|
* blocking code offloading to prevent accidental blocking of the non-blocking
|
||||||
* transport.
|
* transport.
|
||||||
* @param syncSpec a potentially blocking, synchronous specification.
|
* @param syncSpec a potentially blocking, synchronous specification.
|
||||||
|
* @param immediateExecution when true, do not offload. Do NOT set to true when
|
||||||
|
* using a non-blocking transport.
|
||||||
* @return a specification which is protected from blocking calls specified by the
|
* @return a specification which is protected from blocking calls specified by the
|
||||||
* user.
|
* user.
|
||||||
*/
|
*/
|
||||||
static Async fromSync(Sync syncSpec) {
|
static Async fromSync(Sync syncSpec, boolean immediateExecution) {
|
||||||
List<McpServerFeatures.AsyncToolSpecification> tools = new ArrayList<>();
|
List<McpServerFeatures.AsyncToolSpecification> tools = new ArrayList<>();
|
||||||
for (var tool : syncSpec.tools()) {
|
for (var tool : syncSpec.tools()) {
|
||||||
tools.add(AsyncToolSpecification.fromSync(tool));
|
tools.add(AsyncToolSpecification.fromSync(tool, immediateExecution));
|
||||||
}
|
}
|
||||||
|
|
||||||
Map<String, AsyncResourceSpecification> resources = new HashMap<>();
|
Map<String, AsyncResourceSpecification> resources = new HashMap<>();
|
||||||
syncSpec.resources().forEach((key, resource) -> {
|
syncSpec.resources().forEach((key, resource) -> {
|
||||||
resources.put(key, AsyncResourceSpecification.fromSync(resource));
|
resources.put(key, AsyncResourceSpecification.fromSync(resource, immediateExecution));
|
||||||
});
|
});
|
||||||
|
|
||||||
Map<String, AsyncPromptSpecification> prompts = new HashMap<>();
|
Map<String, AsyncPromptSpecification> prompts = new HashMap<>();
|
||||||
syncSpec.prompts().forEach((key, prompt) -> {
|
syncSpec.prompts().forEach((key, prompt) -> {
|
||||||
prompts.put(key, AsyncPromptSpecification.fromSync(prompt));
|
prompts.put(key, AsyncPromptSpecification.fromSync(prompt, immediateExecution));
|
||||||
});
|
});
|
||||||
|
|
||||||
Map<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
|
Map<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
|
||||||
syncSpec.completions().forEach((key, completion) -> {
|
syncSpec.completions().forEach((key, completion) -> {
|
||||||
completions.put(key, AsyncCompletionSpecification.fromSync(completion));
|
completions.put(key, AsyncCompletionSpecification.fromSync(completion, immediateExecution));
|
||||||
});
|
});
|
||||||
|
|
||||||
List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootChangeConsumers = new ArrayList<>();
|
List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootChangeConsumers = new ArrayList<>();
|
||||||
|
|
@ -239,15 +241,15 @@ public class McpServerFeatures {
|
||||||
public record AsyncToolSpecification(McpSchema.Tool tool,
|
public record AsyncToolSpecification(McpSchema.Tool tool,
|
||||||
BiFunction<McpAsyncServerExchange, Map<String, Object>, Mono<McpSchema.CallToolResult>> call) {
|
BiFunction<McpAsyncServerExchange, Map<String, Object>, Mono<McpSchema.CallToolResult>> call) {
|
||||||
|
|
||||||
static AsyncToolSpecification fromSync(SyncToolSpecification tool) {
|
static AsyncToolSpecification fromSync(SyncToolSpecification tool, boolean immediate) {
|
||||||
// FIXME: This is temporary, proper validation should be implemented
|
// FIXME: This is temporary, proper validation should be implemented
|
||||||
if (tool == null) {
|
if (tool == null) {
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
return new AsyncToolSpecification(tool.tool(),
|
return new AsyncToolSpecification(tool.tool(), (exchange, map) -> {
|
||||||
(exchange, map) -> Mono
|
var toolResult = Mono.fromCallable(() -> tool.call().apply(new McpSyncServerExchange(exchange), map));
|
||||||
.fromCallable(() -> tool.call().apply(new McpSyncServerExchange(exchange), map))
|
return immediate ? toolResult : toolResult.subscribeOn(Schedulers.boundedElastic());
|
||||||
.subscribeOn(Schedulers.boundedElastic()));
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -281,15 +283,16 @@ public class McpServerFeatures {
|
||||||
public record AsyncResourceSpecification(McpSchema.Resource resource,
|
public record AsyncResourceSpecification(McpSchema.Resource resource,
|
||||||
BiFunction<McpAsyncServerExchange, McpSchema.ReadResourceRequest, Mono<McpSchema.ReadResourceResult>> readHandler) {
|
BiFunction<McpAsyncServerExchange, McpSchema.ReadResourceRequest, Mono<McpSchema.ReadResourceResult>> readHandler) {
|
||||||
|
|
||||||
static AsyncResourceSpecification fromSync(SyncResourceSpecification resource) {
|
static AsyncResourceSpecification fromSync(SyncResourceSpecification resource, boolean immediateExecution) {
|
||||||
// FIXME: This is temporary, proper validation should be implemented
|
// FIXME: This is temporary, proper validation should be implemented
|
||||||
if (resource == null) {
|
if (resource == null) {
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
return new AsyncResourceSpecification(resource.resource(),
|
return new AsyncResourceSpecification(resource.resource(), (exchange, req) -> {
|
||||||
(exchange, req) -> Mono
|
var resourceResult = Mono
|
||||||
.fromCallable(() -> resource.readHandler().apply(new McpSyncServerExchange(exchange), req))
|
.fromCallable(() -> resource.readHandler().apply(new McpSyncServerExchange(exchange), req));
|
||||||
.subscribeOn(Schedulers.boundedElastic()));
|
return immediateExecution ? resourceResult : resourceResult.subscribeOn(Schedulers.boundedElastic());
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -327,15 +330,16 @@ public class McpServerFeatures {
|
||||||
public record AsyncPromptSpecification(McpSchema.Prompt prompt,
|
public record AsyncPromptSpecification(McpSchema.Prompt prompt,
|
||||||
BiFunction<McpAsyncServerExchange, McpSchema.GetPromptRequest, Mono<McpSchema.GetPromptResult>> promptHandler) {
|
BiFunction<McpAsyncServerExchange, McpSchema.GetPromptRequest, Mono<McpSchema.GetPromptResult>> promptHandler) {
|
||||||
|
|
||||||
static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt) {
|
static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt, boolean immediateExecution) {
|
||||||
// FIXME: This is temporary, proper validation should be implemented
|
// FIXME: This is temporary, proper validation should be implemented
|
||||||
if (prompt == null) {
|
if (prompt == null) {
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
return new AsyncPromptSpecification(prompt.prompt(),
|
return new AsyncPromptSpecification(prompt.prompt(), (exchange, req) -> {
|
||||||
(exchange, req) -> Mono
|
var promptResult = Mono
|
||||||
.fromCallable(() -> prompt.promptHandler().apply(new McpSyncServerExchange(exchange), req))
|
.fromCallable(() -> prompt.promptHandler().apply(new McpSyncServerExchange(exchange), req));
|
||||||
.subscribeOn(Schedulers.boundedElastic()));
|
return immediateExecution ? promptResult : promptResult.subscribeOn(Schedulers.boundedElastic());
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -366,14 +370,17 @@ public class McpServerFeatures {
|
||||||
* @return an asynchronous wrapper of the provided sync specification, or
|
* @return an asynchronous wrapper of the provided sync specification, or
|
||||||
* {@code null} if input is null
|
* {@code null} if input is null
|
||||||
*/
|
*/
|
||||||
static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion) {
|
static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion,
|
||||||
|
boolean immediateExecution) {
|
||||||
if (completion == null) {
|
if (completion == null) {
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
return new AsyncCompletionSpecification(completion.referenceKey(),
|
return new AsyncCompletionSpecification(completion.referenceKey(), (exchange, request) -> {
|
||||||
(exchange, request) -> Mono.fromCallable(
|
var completionResult = Mono.fromCallable(
|
||||||
() -> completion.completionHandler().apply(new McpSyncServerExchange(exchange), request))
|
() -> completion.completionHandler().apply(new McpSyncServerExchange(exchange), request));
|
||||||
.subscribeOn(Schedulers.boundedElastic()));
|
return immediateExecution ? completionResult
|
||||||
|
: completionResult.subscribeOn(Schedulers.boundedElastic());
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -54,13 +54,27 @@ public class McpSyncServer {
|
||||||
*/
|
*/
|
||||||
private final McpAsyncServer asyncServer;
|
private final McpAsyncServer asyncServer;
|
||||||
|
|
||||||
|
private final boolean immediateExecution;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Creates a new synchronous server that wraps the provided async server.
|
* Creates a new synchronous server that wraps the provided async server.
|
||||||
* @param asyncServer The async server to wrap
|
* @param asyncServer The async server to wrap
|
||||||
*/
|
*/
|
||||||
public McpSyncServer(McpAsyncServer asyncServer) {
|
public McpSyncServer(McpAsyncServer asyncServer) {
|
||||||
|
this(asyncServer, false);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Creates a new synchronous server that wraps the provided async server.
|
||||||
|
* @param asyncServer The async server to wrap
|
||||||
|
* @param immediateExecution Tools, prompts, and resources handlers execute work
|
||||||
|
* without blocking code offloading. Do NOT set to true if the {@code asyncServer}'s
|
||||||
|
* transport is non-blocking.
|
||||||
|
*/
|
||||||
|
public McpSyncServer(McpAsyncServer asyncServer, boolean immediateExecution) {
|
||||||
Assert.notNull(asyncServer, "Async server must not be null");
|
Assert.notNull(asyncServer, "Async server must not be null");
|
||||||
this.asyncServer = asyncServer;
|
this.asyncServer = asyncServer;
|
||||||
|
this.immediateExecution = immediateExecution;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
@ -68,7 +82,9 @@ public class McpSyncServer {
|
||||||
* @param toolHandler The tool handler to add
|
* @param toolHandler The tool handler to add
|
||||||
*/
|
*/
|
||||||
public void addTool(McpServerFeatures.SyncToolSpecification toolHandler) {
|
public void addTool(McpServerFeatures.SyncToolSpecification toolHandler) {
|
||||||
this.asyncServer.addTool(McpServerFeatures.AsyncToolSpecification.fromSync(toolHandler)).block();
|
this.asyncServer
|
||||||
|
.addTool(McpServerFeatures.AsyncToolSpecification.fromSync(toolHandler, this.immediateExecution))
|
||||||
|
.block();
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
@ -84,7 +100,10 @@ public class McpSyncServer {
|
||||||
* @param resourceHandler The resource handler to add
|
* @param resourceHandler The resource handler to add
|
||||||
*/
|
*/
|
||||||
public void addResource(McpServerFeatures.SyncResourceSpecification resourceHandler) {
|
public void addResource(McpServerFeatures.SyncResourceSpecification resourceHandler) {
|
||||||
this.asyncServer.addResource(McpServerFeatures.AsyncResourceSpecification.fromSync(resourceHandler)).block();
|
this.asyncServer
|
||||||
|
.addResource(
|
||||||
|
McpServerFeatures.AsyncResourceSpecification.fromSync(resourceHandler, this.immediateExecution))
|
||||||
|
.block();
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
@ -100,7 +119,10 @@ public class McpSyncServer {
|
||||||
* @param promptSpecification The prompt specification to add
|
* @param promptSpecification The prompt specification to add
|
||||||
*/
|
*/
|
||||||
public void addPrompt(McpServerFeatures.SyncPromptSpecification promptSpecification) {
|
public void addPrompt(McpServerFeatures.SyncPromptSpecification promptSpecification) {
|
||||||
this.asyncServer.addPrompt(McpServerFeatures.AsyncPromptSpecification.fromSync(promptSpecification)).block();
|
this.asyncServer
|
||||||
|
.addPrompt(
|
||||||
|
McpServerFeatures.AsyncPromptSpecification.fromSync(promptSpecification, this.immediateExecution))
|
||||||
|
.block();
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,5 @@
|
||||||
/*
|
/*
|
||||||
* Copyright 2024 - 2024 the original author or authors.
|
* Copyright 2024 - 2025 the original author or authors.
|
||||||
*/
|
*/
|
||||||
package io.modelcontextprotocol.server.transport;
|
package io.modelcontextprotocol.server.transport;
|
||||||
|
|
||||||
|
|
@ -37,7 +37,6 @@ import org.apache.catalina.LifecycleState;
|
||||||
import org.apache.catalina.startup.Tomcat;
|
import org.apache.catalina.startup.Tomcat;
|
||||||
import org.junit.jupiter.api.AfterEach;
|
import org.junit.jupiter.api.AfterEach;
|
||||||
import org.junit.jupiter.api.BeforeEach;
|
import org.junit.jupiter.api.BeforeEach;
|
||||||
import org.junit.jupiter.api.Disabled;
|
|
||||||
import org.junit.jupiter.api.Test;
|
import org.junit.jupiter.api.Test;
|
||||||
import reactor.core.publisher.Mono;
|
import reactor.core.publisher.Mono;
|
||||||
import reactor.test.StepVerifier;
|
import reactor.test.StepVerifier;
|
||||||
|
|
@ -46,6 +45,7 @@ import org.springframework.web.client.RestClient;
|
||||||
|
|
||||||
import static org.assertj.core.api.Assertions.assertThat;
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
|
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
|
||||||
|
import static org.assertj.core.api.InstanceOfAssertFactories.type;
|
||||||
import static org.awaitility.Awaitility.await;
|
import static org.awaitility.Awaitility.await;
|
||||||
import static org.mockito.Mockito.mock;
|
import static org.mockito.Mockito.mock;
|
||||||
|
|
||||||
|
|
@ -728,6 +728,9 @@ class HttpServletSseServerTransportProviderIntegrationTests {
|
||||||
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
|
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
|
||||||
McpServerFeatures.SyncToolSpecification tool1 = new McpServerFeatures.SyncToolSpecification(
|
McpServerFeatures.SyncToolSpecification tool1 = new McpServerFeatures.SyncToolSpecification(
|
||||||
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
|
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
|
||||||
|
assertThat(McpTestServletFilter.getThreadLocalValue())
|
||||||
|
.as("blocking code exectuion should be offloaded")
|
||||||
|
.isNull();
|
||||||
// perform a blocking call to a remote service
|
// perform a blocking call to a remote service
|
||||||
String response = RestClient.create()
|
String response = RestClient.create()
|
||||||
.get()
|
.get()
|
||||||
|
|
@ -758,6 +761,37 @@ class HttpServletSseServerTransportProviderIntegrationTests {
|
||||||
mcpServer.close();
|
mcpServer.close();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void testToolCallImmediateExecution() {
|
||||||
|
McpServerFeatures.SyncToolSpecification tool1 = new McpServerFeatures.SyncToolSpecification(
|
||||||
|
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
|
||||||
|
var threadLocalValue = McpTestServletFilter.getThreadLocalValue();
|
||||||
|
return CallToolResult.builder()
|
||||||
|
.addTextContent(threadLocalValue != null ? threadLocalValue : "<unset>")
|
||||||
|
.build();
|
||||||
|
});
|
||||||
|
|
||||||
|
var mcpServer = McpServer.sync(mcpServerTransportProvider)
|
||||||
|
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||||
|
.tools(tool1)
|
||||||
|
.immediateExecution(true)
|
||||||
|
.build();
|
||||||
|
|
||||||
|
try (var mcpClient = clientBuilder.build()) {
|
||||||
|
mcpClient.initialize();
|
||||||
|
|
||||||
|
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
|
||||||
|
|
||||||
|
assertThat(response).isNotNull();
|
||||||
|
assertThat(response.content()).first()
|
||||||
|
.asInstanceOf(type(McpSchema.TextContent.class))
|
||||||
|
.extracting(McpSchema.TextContent::text)
|
||||||
|
.isEqualTo(McpTestServletFilter.THREAD_LOCAL_VALUE);
|
||||||
|
}
|
||||||
|
|
||||||
|
mcpServer.close();
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
void testToolListChangeHandlingSuccess() {
|
void testToolListChangeHandlingSuccess() {
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,43 @@
|
||||||
|
/*
|
||||||
|
* Copyright 2025 - 2025 the original author or authors.
|
||||||
|
*/
|
||||||
|
|
||||||
|
package io.modelcontextprotocol.server.transport;
|
||||||
|
|
||||||
|
import java.io.IOException;
|
||||||
|
|
||||||
|
import jakarta.servlet.Filter;
|
||||||
|
import jakarta.servlet.FilterChain;
|
||||||
|
import jakarta.servlet.ServletException;
|
||||||
|
import jakarta.servlet.ServletRequest;
|
||||||
|
import jakarta.servlet.ServletResponse;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Simple {@link Filter} which sets a value in a thread local. Used to verify whether MCP
|
||||||
|
* executions happen on the thread processing the request or are offloaded.
|
||||||
|
*
|
||||||
|
* @author Daniel Garnier-Moiroux
|
||||||
|
*/
|
||||||
|
public class McpTestServletFilter implements Filter {
|
||||||
|
|
||||||
|
public static final String THREAD_LOCAL_VALUE = McpTestServletFilter.class.getName();
|
||||||
|
|
||||||
|
private static final ThreadLocal<String> holder = new ThreadLocal<>();
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain)
|
||||||
|
throws IOException, ServletException {
|
||||||
|
holder.set(THREAD_LOCAL_VALUE);
|
||||||
|
try {
|
||||||
|
filterChain.doFilter(servletRequest, servletResponse);
|
||||||
|
}
|
||||||
|
finally {
|
||||||
|
holder.remove();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
public static String getThreadLocalValue() {
|
||||||
|
return holder.get();
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
@ -10,9 +10,12 @@ import java.net.ServerSocket;
|
||||||
import jakarta.servlet.Servlet;
|
import jakarta.servlet.Servlet;
|
||||||
import org.apache.catalina.Context;
|
import org.apache.catalina.Context;
|
||||||
import org.apache.catalina.startup.Tomcat;
|
import org.apache.catalina.startup.Tomcat;
|
||||||
|
import org.apache.tomcat.util.descriptor.web.FilterDef;
|
||||||
|
import org.apache.tomcat.util.descriptor.web.FilterMap;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* @author Christian Tzolov
|
* @author Christian Tzolov
|
||||||
|
* @author Daniel Garnier-Moiroux
|
||||||
*/
|
*/
|
||||||
public class TomcatTestUtil {
|
public class TomcatTestUtil {
|
||||||
|
|
||||||
|
|
@ -39,6 +42,16 @@ public class TomcatTestUtil {
|
||||||
context.addChild(wrapper);
|
context.addChild(wrapper);
|
||||||
context.addServletMappingDecoded("/*", "mcpServlet");
|
context.addServletMappingDecoded("/*", "mcpServlet");
|
||||||
|
|
||||||
|
var filterDef = new FilterDef();
|
||||||
|
filterDef.setFilterClass(McpTestServletFilter.class.getName());
|
||||||
|
filterDef.setFilterName(McpTestServletFilter.class.getSimpleName());
|
||||||
|
context.addFilterDef(filterDef);
|
||||||
|
|
||||||
|
var filterMap = new FilterMap();
|
||||||
|
filterMap.setFilterName(McpTestServletFilter.class.getSimpleName());
|
||||||
|
filterMap.addURLPattern("/*");
|
||||||
|
context.addFilterMap(filterMap);
|
||||||
|
|
||||||
var connector = tomcat.getConnector();
|
var connector = tomcat.getConnector();
|
||||||
connector.setAsyncTimeout(3000);
|
connector.setAsyncTimeout(3000);
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue