From 4bc5e0e58db56b8d5c41ef1c6b21a69e3a822836 Mon Sep 17 00:00:00 2001 From: Renxia Wang Date: Sat, 29 Mar 2025 20:49:51 -0400 Subject: [PATCH] Pass request builder to constructor to support custom headers --- .../client/transport/FlowSseClient.java | 13 ++++++- .../HttpClientSseClientTransport.java | 39 +++++++++++++++++-- 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/transport/FlowSseClient.java b/mcp/src/main/java/io/modelcontextprotocol/client/transport/FlowSseClient.java index 7fc67993..4bb00e2c 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/transport/FlowSseClient.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/transport/FlowSseClient.java @@ -38,6 +38,7 @@ public class FlowSseClient { private final HttpClient httpClient; + private final HttpRequest.Builder requestBuilder; /** * Pattern to extract the data content from SSE data field lines. Matches lines @@ -92,7 +93,17 @@ public interface SseEventHandler { * @param httpClient the {@link HttpClient} instance to use for SSE connections */ public FlowSseClient(HttpClient httpClient) { + this(httpClient, HttpRequest.newBuilder()); + } + + /** + * Creates a new FlowSseClient with the specified HTTP client and request builder. + * @param httpClient the {@link HttpClient} instance to use for SSE connections + * @param requestBuilder the {@link HttpRequest.Builder} to use for SSE requests + */ + public FlowSseClient(HttpClient httpClient, HttpRequest.Builder requestBuilder) { this.httpClient = httpClient; + this.requestBuilder = requestBuilder; } /** @@ -109,7 +120,7 @@ public FlowSseClient(HttpClient httpClient) { * @throws RuntimeException if the connection fails with a non-200 status code */ public void subscribe(String url, SseEventHandler eventHandler) { - HttpRequest request = HttpRequest.newBuilder() + HttpRequest request = requestBuilder .uri(URI.create(url)) .header("Accept", "text/event-stream") .header("Cache-Control", "no-cache") diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java b/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java index 696efdff..95573bfe 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java @@ -82,6 +82,9 @@ public class HttpClientSseClientTransport implements McpClientTransport { */ private final HttpClient httpClient; + /** HTTP request builder for building requests to send messages to the server */ + private final HttpRequest.Builder requestBuilder; + /** JSON object mapper for message serialization/deserialization */ protected ObjectMapper objectMapper; @@ -126,15 +129,32 @@ public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, String bas */ public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, String baseUri, String sseEndpoint, ObjectMapper objectMapper) { + this(clientBuilder, HttpRequest.newBuilder(), baseUri, sseEndpoint, objectMapper); + } + + /** + * Creates a new transport instance with custom HTTP client builder, object mapper, and headers. + * @param clientBuilder the HTTP client builder to use + * @param requestBuilder the HTTP request builder to use + * @param baseUri the base URI of the MCP server + * @param sseEndpoint the SSE endpoint path + * @param objectMapper the object mapper for JSON serialization/deserialization + * @throws IllegalArgumentException if objectMapper, clientBuilder, or headers is null + */ + public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, HttpRequest.Builder requestBuilder, + String baseUri, String sseEndpoint, ObjectMapper objectMapper) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); Assert.hasText(baseUri, "baseUri must not be empty"); Assert.hasText(sseEndpoint, "sseEndpoint must not be empty"); Assert.notNull(clientBuilder, "clientBuilder must not be null"); + Assert.notNull(requestBuilder, "requestBuilder must not be null"); this.baseUri = baseUri; this.sseEndpoint = sseEndpoint; this.objectMapper = objectMapper; this.httpClient = clientBuilder.connectTimeout(Duration.ofSeconds(10)).build(); - this.sseClient = new FlowSseClient(this.httpClient); + this.requestBuilder = requestBuilder; + + this.sseClient = new FlowSseClient(this.httpClient, requestBuilder); } /** @@ -159,6 +179,8 @@ public static class Builder { private ObjectMapper objectMapper = new ObjectMapper(); + private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder(); + /** * Creates a new builder with the specified base URI. * @param baseUri the base URI of the MCP server @@ -190,6 +212,17 @@ public Builder clientBuilder(HttpClient.Builder clientBuilder) { return this; } + /** + * Sets the HTTP request builder. + * @param requestBuilder the HTTP request builder + * @return this builder + */ + public Builder requestBuilder(HttpRequest.Builder requestBuilder) { + Assert.notNull(requestBuilder, "requestBuilder must not be null"); + this.requestBuilder = requestBuilder; + return this; + } + /** * Sets the object mapper for JSON serialization/deserialization. * @param objectMapper the object mapper @@ -206,7 +239,7 @@ public Builder objectMapper(ObjectMapper objectMapper) { * @return a new transport instance */ public HttpClientSseClientTransport build() { - return new HttpClientSseClientTransport(clientBuilder, baseUri, sseEndpoint, objectMapper); + return new HttpClientSseClientTransport(clientBuilder, requestBuilder, baseUri, sseEndpoint, objectMapper); } } @@ -301,7 +334,7 @@ public Mono sendMessage(JSONRPCMessage message) { try { String jsonText = this.objectMapper.writeValueAsString(message); - HttpRequest request = HttpRequest.newBuilder() + HttpRequest request = this.requestBuilder .uri(URI.create(this.baseUri + endpoint)) .header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.ofString(jsonText))