diff --git a/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java b/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java index 1e951dbda..83690a94c 100644 --- a/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java +++ b/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java @@ -18,7 +18,6 @@ import static com.google.common.base.Strings.isNullOrEmpty; -import com.google.common.collect.ImmutableMap; import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport; import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; import io.modelcontextprotocol.client.transport.ServerParameters; @@ -29,7 +28,9 @@ import java.net.URI; import java.net.URISyntaxException; import java.util.Collection; +import java.util.Map; import java.util.Optional; +import java.util.function.Supplier; import reactor.core.publisher.Mono; /** @@ -50,18 +51,24 @@ public McpClientTransport build(Object connectionParams) { .sseEndpoint( sseServerParams.sseEndpoint() == null ? "sse" : sseServerParams.sseEndpoint()) .customizeRequest( - builder -> - Optional.ofNullable(sseServerParams.headers()) - .map(ImmutableMap::entrySet) - .stream() - .flatMap(Collection::stream) - .forEach( - entry -> - builder.header( - entry.getKey(), - Optional.ofNullable(entry.getValue()) - .map(Object::toString) - .orElse("")))) + builder -> { + Map effectiveHeaders; + try { + Supplier> provider = sseServerParams.headersProvider(); + effectiveHeaders = provider != null ? provider.get() : sseServerParams.headers(); + } catch (Exception e) { + throw new RuntimeException("Failed to retrieve headers from provider", e); + } + Optional.ofNullable(effectiveHeaders).map(Map::entrySet).stream() + .flatMap(Collection::stream) + .forEach( + entry -> + builder.header( + entry.getKey(), + Optional.ofNullable(entry.getValue()) + .map(Object::toString) + .orElse(""))); + }) .build(); } else if (connectionParams instanceof StreamableHttpServerParameters streamableParams) { // Split the URL so the transport's URI.resolve does not drop a custom path (b/513186321). @@ -72,10 +79,17 @@ public McpClientTransport build(Object connectionParams) { .jsonMapper(jsonMapper) .asyncHttpRequestCustomizer( (requestBuilder, method, uri, body, context) -> { - streamableParams - .headers() - .forEach((key, value) -> requestBuilder.header(key, value)); - return Mono.just(requestBuilder); + try { + Supplier> provider = streamableParams.headersProvider(); + Map effectiveHeaders = + provider != null ? provider.get() : streamableParams.headers(); + if (effectiveHeaders != null) { + effectiveHeaders.forEach((key, value) -> requestBuilder.header(key, value)); + } + return Mono.just(requestBuilder); + } catch (Exception e) { + return Mono.error(e); + } }); if (split.endpoint() != null) { builder.endpoint(split.endpoint()); diff --git a/core/src/main/java/com/google/adk/tools/mcp/SseServerParameters.java b/core/src/main/java/com/google/adk/tools/mcp/SseServerParameters.java index 3b12f7064..3f38054fd 100644 --- a/core/src/main/java/com/google/adk/tools/mcp/SseServerParameters.java +++ b/core/src/main/java/com/google/adk/tools/mcp/SseServerParameters.java @@ -17,12 +17,14 @@ package com.google.adk.tools.mcp; import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder; import com.google.auto.value.AutoValue; import com.google.common.collect.ImmutableMap; import java.time.Duration; import java.util.Map; +import java.util.function.Supplier; import org.jspecify.annotations.Nullable; /** Parameters for establishing a MCP Server-Sent Events (SSE) connection. */ @@ -41,6 +43,11 @@ public abstract class SseServerParameters { @Nullable public abstract ImmutableMap headers(); + /** Optional supplier of headers, evaluated per-request for dynamic authentication. */ + @JsonIgnore + @Nullable + public abstract Supplier> headersProvider(); + /** The timeout for the initial connection attempt. */ @Nullable public abstract Duration timeout(); @@ -75,6 +82,11 @@ static SseServerParameters.Builder jacksonBuilder() { /** Sets the headers for the SSE connection request. */ public abstract Builder headers(@Nullable Map headers); + /** Sets a dynamic supplier of headers, evaluated per-request. */ + @JsonIgnore + public abstract Builder headersProvider( + @Nullable Supplier> headersProvider); + /** Sets the timeout for the initial connection attempt. */ public abstract Builder timeout(@Nullable Duration timeout); diff --git a/core/src/main/java/com/google/adk/tools/mcp/StreamableHttpServerParameters.java b/core/src/main/java/com/google/adk/tools/mcp/StreamableHttpServerParameters.java index e9f8a3ac8..b58ae6176 100644 --- a/core/src/main/java/com/google/adk/tools/mcp/StreamableHttpServerParameters.java +++ b/core/src/main/java/com/google/adk/tools/mcp/StreamableHttpServerParameters.java @@ -21,12 +21,14 @@ import java.time.Duration; import java.util.Collections; import java.util.Map; +import java.util.function.Supplier; import org.jspecify.annotations.Nullable; /** Server parameters for Streamable HTTP client transport. */ public class StreamableHttpServerParameters { private final String url; private final Map headers; + private final @Nullable Supplier> headersProvider; private final Duration timeout; private final Duration readTimeout; private final boolean terminateOnClose; @@ -36,6 +38,7 @@ public class StreamableHttpServerParameters { * * @param url The base URL for the MCP Streamable HTTP server. * @param headers Optional headers to include in requests. + * @param headersProvider Optional supplier of headers, evaluated per-request for dynamic auth. * @param timeout Timeout for HTTP operations (default: 30 seconds). * @param readTimeout Timeout for reading data from the streamed http events(default: 5 minutes). * @param terminateOnClose Whether to terminate the session on close (default: true). @@ -43,12 +46,14 @@ public class StreamableHttpServerParameters { public StreamableHttpServerParameters( String url, Map headers, + @Nullable Supplier> headersProvider, @Nullable Duration timeout, @Nullable Duration readTimeout, @Nullable Boolean terminateOnClose) { Assert.hasText(url, "url must not be empty"); this.url = url; this.headers = headers == null ? Collections.emptyMap() : headers; + this.headersProvider = headersProvider; this.timeout = timeout == null ? Duration.ofSeconds(30) : timeout; this.readTimeout = readTimeout == null ? Duration.ofMinutes(5) : readTimeout; this.terminateOnClose = terminateOnClose == null || terminateOnClose; @@ -62,6 +67,11 @@ public Map headers() { return headers; } + @Nullable + public Supplier> headersProvider() { + return headersProvider; + } + public Duration timeout() { return timeout; } @@ -82,6 +92,7 @@ public static Builder builder() { public static class Builder { private String url; private Map headers = Collections.emptyMap(); + private @Nullable Supplier> headersProvider; private Duration timeout = Duration.ofSeconds(30); private Duration readTimeout = Duration.ofMinutes(5); private boolean terminateOnClose = true; @@ -101,6 +112,12 @@ public Builder headers(Map headers) { return this; } + @CanIgnoreReturnValue + public Builder headersProvider(@Nullable Supplier> headersProvider) { + this.headersProvider = headersProvider; + return this; + } + @CanIgnoreReturnValue public Builder timeout(Duration timeout) { this.timeout = timeout; @@ -121,7 +138,7 @@ public Builder terminateOnClose(boolean terminateOnClose) { public StreamableHttpServerParameters build() { return new StreamableHttpServerParameters( - url, headers, timeout, readTimeout, terminateOnClose); + url, headers, headersProvider, timeout, readTimeout, terminateOnClose); } } } diff --git a/core/src/test/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilderTest.java b/core/src/test/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilderTest.java index d2c263ac4..787bd9d3e 100644 --- a/core/src/test/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilderTest.java +++ b/core/src/test/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilderTest.java @@ -264,6 +264,145 @@ private static McpAsyncHttpClientRequestCustomizer getCustomizer( return (McpAsyncHttpClientRequestCustomizer) field.get(transport); } + @Test + public void build_withStreamableHttpHeadersProvider_providerCalledPerRequest() throws Exception { + java.util.concurrent.atomic.AtomicInteger callCount = + new java.util.concurrent.atomic.AtomicInteger(); + StreamableHttpServerParameters params = + StreamableHttpServerParameters.builder() + .url("http://localhost:8080/mcp/stream") + .headersProvider( + () -> { + callCount.incrementAndGet(); + return ImmutableMap.of("Authorization", "Bearer dynamic-token"); + }) + .build(); + + HttpClientStreamableHttpTransport transport = + (HttpClientStreamableHttpTransport) transportBuilder.build(params); + McpAsyncHttpClientRequestCustomizer customizer = getCustomizer(transport); + + // Call customize twice — provider must be invoked each time + for (int i = 0; i < 2; i++) { + HttpRequest.Builder requestBuilder = HttpRequest.newBuilder().uri(URI.create("http://x/")); + Mono.from( + customizer.customize( + requestBuilder, "POST", URI.create("http://x/"), null, McpTransportContext.EMPTY)) + .block(); + assertThat(collectHeaders(requestBuilder)) + .containsEntry("Authorization", "Bearer dynamic-token"); + } + assertThat(callCount.get()).isEqualTo(2); + } + + @Test + public void build_withSseHeadersProvider_providerCalledPerRequest() throws Exception { + java.util.concurrent.atomic.AtomicInteger callCount = + new java.util.concurrent.atomic.AtomicInteger(); + SseServerParameters params = + SseServerParameters.builder() + .url("http://localhost:8080") + .headersProvider( + () -> { + callCount.incrementAndGet(); + return ImmutableMap.of("Authorization", "Bearer sse-dynamic"); + }) + .build(); + + HttpClientSseClientTransport transport = + (HttpClientSseClientTransport) transportBuilder.build(params); + + // Verify the provider was wired — we can't easily extract the SSE customizer via reflection + // so we verify the parameter object accepted the provider + assertThat(params.headersProvider()).isNotNull(); + assertThat(params.headersProvider().get()).containsEntry("Authorization", "Bearer sse-dynamic"); + assertThat(callCount.get()).isAtLeast(1); + } + + @Test + public void build_withStreamableBothHeadersAndProvider_providerTakesPrecedence() + throws Exception { + StreamableHttpServerParameters params = + StreamableHttpServerParameters.builder() + .url("http://localhost:8080/mcp/stream") + .headers(ImmutableMap.of("Authorization", "Bearer static-token")) + .headersProvider(() -> ImmutableMap.of("Authorization", "Bearer dynamic-token")) + .build(); + + HttpClientStreamableHttpTransport transport = + (HttpClientStreamableHttpTransport) transportBuilder.build(params); + McpAsyncHttpClientRequestCustomizer customizer = getCustomizer(transport); + HttpRequest.Builder requestBuilder = HttpRequest.newBuilder().uri(URI.create("http://x/")); + + Mono.from( + customizer.customize( + requestBuilder, "POST", URI.create("http://x/"), null, McpTransportContext.EMPTY)) + .block(); + + Map headers = collectHeaders(requestBuilder); + assertThat(headers).containsEntry("Authorization", "Bearer dynamic-token"); + } + + @Test + public void build_withStreamableHttpHeadersProviderReturningNull_doesNotThrow() throws Exception { + StreamableHttpServerParameters params = + StreamableHttpServerParameters.builder() + .url("http://localhost:8080/mcp/stream") + .headersProvider(() -> null) + .build(); + + HttpClientStreamableHttpTransport transport = + (HttpClientStreamableHttpTransport) transportBuilder.build(params); + McpAsyncHttpClientRequestCustomizer customizer = getCustomizer(transport); + HttpRequest.Builder requestBuilder = HttpRequest.newBuilder().uri(URI.create("http://x/")); + + HttpRequest.Builder returned = + Mono.from( + customizer.customize( + requestBuilder, + "POST", + URI.create("http://x/"), + null, + McpTransportContext.EMPTY)) + .block(); + + assertThat(returned).isSameInstanceAs(requestBuilder); + assertThat(collectHeaders(requestBuilder)).isEmpty(); + } + + @Test + public void build_withStreamableHttpHeadersProviderThrowing_returnsMono_error() throws Exception { + StreamableHttpServerParameters params = + StreamableHttpServerParameters.builder() + .url("http://localhost:8080/mcp/stream") + .headersProvider( + () -> { + throw new RuntimeException("token refresh failed"); + }) + .build(); + + HttpClientStreamableHttpTransport transport = + (HttpClientStreamableHttpTransport) transportBuilder.build(params); + + McpAsyncHttpClientRequestCustomizer customizer = getCustomizer(transport); + HttpRequest.Builder requestBuilder = HttpRequest.newBuilder().uri(URI.create("http://x/")); + + Exception ex = + assertThrows( + Exception.class, + () -> + Mono.from( + customizer.customize( + requestBuilder, + "POST", + URI.create("http://x/"), + null, + McpTransportContext.EMPTY)) + .block()); + + assertThat(ex).hasMessageThat().contains("token refresh failed"); + } + /** Reads back the headers set on a builder by building a throwaway request. */ private static Map collectHeaders(HttpRequest.Builder builder) { HttpRequest request = builder.GET().build(); diff --git a/core/src/test/java/com/google/adk/tools/mcp/SseServerParametersTest.java b/core/src/test/java/com/google/adk/tools/mcp/SseServerParametersTest.java new file mode 100644 index 000000000..09db84642 --- /dev/null +++ b/core/src/test/java/com/google/adk/tools/mcp/SseServerParametersTest.java @@ -0,0 +1,52 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.tools.mcp; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.common.collect.ImmutableMap; +import java.util.Map; +import java.util.function.Supplier; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link SseServerParameters}. */ +@RunWith(JUnit4.class) +public final class SseServerParametersTest { + + @Test + public void builder_withHeadersProvider_returnsSupplier() { + Supplier> provider = () -> ImmutableMap.of("Authorization", "Bearer token"); + + SseServerParameters params = + SseServerParameters.builder() + .url("http://localhost:8080") + .headersProvider(provider) + .build(); + + assertThat(params.headersProvider()).isSameInstanceAs(provider); + assertThat(params.headersProvider().get()).containsEntry("Authorization", "Bearer token"); + } + + @Test + public void builder_withHeadersProvider_null_returnsNull() { + SseServerParameters params = SseServerParameters.builder().url("http://localhost:8080").build(); + + assertThat(params.headersProvider()).isNull(); + } +} diff --git a/core/src/test/java/com/google/adk/tools/mcp/StreamableHttpServerParametersTest.java b/core/src/test/java/com/google/adk/tools/mcp/StreamableHttpServerParametersTest.java new file mode 100644 index 000000000..bc51dbc73 --- /dev/null +++ b/core/src/test/java/com/google/adk/tools/mcp/StreamableHttpServerParametersTest.java @@ -0,0 +1,65 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.tools.mcp; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.common.collect.ImmutableMap; +import java.util.Map; +import java.util.function.Supplier; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link StreamableHttpServerParameters}. */ +@RunWith(JUnit4.class) +public final class StreamableHttpServerParametersTest { + + @Test + public void builder_withHeadersProvider_returnsSupplier() { + Supplier> provider = () -> ImmutableMap.of("Authorization", "Bearer token"); + + StreamableHttpServerParameters params = + StreamableHttpServerParameters.builder() + .url("http://localhost:8080") + .headersProvider(provider) + .build(); + + assertThat(params.headersProvider()).isSameInstanceAs(provider); + assertThat(params.headersProvider().get()).containsEntry("Authorization", "Bearer token"); + } + + @Test + public void builder_defaultHeadersProvider_isNull() { + StreamableHttpServerParameters params = + StreamableHttpServerParameters.builder().url("http://localhost:8080").build(); + + assertThat(params.headersProvider()).isNull(); + } + + @Test + public void builder_withStaticHeaders_headersProviderIsNull() { + StreamableHttpServerParameters params = + StreamableHttpServerParameters.builder() + .url("http://localhost:8080") + .headers(ImmutableMap.of("X-Key", "value")) + .build(); + + assertThat(params.headers()).containsEntry("X-Key", "value"); + assertThat(params.headersProvider()).isNull(); + } +}