From d01c6ddcd3bd7e9229f86f419c6ab1ccddf015f9 Mon Sep 17 00:00:00 2001 From: Tarun Agarwal <45505483+tarun-agarwal1044@users.noreply.github.com> Date: Wed, 21 Oct 2026 18:11:20 +0530 Subject: [PATCH] Add request-scoped notifier factory to modern server Allow applications to wrap the default async notifier using the request context. Preserve permissive defaults and apply customization to sync and async streaming handlers. Report factory failures through the existing streaming error path. --- .../modern/server/McpAsyncNotifier.java | 3 +- .../modern/server/McpServer.java | 39 +++++- .../modern/server/McpServerTests.java | 126 +++++++++++++++++- .../server/feature/ToolsFeatureTests.java | 28 ++++ 4 files changed, 187 insertions(+), 9 deletions(-) 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() {