Skip to content

Commit a5a866f

Browse files
KehrlannKehrlann
committed
Expose request URI in McpHttpClientAuthorizationErrorHandler
Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
1 parent 6b71136 commit a5a866f

5 files changed

Lines changed: 57 additions & 30 deletions

File tree

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -298,7 +298,7 @@ private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
298298
return Mono.<McpSchema.JSONRPCMessage>error(
299299
new McpHttpClientTransportAuthorizationException(
300300
"Authorization error connecting to SSE stream",
301-
responseEvent.responseInfo()));
301+
responseEvent.responseInfo(), uri));
302302
}
303303
else if (statusCode == METHOD_NOT_ALLOWED) {
304304
logger.debug("The server does not support SSE streams, using request-response mode.");
@@ -417,7 +417,8 @@ private Retry authorizationErrorRetrySpec() {
417417
return Mono.deferContextual(ctx -> {
418418
var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
419419
return Mono
420-
.from(this.authorizationErrorHandler.handle(authException.getResponseInfo(), transportContext))
420+
.from(this.authorizationErrorHandler.handle(authException.getResponseInfo(),
421+
authException.getRequestUri(), transportContext))
421422
.switchIfEmpty(Mono.just(false))
422423
.flatMap(shouldRetry -> shouldRetry ? Mono.just(retrySignal.totalRetries())
423424
: Mono.error(retrySignal.failure()));
@@ -507,7 +508,7 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
507508
if (statusCode == 401 || statusCode == 403) {
508509
logger.debug("Authorization error in sendMessage with code {}", statusCode);
509510
return Mono.<McpSchema.JSONRPCMessage>error(new McpHttpClientTransportAuthorizationException(
510-
"Authorization error when sending message", responseEvent.responseInfo()));
511+
"Authorization error when sending message", responseEvent.responseInfo(), uri));
511512
}
512513

513514
if (transportSession.markInitialized(

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
package io.modelcontextprotocol.client.transport;
66

7+
import java.net.URI;
78
import java.net.http.HttpResponse;
89

910
import io.modelcontextprotocol.spec.McpTransportException;
@@ -19,13 +20,21 @@ public class McpHttpClientTransportAuthorizationException extends McpTransportEx
1920

2021
private final HttpResponse.ResponseInfo responseInfo;
2122

22-
public McpHttpClientTransportAuthorizationException(String message, HttpResponse.ResponseInfo responseInfo) {
23+
private final URI requestUri;
24+
25+
public McpHttpClientTransportAuthorizationException(String message, HttpResponse.ResponseInfo responseInfo,
26+
URI requestUri) {
2327
super(message);
2428
this.responseInfo = responseInfo;
29+
this.requestUri = requestUri;
2530
}
2631

2732
public HttpResponse.ResponseInfo getResponseInfo() {
2833
return responseInfo;
2934
}
3035

36+
public URI getRequestUri() {
37+
return requestUri;
38+
}
39+
3140
}

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
package io.modelcontextprotocol.client.transport.customizer;
66

7+
import java.net.URI;
78
import java.net.http.HttpResponse;
89

910
import io.modelcontextprotocol.client.transport.McpHttpClientTransportAuthorizationException;
@@ -35,11 +36,12 @@ public interface McpHttpClientAuthorizationErrorHandler {
3536
* <p>
3637
* The number of retries is bounded by {@link #maxRetries()}.
3738
* @param responseInfo the HTTP response information
39+
* @param requestUri the HTTP request URI that failed authorization
3840
* @param context the MCP client transport context
3941
* @return {@link Publisher} emitting true if the original request should be replayed,
4042
* false otherwise.
4143
*/
42-
Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo, McpTransportContext context);
44+
Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo, URI requestUri, McpTransportContext context);
4345

4446
/**
4547
* Maximum number of authorization error retries the transport will attempt. When the
@@ -68,7 +70,7 @@ default int maxRetries() {
6870
* @return an async handler
6971
*/
7072
static McpHttpClientAuthorizationErrorHandler fromSync(Sync handler) {
71-
return (info, context) -> Mono.fromCallable(() -> handler.handle(info, context))
73+
return (info, uri, context) -> Mono.fromCallable(() -> handler.handle(info, uri, context))
7274
.subscribeOn(Schedulers.boundedElastic());
7375
}
7476

@@ -85,17 +87,19 @@ interface Sync {
8587
* {@link McpHttpClientTransportAuthorizationException}, indicating the error
8688
* status.
8789
* @param responseInfo the HTTP response information
90+
* @param requestUri the HTTP request URI that failed authorization
8891
* @param context the MCP client transport context
8992
* @return true if the original request should be replayed, false otherwise.
9093
*/
91-
boolean handle(HttpResponse.ResponseInfo responseInfo, McpTransportContext context);
94+
boolean handle(HttpResponse.ResponseInfo responseInfo, URI requestUri, McpTransportContext context);
9295

9396
}
9497

9598
class Noop implements McpHttpClientAuthorizationErrorHandler {
9699

97100
@Override
98-
public Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo, McpTransportContext context) {
101+
public Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo, URI requestUri,
102+
McpTransportContext context) {
99103
return Mono.just(false);
100104
}
101105

‎mcp-core/src/test/java/io/modelcontextprotocol/client/transport/customizer/McpHttpClientAuthorizationErrorHandlerTest.java‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
*/
44
package io.modelcontextprotocol.client.transport.customizer;
55

6+
import java.net.URI;
67
import java.net.http.HttpResponse;
78

89
import io.modelcontextprotocol.common.McpTransportContext;
@@ -18,29 +19,31 @@ class McpHttpClientAuthorizationErrorHandlerTest {
1819

1920
private final HttpResponse.ResponseInfo responseInfo = mock(HttpResponse.ResponseInfo.class);
2021

22+
private final URI requestUri = URI.create("http://localhost/mcp");
23+
2124
private final McpTransportContext context = McpTransportContext.EMPTY;
2225

2326
@Test
2427
void whenTrueThenRetry() {
2528
McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
26-
.fromSync((info, ctx) -> true);
27-
StepVerifier.create(handler.handle(responseInfo, context)).expectNext(true).verifyComplete();
29+
.fromSync((info, uri, ctx) -> true);
30+
StepVerifier.create(handler.handle(responseInfo, requestUri, context)).expectNext(true).verifyComplete();
2831
}
2932

3033
@Test
3134
void whenFalseThenError() {
3235
McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
33-
.fromSync((info, ctx) -> false);
34-
StepVerifier.create(handler.handle(responseInfo, context)).expectNext(false).verifyComplete();
36+
.fromSync((info, uri, ctx) -> false);
37+
StepVerifier.create(handler.handle(responseInfo, requestUri, context)).expectNext(false).verifyComplete();
3538
}
3639

3740
@Test
3841
void whenExceptionThenPropagate() {
3942
McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
40-
.fromSync((info, ctx) -> {
43+
.fromSync((info, uri, ctx) -> {
4144
throw new IllegalStateException("sync handler error");
4245
});
43-
StepVerifier.create(handler.handle(responseInfo, context))
46+
StepVerifier.create(handler.handle(responseInfo, requestUri, context))
4447
.expectErrorMatches(t -> t instanceof IllegalStateException && t.getMessage().equals("sync handler error"))
4548
.verify();
4649
}

‎mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportErrorHandlingTest.java‎

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
import java.io.IOException;
88
import java.net.InetSocketAddress;
9+
import java.net.URI;
910
import java.net.http.HttpResponse;
1011
import java.time.Duration;
1112
import java.util.ArrayList;
@@ -403,9 +404,12 @@ void invokeHandler(int httpStatus) {
403404
AtomicReference<HttpResponse.ResponseInfo> capturedResponseInfo = new AtomicReference<>();
404405
AtomicReference<McpTransportContext> capturedContext = new AtomicReference<>();
405406

407+
AtomicReference<URI> capturedUri = new AtomicReference<>();
408+
406409
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
407-
.authorizationErrorHandler((responseInfo, context) -> {
410+
.authorizationErrorHandler((responseInfo, requestUri, context) -> {
408411
capturedResponseInfo.set(responseInfo);
412+
capturedUri.set(requestUri);
409413
capturedContext.set(context);
410414
return Mono.just(false);
411415
})
@@ -417,6 +421,8 @@ void invokeHandler(int httpStatus) {
417421
assertThat(processedMessagesCount.get()).isEqualTo(1);
418422
assertThat(capturedResponseInfo.get()).isNotNull();
419423
assertThat(capturedResponseInfo.get().statusCode()).isEqualTo(httpStatus);
424+
assertThat(capturedUri.get()).isNotNull();
425+
assertThat(capturedUri.get().toString()).isEqualTo(HOST + "/mcp");
420426
assertThat(capturedContext.get()).isNotNull();
421427

422428
StepVerifier.create(authTransport.closeGracefully()).verifyComplete();
@@ -440,7 +446,7 @@ void defaultHandler() {
440446
void retry() {
441447
serverResponseStatus.set(401);
442448
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
443-
.authorizationErrorHandler((responseInfo, context) -> {
449+
.authorizationErrorHandler((responseInfo, requestUri, context) -> {
444450
serverResponseStatus.set(200);
445451
return Mono.just(true);
446452
})
@@ -456,7 +462,7 @@ void retry() {
456462
void retryAtMostOnce() {
457463
serverResponseStatus.set(401);
458464
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
459-
.authorizationErrorHandler((responseInfo, context) -> Mono.just(true))
465+
.authorizationErrorHandler((responseInfo, requestUri, context) -> Mono.just(true))
460466
.build();
461467
StepVerifier.create(authTransport.sendMessage(createTestRequestMessage()))
462468
.expectErrorMatches(authorizationError(401))
@@ -473,7 +479,7 @@ void customMaxRetries() {
473479
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
474480
.authorizationErrorHandler(new McpHttpClientAuthorizationErrorHandler() {
475481
@Override
476-
public Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo,
482+
public Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo, URI requestUri,
477483
McpTransportContext context) {
478484
return Mono.just(true);
479485
}
@@ -498,7 +504,7 @@ void noRetry() {
498504
serverResponseStatus.set(401);
499505

500506
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
501-
.authorizationErrorHandler((responseInfo, context) -> Mono.just(false))
507+
.authorizationErrorHandler((responseInfo, requestUri, context) -> Mono.just(false))
502508
.build();
503509

504510
StepVerifier.create(authTransport.sendMessage(createTestRequestMessage()))
@@ -513,8 +519,8 @@ void noRetry() {
513519
void propagateHandlerError() {
514520
serverResponseStatus.set(401);
515521
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
516-
.authorizationErrorHandler(
517-
(responseInfo, context) -> Mono.error(new IllegalStateException("handler error")))
522+
.authorizationErrorHandler((responseInfo, requestUri, context) -> Mono
523+
.error(new IllegalStateException("handler error")))
518524
.build();
519525

520526
StepVerifier.create(authTransport.sendMessage(createTestRequestMessage()))
@@ -529,7 +535,7 @@ void propagateHandlerError() {
529535
void emptyHandler() {
530536
serverResponseStatus.set(401);
531537
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
532-
.authorizationErrorHandler((responseInfo, context) -> Mono.empty())
538+
.authorizationErrorHandler((responseInfo, requestUri, context) -> Mono.empty())
533539
.build();
534540

535541
StepVerifier.create(authTransport.sendMessage(createTestRequestMessage()))
@@ -552,11 +558,13 @@ void invokeHandler(int httpStatus) {
552558
AtomicReference<Throwable> capturedException = new AtomicReference<>();
553559

554560
AtomicReference<HttpResponse.ResponseInfo> capturedResponseInfo = new AtomicReference<>();
561+
AtomicReference<URI> capturedUri = new AtomicReference<>();
555562
AtomicReference<McpTransportContext> capturedContext = new AtomicReference<>();
556563

557564
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
558-
.authorizationErrorHandler((responseInfo, context) -> {
565+
.authorizationErrorHandler((responseInfo, requestUri, context) -> {
559566
capturedResponseInfo.set(responseInfo);
567+
capturedUri.set(requestUri);
560568
capturedContext.set(context);
561569
return Mono.just(false);
562570
})
@@ -572,6 +580,8 @@ void invokeHandler(int httpStatus) {
572580
assertThat(messages).isEmpty();
573581
assertThat(capturedResponseInfo.get()).isNotNull();
574582
assertThat(capturedResponseInfo.get().statusCode()).isEqualTo(httpStatus);
583+
assertThat(capturedUri.get()).isNotNull();
584+
assertThat(capturedUri.get().toString()).isEqualTo(HOST + "/mcp");
575585
assertThat(capturedContext.get()).isNotNull();
576586
assertThat(capturedException.get()).hasMessage("Authorization error connecting to SSE stream")
577587
.asInstanceOf(type(McpHttpClientTransportAuthorizationException.class))
@@ -606,7 +616,7 @@ void retry() {
606616
AtomicReference<Throwable> capturedException = new AtomicReference<>();
607617
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
608618
.openConnectionOnStartup(true)
609-
.authorizationErrorHandler((responseInfo, context) -> {
619+
.authorizationErrorHandler((responseInfo, requestUri, context) -> {
610620
serverSseResponseStatus.set(200);
611621
return Mono.just(true);
612622
})
@@ -636,7 +646,7 @@ void retryAtMostOnce() {
636646
AtomicReference<Throwable> capturedException = new AtomicReference<>();
637647
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
638648
.openConnectionOnStartup(true)
639-
.authorizationErrorHandler((responseInfo, context) -> {
649+
.authorizationErrorHandler((responseInfo, requestUri, context) -> {
640650
return Mono.just(true);
641651
})
642652
.build();
@@ -663,7 +673,7 @@ void customMaxRetries() {
663673
.openConnectionOnStartup(true)
664674
.authorizationErrorHandler(new McpHttpClientAuthorizationErrorHandler() {
665675
@Override
666-
public Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo,
676+
public Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo, URI requestUri,
667677
McpTransportContext context) {
668678
return Mono.just(true);
669679
}
@@ -695,7 +705,7 @@ void noRetry() {
695705
AtomicReference<Throwable> capturedException = new AtomicReference<>();
696706
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
697707
.openConnectionOnStartup(true)
698-
.authorizationErrorHandler((responseInfo, context) -> {
708+
.authorizationErrorHandler((responseInfo, requestUri, context) -> {
699709
// if there was a retry, the request would succeed.
700710
serverSseResponseStatus.set(200);
701711
return Mono.just(false);
@@ -720,7 +730,7 @@ void emptyHandler() {
720730
AtomicReference<Throwable> capturedException = new AtomicReference<>();
721731
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
722732
.openConnectionOnStartup(true)
723-
.authorizationErrorHandler((responseInfo, context) -> Mono.empty())
733+
.authorizationErrorHandler((responseInfo, requestUri, context) -> Mono.empty())
724734
.build();
725735
authTransport.setExceptionHandler(capturedException::set);
726736

@@ -741,8 +751,8 @@ void propagateHandlerError() {
741751
AtomicReference<Throwable> capturedException = new AtomicReference<>();
742752
var authTransport = HttpClientStreamableHttpTransport.builder(HOST)
743753
.openConnectionOnStartup(true)
744-
.authorizationErrorHandler(
745-
(responseInfo, context) -> Mono.error(new IllegalStateException("handler error")))
754+
.authorizationErrorHandler((responseInfo, requestUri, context) -> Mono
755+
.error(new IllegalStateException("handler error")))
746756
.build();
747757
authTransport.setExceptionHandler(capturedException::set);
748758

0 commit comments

Comments
 (0)