Add request-scoped notifier factory to modern server

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
  */

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,16 +73,20 @@ public final class McpServer implements McpRequestManager {
 
 	private final SubscriptionsFeature subscriptionsFeature;
 
+	private final BiFunction<McpRequestContext, McpAsyncNotifier, McpAsyncNotifier> notifierFactory;
+
 	private McpServer(Implementation serverInfo, List<String> supportedVersions, McpJsonMapper jsonMapper,
 			Map<String, McpFeature> features, Set<String> inputRequiredMethods, RequestStateCodec requestStateCodec,
-			SubscriptionsFeature subscriptionsFeature) {
+			SubscriptionsFeature subscriptionsFeature,
+			BiFunction<McpRequestContext, McpAsyncNotifier, McpAsyncNotifier> notifierFactory) {
 		this.serverInfo = serverInfo;
 		this.supportedVersions = supportedVersions;
 		this.jsonMapper = jsonMapper;
 		this.features = features;
 		this.inputRequiredMethods = inputRequiredMethods;
 		this.requestStateCodec = requestStateCodec;
 		this.subscriptionsFeature = subscriptionsFeature;
+		this.notifierFactory = notifierFactory;
 	}
 
 	public static Builder builder() {
@@ -218,8 +223,11 @@ private Flux<JSONRPCMessage> stream(McpRequestContext ctx, McpAsyncResponse.Body
 		// Deferred: the body must not run before the transport has started the stream.
 		return Flux.defer(() -> {
 			Sinks.Many<JSONRPCNotification> sink = Sinks.many().unicast().onBackpressureBuffer();
-			DefaultAsyncNotifier notifier = new DefaultAsyncNotifier(ctx, sink);
-			Mono<JSONRPCMessage> terminal = Mono.defer(() -> body.run(notifier))
+			Mono<JSONRPCMessage> 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<McpRequestContext, McpAsyncNotifier, McpAsyncNotifier> 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.
+		 * <p>
+		 * 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<McpRequestContext, McpAsyncNotifier, McpAsyncNotifier> 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);
 		}
 
 	}

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<Void> progress(double progress, Double total, String message) {
+					return delegate.progress(progress, total, message);
+				}
+
+				@Override
+				public Mono<Void> 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<Void> progress(double progress, Double total, String message) {
+				return delegate.progress(progress, total, message);
+			}
+
+			@Override
+			public Mono<Void> 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()

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<String> handlerThread = new AtomicReference<>();
+		List<JSONRPCMessage> delivered = new CopyOnWriteArrayList<>();
+		McpServer server = baseBuilder().notifierFactory((ctx, delegate) -> new McpAsyncNotifier() {
+			@Override
+			public Mono<Void> progress(double progress, Double total, String message) {
+				return delegate.notify("extension/progress", Map.of("requestId", ctx.requestId()));
+			}
+
+			@Override
+			public Mono<Void> 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() {