Skip to content

Commit 6f1dbf4

Browse files
stenalpjollystenalpjolly
committed
feat: add pre and post processing hooks for tools
1 parent fb8cb3c commit 6f1dbf4

7 files changed

Lines changed: 374 additions & 47 deletions

File tree

‎src/main/java/com/google/cloud/mcp/HttpMcpToolboxClient.java‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ public class HttpMcpToolboxClient implements McpToolboxClient {
4444
private final ObjectMapper objectMapper;
4545
private boolean initialized = false;
4646
private final String protocolVersion = "2025-11-25";
47+
private final List<ToolPreProcessor> preProcessors;
48+
private final List<ToolPostProcessor> postProcessors;
4749

4850
/**
4951
* Constructs a new HttpMcpToolboxClient.
@@ -56,7 +58,9 @@ public HttpMcpToolboxClient(String baseUrl, String apiKey) {
5658
baseUrl,
5759
apiKey != null && !apiKey.isEmpty()
5860
? Map.of("Authorization", apiKey.startsWith("Bearer ") ? apiKey : "Bearer " + apiKey)
59-
: Collections.emptyMap());
61+
: Collections.emptyMap(),
62+
Collections.emptyList(),
63+
Collections.emptyList());
6064
}
6165

6266
/**
@@ -66,13 +70,31 @@ public HttpMcpToolboxClient(String baseUrl, String apiKey) {
6670
* @param headers The HTTP headers to include in requests.
6771
*/
6872
public HttpMcpToolboxClient(String baseUrl, Map<String, String> headers) {
73+
this(baseUrl, headers, Collections.emptyList(), Collections.emptyList());
74+
}
75+
76+
/**
77+
* Constructs a new HttpMcpToolboxClient with generic headers and processors.
78+
*
79+
* @param baseUrl The base URL of the MCP Toolbox Server.
80+
* @param headers The HTTP headers to include in requests.
81+
* @param preProcessors The pre-processors to apply.
82+
* @param postProcessors The post-processors to apply.
83+
*/
84+
public HttpMcpToolboxClient(
85+
String baseUrl,
86+
Map<String, String> headers,
87+
List<ToolPreProcessor> preProcessors,
88+
List<ToolPostProcessor> postProcessors) {
6989
this.baseUrl = baseUrl.endsWith("/") ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl;
7090
this.headers =
7191
headers != null
7292
? java.util.Collections.unmodifiableMap(new java.util.HashMap<>(headers))
7393
: java.util.Collections.emptyMap();
7494
this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build();
7595
this.objectMapper = new ObjectMapper();
96+
this.preProcessors = preProcessors != null ? new ArrayList<>(preProcessors) : new ArrayList<>();
97+
this.postProcessors = postProcessors != null ? new ArrayList<>(postProcessors) : new ArrayList<>();
7698
}
7799

78100
private synchronized CompletableFuture<Void> ensureInitialized(String authHeader) {
@@ -206,6 +228,12 @@ public CompletableFuture<Map<String, Tool>> loadToolset(
206228
if (authBinds != null && authBinds.containsKey(toolName)) {
207229
authBinds.get(toolName).forEach(tool::addAuthTokenGetter);
208230
}
231+
for (ToolPreProcessor preProcessor : this.preProcessors) {
232+
tool.addPreProcessor(preProcessor);
233+
}
234+
for (ToolPostProcessor postProcessor : this.postProcessors) {
235+
tool.addPostProcessor(postProcessor);
236+
}
209237
tools.put(toolName, tool);
210238
}
211239
return tools;
@@ -230,6 +258,12 @@ public CompletableFuture<Tool> loadTool(
230258
if (authTokenGetters != null) {
231259
authTokenGetters.forEach(tool::addAuthTokenGetter);
232260
}
261+
for (ToolPreProcessor preProcessor : this.preProcessors) {
262+
tool.addPreProcessor(preProcessor);
263+
}
264+
for (ToolPostProcessor postProcessor : this.postProcessors) {
265+
tool.addPostProcessor(postProcessor);
266+
}
233267
return tool;
234268
});
235269
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,22 @@ interface Builder {
135135
*/
136136
Builder headers(Map<String, String> headers);
137137

138+
/**
139+
* Adds a global pre-processor that will be applied to all tools loaded by this client.
140+
*
141+
* @param preProcessor The pre-processor to add.
142+
* @return The builder instance.
143+
*/
144+
Builder preProcessor(ToolPreProcessor preProcessor);
145+
146+
/**
147+
* Adds a global post-processor that will be applied to all tools loaded by this client.
148+
*
149+
* @param postProcessor The post-processor to add.
150+
* @return The builder instance.
151+
*/
152+
Builder postProcessor(ToolPostProcessor postProcessor);
153+
138154
/**
139155
* Builds and returns a new {@link McpToolboxClient} instance.
140156
*

‎src/main/java/com/google/cloud/mcp/McpToolboxClientBuilder.java‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,18 @@
1616

1717
package com.google.cloud.mcp;
1818

19+
import java.util.ArrayList;
1920
import java.util.HashMap;
21+
import java.util.List;
2022
import java.util.Map;
2123

2224
/** Implementation of the {@link McpToolboxClient.Builder} interface. */
2325
public class McpToolboxClientBuilder implements McpToolboxClient.Builder {
2426
private String baseUrl;
2527
private String apiKey;
2628
private Map<String, String> headers = new HashMap<>();
29+
private final List<ToolPreProcessor> preProcessors = new ArrayList<>();
30+
private final List<ToolPostProcessor> postProcessors = new ArrayList<>();
2731

2832
/** Constructs a new McpToolboxClientBuilder. */
2933
public McpToolboxClientBuilder() {}
@@ -48,6 +52,22 @@ public McpToolboxClient.Builder headers(Map<String, String> headers) {
4852
return this;
4953
}
5054

55+
@Override
56+
public McpToolboxClient.Builder preProcessor(ToolPreProcessor preProcessor) {
57+
if (preProcessor != null) {
58+
this.preProcessors.add(preProcessor);
59+
}
60+
return this;
61+
}
62+
63+
@Override
64+
public McpToolboxClient.Builder postProcessor(ToolPostProcessor postProcessor) {
65+
if (postProcessor != null) {
66+
this.postProcessors.add(postProcessor);
67+
}
68+
return this;
69+
}
70+
5171
@Override
5272
public McpToolboxClient build() {
5373
if (baseUrl == null || baseUrl.isEmpty()) {
@@ -66,6 +86,6 @@ public McpToolboxClient build() {
6686
}
6787
}
6888

69-
return new HttpMcpToolboxClient(baseUrl, finalHeaders);
89+
return new HttpMcpToolboxClient(baseUrl, finalHeaders, preProcessors, postProcessors);
7090
}
7191
}

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

Lines changed: 97 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616

1717
package com.google.cloud.mcp;
1818

19+
import java.util.ArrayList;
1920
import java.util.HashMap;
21+
import java.util.List;
2022
import java.util.Map;
2123
import java.util.Objects;
2224
import java.util.concurrent.CompletableFuture;
@@ -33,6 +35,8 @@ public class Tool {
3335

3436
private final Map<String, Object> boundParameters = new HashMap<>();
3537
private final Map<String, AuthTokenGetter> authGetters = new HashMap<>();
38+
private final List<ToolPreProcessor> preProcessors = new ArrayList<>();
39+
private final List<ToolPostProcessor> postProcessors = new ArrayList<>();
3640

3741
/**
3842
* Constructs a new Tool.
@@ -101,6 +105,28 @@ public Tool addAuthTokenGetter(String serviceName, AuthTokenGetter getter) {
101105
return this;
102106
}
103107

108+
/**
109+
* Adds a pre-processor to the tool.
110+
*
111+
* @param processor The pre-processor to add.
112+
* @return The tool instance.
113+
*/
114+
public Tool addPreProcessor(ToolPreProcessor processor) {
115+
this.preProcessors.add(processor);
116+
return this;
117+
}
118+
119+
/**
120+
* Adds a post-processor to the tool.
121+
*
122+
* @param processor The post-processor to add.
123+
* @return The tool instance.
124+
*/
125+
public Tool addPostProcessor(ToolPostProcessor processor) {
126+
this.postProcessors.add(processor);
127+
return this;
128+
}
129+
104130
/**
105131
* Executes the tool with the provided arguments, applying any bound parameters and resolving
106132
* authentication tokens.
@@ -109,55 +135,81 @@ public Tool addAuthTokenGetter(String serviceName, AuthTokenGetter getter) {
109135
* @return A CompletableFuture containing the result of the tool execution.
110136
*/
111137
public CompletableFuture<ToolResult> execute(Map<String, Object> args) {
112-
Map<String, Object> finalArgs = new HashMap<>(args);
113-
Map<String, String> extraHeaders = new HashMap<>();
114-
115-
// 1. Apply Bound Parameters
116-
for (Map.Entry<String, Object> entry : boundParameters.entrySet()) {
117-
Object val = entry.getValue();
118-
if (val instanceof Supplier) {
119-
finalArgs.put(entry.getKey(), ((Supplier<?>) val).get());
120-
} else {
121-
finalArgs.put(entry.getKey(), val);
122-
}
138+
CompletableFuture<Map<String, Object>> argsFuture =
139+
CompletableFuture.completedFuture(new HashMap<>(args));
140+
141+
for (ToolPreProcessor preProcessor : preProcessors) {
142+
argsFuture = argsFuture.thenCompose(currentArgs -> preProcessor.process(name, currentArgs));
123143
}
124144

125-
// 2. Resolve Auth Tokens
126-
return CompletableFuture.allOf(
127-
authGetters.entrySet().stream()
128-
.map(
129-
entry -> {
130-
String serviceName = entry.getKey();
131-
return entry
132-
.getValue()
133-
.getToken()
134-
.thenAccept(
135-
token -> {
136-
// A. Check if mapped to a Parameter (Authenticated Parameters)
137-
String paramName = findParameterForService(serviceName);
138-
if (paramName != null) {
139-
finalArgs.put(paramName, token);
140-
}
141-
142-
// B. Always add to Headers to support Authorized Invocation
143-
// 1. Standard OIDC Header (Cloud Run)
144-
extraHeaders.put("Authorization", "Bearer " + token);
145-
146-
// 2. SDK Convention Header (Framework Compatibility)
147-
extraHeaders.put(serviceName + "_token", token);
148-
});
149-
})
150-
.toArray(CompletableFuture[]::new))
151-
.thenCompose(
152-
v -> {
153-
try {
154-
// 3. Validation & Cleanup
155-
validateAndSanitizeArgs(finalArgs);
156-
return client.invokeTool(name, finalArgs, extraHeaders);
157-
} catch (Exception e) {
158-
return CompletableFuture.failedFuture(e);
145+
CompletableFuture<ToolResult> resultFuture =
146+
argsFuture.thenCompose(
147+
processedArgs -> {
148+
Map<String, Object> finalArgs = new HashMap<>(processedArgs);
149+
Map<String, String> extraHeaders = new HashMap<>();
150+
151+
// 1. Apply Bound Parameters
152+
for (Map.Entry<String, Object> entry : boundParameters.entrySet()) {
153+
Object val = entry.getValue();
154+
if (val instanceof Supplier) {
155+
finalArgs.put(entry.getKey(), ((Supplier<?>) val).get());
156+
} else {
157+
finalArgs.put(entry.getKey(), val);
158+
}
159159
}
160+
161+
// 2. Resolve Auth Tokens
162+
return CompletableFuture.allOf(
163+
authGetters.entrySet().stream()
164+
.map(
165+
entry -> {
166+
String serviceName = entry.getKey();
167+
return entry
168+
.getValue()
169+
.getToken()
170+
.thenAccept(
171+
token -> {
172+
// A. Check if mapped to a Parameter (Authenticated Parameters)
173+
String paramName = findParameterForService(serviceName);
174+
if (paramName != null) {
175+
finalArgs.put(paramName, token);
176+
}
177+
178+
// B. Always add to Headers to support Authorized Invocation
179+
// 1. Standard OIDC Header (Cloud Run)
180+
extraHeaders.put("Authorization", "Bearer " + token);
181+
182+
// 2. SDK Convention Header (Framework Compatibility)
183+
extraHeaders.put(serviceName + "_token", token);
184+
});
185+
})
186+
.toArray(CompletableFuture[]::new))
187+
.thenCompose(
188+
v -> {
189+
try {
190+
// 3. Validation & Cleanup
191+
validateAndSanitizeArgs(finalArgs);
192+
return client.invokeTool(name, finalArgs, extraHeaders);
193+
} catch (Exception e) {
194+
return CompletableFuture.failedFuture(e);
195+
}
196+
});
160197
});
198+
199+
for (ToolPostProcessor postProcessor : postProcessors) {
200+
resultFuture =
201+
resultFuture
202+
.handle(
203+
(res, err) -> {
204+
if (err != null) {
205+
return CompletableFuture.<ToolResult>failedFuture(err);
206+
}
207+
return postProcessor.process(name, res);
208+
})
209+
.thenCompose(f -> f);
210+
}
211+
212+
return resultFuture;
161213
}
162214

163215
private String findParameterForService(String serviceName) {

‎src/main/java/com/google/cloud/mcp/ToolPostProcessor.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.mcp;
18+
19+
import java.util.concurrent.CompletableFuture;
20+
21+
/** A functional interface for post-processing tool results after invocation. */
22+
@FunctionalInterface
23+
public interface ToolPostProcessor {
24+
25+
/**
26+
* Processes the result of a tool after it has been invoked.
27+
*
28+
* @param toolName The name of the tool that was invoked.
29+
* @param result The original tool result.
30+
* @return A CompletableFuture containing the processed tool result.
31+
*/
32+
CompletableFuture<ToolResult> process(String toolName, ToolResult result);
33+
}

‎src/main/java/com/google/cloud/mcp/ToolPreProcessor.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.mcp;
18+
19+
import java.util.Map;
20+
import java.util.concurrent.CompletableFuture;
21+
22+
/** A functional interface for pre-processing tool inputs before invocation. */
23+
@FunctionalInterface
24+
public interface ToolPreProcessor {
25+
26+
/**
27+
* Processes the input arguments for a tool before it is invoked.
28+
*
29+
* @param toolName The name of the tool being invoked.
30+
* @param arguments The original arguments provided to the tool.
31+
* @return A CompletableFuture containing the processed arguments.
32+
*/
33+
CompletableFuture<Map<String, Object>> process(String toolName, Map<String, Object> arguments);
34+
}

0 commit comments

Comments
 (0)