From 6dc5ddc967fce9a341040766d1ba3b08162efc21 Mon Sep 17 00:00:00 2001 From: Stenal P Jolly Date: Wed, 24 Jun 2026 14:04:06 +0530 Subject: [PATCH] feat: prune bound parameters from exposed tool definition --- .../cymbal/web/CymbalTransitController.java | 7 +- .../cloudcode/helloworld/ExampleUsage.java | 4 +- .../mcp/client/McpToolboxClientImpl.java | 20 ++- .../java/com/google/cloud/mcp/tool/Tool.java | 170 ++++++++++++++---- .../cloud/mcp/e2e/McpToolboxClientTest.java | 37 ++-- .../com/google/cloud/mcp/tool/ToolTest.java | 93 ++++++++-- 6 files changed, 249 insertions(+), 82 deletions(-) diff --git a/demo-applications/cymbal-transit/src/main/java/cloudcode/cymbal/web/CymbalTransitController.java b/demo-applications/cymbal-transit/src/main/java/cloudcode/cymbal/web/CymbalTransitController.java index 32a8255..db75065 100644 --- a/demo-applications/cymbal-transit/src/main/java/cloudcode/cymbal/web/CymbalTransitController.java +++ b/demo-applications/cymbal-transit/src/main/java/cloudcode/cymbal/web/CymbalTransitController.java @@ -230,10 +230,9 @@ public CompletableFuture bookTicket(String tripId, String passengerName) return mcpClient .loadTool("book-ticket", Collections.singletonMap("google_auth", toolAuthGetter)) .thenCompose( - tool -> { - tool.bindParam("passenger_name", passengerName); - return tool.execute(Collections.singletonMap("trip_id", tripId)); - }) + tool -> + tool.bindParam("passenger_name", passengerName) + .execute(Collections.singletonMap("trip_id", tripId))) .thenApply( result -> { if (result.isError() || result.content() == null || result.content().isEmpty()) { diff --git a/example/src/main/java/cloudcode/helloworld/ExampleUsage.java b/example/src/main/java/cloudcode/helloworld/ExampleUsage.java index 5570ad2..f22efba 100644 --- a/example/src/main/java/cloudcode/helloworld/ExampleUsage.java +++ b/example/src/main/java/cloudcode/helloworld/ExampleUsage.java @@ -144,14 +144,14 @@ public static void main(String[] args) { // NOW bind the parameter System.out.println("\n[B] Binding 'description' to 'soft toy'..."); - tool.bindParam("description", "soft toy"); System.out.println( " -> Executing BOUND (Runtime arg: 'barbie' - should be" + " IGNORED)..."); // We pass 'barbie', but expecting 'soft toy' price because of binding // override - return tool.execute(Map.of("description", "barbie")); + return tool.bindParam("description", "soft toy") + .execute(Map.of("description", "barbie")); }); }) .thenAccept( diff --git a/src/main/java/com/google/cloud/mcp/client/McpToolboxClientImpl.java b/src/main/java/com/google/cloud/mcp/client/McpToolboxClientImpl.java index d8ebdd2..b3b8af6 100644 --- a/src/main/java/com/google/cloud/mcp/client/McpToolboxClientImpl.java +++ b/src/main/java/com/google/cloud/mcp/client/McpToolboxClientImpl.java @@ -279,16 +279,20 @@ public CompletableFuture> loadToolset( String toolName = entry.getKey(); Tool tool = new Tool(toolName, entry.getValue(), this); if (paramBinds != null && paramBinds.containsKey(toolName)) { - paramBinds.get(toolName).forEach(tool::bindParam); + for (Map.Entry bind : paramBinds.get(toolName).entrySet()) { + tool = tool.bindParam(bind.getKey(), bind.getValue()); + } } if (authBinds != null && authBinds.containsKey(toolName)) { - authBinds.get(toolName).forEach(tool::addAuthTokenGetter); + for (Map.Entry bind : authBinds.get(toolName).entrySet()) { + tool = tool.addAuthTokenGetter(bind.getKey(), bind.getValue()); + } } for (ToolPreProcessor preProcessor : this.preProcessors) { - tool.addPreProcessor(preProcessor); + tool = tool.addPreProcessor(preProcessor); } for (ToolPostProcessor postProcessor : this.postProcessors) { - tool.addPostProcessor(postProcessor); + tool = tool.addPostProcessor(postProcessor); } tools.put(toolName, tool); } @@ -317,13 +321,15 @@ public CompletableFuture loadTool( } Tool tool = new Tool(toolName, tools.get(toolName), this); if (authTokenGetters != null) { - authTokenGetters.forEach(tool::addAuthTokenGetter); + for (Map.Entry entry : authTokenGetters.entrySet()) { + tool = tool.addAuthTokenGetter(entry.getKey(), entry.getValue()); + } } for (ToolPreProcessor preProcessor : this.preProcessors) { - tool.addPreProcessor(preProcessor); + tool = tool.addPreProcessor(preProcessor); } for (ToolPostProcessor postProcessor : this.postProcessors) { - tool.addPostProcessor(postProcessor); + tool = tool.addPostProcessor(postProcessor); } return tool; }); diff --git a/src/main/java/com/google/cloud/mcp/tool/Tool.java b/src/main/java/com/google/cloud/mcp/tool/Tool.java index 4fb0229..f479380 100644 --- a/src/main/java/com/google/cloud/mcp/tool/Tool.java +++ b/src/main/java/com/google/cloud/mcp/tool/Tool.java @@ -20,6 +20,7 @@ import com.google.cloud.mcp.auth.AuthResolver; import com.google.cloud.mcp.auth.AuthTokenGetter; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -36,22 +37,46 @@ public class Tool { private final ToolDefinition definition; private final McpToolboxClient client; - private final Map boundParameters = new HashMap<>(); - private final Map authGetters = new HashMap<>(); - private final List preProcessors = new ArrayList<>(); - private final List postProcessors = new ArrayList<>(); + private final Map boundParameters; + private final Map authGetters; + private final List preProcessors; + private final List postProcessors; /** * Constructs a new Tool. * - * @param name The name of the tool. - * @param definition The definition of the tool. - * @param client The client used to invoke the tool. + * @param toolName The name of the tool. + * @param toolDefinition The definition of the tool. + * @param toolboxClient The client used to invoke the tool. */ - public Tool(String name, ToolDefinition definition, McpToolboxClient client) { - this.name = name; - this.definition = definition; - this.client = client; + public Tool( + final String toolName, + final ToolDefinition toolDefinition, + final McpToolboxClient toolboxClient) { + this.name = toolName; + this.definition = toolDefinition; + this.client = toolboxClient; + this.boundParameters = Collections.emptyMap(); + this.authGetters = Collections.emptyMap(); + this.preProcessors = Collections.emptyList(); + this.postProcessors = Collections.emptyList(); + } + + private Tool( + final String toolName, + final ToolDefinition toolDefinition, + final McpToolboxClient toolboxClient, + final Map initialBoundParameters, + final Map initialAuthGetters, + final List initialPreProcessors, + final List initialPostProcessors) { + this.name = toolName; + this.definition = toolDefinition; + this.client = toolboxClient; + this.boundParameters = Collections.unmodifiableMap(new HashMap<>(initialBoundParameters)); + this.authGetters = Collections.unmodifiableMap(new HashMap<>(initialAuthGetters)); + this.preProcessors = Collections.unmodifiableList(new ArrayList<>(initialPreProcessors)); + this.postProcessors = Collections.unmodifiableList(new ArrayList<>(initialPostProcessors)); } /** @@ -77,11 +102,20 @@ public ToolDefinition definition() { * * @param key The parameter name. * @param value The value to bind. - * @return The tool instance. + * @return A new tool instance with the parameter bound and pruned from definition. */ - public Tool bindParam(String key, Object value) { - this.boundParameters.put(key, value); - return this; + public Tool bindParam(final String key, final Object value) { + Map newBound = new HashMap<>(this.boundParameters); + newBound.put(key, value); + ToolDefinition newDef = pruneParameter(this.definition, key); + return new Tool( + this.name, + newDef, + this.client, + newBound, + this.authGetters, + this.preProcessors, + this.postProcessors); } /** @@ -89,11 +123,20 @@ public Tool bindParam(String key, Object value) { * * @param key The parameter name. * @param valueSupplier The supplier that provides the value at execution time. - * @return The tool instance. + * @return A new tool instance with the parameter bound and pruned from definition. */ - public Tool bindParam(String key, Supplier valueSupplier) { - this.boundParameters.put(key, valueSupplier); - return this; + public Tool bindParam(final String key, final Supplier valueSupplier) { + Map newBound = new HashMap<>(this.boundParameters); + newBound.put(key, valueSupplier); + ToolDefinition newDef = pruneParameter(this.definition, key); + return new Tool( + this.name, + newDef, + this.client, + newBound, + this.authGetters, + this.preProcessors, + this.postProcessors); } /** @@ -101,11 +144,38 @@ public Tool bindParam(String key, Supplier valueSupplier) { * * @param serviceName The name of the service. * @param getter The token getter. - * @return The tool instance. + * @return A new tool instance with the token getter registered. */ - public Tool addAuthTokenGetter(String serviceName, AuthTokenGetter getter) { - this.authGetters.put(serviceName, getter); - return this; + public Tool addAuthTokenGetter(final String serviceName, final AuthTokenGetter getter) { + Map newAuth = new HashMap<>(this.authGetters); + newAuth.put(serviceName, getter); + return new Tool( + this.name, + this.definition, + this.client, + this.boundParameters, + newAuth, + this.preProcessors, + this.postProcessors); + } + + private static ToolDefinition pruneParameter( + final ToolDefinition original, final String paramName) { + if (original.parameters() == null) { + return original; + } + List newParams = new ArrayList<>(); + for (ToolDefinition.Parameter param : original.parameters()) { + if (!param.name().equals(paramName)) { + newParams.add(param); + } + } + return new ToolDefinition( + original.description(), + newParams, + original.authRequired(), + original.readOnlyHint(), + original.destructiveHint()); } /** @@ -114,20 +184,39 @@ public Tool addAuthTokenGetter(String serviceName, AuthTokenGetter getter) { * @param processor The pre-processor to add. * @return The tool instance. */ - public Tool addPreProcessor(ToolPreProcessor processor) { - this.preProcessors.add(processor); - return this; + public Tool addPreProcessor(final ToolPreProcessor processor) { + List newPre = + new ArrayList<>(this.preProcessors != null ? this.preProcessors : Collections.emptyList()); + newPre.add(processor); + return new Tool( + this.name, + this.definition, + this.client, + this.boundParameters, + this.authGetters, + newPre, + this.postProcessors); } /** * Adds a post-processor to the tool. * * @param processor The post-processor to add. - * @return The tool instance. + * @return A new tool instance with the post-processor added. */ - public Tool addPostProcessor(ToolPostProcessor processor) { - this.postProcessors.add(processor); - return this; + public Tool addPostProcessor(final ToolPostProcessor processor) { + List newPost = + new ArrayList<>( + this.postProcessors != null ? this.postProcessors : Collections.emptyList()); + newPost.add(processor); + return new Tool( + this.name, + this.definition, + this.client, + this.boundParameters, + this.authGetters, + this.preProcessors, + newPost); } /** @@ -137,7 +226,7 @@ public Tool addPostProcessor(ToolPostProcessor processor) { * @param args The arguments for the tool invocation. * @return A CompletableFuture containing the result of the tool execution. */ - public CompletableFuture execute(Map args) { + public CompletableFuture execute(final Map args) { CompletableFuture> argsFuture = CompletableFuture.completedFuture(new HashMap<>(args)); @@ -168,7 +257,8 @@ public CompletableFuture execute(Map args) { .thenCompose( resolvedAuth -> { try { - // Apply credential parameter bindings and extra headers + // Apply credential parameter bindings and extra + // headers resolvedAuth.applyTo(finalArgs, extraHeaders, definition); // Validation & Cleanup @@ -187,12 +277,18 @@ public CompletableFuture execute(Map args) { return resultFuture; } - /** Validates arguments against the tool definition and removes null values. */ - private void validateAndSanitizeArgs(Map args) { + /** + * Validates arguments against the tool definition and removes null values. + * + * @param args The arguments to validate. + */ + private void validateAndSanitizeArgs(final Map args) { // Remove nulls first (filtering none values) args.values().removeIf(Objects::isNull); - if (definition.parameters() == null) return; + if (definition.parameters() == null) { + return; + } for (ToolDefinition.Parameter param : definition.parameters()) { Object value = args.get(param.name()); @@ -221,7 +317,7 @@ private void validateAndSanitizeArgs(Map args) { } } - private Object deepCopy(Object value) { + private Object deepCopy(final Object value) { if (value instanceof Map) { Map map = (Map) value; Map copy = new HashMap<>(); @@ -240,7 +336,7 @@ private Object deepCopy(Object value) { return value; } - private boolean isTypeMatch(Object value, String type) { + private boolean isTypeMatch(final Object value, final String type) { switch (type.toLowerCase()) { case "string": return value instanceof String; diff --git a/src/test/java/com/google/cloud/mcp/e2e/McpToolboxClientTest.java b/src/test/java/com/google/cloud/mcp/e2e/McpToolboxClientTest.java index 77f9823..883557b 100644 --- a/src/test/java/com/google/cloud/mcp/e2e/McpToolboxClientTest.java +++ b/src/test/java/com/google/cloud/mcp/e2e/McpToolboxClientTest.java @@ -119,9 +119,12 @@ void testBindParamsCallable() { @Test void testRunToolAuth() { - Tool tool = client.loadTool("get-row-by-id-auth").join(); - tool.addAuthTokenGetter( - "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken1())); + Tool tool = + client + .loadTool("get-row-by-id-auth") + .join() + .addAuthTokenGetter( + "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken1())); ToolResult result = tool.execute(Map.of("id", "2")).join(); assertFalse(result.isError()); @@ -131,10 +134,12 @@ void testRunToolAuth() { @Test void testRunToolWrongAuth() { - Tool tool = client.loadTool("get-row-by-id-auth").join(); - - tool.addAuthTokenGetter( - "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken2())); + Tool tool = + client + .loadTool("get-row-by-id-auth") + .join() + .addAuthTokenGetter( + "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken2())); ToolResult result = tool.execute(Map.of("id", "2")).join(); assertTrue( @@ -147,9 +152,12 @@ void testRunToolWrongAuth() { @Test void testRunToolParamAuth() { - Tool tool = client.loadTool("get-row-by-email-auth").join(); - tool.addAuthTokenGetter( - "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken1())); + Tool tool = + client + .loadTool("get-row-by-email-auth") + .join() + .addAuthTokenGetter( + "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken1())); ToolResult result = tool.execute(Map.of()).join(); assertFalse(result.isError(), "Expected success but got error: " + getTextContent(result)); @@ -161,9 +169,12 @@ void testRunToolParamAuth() { @Test void testRunToolParamAuthNoField() { - Tool tool = client.loadTool("get-row-by-content-auth").join(); - tool.addAuthTokenGetter( - "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken1())); + Tool tool = + client + .loadTool("get-row-by-content-auth") + .join() + .addAuthTokenGetter( + "my-test-auth", () -> CompletableFuture.completedFuture(server.getAuthToken1())); ToolResult result = tool.execute(Map.of()).join(); assertTrue(result.isError()); diff --git a/src/test/java/com/google/cloud/mcp/tool/ToolTest.java b/src/test/java/com/google/cloud/mcp/tool/ToolTest.java index 9ba2c0d..4543e6e 100644 --- a/src/test/java/com/google/cloud/mcp/tool/ToolTest.java +++ b/src/test/java/com/google/cloud/mcp/tool/ToolTest.java @@ -101,18 +101,19 @@ void execute_withManyConcurrentAuthGetters_doesNotDropCredentials() { Tool raceTool = new Tool("race-tool", def, client); for (int i = 0; i < services; i++) { final String token = "tok-" + i; - raceTool.addAuthTokenGetter( - "svc" + i, - () -> - CompletableFuture.supplyAsync( - () -> { - int spins = ThreadLocalRandom.current().nextInt(50); - for (int s = 0; s < spins; s++) { - Thread.onSpinWait(); - } - return token; - }, - pool)); + raceTool = + raceTool.addAuthTokenGetter( + "svc" + i, + () -> + CompletableFuture.supplyAsync( + () -> { + int spins = ThreadLocalRandom.current().nextInt(50); + for (int s = 0; s < spins; s++) { + Thread.onSpinWait(); + } + return token; + }, + pool)); } for (int iter = 0; iter < iterations; iter++) { @@ -210,9 +211,10 @@ void testBindParamStaticAndSupplier() throws Exception { return CompletableFuture.completedFuture(new ToolResult(List.of(), false)); }); - Tool tool = new Tool("test-tool", def, client); - tool.bindParam("p-static", "static-value"); - tool.bindParam("p-supplier", () -> "supplier-value"); + Tool tool = + new Tool("test-tool", def, client) + .bindParam("p-static", "static-value") + .bindParam("p-supplier", () -> "supplier-value"); tool.execute(Map.of()).join(); @@ -315,9 +317,10 @@ void testExecute_withPreAndPostProcessors_modifiesArgsAndResult() throws Excepti return CompletableFuture.completedFuture(result); }; - tool.addPreProcessor(preProcessor1); - tool.addPreProcessor(preProcessor2); - tool.addPostProcessor(postProcessor); + tool = + tool.addPreProcessor(preProcessor1) + .addPreProcessor(preProcessor2) + .addPostProcessor(postProcessor); when(mockClient.invokeTool(eq("test_tool"), anyMap(), anyMap())) .thenReturn(CompletableFuture.completedFuture(originalResult)); @@ -347,7 +350,7 @@ void testExecute_preProcessorException_failsFutureWithoutInvokingClient() { ToolPreProcessor preProcessor = (name, args) -> CompletableFuture.failedFuture(new RuntimeException("PreProcessor failed")); - tool.addPreProcessor(preProcessor); + tool = tool.addPreProcessor(preProcessor); // Act CompletableFuture futureResult = tool.execute(initialArgs); @@ -367,4 +370,56 @@ void testExecute_preProcessorException_failsFutureWithoutInvokingClient() { verify(mockClient, never()).invokeTool(eq("test_tool"), anyMap(), anyMap()); verify(mockClient, never()).invokeTool(eq("test_tool"), anyMap()); } + + @Test + void testBoundParameterPruning() { + List params = + List.of( + new ToolDefinition.Parameter("p1", "string", true, "param 1", List.of()), + new ToolDefinition.Parameter("p2", "integer", false, "param 2", List.of())); + ToolDefinition def = new ToolDefinition("test-tool", params, List.of()); + McpToolboxClient client = mock(McpToolboxClient.class); + + List> capturedArgs = new ArrayList<>(); + when(client.invokeTool(anyString(), anyMap(), anyMap())) + .thenAnswer( + inv -> { + capturedArgs.add(new HashMap<>(inv.getArgument(1))); + return CompletableFuture.completedFuture(new ToolResult(List.of(), false)); + }); + + Tool tool = new Tool("test-tool", def, client); + + // Initial check: both parameters present + assertEquals(2, tool.definition().parameters().size()); + + // Bind p1 statically + Tool boundTool = tool.bindParam("p1", "bound-value"); + + // Check that p1 is pruned from boundTool definition but remains in original tool definition + assertEquals(2, tool.definition().parameters().size()); + assertEquals(1, boundTool.definition().parameters().size()); + assertEquals("p2", boundTool.definition().parameters().get(0).name()); + + // Bind p2 with supplier + Tool boundTool2 = boundTool.bindParam("p2", () -> 42); + + // Check that p2 is pruned from boundTool2 definition + assertEquals(0, boundTool2.definition().parameters().size()); + + // Execute boundTool2. Since both parameters are bound, we don't need to supply them at runtime, + // and they should be passed to the client. + boundTool2.execute(Map.of()).join(); + + assertEquals(1, capturedArgs.size()); + Map args = capturedArgs.get(0); + assertEquals("bound-value", args.get("p1")); + assertEquals(42, args.get("p2")); + + // Test pruning with null parameters + ToolDefinition nullParamsDef = new ToolDefinition("test-tool", null, List.of()); + Tool nullParamsTool = new Tool("test-tool", nullParamsDef, client); + Tool boundNullParamsTool = nullParamsTool.bindParam("p1", "val"); + assertEquals(null, boundNullParamsTool.definition().parameters()); + } }