feat: add pre & post processing hooks for tools (#58)

Author: stenalpjolly

src/main/java/com/google/cloud/mcp/McpToolboxClient.java

@@ -143,6 +143,22 @@ interface Builder {
      */
     Builder credentialsProvider(CredentialsProvider credentialsProvider);
 
+    /**
+     * 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.
      *

src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.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.concurrent.CompletableFuture;
 
@@ -26,6 +28,8 @@ public final class McpToolboxClientBuilder implements McpToolboxClient.Builder {
   private String apiKey;
   private Map<String, String> headers = new HashMap<>();
   private CredentialsProvider credentialsProvider;
+  private final List<ToolPreProcessor> preProcessors = new ArrayList<>();
+  private final List<ToolPostProcessor> postProcessors = new ArrayList<>();
 
   /** Constructs a new McpToolboxClientBuilder. */
   public McpToolboxClientBuilder() {}
@@ -56,6 +60,22 @@ public McpToolboxClient.Builder credentialsProvider(CredentialsProvider credenti
     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()) {
@@ -83,6 +103,7 @@ public McpToolboxClient build() {
     }
 
     Transport transport = new HttpMcpTransport(baseUrl, this.headers, resolvedProvider);
-    return new McpToolboxClientImpl(transport, this.headers, resolvedProvider);
+    return new McpToolboxClientImpl(
+        transport, this.headers, resolvedProvider, preProcessors, postProcessors);
   }
 }

src/main/java/com/google/cloud/mcp/McpToolboxClientImpl.java

@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
@@ -49,6 +50,9 @@ public final class McpToolboxClientImpl implements McpToolboxClient {
   /** Jackson ObjectMapper for JSON parsing. */
   private final ObjectMapper objectMapper;
 
+  private final List<ToolPreProcessor> preProcessors;
+  private final List<ToolPostProcessor> postProcessors;
+
   /**
    * Constructs a new McpToolboxClientImpl.
    *
@@ -61,22 +65,14 @@ public McpToolboxClientImpl(final Transport clientTransport) {
   /**
    * Constructs a new McpToolboxClientImpl.
    *
-   * @param clientTransport The underlying MCP transport layer.
-   * @param clientHeaders Fallback headers for deprecated constructor compatibility.
-   * @param provider Fallback provider for deprecated constructor compatibility.
+   * @param transport The underlying MCP transport layer.
+   * @param headers Fallback headers for deprecated constructor compatibility.
+   * @param credentialsProvider Fallback provider for deprecated constructor compatibility.
    */
   @Deprecated
   public McpToolboxClientImpl(
-      final Transport clientTransport,
-      final Map<String, String> clientHeaders,
-      final CredentialsProvider provider) {
-    this.transport = clientTransport;
-    this.headers =
-        clientHeaders != null
-            ? java.util.Collections.unmodifiableMap(new java.util.HashMap<>(clientHeaders))
-            : java.util.Collections.emptyMap();
-    this.credentialsProvider = provider;
-    this.objectMapper = new ObjectMapper();
+      Transport transport, Map<String, String> headers, CredentialsProvider credentialsProvider) {
+    this(transport, headers, credentialsProvider, null, null);
   }
 
   /**
@@ -152,6 +148,32 @@ private static CredentialsProvider apiKeyToProvider(final String apiKey) {
     return () -> CompletableFuture.completedFuture(bearerKey);
   }
 
+  /**
+   * Primary constructor for McpToolboxClientImpl.
+   *
+   * @param transport The underlying MCP transport layer.
+   * @param headers Default HTTP headers.
+   * @param credentialsProvider Provider for credentials.
+   * @param preProcessors List of pre-processors.
+   * @param postProcessors List of post-processors.
+   */
+  public McpToolboxClientImpl(
+      Transport transport,
+      Map<String, String> headers,
+      CredentialsProvider credentialsProvider,
+      List<ToolPreProcessor> preProcessors,
+      List<ToolPostProcessor> postProcessors) {
+    this.transport = transport;
+    this.headers =
+        headers != null
+            ? java.util.Collections.unmodifiableMap(new java.util.HashMap<>(headers))
+            : java.util.Collections.emptyMap();
+    this.credentialsProvider = credentialsProvider;
+    this.preProcessors = preProcessors != null ? List.copyOf(preProcessors) : List.of();
+    this.postProcessors = postProcessors != null ? List.copyOf(postProcessors) : List.of();
+    this.objectMapper = new ObjectMapper();
+  }
+
   private CompletableFuture<Map<String, String>> getMergedMetadata(
       final Map<String, String> extraMetadata) {
     if (this.transport instanceof HttpMcpTransport) {
@@ -250,6 +272,12 @@ public CompletableFuture<Map<String, Tool>> 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;
@@ -279,6 +307,12 @@ public CompletableFuture<Tool> 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;
             });
   }

src/main/java/com/google/cloud/mcp/ResolvedAuth.java

@@ -22,6 +22,11 @@
 public final class ResolvedAuth {
   private final Map<String, String> tokens;
 
+  /**
+   * Constructs a new ResolvedAuth.
+   *
+   * @param tokens The map of resolved auth tokens.
+   */
   public ResolvedAuth(Map<String, String> tokens) {
     Map<String, String> copy = new java.util.HashMap<>();
     if (tokens != null) {

src/main/java/com/google/cloud/mcp/Tool.java

@@ -35,6 +35,8 @@ public class Tool {
 
   private final Map<String, Object> boundParameters = new HashMap<>();
   private final Map<String, AuthTokenGetter> authGetters = new HashMap<>();
+  private final List<ToolPreProcessor> preProcessors = new ArrayList<>();
+  private final List<ToolPostProcessor> postProcessors = new ArrayList<>();
 
   /**
    * Constructs a new Tool.
@@ -103,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.
@@ -111,34 +135,53 @@ public Tool addAuthTokenGetter(String serviceName, AuthTokenGetter getter) {
    * @return A CompletableFuture containing the result of the tool execution.
    */
   public CompletableFuture<ToolResult> execute(Map<String, Object> args) {
-    Map<String, Object> finalArgs = new HashMap<>(args);
-    Map<String, String> extraHeaders = new HashMap<>();
-
-    // 1. Apply Bound Parameters
-    for (Map.Entry<String, Object> 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<Map<String, Object>> argsFuture =
+        CompletableFuture.completedFuture(new HashMap<>(args));
+
+    for (ToolPreProcessor preProcessor : preProcessors) {
+      argsFuture = argsFuture.thenCompose(currentArgs -> preProcessor.process(name, currentArgs));
     }
 
-    // 2. Resolve Auth & Execute
-    return AuthResolver.resolve(authGetters)
-        .thenCompose(
-            resolvedAuth -> {
-              try {
-                // Apply credential parameter bindings and extra headers
-                resolvedAuth.applyTo(finalArgs, extraHeaders, definition);
-
-                // 3. Validation & Cleanup
-                validateAndSanitizeArgs(finalArgs);
-                return client.invokeTool(name, finalArgs, extraHeaders);
-              } catch (Exception e) {
-                return CompletableFuture.failedFuture(e);
+    CompletableFuture<ToolResult> resultFuture =
+        argsFuture.thenCompose(
+            processedArgs -> {
+              Map<String, Object> finalArgs =
+                  java.util.Collections.synchronizedMap(new HashMap<>(processedArgs));
+              Map<String, String> extraHeaders =
+                  java.util.Collections.synchronizedMap(new HashMap<>());
+
+              // 1. Apply Bound Parameters
+              for (Map.Entry<String, Object> 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 & Execute
+              return AuthResolver.resolve(authGetters)
+                  .thenCompose(
+                      resolvedAuth -> {
+                        try {
+                          // Apply credential parameter bindings and extra headers
+                          resolvedAuth.applyTo(finalArgs, extraHeaders, definition);
+
+                          // Validation & Cleanup
+                          validateAndSanitizeArgs(finalArgs);
+                          return client.invokeTool(name, finalArgs, extraHeaders);
+                        } catch (Exception e) {
+                          return CompletableFuture.failedFuture(e);
+                        }
+                      });
             });
+
+    for (ToolPostProcessor postProcessor : postProcessors) {
+      resultFuture = resultFuture.thenCompose(res -> postProcessor.process(name, res));
+    }
+
+    return resultFuture;
   }
 
   /** Validates arguments against the tool definition and removes null values. */

src/main/java/com/google/cloud/mcp/ToolDefinition.java

@@ -25,6 +25,9 @@
  *
  * @param description A description of what the tool does.
  * @param parameters A list of parameters the tool accepts.
+ * @param authRequired List of auth services required by the tool.
+ * @param readOnlyHint Hint indicating whether the tool is read-only.
+ * @param destructiveHint Hint indicating whether the tool is destructive.
  */
 @JsonIgnoreProperties(ignoreUnknown = true)
 public record ToolDefinition(
@@ -34,7 +37,13 @@ public record ToolDefinition(
     Boolean readOnlyHint,
     Boolean destructiveHint) {
 
-  /** Backward-compatible constructor. */
+  /**
+   * Backward-compatible constructor.
+   *
+   * @param description A description of what the tool does.
+   * @param parameters A list of parameters the tool accepts.
+   * @param authRequired List of auth services required.
+   */
   public ToolDefinition(String description, List<Parameter> parameters, List<String> authRequired) {
     this(description, parameters, authRequired, null, null);
   }
@@ -58,7 +67,15 @@ public record Parameter(
       List<String> authSources, // Maps services to parameters
       @JsonProperty("default") Object defaultValue) {
 
-    /** Backward-compatible constructor. */
+    /**
+     * Backward-compatible constructor.
+     *
+     * @param name The name of the parameter.
+     * @param type The type of the parameter.
+     * @param required Whether the parameter is required.
+     * @param description A description of the parameter.
+     * @param authSources Authentication sources list.
+     */
     public Parameter(
         String name, String type, boolean required, String description, List<String> authSources) {
       this(name, type, required, description, authSources, null);

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<ToolResult> process(String toolName, ToolResult result);
+}