fix(client): enforce pagination timeout across all pages

Author: ZYZ666-RGB

mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

@@ -1055,14 +1055,20 @@ private <R, A> Mono<R> paginate(Function<String, Mono<R>> pageFetcher, Function<
 			Supplier<A> initialAccumulator, BiFunction<A, R, A> accumulate, Function<A, R> finalize) {
 		return Mono.defer(() -> {
 			PaginationGuard guard = new PaginationGuard(this.paginationConfig);
-			return pageFetcher.apply(McpSchema.FIRST_PAGE).expand(page -> {
+			Mono<R> result = pageFetcher.apply(McpSchema.FIRST_PAGE).expand(page -> {
 				String next = nextCursorOf.apply(page);
 				if (next == null || next.isEmpty()) {
 					return Mono.empty();
 				}
 				guard.beforeNextPage(next);
 				return pageFetcher.apply(next);
 			}).reduce(initialAccumulator.get(), accumulate).map(finalize);
+			if (this.paginationConfig.timeout() != null) {
+				return result.timeout(this.paginationConfig.timeout(), Mono.error(new McpPaginationException(
+						"Pagination timed out after " + this.paginationConfig.timeout()
+								+ ". Increase paginationTimeout if this is expected for the server.")));
+			}
+			return result;
 		});
 	}
 
@@ -1077,8 +1083,6 @@ private static final class PaginationGuard {
 
 		private final PaginationConfig config;
 
-		private final long startNanos = System.nanoTime();
-
 		private int pagesFetched = 1;
 
 		PaginationGuard(PaginationConfig config) {
@@ -1100,11 +1104,6 @@ void beforeNextPage(String cursor) {
 						"Pagination limit exceeded: the server returned more than " + this.config.maxPages()
 								+ " pages. Increase maxPaginationPages if this is expected for the server.");
 			}
-			if (this.config.timeout() != null
-					&& Duration.ofNanos(System.nanoTime() - this.startNanos).compareTo(this.config.timeout()) > 0) {
-				throw new McpPaginationException("Pagination timed out after " + this.config.timeout()
-						+ ". Increase paginationTimeout if this is expected for the server.");
-			}
 			this.pagesFetched++;
 		}
 

mcp-test/src/test/java/io/modelcontextprotocol/client/McpAsyncClientPaginationTests.java

@@ -65,6 +65,10 @@ default Duration pageDelay() {
 			return Duration.ZERO;
 		}
 
+		default Duration pageDelay(String cursor) {
+			return pageDelay();
+		}
+
 	}
 
 	private McpClientTransport createPaginatedTransport(PaginatedServer server, AtomicInteger toolsRequests,
@@ -119,8 +123,9 @@ else if (McpSchema.METHOD_RESOURCES_LIST.equals(request.method())) {
 				}
 
 				Mono<McpSchema.JSONRPCMessage> responseMono = Mono.just(response);
-				if (!server.pageDelay().isZero()) {
-					responseMono = responseMono.delayElement(server.pageDelay());
+				Duration delay = server.pageDelay(cursorOf(request));
+				if (!delay.isZero()) {
+					responseMono = responseMono.delayElement(delay);
 				}
 				return responseMono.flatMap(r -> handler.apply(Mono.just(r))).then();
 			}
@@ -235,6 +240,34 @@ public Duration pageDelay() {
 		}).verify();
 	}
 
+	@Test
+	void listToolsTimesOutWhileWaitingForFinalPage() {
+		AtomicInteger toolsRequests = new AtomicInteger();
+		AtomicInteger resourcesRequests = new AtomicInteger();
+		PaginatedServer server = new PaginatedServer() {
+			@Override
+			public String nextCursorFor(String cursor) {
+				return cursor == null ? "last-page" : null;
+			}
+
+			@Override
+			public Duration pageDelay(String cursor) {
+				return cursor == null ? Duration.ofMillis(20) : Duration.ofMillis(800);
+			}
+		};
+		McpAsyncClient client = McpClient.async(createPaginatedTransport(server, toolsRequests, resourcesRequests))
+			.paginationTimeout(Duration.ofMillis(300))
+			.build();
+
+		client.initialize().block();
+
+		StepVerifier.create(client.listTools()).expectErrorSatisfies(error -> {
+			assertThat(error).isInstanceOf(McpPaginationException.class);
+			assertThat(error.getMessage()).contains("timed out");
+		}).verify();
+		assertThat(toolsRequests.get()).isEqualTo(2);
+	}
+
 	@Test
 	void listResourcesAggregatesAllPagesUntilNullCursor() {
 		AtomicInteger toolsRequests = new AtomicInteger();