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>
This commit is contained in:
JiHwan Oh 2025-07-17 10:31:34 +09:00 • committed by Christian Tzolov
commit b04c30b60e
3 changed files with 535 additions and 8 deletions

View file

@ -12,8 +12,10 @@ import static org.junit.jupiter.api.Assertions.assertInstanceOf;
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.McpSchema.UnsubscribeRequest;
import io.modelcontextprotocol.spec.McpTransport;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.test.StepVerifier;
/**
@ -793,4 +796,39 @@ public abstract class AbstractMcpAsyncClientTests {
});
}
// ---------------------------------------
// 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");
});
}
}

View file

@ -13,6 +13,9 @@ import static org.junit.jupiter.api.Assertions.assertInstanceOf;
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;
@ -649,4 +652,48 @@ public abstract class AbstractMcpSyncClientTests {
});
}
// ---------------------------------------
// 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");
});
}
}

View file

@ -7,6 +7,7 @@ import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
@ -907,7 +908,7 @@ class HttpServletSseServerTransportProviderIntegrationTests {
@Test
void testLoggingNotification() {
// Create a list to store received logging notifications
List<McpSchema.LoggingMessageNotification> receivedNotifications = new ArrayList<>();
List<McpSchema.LoggingMessageNotification> receivedNotifications = new CopyOnWriteArrayList<>();
// Create server with a tool that sends logging notifications
McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
@ -1020,6 +1021,80 @@ class HttpServletSseServerTransportProviderIntegrationTests {
mcpServer.close();
}
// ---------------------------------------
// Progress Tests
// ---------------------------------------
@Test
void testProgressNotification() {
// Create a list to store received progress notifications
List<McpSchema.ProgressNotification> receivedNotifications = new CopyOnWriteArrayList<>();
// Create server with a tool that sends progress notifications
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
McpSchema.Tool.builder()
.name("progress-test")
.description("Test progress notifications")
.inputSchema(emptyJsonSchema)
.build(),
null, (exchange, request) -> {
var progressToken = request.progressToken();
exchange
.progressNotification(
new McpSchema.ProgressNotification(progressToken, 0.1, 1.0, "Test progress 1/10"))
.block();
exchange
.progressNotification(
new McpSchema.ProgressNotification(progressToken, 0.5, 1.0, "Test progress 5/10"))
.block();
exchange
.progressNotification(
new McpSchema.ProgressNotification(progressToken, 1.0, 1.0, "Test progress 10/10"))
.block();
return Mono.just(new CallToolResult("Progress test completed", false));
});
var mcpServer = McpServer.async(mcpServerTransportProvider)
.serverInfo("test-server", "1.0.0")
.capabilities(ServerCapabilities.builder().logging().tools(true).build())
.tools(tool)
.build();
// Create client with progress notification handler
try (var mcpClient = clientBuilder.progressConsumer(receivedNotifications::add).build()) {
// Initialize client
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
// Call the tool that sends progress notifications
CallToolResult result = mcpClient.callTool(
new McpSchema.CallToolRequest("progress-test", Map.of(), Map.of("progressToken", "test-token")));
assertThat(result).isNotNull();
assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class);
assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Progress test completed");
// Wait for notifications to be processed
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
// Should have received 3 notifications
assertThat(receivedNotifications).hasSize(3);
// Check the progress notifications
assertThat(receivedNotifications.stream().map(McpSchema.ProgressNotification::progressToken))
.containsExactlyInAnyOrder("test-token", "test-token", "test-token");
assertThat(receivedNotifications.stream().map(McpSchema.ProgressNotification::progress))
.containsExactlyInAnyOrder(0.1, 0.5, 1.0);
});
}
finally {
mcpServer.close();
}
}
// ---------------------------------------
// Ping Tests
// ---------------------------------------