Introduce retry handling

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>

Author: Kehrlann

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

@@ -51,6 +51,7 @@
 import reactor.core.publisher.Mono;
 import reactor.util.function.Tuple2;
 import reactor.util.function.Tuples;
+import reactor.util.retry.Retry;
 
 /**
  * An implementation of the Streamable HTTP protocol as defined by the
@@ -291,26 +292,17 @@ private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
 					})).flatMap(responseEvent -> {
 						int statusCode = responseEvent.responseInfo().statusCode();
 						if (statusCode == 401 || statusCode == 403) {
-							logger.debug("Authorization error in sendMessage with code {}", statusCode);
-							return Mono.deferContextual(innerCtx -> {
-								var transportContext = innerCtx.getOrDefault(McpTransportContext.KEY,
-										McpTransportContext.EMPTY);
-								return Mono.from(this.authorizationErrorHandler.onAuthorizationError(
-										responseEvent.responseInfo(), transportContext, Mono.defer(() -> {
-											logger.debug("Authorization error handled, retrying original request");
-											return this.reconnect(stream).then();
-										}),
-										Mono.error(new McpHttpClientTransportException(
-												"Authorization error connecting to SSE stream",
-												responseEvent.responseInfo()))))
-									.then(Mono.empty());
-							});
+							logger.debug("Authorization error in reconnect with code {}", statusCode);
+							return Mono.<McpSchema.JSONRPCMessage>error(
+									new McpHttpClientTransportAuthorizationException(
+											"Authorization error connecting to SSE stream",
+											responseEvent.responseInfo()));
 						}
 
 						if (!(responseEvent instanceof ResponseSubscribers.SseResponseEvent sseResponseEvent)) {
-							return Flux.<McpSchema.JSONRPCMessage>error(new McpHttpClientTransportException(
-									"Unrecognized server error when connecting to SSE stream",
-									responseEvent.responseInfo()));
+							return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
+									"Unrecognized server error when connecting to SSE stream, status code: "
+											+ statusCode));
 						}
 						else if (statusCode >= 200 && statusCode < 300) {
 							if (MESSAGE_EVENT_TYPE.equals(sseResponseEvent.sseEvent().event())) {
@@ -389,6 +381,7 @@ else if (statusCode == BAD_REQUEST) {
 						return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
 								"Received unrecognized SSE event type: " + sseResponseEvent.sseEvent().event()));
 					})
+					.retryWhen(authorizationErrorRetrySpec())
 					.flatMap(jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
 					.onErrorMap(CompletionException.class, t -> t.getCause())
 					.onErrorComplete(t -> {
@@ -411,6 +404,25 @@ else if (statusCode == BAD_REQUEST) {
 
 	}
 
+	private Retry authorizationErrorRetrySpec() {
+		return Retry.from(companion -> companion.flatMap(retrySignal -> {
+			if (!(retrySignal.failure() instanceof McpHttpClientTransportAuthorizationException authException)) {
+				return Mono.error(retrySignal.failure());
+			}
+			if (retrySignal.totalRetriesInARow() >= this.authorizationErrorHandler.maxRetries()) {
+				return Mono.error(retrySignal.failure());
+			}
+			return Mono.deferContextual(ctx -> {
+				var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
+				return Mono
+					.from(this.authorizationErrorHandler.handle(authException.getResponseInfo(), transportContext))
+					.switchIfEmpty(Mono.just(false))
+					.flatMap(shouldRetry -> shouldRetry ? Mono.just(retrySignal.totalRetries())
+							: Mono.error(retrySignal.failure()));
+			});
+		}));
+	}
+
 	private BodyHandler<Void> toSendMessageBodySubscriber(FluxSink<ResponseEvent> sink) {
 
 		BodyHandler<Void> responseBodyHandler = responseInfo -> {
@@ -492,17 +504,8 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
 				int statusCode = responseEvent.responseInfo().statusCode();
 				if (statusCode == 401 || statusCode == 403) {
 					logger.debug("Authorization error in sendMessage with code {}", statusCode);
-					return Mono.deferContextual(ctx -> {
-						var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
-						return Mono.from(this.authorizationErrorHandler
-							.onAuthorizationError(responseEvent.responseInfo(), transportContext, Mono.defer(() -> {
-								logger.debug("Authorization error handled, retrying original request");
-								return this.sendMessage(sentMessage);
-							}), Mono.error(new McpHttpClientTransportException(
-									"Authorization error when sending message", responseEvent.responseInfo()))))
-							.doOnSuccess(s -> deliveredSink.success())
-							.then(Mono.empty());
-					});
+					return Mono.<McpSchema.JSONRPCMessage>error(new McpHttpClientTransportAuthorizationException(
+							"Authorization error when sending message", responseEvent.responseInfo()));
 				}
 
 				if (transportSession.markInitialized(
@@ -630,6 +633,7 @@ else if (statusCode == BAD_REQUEST) {
 				return Flux.<McpSchema.JSONRPCMessage>error(
 						new RuntimeException("Failed to send message: " + responseEvent));
 			})
+				.retryWhen(authorizationErrorRetrySpec())
 				.flatMap(jsonRpcMessage -> this.handler.get().apply(Mono.just(jsonRpcMessage)))
 				.onErrorMap(CompletionException.class, t -> t.getCause())
 				.onErrorComplete(t -> {

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2026-2026 the original author or authors.
+ */
+
+package io.modelcontextprotocol.client.transport;
+
+import java.net.http.HttpResponse;
+
+import io.modelcontextprotocol.spec.McpTransportException;
+
+/**
+ * Thrown when the MCP server responds with an authorization error (HTTP 401 or HTTP 403).
+ * Subclass of {@link McpTransportException} for targeted retry handling in
+ * {@link HttpClientStreamableHttpTransport}.
+ *
+ * @author Daniel Garnier-Moiroux
+ */
+public class McpHttpClientTransportAuthorizationException extends McpTransportException {
+
+	private final HttpResponse.ResponseInfo responseInfo;
+
+	public McpHttpClientTransportAuthorizationException(String message, HttpResponse.ResponseInfo responseInfo) {
+		super(message);
+		this.responseInfo = responseInfo;
+	}
+
+	public HttpResponse.ResponseInfo getResponseInfo() {
+		return responseInfo;
+	}
+
+}

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

@@ -6,7 +6,7 @@
 
 import java.net.http.HttpResponse;
 
-import io.modelcontextprotocol.client.transport.McpHttpClientTransportException;
+import io.modelcontextprotocol.client.transport.McpHttpClientTransportAuthorizationException;
 import io.modelcontextprotocol.common.McpTransportContext;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Mono;
@@ -27,13 +27,13 @@ public interface McpHttpClientAuthorizationErrorHandler {
 	 * Handle authorization error (HTTP 401 or 403), and signal whether the HTTP request
 	 * should be retried or not. If the publisher returns true, the original transport
 	 * method (connect, sendMessage) will be replayed with the original arguments.
-	 * Otherwise, the transport will throw an {@link McpHttpClientTransportException},
-	 * indicating the error status.
+	 * Otherwise, the transport will throw an
+	 * {@link McpHttpClientTransportAuthorizationException}, indicating the error status.
 	 * <p>
 	 * If the returned {@link Publisher} errors, the error will be propagated to the
 	 * calling method, to be handled by the caller.
 	 * <p>
-	 * The caller is responsible for bounding the number of retries.
+	 * The number of retries is bounded by {@link #maxRetries()}.
 	 * @param responseInfo the HTTP response information
 	 * @param context the MCP client transport context
 	 * @return {@link Publisher} emitting true if the original request should be replayed,
@@ -42,36 +42,23 @@ public interface McpHttpClientAuthorizationErrorHandler {
 	Publisher<Boolean> handle(HttpResponse.ResponseInfo responseInfo, McpTransportContext context);
 
 	/**
-	 * A no-op handler, used in the default use-case.
+	 * Maximum number of authorization error retries the transport will attempt. When the
+	 * handler signals a retry via {@link #handle}, the transport will replay the original
+	 * request at most this many times. If the authorization error persists after
+	 * exhausting all retries, the transport will propagate the
+	 * {@link McpHttpClientTransportAuthorizationException}.
+	 * <p>
+	 * Defaults to {@code 1}.
+	 * @return the maximum number of retries
 	 */
-	McpHttpClientAuthorizationErrorHandler NOOP = new Noop();
+	default int maxRetries() {
+		return 1;
+	}
 
 	/**
-	 * Handle authorization error (HTTP 401 or 403), and optionally retry the HTTP
-	 * request, or trigger a transport error. To retry, use the {@code retryAction}
-	 * publisher. To emit the default transport error, use the {@code defaultError}
-	 * publisher.
-	 * <p>
-	 * Optionally, the returned {@link Publisher} may error to trigger an out-of-band
-	 * action. In that case, the error will be propagated to the calling method, to be
-	 * handled by the caller.
-	 * <p>
-	 * Defaults to {@link #handle(HttpResponse.ResponseInfo, McpTransportContext)}, and
-	 * uses the boolean from the return value to decide whether it should retry the
-	 * request.
-	 * @param responseInfo the HTTP response information
-	 * @param context the MCP client transport context
-	 * @param retryAction handler to retry the original request
-	 * @param defaultError handler to emit an error
-	 * @return a {@link Publisher} to signal either an error or a retry
+	 * A no-op handler, used in the default use-case.
 	 */
-	default Publisher<Void> onAuthorizationError(HttpResponse.ResponseInfo responseInfo, McpTransportContext context,
-			Publisher<Void> retryAction, Publisher<Void> defaultError) {
-		return Mono.from(this.handle(responseInfo, context))
-			.switchIfEmpty(Mono.just(false))
-			.flatMap(shouldRetry -> shouldRetry != null && shouldRetry ? Mono.from(retryAction)
-					: Mono.from(defaultError));
-	}
+	McpHttpClientAuthorizationErrorHandler NOOP = new Noop();
 
 	/**
 	 * Create a {@link McpHttpClientAuthorizationErrorHandler} from a synchronous handler.
@@ -95,7 +82,8 @@ interface Sync {
 		 * request should be retried or not. If the return value is true, the original
 		 * transport method (connect, sendMessage) will be replayed with the original
 		 * arguments. Otherwise, the transport will throw an
-		 * {@link McpHttpClientTransportException}, indicating the error status.
+		 * {@link McpHttpClientTransportAuthorizationException}, indicating the error
+		 * status.
 		 * @param responseInfo the HTTP response information
 		 * @param context the MCP client transport context
 		 * @return true if the original request should be replayed, false otherwise.

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

@@ -6,9 +6,7 @@
 import java.net.http.HttpResponse;
 
 import io.modelcontextprotocol.common.McpTransportContext;
-import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
-import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import static org.mockito.Mockito.mock;
@@ -22,85 +20,29 @@ class McpHttpClientAuthorizationErrorHandlerTest {
 
 	private final McpTransportContext context = McpTransportContext.EMPTY;
 
-	@Nested
-	class OnAuthorizationError {
-
-		@Test
-		void whenTrueThenRetry() {
-			McpHttpClientAuthorizationErrorHandler handler = (info, ctx) -> Mono.just(true);
-			Mono<Void> retryAction = Mono.empty();
-			Mono<Void> defaultError = Mono.error(new RuntimeException("should not be called"));
-
-			StepVerifier.create(handler.onAuthorizationError(responseInfo, context, retryAction, defaultError))
-				.verifyComplete();
-		}
-
-		@Test
-		void whenFalseThenError() {
-			McpHttpClientAuthorizationErrorHandler handler = (info, ctx) -> Mono.just(false);
-			Mono<Void> retryAction = Mono.error(new RuntimeException("should not be called"));
-			Mono<Void> defaultError = Mono.error(new RuntimeException("authorization error"));
-
-			StepVerifier.create(handler.onAuthorizationError(responseInfo, context, retryAction, defaultError))
-				.expectErrorMatches(t -> t instanceof RuntimeException && t.getMessage().equals("authorization error"))
-				.verify();
-		}
-
-		@Test
-		void whenErrorThenPropagate() {
-			McpHttpClientAuthorizationErrorHandler handler = (info, ctx) -> Mono
-				.error(new IllegalStateException("handler error"));
-			Mono<Void> retryAction = Mono.error(new RuntimeException("should not be called"));
-			Mono<Void> defaultError = Mono.error(new RuntimeException("should not be called"));
-
-			StepVerifier.create(handler.onAuthorizationError(responseInfo, context, retryAction, defaultError))
-				.expectErrorMatches(t -> t instanceof IllegalStateException && t.getMessage().equals("handler error"))
-				.verify();
-		}
-
+	@Test
+	void whenTrueThenRetry() {
+		McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
+			.fromSync((info, ctx) -> true);
+		StepVerifier.create(handler.handle(responseInfo, context)).expectNext(true).verifyComplete();
 	}
 
-	@Nested
-	class FromSync {
-
-		@Test
-		void whenTrueThenRetry() {
-			McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
-				.fromSync((info, ctx) -> true);
-			Mono<Void> retryAction = Mono.empty();
-			Mono<Void> defaultError = Mono.error(new RuntimeException("should not be called"));
-
-			StepVerifier.create(handler.onAuthorizationError(responseInfo, context, retryAction, defaultError))
-				.verifyComplete();
-		}
-
-		@Test
-		void whenFalseThenError() {
-			McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
-				.fromSync((info, ctx) -> false);
-			Mono<Void> retryAction = Mono.error(new RuntimeException("should not be called"));
-			Mono<Void> defaultError = Mono.error(new RuntimeException("authorization error"));
-
-			StepVerifier.create(handler.onAuthorizationError(responseInfo, context, retryAction, defaultError))
-				.expectErrorMatches(t -> t instanceof RuntimeException && t.getMessage().equals("authorization error"))
-				.verify();
-		}
-
-		@Test
-		void whenExceptionThenPropagate() {
-			McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
-				.fromSync((info, ctx) -> {
-					throw new IllegalStateException("sync handler error");
-				});
-			Mono<Void> retryAction = Mono.error(new RuntimeException("should not be called"));
-			Mono<Void> defaultError = Mono.error(new RuntimeException("should not be called"));
-
-			StepVerifier.create(handler.onAuthorizationError(responseInfo, context, retryAction, defaultError))
-				.expectErrorMatches(
-						t -> t instanceof IllegalStateException && t.getMessage().equals("sync handler error"))
-				.verify();
-		}
+	@Test
+	void whenFalseThenError() {
+		McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
+			.fromSync((info, ctx) -> false);
+		StepVerifier.create(handler.handle(responseInfo, context)).expectNext(false).verifyComplete();
+	}
 
+	@Test
+	void whenExceptionThenPropagate() {
+		McpHttpClientAuthorizationErrorHandler handler = McpHttpClientAuthorizationErrorHandler
+			.fromSync((info, ctx) -> {
+				throw new IllegalStateException("sync handler error");
+			});
+		StepVerifier.create(handler.handle(responseInfo, context))
+			.expectErrorMatches(t -> t instanceof IllegalStateException && t.getMessage().equals("sync handler error"))
+			.verify();
 	}
 
 }