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
3 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.transport;
|
||||
|
||||
|
|
@ -37,7 +37,6 @@ import org.apache.catalina.LifecycleState;
|
|||
import org.apache.catalina.startup.Tomcat;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Mono;
|
||||
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.assertThatExceptionOfType;
|
||||
import static org.assertj.core.api.InstanceOfAssertFactories.type;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
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);
|
||||
McpServerFeatures.SyncToolSpecification tool1 = new McpServerFeatures.SyncToolSpecification(
|
||||
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
|
||||
String response = RestClient.create()
|
||||
.get()
|
||||
|
|
@ -758,6 +761,37 @@ class HttpServletSseServerTransportProviderIntegrationTests {
|
|||
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
|
||||
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 org.apache.catalina.Context;
|
||||
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 Daniel Garnier-Moiroux
|
||||
*/
|
||||
public class TomcatTestUtil {
|
||||
|
||||
|
|
@ -39,6 +42,16 @@ public class TomcatTestUtil {
|
|||
context.addChild(wrapper);
|
||||
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();
|
||||
connector.setAsyncTimeout(3000);
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue