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:
parent
87cdaf883a
commit
b04c30b60e
3 changed files with 535 additions and 8 deletions
|
|
@ -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");
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
// ---------------------------------------
|
||||
|
|
|
|||
Loading…
Reference in a new issue