mcp-client: introduce support for chat streaming in AuthenticationMcpTransportContextProvider

- Backfilled tests that interact with MCP through the lens of the Spring
  AI ChatClient.
- Closes #19

Signed-off-by: Daniel Garnier-Moiroux <[email protected]>

Author: Kehrlann

mcp-client-security/pom.xml

@@ -83,6 +83,11 @@
             <artifactId>spring-webflux</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-model</artifactId>
+            <scope>provided</scope>
+        </dependency>
         <dependency>
             <groupId>jakarta.servlet</groupId>
             <artifactId>jakarta.servlet-api</artifactId>

mcp-client-security/src/main/java/org/springaicommunity/mcp/security/client/sync/AuthenticationMcpTransportContextProvider.java

@@ -16,26 +16,120 @@
 
 package org.springaicommunity.mcp.security.client.sync;
 
-import io.modelcontextprotocol.common.McpTransportContext;
+import java.net.http.HttpClient;
 import java.util.HashMap;
 import java.util.function.Supplier;
 
+import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer;
+import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer;
+import io.modelcontextprotocol.common.McpTransportContext;
+import org.springaicommunity.mcp.security.client.sync.oauth2.http.client.OAuth2AuthorizationCodeSyncHttpRequestCustomizer;
+import org.springaicommunity.mcp.security.client.sync.oauth2.http.client.OAuth2HybridSyncHttpRequestCustomizer;
+import org.springaicommunity.mcp.security.client.sync.oauth2.webclient.McpOAuth2AuthorizationCodeExchangeFilterFunction;
+import org.springaicommunity.mcp.security.client.sync.oauth2.webclient.McpOAuth2HybridExchangeFilterFunction;
+import reactor.util.context.Context;
+import reactor.util.context.ContextView;
+
+import org.springframework.ai.model.tool.internal.ToolCallReactiveContextHolder;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.context.SecurityContextHolder;
 import org.springframework.web.context.request.RequestAttributes;
 import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
+import org.springframework.web.reactive.function.client.WebClient;
 
 /**
+ * A supplier that extracts security-related information from the "context", and make it
+ * available to MCP clients when they send requests to MCP servers. It extracts request
+ * attributes and the current authentication object. In Servlet application, this is
+ * achieved with {@link SecurityContextHolder} and {@link RequestContextHolder}.
+ * <p>
+ * This can be used in conjunction with {@link McpSyncHttpClientRequestCustomizer} and
+ * {@link McpAsyncHttpClientRequestCustomizer} for {@link HttpClient}-based transports,
+ * and with {@link ExchangeFilterFunction} for {@link WebClient}-based transports.
+ * <p>
+ * This is usually used through a Spring AI {@code McpSyncClientCustomizer} or
+ * {@code McpAsyncClientCustomizer}, like so:
+ *
+ * <pre>
+ * &#x40;Bean
+ * McpSyncClientCustomizer syncClientCustomizer() {
+ *   return (name, syncSpec) -> syncSpec
+ *     .transportContextProvider(
+ *       new AuthenticationMcpTransportContextProvider()
+ *     );
+ * }
+ * </pre>
+ *
+ * <p>
+ * When using Spring's {@code ChatClient} "streaming" capabilities, you must also use
+ * {@link #writeToReactorContext()} to make thread-locals available in the stream's
+ * reactor context:
+ *
+ * <pre>
+ * chatClientSupplier.get()
+ *     .prompt("your LLM prompt")
+ *     .stream()
+ *     .content()
+ *     .contextWrite(AuthenticationMcpTransportContextProvider.writeToReactorContext())
+ *     // ...
+ * </pre>
+ *
  * @author Daniel Garnier-Moiroux
+ * @see OAuth2AuthorizationCodeSyncHttpRequestCustomizer
+ * @see OAuth2HybridSyncHttpRequestCustomizer
+ * @see McpOAuth2AuthorizationCodeExchangeFilterFunction
+ * @see McpOAuth2HybridExchangeFilterFunction
  */
 public class AuthenticationMcpTransportContextProvider implements Supplier<McpTransportContext> {
 
 	public static final String AUTHENTICATION_KEY = Authentication.class.getName();
 
 	public static final String REQUEST_ATTRIBUTES_KEY = RequestAttributes.class.getName();
 
+	public static final String REACTOR_CONTEXT_KEY = "org.springaicommunity.mcp.security.client.sync.REACTOR_CONTEXT";
+
+	private final boolean reactiveContextHolderAvailable;
+
+	public AuthenticationMcpTransportContextProvider() {
+		boolean reactiveContextHolderAvailable = false;
+		try {
+			Class.forName("org.springframework.ai.model.tool.internal.ToolCallReactiveContextHolder");
+			reactiveContextHolderAvailable = true;
+		}
+		catch (ClassNotFoundException ignored) {
+		}
+		this.reactiveContextHolderAvailable = reactiveContextHolderAvailable;
+	}
+
+	/**
+	 * Helper function to write to thread-locals to the reactor context. Use it on your
+	 * reactive {@code ChatClient} operations, such as
+	 * {@code chatClient.prompt("...").stream().content()}.
+	 * <p>
+	 * Do NOT use if Reactor is not on the classpath.
+	 */
+	public static ContextView writeToReactorContext() {
+		return Context.empty().put(REACTOR_CONTEXT_KEY, fromThreadLocals());
+	}
+
+	/**
+	 * Read authentication and request data from thread-locals. If they are not available,
+	 * and a Spring AI {@code ToolCallReactiveContextHolder} is available on the
+	 * classpath, it will try to access the values there.
+	 */
 	@Override
 	public McpTransportContext get() {
+		var transportContext = fromThreadLocals();
+
+		if (this.reactiveContextHolderAvailable && transportContext == McpTransportContext.EMPTY) {
+			transportContext = fromToolCallReactiveContextHolder();
+		}
+
+		return transportContext;
+	}
+
+	private static McpTransportContext fromThreadLocals() {
 		var data = new HashMap<String, Object>();
 
 		var securityContext = SecurityContextHolder.getContext();
@@ -48,7 +142,19 @@ public McpTransportContext get() {
 			data.put(REQUEST_ATTRIBUTES_KEY, requestAttributes);
 		}
 
+		if (data.isEmpty()) {
+			return McpTransportContext.EMPTY;
+		}
+
 		return McpTransportContext.create(data);
 	}
 
+	private static McpTransportContext fromToolCallReactiveContextHolder() {
+		var reactorContext = ToolCallReactiveContextHolder.getContext();
+		if (reactorContext == Context.empty()) {
+			return McpTransportContext.EMPTY;
+		}
+		return reactorContext.getOrDefault(REACTOR_CONTEXT_KEY, McpTransportContext.EMPTY);
+	}
+
 }

samples/integration-tests/pom.xml

@@ -48,7 +48,10 @@
             <artifactId>mcp-client-security</artifactId>
             <version>0.0.4-SNAPSHOT</version>
         </dependency>
-
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-starter-model-anthropic</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.springframework.experimental.boot</groupId>
             <artifactId>spring-boot-testjars</artifactId>

samples/integration-tests/src/main/java/org/springaicommunity/mcp/security/tests/McpController.java

@@ -16,8 +16,15 @@
 
 package org.springaicommunity.mcp.security.tests;
 
+import java.util.Optional;
+
 import io.modelcontextprotocol.spec.McpSchema;
+import org.springaicommunity.mcp.security.client.sync.AuthenticationMcpTransportContextProvider;
 
+import org.springframework.ai.chat.client.ChatClient;
+import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.util.function.SingletonSupplier;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RestController;
 
@@ -29,8 +36,15 @@ public class McpController {
 
 	private final InMemoryMcpClientRepository repository;
 
-	McpController(InMemoryMcpClientRepository repository) {
+	// Not all tests have ChatClient support, so we wrap it in a supplier that will
+	// only be used in LLM-powered tests.
+	private final SingletonSupplier<ChatClient> chatClientSupplier;
+
+	McpController(InMemoryMcpClientRepository repository, ObjectProvider<ChatClient.Builder> chatClientBuilder,
+			Optional<SyncMcpToolCallbackProvider> mcpTools) {
 		this.repository = repository;
+		this.chatClientSupplier = SingletonSupplier
+			.of(() -> chatClientBuilder.getIfUnique().defaultToolCallbacks(mcpTools.get()).build());
 	}
 
 	@GetMapping("/tool/call")
@@ -42,4 +56,24 @@ public String callTool(String clientName, String toolName) {
 		return "Called [client: %s, tool: %s], got response [%s]".formatted(clientName, toolName, toolResponse);
 	}
 
+	@GetMapping("/chat")
+	public String chat(String question) {
+		return chatClientSupplier.get().prompt(question).call().content();
+	}
+
+	@GetMapping("/stream")
+	public String stream(String question) {
+		return chatClientSupplier.get()
+			.prompt(question)
+			.stream()
+			.content()
+			.contextWrite(AuthenticationMcpTransportContextProvider.writeToReactorContext())
+			.blockLast();
+	}
+
+	@GetMapping("/stream-no-context")
+	public String streamNoContext(String question) {
+		return chatClientSupplier.get().prompt(question).stream().content().blockLast();
+	}
+
 }