Skip to content

Commit e5460c9

Browse files
KehrlannKehrlann
committed
Compose request customizers on the HTTP client transport builders
Fixes #1073 Signed-off-by: Daniel Garnier-Moiroux <[email protected]>
1 parent f856940 commit e5460c9

5 files changed

Lines changed: 225 additions & 22 deletions

File tree

‎mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java‎

Lines changed: 76 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,15 @@
88
import java.net.URI;
99
import java.net.http.HttpClient;
1010
import java.net.http.HttpRequest;
11-
import java.net.http.HttpResponse;
1211
import java.time.Duration;
12+
import java.util.ArrayList;
1313
import java.util.List;
1414
import java.util.Optional;
1515
import java.util.concurrent.atomic.AtomicReference;
1616
import java.util.function.Consumer;
1717
import java.util.function.Function;
1818

19+
import io.modelcontextprotocol.client.transport.customizer.DelegatingMcpAsyncHttpClientRequestCustomizer;
1920
import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer;
2021
import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer;
2122
import io.modelcontextprotocol.common.McpTransportContext;
@@ -203,7 +204,7 @@ public static class Builder {
203204

204205
private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder();
205206

206-
private McpAsyncHttpClientRequestCustomizer httpRequestCustomizer = McpAsyncHttpClientRequestCustomizer.NOOP;
207+
private final List<McpAsyncHttpClientRequestCustomizer> httpRequestCustomizers = new ArrayList<>();
207208

208209
private Duration connectTimeout = Duration.ofSeconds(10);
209210

@@ -286,36 +287,92 @@ public Builder jsonMapper(McpJsonMapper jsonMapper) {
286287

287288
/**
288289
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
289-
* executing them.
290+
* executing them. The customizer is internally converted to
291+
* {@link McpAsyncHttpClientRequestCustomizer}.
290292
* <p>
291-
* This overrides the customizer from
292-
* {@link #asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}.
293+
* This replaces all customizers previously registered on this builder.
293294
* <p>
294295
* Do NOT use a blocking {@link McpSyncHttpClientRequestCustomizer} in a
295296
* non-blocking context. Use
296297
* {@link #asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}
297298
* instead.
298299
* @param syncHttpRequestCustomizer the request customizer
299300
* @return this builder
301+
* @deprecated Use {@link #addHttpRequestCustomizer} instead.
300302
*/
303+
@Deprecated
301304
public Builder httpRequestCustomizer(McpSyncHttpClientRequestCustomizer syncHttpRequestCustomizer) {
302-
this.httpRequestCustomizer = McpAsyncHttpClientRequestCustomizer.fromSync(syncHttpRequestCustomizer);
303-
return this;
305+
Assert.notNull(syncHttpRequestCustomizer, "syncHttpRequestCustomizer must not be null");
306+
this.httpRequestCustomizers.clear();
307+
return addHttpRequestCustomizer(syncHttpRequestCustomizer);
304308
}
305309

306310
/**
307311
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
308312
* executing them.
309313
* <p>
310-
* This overrides the customizer from
311-
* {@link #httpRequestCustomizer(McpSyncHttpClientRequestCustomizer)}.
314+
* This replaces all customizers previously registered on this builder.
312315
* <p>
313316
* Do NOT use a blocking implementation in a non-blocking context.
314317
* @param asyncHttpRequestCustomizer the request customizer
315318
* @return this builder
319+
* @deprecated Use {@link #addAsyncHttpRequestCustomizer} instead.
316320
*/
321+
@Deprecated
317322
public Builder asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer asyncHttpRequestCustomizer) {
318-
this.httpRequestCustomizer = asyncHttpRequestCustomizer;
323+
Assert.notNull(asyncHttpRequestCustomizer, "asyncHttpRequestCustomizer must not be null");
324+
this.httpRequestCustomizers.clear();
325+
return addAsyncHttpRequestCustomizer(asyncHttpRequestCustomizer);
326+
}
327+
328+
/**
329+
* Adds a customizer for {@link HttpRequest.Builder}, to modify requests before
330+
* executing them. Customizers are applied in the order they are added, after
331+
* those already registered on this builder. The customizer is internally
332+
* converted to * {@link McpAsyncHttpClientRequestCustomizer}.
333+
* <p>
334+
* Do NOT use a blocking {@link McpSyncHttpClientRequestCustomizer} in a
335+
* non-blocking context. Use
336+
* {@link #addAsyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}
337+
* instead.
338+
* @param syncHttpRequestCustomizer the request customizer
339+
* @return this builder
340+
*/
341+
public Builder addHttpRequestCustomizer(McpSyncHttpClientRequestCustomizer syncHttpRequestCustomizer) {
342+
Assert.notNull(syncHttpRequestCustomizer, "syncHttpRequestCustomizer must not be null");
343+
this.httpRequestCustomizers.add(McpAsyncHttpClientRequestCustomizer.fromSync(syncHttpRequestCustomizer));
344+
return this;
345+
}
346+
347+
/**
348+
* Adds a customizer for {@link HttpRequest.Builder}, to modify requests before
349+
* executing them. Customizers are applied in the order they are added, after
350+
* those already registered on this builder.
351+
* <p>
352+
* Do NOT use a blocking implementation in a non-blocking context.
353+
* @param asyncHttpRequestCustomizer the request customizer
354+
* @return this builder
355+
*/
356+
public Builder addAsyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer asyncHttpRequestCustomizer) {
357+
Assert.notNull(asyncHttpRequestCustomizer, "asyncHttpRequestCustomizer must not be null");
358+
this.httpRequestCustomizers.add(asyncHttpRequestCustomizer);
359+
return this;
360+
}
361+
362+
/**
363+
* Provides access to the mutable list of request customizers registered on this
364+
* builder, so they can be inspected, reordered, added or removed. Customizers are
365+
* applied in list order. Synchronous customizers registered through
366+
* {@link #httpRequestCustomizer(McpSyncHttpClientRequestCustomizer)} or
367+
* {@link #addHttpRequestCustomizer(McpSyncHttpClientRequestCustomizer)} appear in
368+
* the list wrapped as {@link McpAsyncHttpClientRequestCustomizer}.
369+
* @param customizersConsumer a consumer of the list of customizers
370+
* @return this builder
371+
*/
372+
public Builder asyncHttpRequestCustomizers(
373+
Consumer<List<McpAsyncHttpClientRequestCustomizer>> customizersConsumer) {
374+
Assert.notNull(customizersConsumer, "customizersConsumer must not be null");
375+
customizersConsumer.accept(this.httpRequestCustomizers);
319376
return this;
320377
}
321378

@@ -369,10 +426,18 @@ public Builder maxResponseSize(int maxResponseSize) {
369426
public HttpClientSseClientTransport build() {
370427
HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build();
371428
return new HttpClientSseClientTransport(httpClient, requestBuilder, baseUri, sseEndpoint,
372-
jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, httpRequestCustomizer,
429+
jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, httpRequestCustomizer(),
373430
messageEndpointValidator, maxResponseSize);
374431
}
375432

433+
private McpAsyncHttpClientRequestCustomizer httpRequestCustomizer() {
434+
return switch (this.httpRequestCustomizers.size()) {
435+
case 0 -> McpAsyncHttpClientRequestCustomizer.NOOP;
436+
case 1 -> this.httpRequestCustomizers.get(0);
437+
default -> new DelegatingMcpAsyncHttpClientRequestCustomizer(List.copyOf(this.httpRequestCustomizers));
438+
};
439+
}
440+
376441
}
377442

378443
@Override

‎mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java‎

Lines changed: 81 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import java.net.http.HttpResponse;
1212
import java.nio.ByteBuffer;
1313
import java.time.Duration;
14+
import java.util.ArrayList;
1415
import java.util.Collections;
1516
import java.util.Comparator;
1617
import java.util.List;
@@ -22,6 +23,7 @@
2223
import java.util.function.Function;
2324

2425
import io.modelcontextprotocol.client.McpAsyncClient;
26+
import io.modelcontextprotocol.client.transport.customizer.DelegatingMcpAsyncHttpClientRequestCustomizer;
2527
import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer;
2628
import io.modelcontextprotocol.client.transport.customizer.McpHttpClientAuthorizationErrorHandler;
2729
import io.modelcontextprotocol.client.transport.customizer.McpHttpClientTransportAuthorizationErrorHandler;
@@ -672,7 +674,7 @@ public static class Builder {
672674

673675
private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder();
674676

675-
private McpAsyncHttpClientRequestCustomizer httpRequestCustomizer = McpAsyncHttpClientRequestCustomizer.NOOP;
677+
private final List<McpAsyncHttpClientRequestCustomizer> httpRequestCustomizers = new ArrayList<>();
676678

677679
private Duration connectTimeout = Duration.ofSeconds(10);
678680

@@ -776,36 +778,96 @@ public Builder openConnectionOnStartup(boolean openConnectionOnStartup) {
776778

777779
/**
778780
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
779-
* executing them.
781+
* executing them. The customizer is internally converted to
782+
* {@link McpAsyncHttpClientRequestCustomizer}.
780783
* <p>
781-
* This overrides the customizer from
782-
* {@link #asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}.
784+
* This replaces all customizers previously registered on this builder. Use
785+
* {@link #addHttpRequestCustomizer(McpSyncHttpClientRequestCustomizer)} to keep
786+
* them.
783787
* <p>
784788
* Do NOT use a blocking {@link McpSyncHttpClientRequestCustomizer} in a
785789
* non-blocking context. Use
786790
* {@link #asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}
787791
* instead.
788792
* @param syncHttpRequestCustomizer the request customizer
789793
* @return this builder
794+
* @deprecated use {@link #addHttpRequestCustomizer} instead
790795
*/
796+
@Deprecated
791797
public Builder httpRequestCustomizer(McpSyncHttpClientRequestCustomizer syncHttpRequestCustomizer) {
792-
this.httpRequestCustomizer = McpAsyncHttpClientRequestCustomizer.fromSync(syncHttpRequestCustomizer);
793-
return this;
798+
Assert.notNull(syncHttpRequestCustomizer, "syncHttpRequestCustomizer must not be null");
799+
this.httpRequestCustomizers.clear();
800+
return addHttpRequestCustomizer(syncHttpRequestCustomizer);
794801
}
795802

796803
/**
797804
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
798-
* executing them.
805+
* executing them. The customizer is internally converted to
806+
* {@link McpAsyncHttpClientRequestCustomizer}.
799807
* <p>
800-
* This overrides the customizer from
801-
* {@link #httpRequestCustomizer(McpSyncHttpClientRequestCustomizer)}.
808+
* This replaces all customizers previously registered on this builder. Use
809+
* {@link #addAsyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)} to
810+
* keep them.
802811
* <p>
803812
* Do NOT use a blocking implementation in a non-blocking context.
804813
* @param asyncHttpRequestCustomizer the request customizer
805814
* @return this builder
815+
* @deprecated use {@link #addAsyncHttpRequestCustomizer} instead.
806816
*/
817+
@Deprecated
807818
public Builder asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer asyncHttpRequestCustomizer) {
808-
this.httpRequestCustomizer = asyncHttpRequestCustomizer;
819+
Assert.notNull(asyncHttpRequestCustomizer, "asyncHttpRequestCustomizer must not be null");
820+
this.httpRequestCustomizers.clear();
821+
return addAsyncHttpRequestCustomizer(asyncHttpRequestCustomizer);
822+
}
823+
824+
/**
825+
* Adds a customizer for {@link HttpRequest.Builder}, to modify requests before
826+
* executing them. Customizers are applied in the order they are added, after
827+
* those already registered on this builder.
828+
* <p>
829+
* Do NOT use a blocking {@link McpSyncHttpClientRequestCustomizer} in a
830+
* non-blocking context. Use
831+
* {@link #addAsyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}
832+
* instead.
833+
* @param syncHttpRequestCustomizer the request customizer
834+
* @return this builder
835+
*/
836+
public Builder addHttpRequestCustomizer(McpSyncHttpClientRequestCustomizer syncHttpRequestCustomizer) {
837+
Assert.notNull(syncHttpRequestCustomizer, "syncHttpRequestCustomizer must not be null");
838+
this.httpRequestCustomizers.add(McpAsyncHttpClientRequestCustomizer.fromSync(syncHttpRequestCustomizer));
839+
return this;
840+
}
841+
842+
/**
843+
* Adds a customizer for {@link HttpRequest.Builder}, to modify requests before
844+
* executing them. Customizers are applied in the order they are added, after
845+
* those already registered on this builder.
846+
* <p>
847+
* Do NOT use a blocking implementation in a non-blocking context.
848+
* @param asyncHttpRequestCustomizer the request customizer
849+
* @return this builder
850+
*/
851+
public Builder addAsyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer asyncHttpRequestCustomizer) {
852+
Assert.notNull(asyncHttpRequestCustomizer, "asyncHttpRequestCustomizer must not be null");
853+
this.httpRequestCustomizers.add(asyncHttpRequestCustomizer);
854+
return this;
855+
}
856+
857+
/**
858+
* Provides access to the mutable list of request customizers registered on this
859+
* builder, so they can be inspected, reordered, added or removed. Customizers are
860+
* applied in list order. Synchronous customizers registered through
861+
* {@link #httpRequestCustomizer(McpSyncHttpClientRequestCustomizer)} or
862+
* {@link #addHttpRequestCustomizer(McpSyncHttpClientRequestCustomizer)} appear in
863+
* the list wrapped as {@link McpAsyncHttpClientRequestCustomizer}.
864+
* @param customizersConsumer a consumer of the list of customizers
865+
* @return this builder
866+
*/
867+
public Builder asyncHttpRequestCustomizers(
868+
Consumer<List<McpAsyncHttpClientRequestCustomizer>> customizersConsumer) {
869+
Assert.notNull(customizersConsumer, "customizersConsumer must not be null");
870+
customizersConsumer.accept(this.httpRequestCustomizers);
809871
return this;
810872
}
811873

@@ -910,7 +972,15 @@ public HttpClientStreamableHttpTransport build() {
910972
HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build();
911973
return new HttpClientStreamableHttpTransport(jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper,
912974
httpClient, requestBuilder, baseUri, endpoint, resumableStreams, openConnectionOnStartup,
913-
httpRequestCustomizer, authorizationErrorHandler, supportedProtocolVersions, maxResponseSize);
975+
httpRequestCustomizer(), authorizationErrorHandler, supportedProtocolVersions, maxResponseSize);
976+
}
977+
978+
private McpAsyncHttpClientRequestCustomizer httpRequestCustomizer() {
979+
return switch (this.httpRequestCustomizers.size()) {
980+
case 0 -> McpAsyncHttpClientRequestCustomizer.NOOP;
981+
case 1 -> this.httpRequestCustomizers.get(0);
982+
default -> new DelegatingMcpAsyncHttpClientRequestCustomizer(List.copyOf(this.httpRequestCustomizers));
983+
};
914984
}
915985

916986
}

‎mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,14 @@
2424
import java.util.stream.Collectors;
2525

2626
import io.modelcontextprotocol.client.McpClient;
27+
import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer;
2728
import io.modelcontextprotocol.common.McpTransportContext;
2829
import io.modelcontextprotocol.json.McpJsonDefaults;
2930
import io.modelcontextprotocol.server.McpServer;
3031
import io.modelcontextprotocol.server.McpServerFeatures;
3132
import io.modelcontextprotocol.server.McpSyncServer;
3233
import io.modelcontextprotocol.server.McpSyncServerExchange;
34+
import io.modelcontextprotocol.spec.McpClientTransport;
3335
import io.modelcontextprotocol.spec.McpError;
3436
import io.modelcontextprotocol.spec.McpSchema;
3537
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
@@ -52,6 +54,7 @@
5254
import io.modelcontextprotocol.spec.McpSchema.Tool;
5355
import io.modelcontextprotocol.util.Utils;
5456
import net.javacrumbs.jsonunit.core.Option;
57+
import org.junit.jupiter.api.Test;
5558
import org.junit.jupiter.params.ParameterizedTest;
5659
import org.junit.jupiter.params.provider.MethodSource;
5760
import org.junit.jupiter.params.provider.ValueSource;
@@ -68,6 +71,7 @@
6871
import static org.assertj.core.api.InstanceOfAssertFactories.MAP;
6972
import static org.assertj.core.api.InstanceOfAssertFactories.type;
7073
import static org.awaitility.Awaitility.await;
74+
import static org.junit.jupiter.api.Assumptions.assumeTrue;
7175
import static org.mockito.Mockito.mock;
7276

7377
public abstract class AbstractMcpClientServerIntegrationTests {
@@ -89,6 +93,15 @@ public abstract class AbstractMcpClientServerIntegrationTests {
8993
protected void awaitClientStreamEstablished() {
9094
}
9195

96+
// Creates a client transport that applies request customizers, in the order they
97+
// appear in the list. Subclasses whose transport supports
98+
// McpAsyncHttpClientRequestCustomizer should override this; tests relying on it are
99+
// skipped otherwise.
100+
protected McpClientTransport createClientTransport(
101+
Consumer<List<McpAsyncHttpClientRequestCustomizer>> requestCustomizers) {
102+
return null;
103+
}
104+
92105
@ParameterizedTest(name = "{0} : {displayName} ")
93106
@MethodSource("clientsForTesting")
94107
void simple(String clientType) {
@@ -113,6 +126,35 @@ void simple(String clientType) {
113126
}
114127
}
115128

129+
@Test
130+
void requestCustomizersAreAppliedInOrder() {
131+
var invocations = new CopyOnWriteArrayList<String>();
132+
Function<String, McpAsyncHttpClientRequestCustomizer> recording = name -> (builder, method, endpoint, body,
133+
context) -> {
134+
invocations.add(name);
135+
return Mono.just(builder);
136+
};
137+
138+
var transport = createClientTransport(customizers -> {
139+
customizers.add(recording.apply("second"));
140+
customizers.add(recording.apply("third"));
141+
customizers.add(0, recording.apply("first"));
142+
});
143+
assumeTrue(transport != null, "Transport does not support request customizers");
144+
145+
var server = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build();
146+
try (var client = McpClient.sync(transport).build()) {
147+
assertThat(client.initialize()).isNotNull();
148+
// Later requests, e.g. the GET stream in Streamable HTTP, may run
149+
// concurrently,
150+
// only check the initialize request.
151+
assertThat(invocations).startsWith("first", "second", "third");
152+
}
153+
finally {
154+
server.closeGracefully().block();
155+
}
156+
}
157+
116158
// ---------------------------------------
117159
// Sampling Tests
118160
// ---------------------------------------

0 commit comments

Comments
 (0)