diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpAsyncNotifier.java b/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpAsyncNotifier.java index b3850a771..d4e0893e9 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpAsyncNotifier.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpAsyncNotifier.java @@ -8,7 +8,8 @@ /** * Emits request-scoped notifications from inside the body of a streaming - * {@link McpAsyncResponse}. + * {@link McpAsyncResponse}. Customize notification policy through + * {@link McpServer.Builder#notifierFactory}. * * @author Dariusz Jędrzejczyk */ diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpServer.java b/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpServer.java index de86c2321..c6639c867 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpServer.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpServer.java @@ -12,6 +12,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.BiFunction; import io.modelcontextprotocol.common.McpTransportContext; import io.modelcontextprotocol.json.McpJsonDefaults; @@ -72,9 +73,12 @@ public final class McpServer implements McpRequestManager { private final SubscriptionsFeature subscriptionsFeature; + private final BiFunction notifierFactory; + private McpServer(Implementation serverInfo, List supportedVersions, McpJsonMapper jsonMapper, Map features, Set inputRequiredMethods, RequestStateCodec requestStateCodec, - SubscriptionsFeature subscriptionsFeature) { + SubscriptionsFeature subscriptionsFeature, + BiFunction notifierFactory) { this.serverInfo = serverInfo; this.supportedVersions = supportedVersions; this.jsonMapper = jsonMapper; @@ -82,6 +86,7 @@ private McpServer(Implementation serverInfo, List supportedVersions, Mcp this.inputRequiredMethods = inputRequiredMethods; this.requestStateCodec = requestStateCodec; this.subscriptionsFeature = subscriptionsFeature; + this.notifierFactory = notifierFactory; } public static Builder builder() { @@ -218,8 +223,11 @@ private Flux stream(McpRequestContext ctx, McpAsyncResponse.Body // Deferred: the body must not run before the transport has started the stream. return Flux.defer(() -> { Sinks.Many sink = Sinks.many().unicast().onBackpressureBuffer(); - DefaultAsyncNotifier notifier = new DefaultAsyncNotifier(ctx, sink); - Mono terminal = Mono.defer(() -> body.run(notifier)) + Mono terminal = Mono.defer(() -> { + McpAsyncNotifier notifier = this.notifierFactory.apply(ctx, new DefaultAsyncNotifier(ctx, sink)); + Assert.notNull(notifier, "notifierFactory must not return null"); + return body.run(notifier); + }) .map(result -> (JSONRPCMessage) resultResponse(ctx, result)) .switchIfEmpty(Mono.fromSupplier(() -> bug(ctx, "Streaming body completed without a result"))) // The stream has already started, so a failure can only be reported @@ -346,6 +354,9 @@ public static final class Builder { private McpChangePublisher changePublisher; + private BiFunction notifierFactory = (ctx, + notifier) -> notifier; + private Builder() { } @@ -392,6 +403,26 @@ public Builder subscriptions(McpChangePublisher publisher) { return this; } + /** + * Customizes the notifier passed to streaming bodies, including sync handlers and + * {@code subscriptions/listen}. The factory receives the request context and the + * default notifier, which it may wrap to filter or reject notifications. By + * default, the supplied notifier is used unchanged; no additional notification + * policy is enforced. + *

+ * Invoked once per stream subscription, before the body runs. The factory must be + * thread-safe and return a non-null notifier. Factory failures are reported as + * errors in the stream, whose HTTP status has already been committed. + * @param notifierFactory the factory receiving the context and default notifier + * @return this builder + */ + public Builder notifierFactory( + BiFunction notifierFactory) { + Assert.notNull(notifierFactory, "notifierFactory must not be null"); + this.notifierFactory = notifierFactory; + return this; + } + /** * The codec used to seal/open MRTR {@code requestState}. Defaults to * {@link HmacRequestStateCodec#builder()}{@code .build()}. @@ -459,7 +490,7 @@ public McpServer build() { : HmacRequestStateCodec.builder().jsonMapper(mapper).build(); return new McpServer(this.serverInfo, this.supportedVersions, mapper, Map.copyOf(byMethod), - Set.copyOf(inputRequiredMethods), codec, subscriptionsFeature); + Set.copyOf(inputRequiredMethods), codec, subscriptionsFeature, this.notifierFactory); } } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/McpServerTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/McpServerTests.java index ab62c5aee..f93fdc1f7 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/McpServerTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/McpServerTests.java @@ -9,6 +9,7 @@ import java.util.Map; import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiFunction; import io.modelcontextprotocol.modern.JsonRpc.JSONRPCNotification; @@ -26,6 +27,8 @@ import io.modelcontextprotocol.modern.McpSchema.TextContent; import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -297,10 +300,11 @@ void streamingBodyRunsOnlyOnceTheStreamIsSubscribed() { @Test void streamingBodyNotificationsPrecedeTheResult() { - McpServer server = baseBuilder().feature(feature("tools/call", - (ctx, params) -> Mono.just(McpAsyncResponse.streaming(notifier -> notifier.progress(1.0, 2.0, "half") - .then(notifier.progress(2.0, 2.0, "done")) - .thenReturn(ok()))))) + McpServer server = baseBuilder().feature(feature("tools/call", (ctx, params) -> Mono + .just(McpAsyncResponse.streaming(notifier -> notifier.progress(1.0, 2.0, "half") + .then(notifier.progress(2.0, 2.0, "done")) + .then(notifier.notify(McpSchema.METHOD_NOTIFICATION_MESSAGE, Map.of("level", "info", "data", "hello"))) + .thenReturn(ok()))))) .build(); McpTransportResponse invocation = invoke(server, request("tools/call", meta(MetaKeys.PROGRESS_TOKEN, "tok-1"))) @@ -309,6 +313,8 @@ void streamingBodyNotificationsPrecedeTheResult() { StepVerifier.create(((McpTransportResponse.Streaming) invocation).messages()) .expectNextMatches(msg -> msg instanceof JSONRPCNotification) .expectNextMatches(msg -> msg instanceof JSONRPCNotification) + .expectNextMatches(msg -> msg instanceof JSONRPCNotification notification + && notification.method().equals(McpSchema.METHOD_NOTIFICATION_MESSAGE)) .expectNextMatches(msg -> msg instanceof JSONRPCResponse) .verifyComplete(); } @@ -328,6 +334,118 @@ void progressIsSuppressedWithoutProgressToken() { .verifyComplete(); } + @Test + void notifierFactoryFiltersUsingEachRequestsContextAndPreservesProgress() { + AtomicInteger created = new AtomicInteger(); + McpServer server = baseBuilder().notifierFactory((ctx, delegate) -> { + created.incrementAndGet(); + return new McpAsyncNotifier() { + @Override + public Mono progress(double progress, Double total, String message) { + return delegate.progress(progress, total, message); + } + + @Override + public Mono notify(String method, Object params) { + if (McpSchema.METHOD_NOTIFICATION_MESSAGE.equals(method) + && !ctx.meta().containsKey(MetaKeys.LOG_LEVEL)) { + return Mono.empty(); + } + return delegate.notify(method, params); + } + }; + }) + .feature(feature("tools/call", + (ctx, params) -> Mono.just(McpAsyncResponse.streaming(notifier -> notifier + .notify(McpSchema.METHOD_NOTIFICATION_MESSAGE, Map.of("level", "info", "data", "hello")) + .then(notifier.progress(1.0, 1.0, "done")) + .thenReturn(ok()))))) + .build(); + + var quiet = (McpTransportResponse.Streaming) invoke(server, + request("tools/call", meta(MetaKeys.PROGRESS_TOKEN, "quiet"))) + .block(); + var logging = (McpTransportResponse.Streaming) invoke(server, new JSONRPCRequest("tools/call", 2, + Map.of("_meta", meta(MetaKeys.PROGRESS_TOKEN, "logging", MetaKeys.LOG_LEVEL, "info")))) + .block(); + + assertThat(created).hasValue(0); + StepVerifier.create(quiet.messages()) + .expectNextMatches(msg -> msg instanceof JSONRPCNotification notification + && notification.method().equals(McpSchema.METHOD_NOTIFICATION_PROGRESS)) + .expectNextMatches(msg -> msg instanceof JSONRPCResponse response && response.error() == null) + .verifyComplete(); + StepVerifier.create(logging.messages()) + .expectNextMatches(msg -> msg instanceof JSONRPCNotification notification + && notification.method().equals(McpSchema.METHOD_NOTIFICATION_MESSAGE)) + .expectNextMatches(msg -> msg instanceof JSONRPCNotification notification + && notification.method().equals(McpSchema.METHOD_NOTIFICATION_PROGRESS)) + .expectNextMatches(msg -> msg instanceof JSONRPCResponse response && response.error() == null) + .verifyComplete(); + assertThat(created).hasValue(2); + } + + @ParameterizedTest + @ValueSource(booleans = { false, true }) + void invalidNotifierFactoryIsAnInternalErrorInStream(boolean returnsNull) { + AtomicBoolean ran = new AtomicBoolean(); + McpServer server = baseBuilder().notifierFactory((ctx, delegate) -> { + if (returnsNull) { + return null; + } + throw new IllegalStateException("factory failed"); + }).feature(feature("tools/call", (ctx, params) -> Mono.just(McpAsyncResponse.streaming(notifier -> { + ran.set(true); + return Mono.just(ok()); + })))).build(); + + var invocation = (McpTransportResponse.Streaming) invoke(server, request("tools/call", meta())).block(); + + StepVerifier.create(invocation.messages()) + .expectNextMatches(msg -> msg instanceof JSONRPCResponse response + && response.error().code() == ErrorCodes.INTERNAL_ERROR) + .verifyComplete(); + assertThat(ran).isFalse(); + } + + @Test + void notifierFactoryCanRejectANotificationWithAnMcpError() { + McpServer server = baseBuilder().notifierFactory((ctx, delegate) -> new McpAsyncNotifier() { + @Override + public Mono progress(double progress, Double total, String message) { + return delegate.progress(progress, total, message); + } + + @Override + public Mono notify(String method, Object params) { + return Mono.error(McpException.invalidParams("Notification not allowed")); + } + }) + .feature(feature("tools/call", + (ctx, params) -> Mono.just(McpAsyncResponse + .streaming(notifier -> notifier.notify("extension/forbidden", Map.of()).thenReturn(ok()))))) + .build(); + + var invocation = (McpTransportResponse.Streaming) invoke(server, request("tools/call", meta())).block(); + + StepVerifier.create(invocation.messages()) + .expectNextMatches(msg -> msg instanceof JSONRPCResponse response + && response.error().code() == ErrorCodes.INVALID_PARAMS + && response.error().message().equals("Notification not allowed")) + .verifyComplete(); + } + + @Test + void notifierFactoryIsNotInvokedForSingleResults() { + McpServer server = baseBuilder().notifierFactory((ctx, delegate) -> { + throw new IllegalStateException("No notifier needed"); + }).feature(echoFeature("tools/call")).build(); + + StepVerifier.create(respond(server, request("tools/call", meta()))) + .assertNext(response -> assertThat(response.error()).isNull()) + .verifyComplete(); + } + @Test void streamingBodyExceptionIsAnInternalErrorInStream() { McpServer server = baseBuilder() diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/feature/ToolsFeatureTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/feature/ToolsFeatureTests.java index 3c92dffcb..0eeb56f00 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/feature/ToolsFeatureTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/modern/server/feature/ToolsFeatureTests.java @@ -34,6 +34,7 @@ import io.modelcontextprotocol.modern.server.McpTransportResponse; import io.modelcontextprotocol.modern.server.McpRequestContext; import io.modelcontextprotocol.modern.server.McpAsyncResponse; +import io.modelcontextprotocol.modern.server.McpAsyncNotifier; import io.modelcontextprotocol.modern.server.McpServer; import io.modelcontextprotocol.modern.server.McpSyncResponse; import io.modelcontextprotocol.modern.server.ModernTestFixtures; @@ -322,6 +323,33 @@ void blockingCallerRunsSyncStreamingHandlerOnItsOwnThreadAndReceivesNotification assertThat(greetingText(delivered.get(1))).isEqualTo("hello alice"); } + @Test + void notifierFactoryCustomizesSyncStreamingHandlers() { + AtomicReference handlerThread = new AtomicReference<>(); + List delivered = new CopyOnWriteArrayList<>(); + McpServer server = baseBuilder().notifierFactory((ctx, delegate) -> new McpAsyncNotifier() { + @Override + public Mono progress(double progress, Double total, String message) { + return delegate.notify("extension/progress", Map.of("requestId", ctx.requestId())); + } + + @Override + public Mono notify(String method, Object params) { + return delegate.notify(method, params); + } + }).feature(tools(streamingGreetingRepo(ctx -> "alice", handlerThread, delivered, new AtomicInteger()))).build(); + + callStreaming(server.handleBlocking(McpTransportContext.EMPTY, callEcho()), delivered); + + assertThat(handlerThread.get()).isEqualTo(Thread.currentThread().getName()); + assertThat(delivered).hasSize(2); + assertThat(delivered.get(0)).isInstanceOfSatisfying(JSONRPCNotification.class, notification -> { + assertThat(notification.method()).isEqualTo("extension/progress"); + assertThat(notification.params()).isEqualTo(Map.of("requestId", callEcho().id())); + }); + assertThat(greetingText(delivered.get(1))).isEqualTo("hello alice"); + } + @Test @SuppressWarnings("unchecked") void blockingCallerRunsSyncListOnItsOwnThread() {