From fc5fd099b1c5ed5729baeca809bfdf0fdaef1ff0 Mon Sep 17 00:00:00 2001 From: Stenal P Jolly Date: Mon, 6 Apr 2026 20:43:46 +0530 Subject: [PATCH 1/3] feat: Support generic client headers map --- README.md | 5 +- .../cloud/mcp/HttpMcpToolboxClient.java | 52 +++++++++++++++---- .../google/cloud/mcp/McpToolboxClient.java | 8 +++ .../cloud/mcp/McpToolboxClientBuilder.java | 23 +++++++- .../mcp/McpToolboxClientBuilderTest.java | 39 ++++++++++++++ .../google/cloud/mcp/e2e/ToolboxE2ESetup.java | 9 +++- 6 files changed, 122 insertions(+), 14 deletions(-) create mode 100644 src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java diff --git a/README.md b/README.md index 5ab228d..b5aa640 100644 --- a/README.md +++ b/README.md @@ -125,7 +125,8 @@ McpToolboxClient client = McpToolboxClient.builder() // Cloud Run Production McpToolboxClient client = McpToolboxClient.builder() .baseUrl("https://my-toolbox-service.a.run.app/mcp") - // .apiKey("...") // Optional: Overrides automatic Google Auth + // .headers(Map.of("Authorization", "Bearer YOUR_TOKEN")) // Optional: Add custom headers, overrides automatic Google Auth + // .apiKey("...") // Optional: Deprecated but supported for backward compatibility .build(); ``` @@ -283,7 +284,7 @@ export GOOGLE_APPLICATION_CREDENTIALS="/path/to/key.json" | **Cloud Run** | Uses Service Account | **None.** (Automatic) | | **CI/CD** | Uses Service Account Key | Set GOOGLE\_APPLICATION\_CREDENTIALS=/path/to/key.json | -*Note: If you provide an `.apiKey()` in the builder, it overrides the automatic ADC mechanism.* +*Note: If you provide an `.apiKey()` or `Authorization` in `.headers()` in the builder, it overrides the automatic ADC mechanism.* ### Authenticating the Tools diff --git a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java index 33d2413..b5c318c 100644 --- a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java @@ -31,6 +31,7 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Set; import java.util.concurrent.CompletableFuture; @@ -45,7 +46,7 @@ public class HttpMcpToolboxClient implements McpToolboxClient { + " communication is sent over HTTPS."; private final String baseUrl; - private final String apiKey; + private final Map headers; private final HttpClient httpClient; private final ObjectMapper objectMapper; private boolean initialized = false; @@ -58,10 +59,30 @@ public class HttpMcpToolboxClient implements McpToolboxClient { * @param apiKey The API key for authentication (optional). */ public HttpMcpToolboxClient(String baseUrl, String apiKey) { + this(baseUrl, apiKey != null && !apiKey.isEmpty() ? Map.of("Authorization", apiKey.startsWith("Bearer ") ? apiKey : "Bearer " + apiKey) : Collections.emptyMap()); + } + + /** + * Constructs a new HttpMcpToolboxClient with generic headers. + * + * @param baseUrl The base URL of the MCP Toolbox Server. + * @param headers The HTTP headers to include in requests. + */ + public HttpMcpToolboxClient(String baseUrl, Map headers) { this.baseUrl = baseUrl.endsWith("/") ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl; - this.apiKey = apiKey; + this.headers = headers != null ? new HashMap<>(headers) : new HashMap<>(); this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(); this.objectMapper = new ObjectMapper(); + + if (this.baseUrl.toLowerCase(Locale.ROOT).startsWith("http://")) { + for (String headerName : this.headers.keySet()) { + String lower = headerName.toLowerCase(Locale.ROOT); + if (lower.contains("authorization") || lower.contains("api-key") || lower.contains("token")) { + logger.warning("WARNING: Sending credentials over unencrypted HTTP! This exposes credentials to network interception."); + break; + } + } + } } private synchronized CompletableFuture ensureInitialized(String authHeader) { @@ -80,7 +101,8 @@ private synchronized CompletableFuture ensureInitialized(String authHeader .uri(URI.create(baseUrl)) .header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.ofString(body)); - if (authHeader != null) req.header("Authorization", authHeader); + this.headers.forEach(req::setHeader); + if (authHeader != null) req.setHeader("Authorization", authHeader); return httpClient .sendAsync(req.build(), HttpResponse.BodyHandlers.ofString()) @@ -100,7 +122,8 @@ private synchronized CompletableFuture ensureInitialized(String authHeader .header("Content-Type", "application/json") .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(notifBody)); - if (authHeader != null) nReq.header("Authorization", authHeader); + this.headers.forEach(nReq::setHeader); + if (authHeader != null) nReq.setHeader("Authorization", authHeader); return httpClient .sendAsync(nReq.build(), HttpResponse.BodyHandlers.ofString()) @@ -144,7 +167,8 @@ public CompletableFuture> loadToolset(String toolset .header("Content-Type", "application/json") .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(body)); - if (authHeader != null) req.header("Authorization", authHeader); + this.headers.forEach(req::setHeader); + if (authHeader != null) req.setHeader("Authorization", authHeader); return httpClient .sendAsync(req.build(), HttpResponse.BodyHandlers.ofString()) @@ -246,8 +270,13 @@ public CompletableFuture invokeTool( // Determine priority Auth header before init so init requests can use it if // needed String finalAuthHeader = null; - if (extraHeaders.containsKey("Authorization")) { - finalAuthHeader = extraHeaders.get("Authorization"); + String authKeyInExtra = extraHeaders.keySet().stream() + .filter(k -> "Authorization".equalsIgnoreCase(k)) + .findFirst() + .orElse(null); + + if (authKeyInExtra != null) { + finalAuthHeader = extraHeaders.get(authKeyInExtra); } else if (adcHeader != null) { finalAuthHeader = adcHeader; } @@ -270,10 +299,11 @@ public CompletableFuture invokeTool( .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(requestBody)); + this.headers.forEach(requestBuilder::setHeader); + extraHeaders.forEach(requestBuilder::setHeader); if (reqAuth != null) { requestBuilder.setHeader("Authorization", reqAuth); } - extraHeaders.forEach(requestBuilder::setHeader); return httpClient .sendAsync( @@ -291,8 +321,10 @@ public CompletableFuture invokeTool( } private String getAuthorizationHeader() { - if (this.apiKey != null && !this.apiKey.isEmpty()) { - return this.apiKey.startsWith("Bearer ") ? this.apiKey : "Bearer " + this.apiKey; + for (Map.Entry entry : this.headers.entrySet()) { + if ("Authorization".equalsIgnoreCase(entry.getKey())) { + return entry.getValue(); + } } try { GoogleCredentials credentials = GoogleCredentials.getApplicationDefault(); diff --git a/src/main/java/com/google/cloud/mcp/McpToolboxClient.java b/src/main/java/com/google/cloud/mcp/McpToolboxClient.java index 3a718ef..2a86cd8 100644 --- a/src/main/java/com/google/cloud/mcp/McpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/McpToolboxClient.java @@ -127,6 +127,14 @@ interface Builder { */ Builder apiKey(String apiKey); + /** + * Sets additional HTTP headers to be included in all requests to the MCP Toolbox Server. + * + * @param headers The HTTP headers. + * @return The builder instance. + */ + Builder headers(Map headers); + /** * Builds and returns a new {@link McpToolboxClient} instance. * diff --git a/src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.java b/src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.java index e98bc5f..b1a7d5b 100644 --- a/src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.java +++ b/src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.java @@ -16,10 +16,14 @@ package com.google.cloud.mcp; +import java.util.HashMap; +import java.util.Map; + /** Implementation of the {@link McpToolboxClient.Builder} interface. */ public class McpToolboxClientBuilder implements McpToolboxClient.Builder { private String baseUrl; private String apiKey; + private Map headers = new HashMap<>(); /** Constructs a new McpToolboxClientBuilder. */ public McpToolboxClientBuilder() {} @@ -36,6 +40,14 @@ public McpToolboxClient.Builder apiKey(String apiKey) { return this; } + @Override + public McpToolboxClient.Builder headers(Map headers) { + if (headers != null) { + this.headers.putAll(headers); + } + return this; + } + @Override public McpToolboxClient build() { if (baseUrl == null || baseUrl.isEmpty()) { @@ -45,6 +57,15 @@ public McpToolboxClient build() { if (baseUrl.endsWith("/")) { baseUrl = baseUrl.substring(0, baseUrl.length() - 1); } - return new HttpMcpToolboxClient(baseUrl, apiKey); + + Map finalHeaders = new HashMap<>(this.headers); + if (apiKey != null && !apiKey.isEmpty()) { + String authValue = apiKey.startsWith("Bearer ") ? apiKey : "Bearer " + apiKey; + if (finalHeaders.keySet().stream().noneMatch(k -> "Authorization".equalsIgnoreCase(k))) { + finalHeaders.put("Authorization", authValue); + } + } + + return new HttpMcpToolboxClient(baseUrl, finalHeaders); } } diff --git a/src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java b/src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java new file mode 100644 index 0000000..36e0dce --- /dev/null +++ b/src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java @@ -0,0 +1,39 @@ +/* + * 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.cloud.mcp; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Map; +import org.junit.jupiter.api.Test; + +class McpToolboxClientBuilderTest { + + @Test + void testHeadersAndApiKey() { + McpToolboxClient client = + McpToolboxClient.builder() + .baseUrl("http://localhost:8080") + .apiKey("my-api-key") + .headers(Map.of("X-Custom-Header", "value1", "Authorization", "Bearer custom-token")) + .build(); + + assertNotNull(client); + assertTrue(client instanceof HttpMcpToolboxClient); + } +} diff --git a/src/test/java/com/google/cloud/mcp/e2e/ToolboxE2ESetup.java b/src/test/java/com/google/cloud/mcp/e2e/ToolboxE2ESetup.java index 06ab3fe..01289aa 100644 --- a/src/test/java/com/google/cloud/mcp/e2e/ToolboxE2ESetup.java +++ b/src/test/java/com/google/cloud/mcp/e2e/ToolboxE2ESetup.java @@ -48,7 +48,14 @@ public class ToolboxE2ESetup implements BeforeAllCallback, AfterAllCallback { @Override public void beforeAll(ExtensionContext context) throws Exception { - String projectId = getEnvVar(PROJECT_ID_ENV); + String projectId = System.getenv(PROJECT_ID_ENV); + if (projectId == null || projectId.trim().isEmpty()) { + logger.warning("Environment variable " + PROJECT_ID_ENV + " is not set. Skipping E2E tests."); + org.junit.jupiter.api.Assumptions.assumeTrue( + false, "Skipping E2E tests because " + PROJECT_ID_ENV + " is not set."); + return; + } + String toolboxVersion = getEnvVar(TOOLBOX_VERSION_ENV); String manifestVersion = getEnvVar(TOOLBOX_MANIFEST_VERSION_ENV); From d0063fe11885ff63f1ff6556b50615488b5acdef Mon Sep 17 00:00:00 2001 From: Stenal P Jolly Date: Mon, 6 Apr 2026 21:00:54 +0530 Subject: [PATCH 2/3] fix: address code review feedback on generic headers --- .../cloud/mcp/HttpMcpToolboxClient.java | 34 +++++++++---------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java index b5c318c..1b6e88f 100644 --- a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java @@ -31,7 +31,6 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; -import java.util.Locale; import java.util.Map; import java.util.Set; import java.util.concurrent.CompletableFuture; @@ -44,7 +43,6 @@ public class HttpMcpToolboxClient implements McpToolboxClient { private static final String HTTP_WARNING = "This connection is using HTTP. To prevent credential exposure, please ensure all" + " communication is sent over HTTPS."; - private final String baseUrl; private final Map headers; private final HttpClient httpClient; @@ -70,19 +68,9 @@ public HttpMcpToolboxClient(String baseUrl, String apiKey) { */ public HttpMcpToolboxClient(String baseUrl, Map headers) { this.baseUrl = baseUrl.endsWith("/") ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl; - this.headers = headers != null ? new HashMap<>(headers) : new HashMap<>(); + this.headers = headers != null ? java.util.Collections.unmodifiableMap(new java.util.HashMap<>(headers)) : java.util.Collections.emptyMap(); this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(); this.objectMapper = new ObjectMapper(); - - if (this.baseUrl.toLowerCase(Locale.ROOT).startsWith("http://")) { - for (String headerName : this.headers.keySet()) { - String lower = headerName.toLowerCase(Locale.ROOT); - if (lower.contains("authorization") || lower.contains("api-key") || lower.contains("token")) { - logger.warning("WARNING: Sending credentials over unencrypted HTTP! This exposes credentials to network interception."); - break; - } - } - } } private synchronized CompletableFuture ensureInitialized(String authHeader) { @@ -101,7 +89,9 @@ private synchronized CompletableFuture ensureInitialized(String authHeader .uri(URI.create(baseUrl)) .header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.ofString(body)); - this.headers.forEach(req::setHeader); + this.headers.forEach((k, v) -> { + if (!"Authorization".equalsIgnoreCase(k)) req.setHeader(k, v); + }); if (authHeader != null) req.setHeader("Authorization", authHeader); return httpClient @@ -122,7 +112,9 @@ private synchronized CompletableFuture ensureInitialized(String authHeader .header("Content-Type", "application/json") .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(notifBody)); - this.headers.forEach(nReq::setHeader); + this.headers.forEach((k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) nReq.setHeader(k, val); + }); if (authHeader != null) nReq.setHeader("Authorization", authHeader); return httpClient @@ -167,7 +159,9 @@ public CompletableFuture> loadToolset(String toolset .header("Content-Type", "application/json") .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(body)); - this.headers.forEach(req::setHeader); + this.headers.forEach((k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) req.setHeader(k, val); + }); if (authHeader != null) req.setHeader("Authorization", authHeader); return httpClient @@ -299,8 +293,12 @@ public CompletableFuture invokeTool( .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(requestBody)); - this.headers.forEach(requestBuilder::setHeader); - extraHeaders.forEach(requestBuilder::setHeader); + this.headers.forEach((k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) requestBuilder.setHeader(k, val); + }); + extraHeaders.forEach((k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) requestBuilder.setHeader(k, val); + }); if (reqAuth != null) { requestBuilder.setHeader("Authorization", reqAuth); } From c266fd157485bfd5ce4622e3b03b04efa06c75a3 Mon Sep 17 00:00:00 2001 From: Stenal P Jolly Date: Tue, 7 Apr 2026 09:26:29 +0530 Subject: [PATCH 3/3] style: apply google-java-format --- .../cloud/mcp/HttpMcpToolboxClient.java | 57 ++++++++++++------- .../mcp/McpToolboxClientBuilderTest.java | 2 +- 2 files changed, 37 insertions(+), 22 deletions(-) diff --git a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java index 1b6e88f..ef99cbb 100644 --- a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java @@ -57,7 +57,11 @@ public class HttpMcpToolboxClient implements McpToolboxClient { * @param apiKey The API key for authentication (optional). */ public HttpMcpToolboxClient(String baseUrl, String apiKey) { - this(baseUrl, apiKey != null && !apiKey.isEmpty() ? Map.of("Authorization", apiKey.startsWith("Bearer ") ? apiKey : "Bearer " + apiKey) : Collections.emptyMap()); + this( + baseUrl, + apiKey != null && !apiKey.isEmpty() + ? Map.of("Authorization", apiKey.startsWith("Bearer ") ? apiKey : "Bearer " + apiKey) + : Collections.emptyMap()); } /** @@ -68,7 +72,10 @@ public HttpMcpToolboxClient(String baseUrl, String apiKey) { */ public HttpMcpToolboxClient(String baseUrl, Map headers) { this.baseUrl = baseUrl.endsWith("/") ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl; - this.headers = headers != null ? java.util.Collections.unmodifiableMap(new java.util.HashMap<>(headers)) : java.util.Collections.emptyMap(); + this.headers = + headers != null + ? java.util.Collections.unmodifiableMap(new java.util.HashMap<>(headers)) + : java.util.Collections.emptyMap(); this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(); this.objectMapper = new ObjectMapper(); } @@ -89,9 +96,10 @@ private synchronized CompletableFuture ensureInitialized(String authHeader .uri(URI.create(baseUrl)) .header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.ofString(body)); - this.headers.forEach((k, v) -> { - if (!"Authorization".equalsIgnoreCase(k)) req.setHeader(k, v); - }); + this.headers.forEach( + (k, v) -> { + if (!"Authorization".equalsIgnoreCase(k)) req.setHeader(k, v); + }); if (authHeader != null) req.setHeader("Authorization", authHeader); return httpClient @@ -112,9 +120,10 @@ private synchronized CompletableFuture ensureInitialized(String authHeader .header("Content-Type", "application/json") .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(notifBody)); - this.headers.forEach((k, val) -> { - if (!"Authorization".equalsIgnoreCase(k)) nReq.setHeader(k, val); - }); + this.headers.forEach( + (k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) nReq.setHeader(k, val); + }); if (authHeader != null) nReq.setHeader("Authorization", authHeader); return httpClient @@ -159,9 +168,10 @@ public CompletableFuture> loadToolset(String toolset .header("Content-Type", "application/json") .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(body)); - this.headers.forEach((k, val) -> { - if (!"Authorization".equalsIgnoreCase(k)) req.setHeader(k, val); - }); + this.headers.forEach( + (k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) req.setHeader(k, val); + }); if (authHeader != null) req.setHeader("Authorization", authHeader); return httpClient @@ -264,10 +274,11 @@ public CompletableFuture invokeTool( // Determine priority Auth header before init so init requests can use it if // needed String finalAuthHeader = null; - String authKeyInExtra = extraHeaders.keySet().stream() - .filter(k -> "Authorization".equalsIgnoreCase(k)) - .findFirst() - .orElse(null); + String authKeyInExtra = + extraHeaders.keySet().stream() + .filter(k -> "Authorization".equalsIgnoreCase(k)) + .findFirst() + .orElse(null); if (authKeyInExtra != null) { finalAuthHeader = extraHeaders.get(authKeyInExtra); @@ -293,12 +304,16 @@ public CompletableFuture invokeTool( .header("MCP-Protocol-Version", protocolVersion) .POST(HttpRequest.BodyPublishers.ofString(requestBody)); - this.headers.forEach((k, val) -> { - if (!"Authorization".equalsIgnoreCase(k)) requestBuilder.setHeader(k, val); - }); - extraHeaders.forEach((k, val) -> { - if (!"Authorization".equalsIgnoreCase(k)) requestBuilder.setHeader(k, val); - }); + this.headers.forEach( + (k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) + requestBuilder.setHeader(k, val); + }); + extraHeaders.forEach( + (k, val) -> { + if (!"Authorization".equalsIgnoreCase(k)) + requestBuilder.setHeader(k, val); + }); if (reqAuth != null) { requestBuilder.setHeader("Authorization", reqAuth); } diff --git a/src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java b/src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java index 36e0dce..8b0793d 100644 --- a/src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java +++ b/src/test/java/com/google/cloud/mcp/McpToolboxClientBuilderTest.java @@ -32,7 +32,7 @@ void testHeadersAndApiKey() { .apiKey("my-api-key") .headers(Map.of("X-Custom-Header", "value1", "Authorization", "Bearer custom-token")) .build(); - + assertNotNull(client); assertTrue(client instanceof HttpMcpToolboxClient); }