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
13 changed files with 535 additions and 8 deletions
|
|
@ -1016,6 +1016,113 @@ class WebFluxSseIntegrationTests {
|
|||
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
|
||||
// ---------------------------------------
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
/**
|
||||
|
|
@ -420,7 +423,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
.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 @@ 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;
|
||||
|
|
@ -648,4 +651,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");
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 class McpAsyncClient {
|
|||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* 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.
|
||||
|
|
|
|||
|
|
@ -177,6 +177,8 @@ public interface McpClient {
|
|||
|
||||
private final List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers = new ArrayList<>();
|
||||
|
||||
private final List<Consumer<McpSchema.ProgressNotification>> progressConsumers = new ArrayList<>();
|
||||
|
||||
private Function<CreateMessageRequest, CreateMessageResult> samplingHandler;
|
||||
|
||||
private Function<ElicitRequest, ElicitResult> elicitationHandler;
|
||||
|
|
@ -377,6 +379,36 @@ public interface McpClient {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds a consumer to be notified of progress notifications from the server. This
|
||||
* allows the client to track long-running operations and provide feedback to
|
||||
* users.
|
||||
* @param progressConsumer A consumer that receives progress notifications. Must
|
||||
* not be null.
|
||||
* @return This builder instance for method chaining
|
||||
* @throws IllegalArgumentException if progressConsumer is null
|
||||
*/
|
||||
public SyncSpec progressConsumer(Consumer<McpSchema.ProgressNotification> progressConsumer) {
|
||||
Assert.notNull(progressConsumer, "Progress consumer must not be null");
|
||||
this.progressConsumers.add(progressConsumer);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds a multiple consumers to be notified of progress notifications from the
|
||||
* server. This allows the client to track long-running operations and provide
|
||||
* feedback to users.
|
||||
* @param progressConsumers A list of consumers that receives progress
|
||||
* notifications. Must not be null.
|
||||
* @return This builder instance for method chaining
|
||||
* @throws IllegalArgumentException if progressConsumer is null
|
||||
*/
|
||||
public SyncSpec progressConsumers(List<Consumer<McpSchema.ProgressNotification>> progressConsumers) {
|
||||
Assert.notNull(progressConsumers, "Progress consumers must not be null");
|
||||
this.progressConsumers.addAll(progressConsumers);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an instance of {@link McpSyncClient} with the provided configurations or
|
||||
* sensible defaults.
|
||||
|
|
@ -385,7 +417,8 @@ public interface McpClient {
|
|||
public McpSyncClient build() {
|
||||
McpClientFeatures.Sync syncFeatures = new McpClientFeatures.Sync(this.clientInfo, this.capabilities,
|
||||
this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
|
||||
this.promptsChangeConsumers, this.loggingConsumers, this.samplingHandler, this.elicitationHandler);
|
||||
this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers, this.samplingHandler,
|
||||
this.elicitationHandler);
|
||||
|
||||
McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
|
||||
|
||||
|
|
@ -435,6 +468,8 @@ public interface McpClient {
|
|||
|
||||
private final List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers = new ArrayList<>();
|
||||
|
||||
private final List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers = new ArrayList<>();
|
||||
|
||||
private Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler;
|
||||
|
||||
private Function<ElicitRequest, Mono<ElicitResult>> elicitationHandler;
|
||||
|
|
@ -654,6 +689,37 @@ public interface McpClient {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds a consumer to be notified of progress notifications from the server. This
|
||||
* allows the client to track long-running operations and provide feedback to
|
||||
* users.
|
||||
* @param progressConsumer A consumer that receives progress notifications. Must
|
||||
* not be null.
|
||||
* @return This builder instance for method chaining
|
||||
* @throws IllegalArgumentException if progressConsumer is null
|
||||
*/
|
||||
public AsyncSpec progressConsumer(Function<McpSchema.ProgressNotification, Mono<Void>> progressConsumer) {
|
||||
Assert.notNull(progressConsumer, "Progress consumer must not be null");
|
||||
this.progressConsumers.add(progressConsumer);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds a multiple consumers to be notified of progress notifications from the
|
||||
* server. This allows the client to track long-running operations and provide
|
||||
* feedback to users.
|
||||
* @param progressConsumers A list of consumers that receives progress
|
||||
* notifications. Must not be null.
|
||||
* @return This builder instance for method chaining
|
||||
* @throws IllegalArgumentException if progressConsumer is null
|
||||
*/
|
||||
public AsyncSpec progressConsumers(
|
||||
List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers) {
|
||||
Assert.notNull(progressConsumers, "Progress consumers must not be null");
|
||||
this.progressConsumers.addAll(progressConsumers);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an instance of {@link McpAsyncClient} with the provided configurations
|
||||
* or sensible defaults.
|
||||
|
|
@ -663,8 +729,8 @@ public interface McpClient {
|
|||
return new McpAsyncClient(this.transport, this.requestTimeout, this.initializationTimeout,
|
||||
new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots,
|
||||
this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
|
||||
this.promptsChangeConsumers, this.loggingConsumers, this.samplingHandler,
|
||||
this.elicitationHandler));
|
||||
this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers,
|
||||
this.samplingHandler, this.elicitationHandler));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ class McpClientFeatures {
|
|||
* @param resourcesChangeConsumers the resources change consumers.
|
||||
* @param promptsChangeConsumers the prompts change consumers.
|
||||
* @param loggingConsumers the logging consumers.
|
||||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
*/
|
||||
|
|
@ -68,6 +69,7 @@ class McpClientFeatures {
|
|||
List<Function<List<McpSchema.ResourceContents>, Mono<Void>>> resourcesUpdateConsumers,
|
||||
List<Function<List<McpSchema.Prompt>, Mono<Void>>> promptsChangeConsumers,
|
||||
List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers,
|
||||
List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
|
||||
|
||||
|
|
@ -79,6 +81,7 @@ class McpClientFeatures {
|
|||
* @param resourcesChangeConsumers the resources change consumers.
|
||||
* @param promptsChangeConsumers the prompts change consumers.
|
||||
* @param loggingConsumers the logging consumers.
|
||||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
*/
|
||||
|
|
@ -89,6 +92,7 @@ class McpClientFeatures {
|
|||
List<Function<List<McpSchema.ResourceContents>, Mono<Void>>> resourcesUpdateConsumers,
|
||||
List<Function<List<McpSchema.Prompt>, Mono<Void>>> promptsChangeConsumers,
|
||||
List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers,
|
||||
List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
|
||||
|
||||
|
|
@ -106,10 +110,28 @@ class McpClientFeatures {
|
|||
this.resourcesUpdateConsumers = resourcesUpdateConsumers != null ? resourcesUpdateConsumers : List.of();
|
||||
this.promptsChangeConsumers = promptsChangeConsumers != null ? promptsChangeConsumers : List.of();
|
||||
this.loggingConsumers = loggingConsumers != null ? loggingConsumers : List.of();
|
||||
this.progressConsumers = progressConsumers != null ? progressConsumers : List.of();
|
||||
this.samplingHandler = samplingHandler;
|
||||
this.elicitationHandler = elicitationHandler;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated Only exists for backwards-compatibility purposes.
|
||||
*/
|
||||
public Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
|
||||
Map<String, McpSchema.Root> roots,
|
||||
List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumers,
|
||||
List<Function<List<McpSchema.Resource>, Mono<Void>>> resourcesChangeConsumers,
|
||||
List<Function<List<McpSchema.ResourceContents>, Mono<Void>>> resourcesUpdateConsumers,
|
||||
List<Function<List<McpSchema.Prompt>, Mono<Void>>> promptsChangeConsumers,
|
||||
List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
|
||||
this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers,
|
||||
resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler,
|
||||
elicitationHandler);
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert a synchronous specification into an asynchronous one and provide
|
||||
* blocking code offloading to prevent accidental blocking of the non-blocking
|
||||
|
|
@ -149,6 +171,12 @@ class McpClientFeatures {
|
|||
.subscribeOn(Schedulers.boundedElastic()));
|
||||
}
|
||||
|
||||
List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers = new ArrayList<>();
|
||||
for (Consumer<McpSchema.ProgressNotification> consumer : syncSpec.progressConsumers()) {
|
||||
progressConsumers.add(l -> Mono.<Void>fromRunnable(() -> consumer.accept(l))
|
||||
.subscribeOn(Schedulers.boundedElastic()));
|
||||
}
|
||||
|
||||
Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler = r -> Mono
|
||||
.fromCallable(() -> syncSpec.samplingHandler().apply(r))
|
||||
.subscribeOn(Schedulers.boundedElastic());
|
||||
|
|
@ -159,7 +187,7 @@ class McpClientFeatures {
|
|||
|
||||
return new Async(syncSpec.clientInfo(), syncSpec.clientCapabilities(), syncSpec.roots(),
|
||||
toolsChangeConsumers, resourcesChangeConsumers, resourcesUpdateConsumers, promptsChangeConsumers,
|
||||
loggingConsumers, samplingHandler, elicitationHandler);
|
||||
loggingConsumers, progressConsumers, samplingHandler, elicitationHandler);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -174,6 +202,7 @@ class McpClientFeatures {
|
|||
* @param resourcesChangeConsumers the resources change consumers.
|
||||
* @param promptsChangeConsumers the prompts change consumers.
|
||||
* @param loggingConsumers the logging consumers.
|
||||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
*/
|
||||
|
|
@ -183,6 +212,7 @@ class McpClientFeatures {
|
|||
List<Consumer<List<McpSchema.ResourceContents>>> resourcesUpdateConsumers,
|
||||
List<Consumer<List<McpSchema.Prompt>>> promptsChangeConsumers,
|
||||
List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers,
|
||||
List<Consumer<McpSchema.ProgressNotification>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
|
||||
|
||||
|
|
@ -196,6 +226,7 @@ class McpClientFeatures {
|
|||
* @param resourcesUpdateConsumers the resource update consumers.
|
||||
* @param promptsChangeConsumers the prompts change consumers.
|
||||
* @param loggingConsumers the logging consumers.
|
||||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
*/
|
||||
|
|
@ -205,6 +236,7 @@ class McpClientFeatures {
|
|||
List<Consumer<List<McpSchema.ResourceContents>>> resourcesUpdateConsumers,
|
||||
List<Consumer<List<McpSchema.Prompt>>> promptsChangeConsumers,
|
||||
List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers,
|
||||
List<Consumer<McpSchema.ProgressNotification>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
|
||||
|
||||
|
|
@ -222,9 +254,26 @@ class McpClientFeatures {
|
|||
this.resourcesUpdateConsumers = resourcesUpdateConsumers != null ? resourcesUpdateConsumers : List.of();
|
||||
this.promptsChangeConsumers = promptsChangeConsumers != null ? promptsChangeConsumers : List.of();
|
||||
this.loggingConsumers = loggingConsumers != null ? loggingConsumers : List.of();
|
||||
this.progressConsumers = progressConsumers != null ? progressConsumers : List.of();
|
||||
this.samplingHandler = samplingHandler;
|
||||
this.elicitationHandler = elicitationHandler;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated Only exists for backwards-compatibility purposes.
|
||||
*/
|
||||
public Sync(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
|
||||
Map<String, McpSchema.Root> roots, List<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers,
|
||||
List<Consumer<List<McpSchema.Resource>>> resourcesChangeConsumers,
|
||||
List<Consumer<List<McpSchema.ResourceContents>>> resourcesUpdateConsumers,
|
||||
List<Consumer<List<McpSchema.Prompt>>> promptsChangeConsumers,
|
||||
List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
|
||||
this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers,
|
||||
resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler,
|
||||
elicitationHandler);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -177,6 +177,19 @@ public class McpAsyncServerExchange {
|
|||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a notification to the client that the current progress status has changed for
|
||||
* long-running operations.
|
||||
* @param progressNotification The progress notification to send
|
||||
* @return A Mono that completes when the notification has been sent
|
||||
*/
|
||||
public Mono<Void> progressNotification(McpSchema.ProgressNotification progressNotification) {
|
||||
if (progressNotification == null) {
|
||||
return Mono.error(new McpError("Progress notification must not be null"));
|
||||
}
|
||||
return this.session.sendNotification(McpSchema.METHOD_NOTIFICATION_PROGRESS, progressNotification);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a ping request to the client.
|
||||
* @return A Mono that completes with clients's ping response
|
||||
|
|
|
|||
|
|
@ -209,7 +209,7 @@ public class McpServerFeatures {
|
|||
* represents a specific capability.
|
||||
*
|
||||
* @param tool The tool definition including name, description, and parameter schema
|
||||
* @param call Deprecated. Uset he {@link AsyncToolSpecification#callHandler} instead.
|
||||
* @param call Deprecated. Use the {@link AsyncToolSpecification#callHandler} instead.
|
||||
* @param callHandler The function that implements the tool's logic, receiving a
|
||||
* {@link McpAsyncServerExchange} and a
|
||||
* {@link io.modelcontextprotocol.spec.McpSchema.CallToolRequest} and returning
|
||||
|
|
|
|||
|
|
@ -107,6 +107,15 @@ public class McpSyncServerExchange {
|
|||
this.exchange.loggingNotification(loggingMessageNotification).block();
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a notification to the client that the current progress status has changed for
|
||||
* long-running operations.
|
||||
* @param progressNotification The progress notification to send
|
||||
*/
|
||||
public void progressNotification(McpSchema.ProgressNotification progressNotification) {
|
||||
this.exchange.progressNotification(progressNotification).block();
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a synchronous ping request to the client.
|
||||
* @return
|
||||
|
|
|
|||
|
|
@ -61,6 +61,8 @@ public final class McpSchema {
|
|||
|
||||
public static final String METHOD_PING = "ping";
|
||||
|
||||
public static final String METHOD_NOTIFICATION_PROGRESS = "notifications/progress";
|
||||
|
||||
// Tool Methods
|
||||
public static final String METHOD_TOOLS_LIST = "tools/list";
|
||||
|
||||
|
|
@ -1388,7 +1390,8 @@ public final class McpSchema {
|
|||
* tools/list.
|
||||
* @param arguments Arguments to pass to the tool. These must conform to the tool's
|
||||
* input schema.
|
||||
* @param meta See specification for notes on _meta usage
|
||||
* @param meta Optional metadata about the request. This can include additional
|
||||
* information like `progressToken`
|
||||
*/
|
||||
@JsonInclude(JsonInclude.Include.NON_ABSENT)
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
|
|
|
|||
|
|
@ -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