diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java index 3509b760b..232efd465 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/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>> 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>> loggingConsumersFinal = new ArrayList<>(); + List>> 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. + *

+ * 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> 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. diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java index 7e08f83a0..ee4c05059 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java +++ b/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. + *

+ * 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 loggingConsumer) { + Assert.notNull(loggingConsumer, "Logging consumer must not be null"); + this.delegate.addLoggingConsumer( + l -> Mono.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 diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java index dea7d42e9..73aa0bd1e 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java +++ b/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 {