From 207591b8dae39ad575fb52c6affc6708eb2ad57e Mon Sep 17 00:00:00 2001 From: Stenal P Jolly Date: Tue, 7 Apr 2026 11:19:05 +0530 Subject: [PATCH 1/3] feat: add pre and post processors --- .../cloud/mcp/HttpMcpToolboxClient.java | 32 ++++ .../google/cloud/mcp/McpToolboxClient.java | 16 ++ .../cloud/mcp/McpToolboxClientBuilder.java | 23 ++- src/main/java/com/google/cloud/mcp/Tool.java | 142 ++++++++++++------ .../google/cloud/mcp/ToolPostProcessor.java | 33 ++++ .../google/cloud/mcp/ToolPreProcessor.java | 34 +++++ .../java/com/google/cloud/mcp/ToolTest.java | 138 +++++++++++++++++ .../google/cloud/mcp/e2e/ToolboxE2ESetup.java | 9 +- 8 files changed, 380 insertions(+), 47 deletions(-) create mode 100644 src/main/java/com/google/cloud/mcp/ToolPostProcessor.java create mode 100644 src/main/java/com/google/cloud/mcp/ToolPreProcessor.java create mode 100644 src/test/java/com/google/cloud/mcp/ToolTest.java diff --git a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java index fd699fe..0954795 100644 --- a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java @@ -44,6 +44,8 @@ public class HttpMcpToolboxClient implements McpToolboxClient { private final ObjectMapper objectMapper; private boolean initialized = false; private final String protocolVersion = "2025-11-25"; + private final List preProcessors; + private final List postProcessors; /** * Constructs a new HttpMcpToolboxClient. @@ -52,10 +54,28 @@ public class HttpMcpToolboxClient implements McpToolboxClient { * @param apiKey The API key for authentication (optional). */ public HttpMcpToolboxClient(String baseUrl, String apiKey) { + this(baseUrl, apiKey, Collections.emptyList(), Collections.emptyList()); + } + + /** + * Constructs a new HttpMcpToolboxClient with processors. + * + * @param baseUrl The base URL of the MCP Toolbox Server. + * @param apiKey The API key for authentication (optional). + * @param preProcessors The pre-processors to apply. + * @param postProcessors The post-processors to apply. + */ + public HttpMcpToolboxClient( + String baseUrl, + String apiKey, + List preProcessors, + List postProcessors) { this.baseUrl = baseUrl.endsWith("/") ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl; this.apiKey = apiKey; this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(); this.objectMapper = new ObjectMapper(); + this.preProcessors = preProcessors != null ? new ArrayList<>(preProcessors) : new ArrayList<>(); + this.postProcessors = postProcessors != null ? new ArrayList<>(postProcessors) : new ArrayList<>(); } private synchronized CompletableFuture ensureInitialized(String authHeader) { @@ -177,6 +197,12 @@ public CompletableFuture> loadToolset( if (authBinds != null && authBinds.containsKey(toolName)) { authBinds.get(toolName).forEach(tool::addAuthTokenGetter); } + for (ToolPreProcessor preProcessor : this.preProcessors) { + tool.addPreProcessor(preProcessor); + } + for (ToolPostProcessor postProcessor : this.postProcessors) { + tool.addPostProcessor(postProcessor); + } tools.put(toolName, tool); } return tools; @@ -201,6 +227,12 @@ public CompletableFuture loadTool( if (authTokenGetters != null) { authTokenGetters.forEach(tool::addAuthTokenGetter); } + for (ToolPreProcessor preProcessor : this.preProcessors) { + tool.addPreProcessor(preProcessor); + } + for (ToolPostProcessor postProcessor : this.postProcessors) { + tool.addPostProcessor(postProcessor); + } return tool; }); } diff --git a/src/main/java/com/google/cloud/mcp/McpToolboxClient.java b/src/main/java/com/google/cloud/mcp/McpToolboxClient.java index 3a718ef..014efa4 100644 --- a/src/main/java/com/google/cloud/mcp/McpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/McpToolboxClient.java @@ -127,6 +127,22 @@ interface Builder { */ Builder apiKey(String apiKey); + /** + * Adds a global pre-processor that will be applied to all tools loaded by this client. + * + * @param preProcessor The pre-processor to add. + * @return The builder instance. + */ + Builder preProcessor(ToolPreProcessor preProcessor); + + /** + * Adds a global post-processor that will be applied to all tools loaded by this client. + * + * @param postProcessor The post-processor to add. + * @return The builder instance. + */ + Builder postProcessor(ToolPostProcessor postProcessor); + /** * 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..7bc19d1 100644 --- a/src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.java +++ b/src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.java @@ -16,10 +16,15 @@ package com.google.cloud.mcp; +import java.util.ArrayList; +import java.util.List; + /** Implementation of the {@link McpToolboxClient.Builder} interface. */ public class McpToolboxClientBuilder implements McpToolboxClient.Builder { private String baseUrl; private String apiKey; + private final List preProcessors = new ArrayList<>(); + private final List postProcessors = new ArrayList<>(); /** Constructs a new McpToolboxClientBuilder. */ public McpToolboxClientBuilder() {} @@ -36,6 +41,22 @@ public McpToolboxClient.Builder apiKey(String apiKey) { return this; } + @Override + public McpToolboxClient.Builder preProcessor(ToolPreProcessor preProcessor) { + if (preProcessor != null) { + this.preProcessors.add(preProcessor); + } + return this; + } + + @Override + public McpToolboxClient.Builder postProcessor(ToolPostProcessor postProcessor) { + if (postProcessor != null) { + this.postProcessors.add(postProcessor); + } + return this; + } + @Override public McpToolboxClient build() { if (baseUrl == null || baseUrl.isEmpty()) { @@ -45,6 +66,6 @@ public McpToolboxClient build() { if (baseUrl.endsWith("/")) { baseUrl = baseUrl.substring(0, baseUrl.length() - 1); } - return new HttpMcpToolboxClient(baseUrl, apiKey); + return new HttpMcpToolboxClient(baseUrl, apiKey, preProcessors, postProcessors); } } diff --git a/src/main/java/com/google/cloud/mcp/Tool.java b/src/main/java/com/google/cloud/mcp/Tool.java index e32d87e..ab46690 100644 --- a/src/main/java/com/google/cloud/mcp/Tool.java +++ b/src/main/java/com/google/cloud/mcp/Tool.java @@ -16,7 +16,9 @@ package com.google.cloud.mcp; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.CompletableFuture; @@ -33,6 +35,8 @@ public class Tool { private final Map boundParameters = new HashMap<>(); private final Map authGetters = new HashMap<>(); + private final List preProcessors = new ArrayList<>(); + private final List postProcessors = new ArrayList<>(); /** * Constructs a new Tool. @@ -101,6 +105,28 @@ public Tool addAuthTokenGetter(String serviceName, AuthTokenGetter getter) { return this; } + /** + * Adds a pre-processor to the tool. + * + * @param processor The pre-processor to add. + * @return The tool instance. + */ + public Tool addPreProcessor(ToolPreProcessor processor) { + this.preProcessors.add(processor); + return this; + } + + /** + * Adds a post-processor to the tool. + * + * @param processor The post-processor to add. + * @return The tool instance. + */ + public Tool addPostProcessor(ToolPostProcessor processor) { + this.postProcessors.add(processor); + return this; + } + /** * Executes the tool with the provided arguments, applying any bound parameters and resolving * authentication tokens. @@ -109,55 +135,81 @@ public Tool addAuthTokenGetter(String serviceName, AuthTokenGetter getter) { * @return A CompletableFuture containing the result of the tool execution. */ public CompletableFuture execute(Map args) { - Map finalArgs = new HashMap<>(args); - Map extraHeaders = new HashMap<>(); - - // 1. Apply Bound Parameters - for (Map.Entry entry : boundParameters.entrySet()) { - Object val = entry.getValue(); - if (val instanceof Supplier) { - finalArgs.put(entry.getKey(), ((Supplier) val).get()); - } else { - finalArgs.put(entry.getKey(), val); - } + CompletableFuture> argsFuture = + CompletableFuture.completedFuture(new HashMap<>(args)); + + for (ToolPreProcessor preProcessor : preProcessors) { + argsFuture = argsFuture.thenCompose(currentArgs -> preProcessor.process(name, currentArgs)); } - // 2. Resolve Auth Tokens - return CompletableFuture.allOf( - authGetters.entrySet().stream() - .map( - entry -> { - String serviceName = entry.getKey(); - return entry - .getValue() - .getToken() - .thenAccept( - token -> { - // A. Check if mapped to a Parameter (Authenticated Parameters) - String paramName = findParameterForService(serviceName); - if (paramName != null) { - finalArgs.put(paramName, token); - } - - // B. Always add to Headers to support Authorized Invocation - // 1. Standard OIDC Header (Cloud Run) - extraHeaders.put("Authorization", "Bearer " + token); - - // 2. SDK Convention Header (Framework Compatibility) - extraHeaders.put(serviceName + "_token", token); - }); - }) - .toArray(CompletableFuture[]::new)) - .thenCompose( - v -> { - try { - // 3. Validation & Cleanup - validateAndSanitizeArgs(finalArgs); - return client.invokeTool(name, finalArgs, extraHeaders); - } catch (Exception e) { - return CompletableFuture.failedFuture(e); + CompletableFuture resultFuture = + argsFuture.thenCompose( + processedArgs -> { + Map finalArgs = new HashMap<>(processedArgs); + Map extraHeaders = new HashMap<>(); + + // 1. Apply Bound Parameters + for (Map.Entry entry : boundParameters.entrySet()) { + Object val = entry.getValue(); + if (val instanceof Supplier) { + finalArgs.put(entry.getKey(), ((Supplier) val).get()); + } else { + finalArgs.put(entry.getKey(), val); + } } + + // 2. Resolve Auth Tokens + return CompletableFuture.allOf( + authGetters.entrySet().stream() + .map( + entry -> { + String serviceName = entry.getKey(); + return entry + .getValue() + .getToken() + .thenAccept( + token -> { + // A. Check if mapped to a Parameter (Authenticated Parameters) + String paramName = findParameterForService(serviceName); + if (paramName != null) { + finalArgs.put(paramName, token); + } + + // B. Always add to Headers to support Authorized Invocation + // 1. Standard OIDC Header (Cloud Run) + extraHeaders.put("Authorization", "Bearer " + token); + + // 2. SDK Convention Header (Framework Compatibility) + extraHeaders.put(serviceName + "_token", token); + }); + }) + .toArray(CompletableFuture[]::new)) + .thenCompose( + v -> { + try { + // 3. Validation & Cleanup + validateAndSanitizeArgs(finalArgs); + return client.invokeTool(name, finalArgs, extraHeaders); + } catch (Exception e) { + return CompletableFuture.failedFuture(e); + } + }); }); + + for (ToolPostProcessor postProcessor : postProcessors) { + resultFuture = + resultFuture + .handle( + (res, err) -> { + if (err != null) { + return CompletableFuture.failedFuture(err); + } + return postProcessor.process(name, res); + }) + .thenCompose(f -> f); + } + + return resultFuture; } private String findParameterForService(String serviceName) { diff --git a/src/main/java/com/google/cloud/mcp/ToolPostProcessor.java b/src/main/java/com/google/cloud/mcp/ToolPostProcessor.java new file mode 100644 index 0000000..09000bb --- /dev/null +++ b/src/main/java/com/google/cloud/mcp/ToolPostProcessor.java @@ -0,0 +1,33 @@ +/* + * 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 java.util.concurrent.CompletableFuture; + +/** A functional interface for post-processing tool results after invocation. */ +@FunctionalInterface +public interface ToolPostProcessor { + + /** + * Processes the result of a tool after it has been invoked. + * + * @param toolName The name of the tool that was invoked. + * @param result The original tool result. + * @return A CompletableFuture containing the processed tool result. + */ + CompletableFuture process(String toolName, ToolResult result); +} diff --git a/src/main/java/com/google/cloud/mcp/ToolPreProcessor.java b/src/main/java/com/google/cloud/mcp/ToolPreProcessor.java new file mode 100644 index 0000000..a280a85 --- /dev/null +++ b/src/main/java/com/google/cloud/mcp/ToolPreProcessor.java @@ -0,0 +1,34 @@ +/* + * 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 java.util.Map; +import java.util.concurrent.CompletableFuture; + +/** A functional interface for pre-processing tool inputs before invocation. */ +@FunctionalInterface +public interface ToolPreProcessor { + + /** + * Processes the input arguments for a tool before it is invoked. + * + * @param toolName The name of the tool being invoked. + * @param arguments The original arguments provided to the tool. + * @return A CompletableFuture containing the processed arguments. + */ + CompletableFuture> process(String toolName, Map arguments); +} diff --git a/src/test/java/com/google/cloud/mcp/ToolTest.java b/src/test/java/com/google/cloud/mcp/ToolTest.java new file mode 100644 index 0000000..3e8d1a5 --- /dev/null +++ b/src/test/java/com/google/cloud/mcp/ToolTest.java @@ -0,0 +1,138 @@ +/* + * 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.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +class ToolTest { + + private McpToolboxClient mockClient; + private ToolDefinition toolDefinition; + private Tool tool; + + @BeforeEach + void setUp() { + mockClient = mock(McpToolboxClient.class); + toolDefinition = new ToolDefinition("Test Tool", null, null); + tool = new Tool("test_tool", toolDefinition, mockClient); + } + + @Test + @SuppressWarnings("unchecked") + void testExecute_withPreAndPostProcessors_modifiesArgsAndResult() throws Exception { + // Arrange + Map initialArgs = new HashMap<>(); + initialArgs.put("arg1", "val1"); + + ToolResult originalResult = + new ToolResult(List.of(new ToolResult.Content("text", "original")), false); + ToolResult modifiedResult = + new ToolResult(List.of(new ToolResult.Content("text", "modified")), false); + + ToolPreProcessor preProcessor1 = + (name, args) -> { + Map newArgs = new HashMap<>(args); + newArgs.put("arg2", "val2"); + return CompletableFuture.completedFuture(newArgs); + }; + + ToolPreProcessor preProcessor2 = + (name, args) -> { + Map newArgs = new HashMap<>(args); + newArgs.put("arg3", "val3"); + return CompletableFuture.completedFuture(newArgs); + }; + + ToolPostProcessor postProcessor = + (name, result) -> { + if (result.content().get(0).text().equals("original")) { + return CompletableFuture.completedFuture(modifiedResult); + } + return CompletableFuture.completedFuture(result); + }; + + tool.addPreProcessor(preProcessor1); + tool.addPreProcessor(preProcessor2); + tool.addPostProcessor(postProcessor); + + when(mockClient.invokeTool(eq("test_tool"), anyMap(), anyMap())) + .thenReturn(CompletableFuture.completedFuture(originalResult)); + + // Act + CompletableFuture futureResult = tool.execute(initialArgs); + ToolResult finalResult = futureResult.get(); + + // Assert + ArgumentCaptor> argsCaptor = ArgumentCaptor.forClass(Map.class); + verify(mockClient, times(1)).invokeTool(eq("test_tool"), argsCaptor.capture(), anyMap()); + + Map capturedArgs = argsCaptor.getValue(); + assertEquals(3, capturedArgs.size()); + assertEquals("val1", capturedArgs.get("arg1")); + assertEquals("val2", capturedArgs.get("arg2")); + assertEquals("val3", capturedArgs.get("arg3")); + + assertSame(modifiedResult, finalResult); + } + + @Test + void testExecute_preProcessorException_failsFutureWithoutInvokingClient() { + // Arrange + Map initialArgs = new HashMap<>(); + + ToolPreProcessor preProcessor = + (name, args) -> CompletableFuture.failedFuture(new RuntimeException("PreProcessor failed")); + + tool.addPreProcessor(preProcessor); + + // Act + CompletableFuture futureResult = tool.execute(initialArgs); + + // Assert + assertTrue(futureResult.isCompletedExceptionally()); + + Exception exception = null; + try { + futureResult.get(); + } catch (InterruptedException | ExecutionException e) { + exception = e; + } + assertTrue(exception.getCause() instanceof RuntimeException); + assertEquals("PreProcessor failed", exception.getCause().getMessage()); + + verify(mockClient, never()).invokeTool(eq("test_tool"), anyMap(), anyMap()); + verify(mockClient, never()).invokeTool(eq("test_tool"), anyMap()); + } +} 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 a5044ab..19e55c8 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 b758c0404178c7154e2ca1fd58aec7377300800f Mon Sep 17 00:00:00 2001 From: Stenal P Jolly Date: Tue, 7 Apr 2026 11:55:45 +0530 Subject: [PATCH 2/3] refactor: simplify CompletableFuture chaining for post processors --- src/main/java/com/google/cloud/mcp/Tool.java | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/main/java/com/google/cloud/mcp/Tool.java b/src/main/java/com/google/cloud/mcp/Tool.java index ab46690..10aae99 100644 --- a/src/main/java/com/google/cloud/mcp/Tool.java +++ b/src/main/java/com/google/cloud/mcp/Tool.java @@ -197,16 +197,7 @@ public CompletableFuture execute(Map args) { }); for (ToolPostProcessor postProcessor : postProcessors) { - resultFuture = - resultFuture - .handle( - (res, err) -> { - if (err != null) { - return CompletableFuture.failedFuture(err); - } - return postProcessor.process(name, res); - }) - .thenCompose(f -> f); + resultFuture = resultFuture.thenCompose(res -> postProcessor.process(name, res)); } return resultFuture; From a6145f6b316adeb2c0511737db5e51b14a151ac0 Mon Sep 17 00:00:00 2001 From: Stenal P Jolly Date: Thu, 9 Apr 2026 20:26:00 +0530 Subject: [PATCH 3/3] chore: fix formatting for CI/CD --- .../java/com/google/cloud/mcp/HttpMcpToolboxClient.java | 3 ++- src/main/java/com/google/cloud/mcp/Tool.java | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java index 0954795..7b235e4 100644 --- a/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java +++ b/src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java @@ -75,7 +75,8 @@ public HttpMcpToolboxClient( this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(); this.objectMapper = new ObjectMapper(); this.preProcessors = preProcessors != null ? new ArrayList<>(preProcessors) : new ArrayList<>(); - this.postProcessors = postProcessors != null ? new ArrayList<>(postProcessors) : new ArrayList<>(); + this.postProcessors = + postProcessors != null ? new ArrayList<>(postProcessors) : new ArrayList<>(); } private synchronized CompletableFuture ensureInitialized(String authHeader) { diff --git a/src/main/java/com/google/cloud/mcp/Tool.java b/src/main/java/com/google/cloud/mcp/Tool.java index 10aae99..bfe3278 100644 --- a/src/main/java/com/google/cloud/mcp/Tool.java +++ b/src/main/java/com/google/cloud/mcp/Tool.java @@ -169,13 +169,15 @@ public CompletableFuture execute(Map args) { .getToken() .thenAccept( token -> { - // A. Check if mapped to a Parameter (Authenticated Parameters) + // A. Check if mapped to a Parameter (Authenticated + // Parameters) String paramName = findParameterForService(serviceName); if (paramName != null) { finalArgs.put(paramName, token); } - // B. Always add to Headers to support Authorized Invocation + // B. Always add to Headers to support Authorized + // Invocation // 1. Standard OIDC Header (Cloud Run) extraHeaders.put("Authorization", "Bearer " + token);