Allow registering logging consumers after client initialization

The logging notification handler already iterates the consumer list on
every notification, so consumers can be added after build without
reinstalling the handler. Switch the list to CopyOnWriteArrayList, since
notifications are delivered on transport threads, and expose
addLoggingConsumer on both the async and sync clients.

This lets callers inspect getServerCapabilities() and only attach a
consumer when the server actually advertises logging support.

Fixes #938

Signed-off-by: Armen Petrosyan <liotur.mail@gmail.com>

Author: liotur

mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

@@ -13,6 +13,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Function;
 
 import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
@@ -185,6 +186,13 @@ public class McpAsyncClient {
 
 	private final boolean applyElicitationDefaults;
 
+	/**
+	 * Consumers notified for every logging message received from the server. Held as a
+	 * live reference so that consumers registered after the client was built are picked
+	 * up by the already installed notification handler.
+	 */
+	private final List<Function<LoggingMessageNotification, Mono<Void>>> loggingConsumers;
+
 	/**
 	 * Create a new McpAsyncClient with the given transport and session request-response
 	 * timeout.
@@ -304,11 +312,12 @@ public class McpAsyncClient {
 				asyncPromptsChangeNotificationHandler(promptsChangeConsumersFinal));
 
 		// Utility Logging Notification
-		List<Function<LoggingMessageNotification, Mono<Void>>> loggingConsumersFinal = new ArrayList<>();
+		List<Function<LoggingMessageNotification, Mono<Void>>> loggingConsumersFinal = new CopyOnWriteArrayList<>();
 		loggingConsumersFinal.add((notification) -> Mono.fromRunnable(() -> logger.debug("Logging: {}", notification)));
 		if (!Utils.isEmpty(features.loggingConsumers())) {
 			loggingConsumersFinal.addAll(features.loggingConsumers());
 		}
+		this.loggingConsumers = loggingConsumersFinal;
 		notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_MESSAGE,
 				asyncLoggingNotificationHandler(loggingConsumersFinal));
 
@@ -1103,6 +1112,19 @@ private NotificationHandler asyncLoggingNotificationHandler(
 		};
 	}
 
+	/**
+	 * Registers an additional consumer for logging messages received from the server.
+	 * <p>
+	 * Unlike {@link McpClient.AsyncSpec#loggingConsumer(Function)}, this can be called
+	 * after the client has been built, so a consumer may be registered only once
+	 * {@link #getServerCapabilities()} shows the server supports logging.
+	 * @param loggingConsumer A consumer that receives logging messages. Must not be null.
+	 */
+	public void addLoggingConsumer(Function<LoggingMessageNotification, Mono<Void>> loggingConsumer) {
+		Assert.notNull(loggingConsumer, "Logging consumer must not be null");
+		this.loggingConsumers.add(loggingConsumer);
+	}
+
 	/**
 	 * Sets the minimum logging level for messages received from the server. The client
 	 * will only receive log messages at or above the specified severity level.

mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java

@@ -6,6 +6,7 @@
 
 import java.time.Duration;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import org.slf4j.Logger;
@@ -19,6 +20,7 @@
 import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult;
 import io.modelcontextprotocol.util.Assert;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 
 /**
  * A synchronous client implementation for the Model Context Protocol (MCP) that wraps an
@@ -423,6 +425,20 @@ public GetPromptResult getPrompt(GetPromptRequest getPromptRequest) {
 		return withProvidedContext(this.delegate.getPrompt(getPromptRequest)).block();
 	}
 
+	/**
+	 * Registers an additional consumer for logging messages received from the server.
+	 * <p>
+	 * Unlike {@link McpClient.SyncSpec#loggingConsumer(Consumer)}, this can be called
+	 * after the client has been built, so a consumer may be registered only once
+	 * {@link #getServerCapabilities()} reports that the server supports logging.
+	 * @param loggingConsumer the consumer that receives logging messages
+	 */
+	public void addLoggingConsumer(Consumer<McpSchema.LoggingMessageNotification> loggingConsumer) {
+		Assert.notNull(loggingConsumer, "Logging consumer must not be null");
+		this.delegate.addLoggingConsumer(
+				l -> Mono.<Void>fromRunnable(() -> loggingConsumer.accept(l)).subscribeOn(Schedulers.boundedElastic()));
+	}
+
 	/**
 	 * Client can set the minimum logging level it wants to receive from the server.
 	 * @param loggingLevel the min logging level

mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java

@@ -106,6 +106,31 @@ void bothHandlersPresent() {
 
 		}
 
+		@Nested
+		class LoggingConsumers {
+
+			@Test
+			void addLoggingConsumerAfterBuild() {
+				McpClientTransport transport = mock(McpClientTransport.class);
+				when(transport.protocolVersions()).thenReturn(List.of("2024-11-05"));
+				var client = McpClient.async(transport).jsonSchemaValidator(mock(JsonSchemaValidator.class)).build();
+
+				assertThatCode(() -> client.addLoggingConsumer(notification -> Mono.empty()))
+					.doesNotThrowAnyException();
+			}
+
+			@Test
+			void addLoggingConsumerRejectsNull() {
+				McpClientTransport transport = mock(McpClientTransport.class);
+				when(transport.protocolVersions()).thenReturn(List.of("2024-11-05"));
+				var client = McpClient.async(transport).jsonSchemaValidator(mock(JsonSchemaValidator.class)).build();
+
+				assertThatThrownBy(() -> client.addLoggingConsumer(null)).isInstanceOf(IllegalArgumentException.class)
+					.hasMessage("Logging consumer must not be null");
+			}
+
+		}
+
 		@Nested
 		class ClientCapabilities {