Skip to content

Add context extractor for HTTP headers in WebFlux and WebMvc configurations #4723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add context extractor for HTTP headers in WebFlux and WebMvc configurations #4723

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions ...ramework/ai/mcp/server/autoconfigure/McpServerStreamableHttpWebFluxAutoConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,11 @@

package org.springframework.ai.mcp.server.autoconfigure;

import java.util.HashMap;
import java.util.Map;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.common.McpTransportContext;
import io.modelcontextprotocol.json.jackson.JacksonMcpJsonMapper;
import io.modelcontextprotocol.server.transport.WebFluxStreamableServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
Expand All @@ -33,6 +37,7 @@
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Conditional;
import org.springframework.web.reactive.function.server.RouterFunction;
import org.springframework.web.reactive.function.server.ServerRequest;

/**
* @author Christian Tzolov
Expand All @@ -57,9 +62,20 @@ public WebFluxStreamableServerTransportProvider webFluxStreamableServerTransport
.messageEndpoint(serverProperties.getMcpEndpoint())
.keepAliveInterval(serverProperties.getKeepAliveInterval())
.disallowDelete(serverProperties.isDisallowDelete())
.contextExtractor(this::extractContextFromRequest)
.build();
}

private McpTransportContext extractContextFromRequest(ServerRequest serverRequest) {
Map<String, Object> headersMap = new HashMap<>();
serverRequest.headers().asHttpHeaders().forEach((headerName, headerValues) -> {
if (!headerValues.isEmpty()) {
headersMap.put(headerName, headerValues.get(0));
}
});
return McpTransportContext.create(headersMap);
}

// Router function for streamable http transport used by Spring WebFlux to start an
// HTTP server.
@Bean
Expand Down
16 changes: 16 additions & 0 deletions ...mework/ai/mcp/server/autoconfigure/McpServerStreamableHttpWebFluxAutoConfigurationIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,4 +192,20 @@ void enabledPropertyExplicitlyTrue() {
});
}

@Test
void contextExtractorExtractsHeaders() {
this.contextRunner.run(context -> {
WebFluxStreamableServerTransportProvider provider = context
.getBean(WebFluxStreamableServerTransportProvider.class);

// Verify the provider is properly configured with context extractor
assertThat(provider).isNotNull();

// Note: Testing the actual header extraction requires a live request context
// which is better tested through integration tests with a running server.
// This test verifies that the bean is properly configured with the context
// extractor.
});
}

}
34 changes: 30 additions & 4 deletions ...framework/ai/mcp/server/autoconfigure/McpServerStreamableHttpWebMvcAutoConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,11 @@

package org.springframework.ai.mcp.server.autoconfigure;

import java.util.HashMap;
import java.util.Map;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.common.McpTransportContext;
import io.modelcontextprotocol.json.jackson.JacksonMcpJsonMapper;
import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
Expand All @@ -33,6 +37,7 @@
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Conditional;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerRequest;
import org.springframework.web.servlet.function.ServerResponse;

/**
Expand All @@ -46,10 +51,17 @@
McpServerAutoConfiguration.EnabledStreamableServerCondition.class })
public class McpServerStreamableHttpWebMvcAutoConfiguration {

/**
* Creates a WebMvc streamable server transport provider.
* @param objectMapperProvider the object mapper provider
* @param serverProperties the server properties
* @return the transport provider
*/
@Bean
@ConditionalOnMissingBean
public WebMvcStreamableServerTransportProvider webMvcStreamableServerTransportProvider(
ObjectProvider<ObjectMapper> objectMapperProvider, McpServerStreamableHttpProperties serverProperties) {
final ObjectProvider<ObjectMapper> objectMapperProvider,
final McpServerStreamableHttpProperties serverProperties) {

ObjectMapper objectMapper = objectMapperProvider.getIfAvailable(ObjectMapper::new);

Expand All @@ -58,15 +70,29 @@ public WebMvcStreamableServerTransportProvider webMvcStreamableServerTransportPr
.mcpEndpoint(serverProperties.getMcpEndpoint())
.keepAliveInterval(serverProperties.getKeepAliveInterval())
.disallowDelete(serverProperties.isDisallowDelete())
.contextExtractor(this::extractContextFromRequest)
.build();
}

// Router function for streamable http transport used by Spring WebFlux to start an
// HTTP server.
private McpTransportContext extractContextFromRequest(final ServerRequest serverRequest) {
Map<String, Object> headersMap = new HashMap<>();
serverRequest.headers().asHttpHeaders().forEach((headerName, headerValues) -> {
if (!headerValues.isEmpty()) {
headersMap.put(headerName, headerValues.get(0));
}
});
return McpTransportContext.create(headersMap);
}

/**
* Creates a router function for the streamable server transport.
* @param webMvcProvider the transport provider
* @return the router function
*/
@Bean
@ConditionalOnMissingBean(name = "webMvcStreamableServerRouterFunction")
public RouterFunction<ServerResponse> webMvcStreamableServerRouterFunction(
WebMvcStreamableServerTransportProvider webMvcProvider) {
final WebMvcStreamableServerTransportProvider webMvcProvider) {
return webMvcProvider.getRouterFunction();
}

Expand Down
26 changes: 26 additions & 0 deletions ...amework/ai/mcp/server/autoconfigure/McpServerStreamableHttpWebMvcAutoConfigurationIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@

import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerRequest;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.mock;
Expand Down Expand Up @@ -192,4 +194,28 @@ void enabledPropertyExplicitlyTrue() {
});
}

@Test
void contextExtractorExtractsHeaders() {
this.contextRunner.run(context -> {
WebMvcStreamableServerTransportProvider provider = context
.getBean(WebMvcStreamableServerTransportProvider.class);

// Create a mock ServerRequest with headers
MockHttpServletRequest mockRequest = new MockHttpServletRequest();
mockRequest.addHeader("xxxx", "123456");
mockRequest.addHeader("Authorization", "Bearer token123");
mockRequest.addHeader("Content-Type", "application/json");

ServerRequest serverRequest = ServerRequest.create(mockRequest, java.util.Collections.emptyList());

// Verify the provider is properly configured
assertThat(provider).isNotNull();

// Verify headers are accessible from the ServerRequest
assertThat(serverRequest.headers().firstHeader("xxxx")).isEqualTo("123456");
assertThat(serverRequest.headers().firstHeader("Authorization")).isEqualTo("Bearer token123");
assertThat(serverRequest.headers().firstHeader("Content-Type")).isEqualTo("application/json");
});
}

}