feat: Add progress notification support for MCP operations (#407)

Implement progress tracking for long-running operations with:
- New ProgressNotification schema and client/server support
- Progress consumer handlers in sync/async client builders
- Server exchange methods for sending progress updates
- Comprehensive integration tests
- Backwards compatibility maintained
- Add additional tests

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>

Author: 134130

mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java

@@ -1016,6 +1016,113 @@ void testLoggingNotification(String clientType) throws InterruptedException {
 		mcpServer.close();
 	}
 
+	// ---------------------------------------
+	// Progress Tests
+	// ---------------------------------------
+	@ParameterizedTest(name = "{0} : {displayName} ")
+	@ValueSource(strings = { "httpclient", "webflux" })
+	void testProgressNotification(String clientType) throws InterruptedException {
+		int expectedNotificationsCount = 4; // 3 notifications + 1 for another progress
+											// token
+		CountDownLatch latch = new CountDownLatch(expectedNotificationsCount);
+		// Create a list to store received logging notifications
+		List<McpSchema.ProgressNotification> receivedNotifications = new CopyOnWriteArrayList<>();
+
+		var clientBuilder = clientBuilders.get(clientType);
+
+		// Create server with a tool that sends logging notifications
+		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
+			.tool(McpSchema.Tool.builder()
+				.name("progress-test")
+				.description("Test progress notifications")
+				.inputSchema(emptyJsonSchema)
+				.build())
+			.callHandler((exchange, request) -> {
+
+				// Create and send notifications
+				var progressToken = (String) request.meta().get("progressToken");
+
+				return exchange
+					.progressNotification(
+							new McpSchema.ProgressNotification(progressToken, 0.0, 1.0, "Processing started"))
+					.then(exchange.progressNotification(
+							new McpSchema.ProgressNotification(progressToken, 0.5, 1.0, "Processing data")))
+					.then(// Send a progress notification with another progress value
+							// should
+							exchange.progressNotification(new McpSchema.ProgressNotification("another-progress-token",
+									0.0, 1.0, "Another processing started")))
+					.then(exchange.progressNotification(
+							new McpSchema.ProgressNotification(progressToken, 1.0, 1.0, "Processing completed")))
+					.thenReturn(new CallToolResult(("Progress test completed"), false));
+			})
+			.build();
+
+		var mcpServer = McpServer.async(mcpServerTransportProvider)
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.tools(tool)
+			.build();
+
+		try (
+				// Create client with progress notification handler
+				var mcpClient = clientBuilder.progressConsumer(notification -> {
+					receivedNotifications.add(notification);
+					latch.countDown();
+				}).build()) {
+
+			// Initialize client
+			InitializeResult initResult = mcpClient.initialize();
+			assertThat(initResult).isNotNull();
+
+			// Call the tool that sends progress notifications
+			McpSchema.CallToolRequest callToolRequest = McpSchema.CallToolRequest.builder()
+				.name("progress-test")
+				.meta(Map.of("progressToken", "test-progress-token"))
+				.build();
+			CallToolResult result = mcpClient.callTool(callToolRequest);
+			assertThat(result).isNotNull();
+			assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class);
+			assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Progress test completed");
+
+			assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue();
+
+			// Should have received 3 notifications
+			assertThat(receivedNotifications).hasSize(expectedNotificationsCount);
+
+			Map<String, McpSchema.ProgressNotification> notificationMap = receivedNotifications.stream()
+				.collect(Collectors.toMap(n -> n.message(), n -> n));
+
+			// First notification should be 0.0/1.0 progress
+			assertThat(notificationMap.get("Processing started").progressToken()).isEqualTo("test-progress-token");
+			assertThat(notificationMap.get("Processing started").progress()).isEqualTo(0.0);
+			assertThat(notificationMap.get("Processing started").total()).isEqualTo(1.0);
+			assertThat(notificationMap.get("Processing started").message()).isEqualTo("Processing started");
+
+			// Second notification should be 0.5/1.0 progress
+			assertThat(notificationMap.get("Processing data").progressToken()).isEqualTo("test-progress-token");
+			assertThat(notificationMap.get("Processing data").progress()).isEqualTo(0.5);
+			assertThat(notificationMap.get("Processing data").total()).isEqualTo(1.0);
+			assertThat(notificationMap.get("Processing data").message()).isEqualTo("Processing data");
+
+			// Third notification should be another progress token with 0.0/1.0 progress
+			assertThat(notificationMap.get("Another processing started").progressToken())
+				.isEqualTo("another-progress-token");
+			assertThat(notificationMap.get("Another processing started").progress()).isEqualTo(0.0);
+			assertThat(notificationMap.get("Another processing started").total()).isEqualTo(1.0);
+			assertThat(notificationMap.get("Another processing started").message())
+				.isEqualTo("Another processing started");
+
+			// Fourth notification should be 1.0/1.0 progress
+			assertThat(notificationMap.get("Processing completed").progressToken()).isEqualTo("test-progress-token");
+			assertThat(notificationMap.get("Processing completed").progress()).isEqualTo(1.0);
+			assertThat(notificationMap.get("Processing completed").total()).isEqualTo(1.0);
+			assertThat(notificationMap.get("Processing completed").message()).isEqualTo("Processing completed");
+		}
+		finally {
+			mcpServer.close();
+		}
+	}
+
 	// ---------------------------------------
 	// Completion Tests
 	// ---------------------------------------

mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java

@@ -12,8 +12,10 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -49,6 +51,7 @@
 import io.modelcontextprotocol.spec.McpTransport;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
 import reactor.test.StepVerifier;
 
 /**
@@ -420,7 +423,7 @@ void testListAllPromptsReturnsImmutableList() {
 				.consumeNextWith(result -> {
 					assertThat(result.prompts()).isNotNull();
 					// Verify that the returned list is immutable
-					assertThatThrownBy(() -> result.prompts().add(new Prompt("test", "Test", "test", null)))
+					assertThatThrownBy(() -> result.prompts().add(new Prompt("test", "test", "test", null)))
 						.isInstanceOf(UnsupportedOperationException.class);
 				})
 				.verifyComplete();
@@ -792,4 +795,39 @@ void testSampling() {
 			});
 	}
 
+	// ---------------------------------------
+	// Progress Notification Tests
+	// ---------------------------------------
+
+	@Test
+	void testProgressConsumer() {
+		Sinks.Many<McpSchema.ProgressNotification> sink = Sinks.many().unicast().onBackpressureBuffer();
+		List<McpSchema.ProgressNotification> receivedNotifications = new CopyOnWriteArrayList<>();
+
+		withClient(createMcpTransport(), builder -> builder.progressConsumer(notification -> {
+			receivedNotifications.add(notification);
+			sink.tryEmitNext(notification);
+			return Mono.empty();
+		}), client -> {
+			StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
+
+			// Call a tool that sends progress notifications
+			CallToolRequest request = CallToolRequest.builder()
+				.name("longRunningOperation")
+				.arguments(Map.of("duration", 1, "steps", 2))
+				.progressToken("test-token")
+				.build();
+
+			StepVerifier.create(client.callTool(request)).consumeNextWith(result -> {
+				assertThat(result).isNotNull();
+			}).verifyComplete();
+
+			// Use StepVerifier to verify the progress notifications via the sink
+			StepVerifier.create(sink.asFlux()).expectNextCount(2).thenCancel().verify(Duration.ofSeconds(3));
+
+			assertThat(receivedNotifications).hasSize(2);
+			assertThat(receivedNotifications.get(0).progressToken()).isEqualTo("test-token");
+		});
+	}
+
 }

mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java

@@ -13,6 +13,9 @@
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -648,4 +651,48 @@ void testSampling() {
 			});
 	}
 
+	// ---------------------------------------
+	// Progress Notification Tests
+	// ---------------------------------------
+
+	@Test
+	void testProgressConsumer() {
+		AtomicInteger progressNotificationCount = new AtomicInteger(0);
+		List<McpSchema.ProgressNotification> receivedNotifications = new CopyOnWriteArrayList<>();
+		CountDownLatch latch = new CountDownLatch(2);
+
+		withClient(createMcpTransport(), builder -> builder.progressConsumer(notification -> {
+			System.out.println("Received progress notification: " + notification);
+			receivedNotifications.add(notification);
+			progressNotificationCount.incrementAndGet();
+			latch.countDown();
+		}), client -> {
+			client.initialize();
+
+			// Call a tool that sends progress notifications
+			CallToolRequest request = CallToolRequest.builder()
+				.name("longRunningOperation")
+				.arguments(Map.of("duration", 1, "steps", 2))
+				.progressToken("test-token")
+				.build();
+
+			CallToolResult result = client.callTool(request);
+
+			assertThat(result).isNotNull();
+
+			try {
+				// Wait for progress notifications to be processed
+				latch.await(3, TimeUnit.SECONDS);
+			}
+			catch (InterruptedException e) {
+				e.printStackTrace();
+			}
+
+			assertThat(progressNotificationCount.get()).isEqualTo(2);
+
+			assertThat(receivedNotifications).isNotEmpty();
+			assertThat(receivedNotifications.get(0).progressToken()).isEqualTo("test-token");
+		});
+	}
+
 }

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

@@ -100,6 +100,9 @@ public class McpAsyncClient {
 	public static final TypeReference<LoggingMessageNotification> LOGGING_MESSAGE_NOTIFICATION_TYPE_REF = new TypeReference<>() {
 	};
 
+	public static final TypeReference<McpSchema.ProgressNotification> PROGRESS_NOTIFICATION_TYPE_REF = new TypeReference<>() {
+	};
+
 	/**
 	 * Client capabilities.
 	 */
@@ -253,6 +256,16 @@ public class McpAsyncClient {
 		notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_MESSAGE,
 				asyncLoggingNotificationHandler(loggingConsumersFinal));
 
+		// Utility Progress Notification
+		List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumersFinal = new ArrayList<>();
+		progressConsumersFinal
+			.add((notification) -> Mono.fromRunnable(() -> logger.debug("Progress: {}", notification)));
+		if (!Utils.isEmpty(features.progressConsumers())) {
+			progressConsumersFinal.addAll(features.progressConsumers());
+		}
+		notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_PROGRESS,
+				asyncProgressNotificationHandler(progressConsumersFinal));
+
 		this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo,
 				List.of(McpSchema.LATEST_PROTOCOL_VERSION), initializationTimeout,
 				ctx -> new McpClientSession(requestTimeout, transport, requestHandlers, notificationHandlers,
@@ -846,6 +859,28 @@ public Mono<Void> setLoggingLevel(LoggingLevel loggingLevel) {
 		});
 	}
 
+	/**
+	 * Create a notification handler for progress notifications from the server. This
+	 * handler automatically distributes progress notifications to all registered
+	 * consumers.
+	 * @param progressConsumers List of consumers that will be notified when a progress
+	 * message is received. Each consumer receives the progress notification.
+	 * @return A NotificationHandler that processes progress notifications by distributing
+	 * the message to all registered consumers
+	 */
+	private NotificationHandler asyncProgressNotificationHandler(
+			List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers) {
+
+		return params -> {
+			McpSchema.ProgressNotification progressNotification = transport.unmarshalFrom(params,
+					PROGRESS_NOTIFICATION_TYPE_REF);
+
+			return Flux.fromIterable(progressConsumers)
+				.flatMap(consumer -> consumer.apply(progressNotification))
+				.then();
+		};
+	}
+
 	/**
 	 * This method is package-private and used for test only. Should not be called by user
 	 * code.