options) {
+ return mcpCheckOutputAsync(connectorType, responseData, options);
+ }
+
// ========================================================================
// Policy CRUD - Static Policies
// ========================================================================
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/GovernedTool.java b/src/main/java/com/getaxonflow/sdk/adapters/GovernedTool.java
new file mode 100644
index 0000000..1180ba6
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/GovernedTool.java
@@ -0,0 +1,277 @@
+/*
+ * Copyright 2026 AxonFlow
+ *
+ * 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.getaxonflow.sdk.adapters;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.getaxonflow.sdk.AxonFlow;
+import com.getaxonflow.sdk.exceptions.PolicyViolationException;
+import com.getaxonflow.sdk.types.MCPCheckInputResponse;
+import com.getaxonflow.sdk.types.MCPCheckOutputResponse;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.function.Function;
+
+/**
+ * Wraps a {@link Tool} with AxonFlow input/output policy enforcement.
+ *
+ * Every {@link #invoke} call runs through two policy checks:
+ *
+ *
+ * - Input check ({@code mcpCheckInput}): evaluates tool arguments before
+ * execution. Blocked calls throw {@link PolicyViolationException} and the tool never runs.
+ *
- Output check ({@code mcpCheckOutput}): evaluates tool results after
+ * execution. Can block (throw), redact (return cleaned data), or allow.
+ *
+ *
+ * Example usage:
+ *
+ *
{@code
+ * // Simple wrapping
+ * GovernedTool governed = GovernedTool.wrap(myTool, axonflow);
+ * Object result = governed.invoke("SELECT * FROM users");
+ *
+ * // Builder pattern with custom connector type
+ * GovernedTool governed = GovernedTool.builder(myTool, axonflow)
+ * .connectorTypeFn(name -> "custom." + name)
+ * .operation("query")
+ * .build();
+ *
+ * // Batch wrapping
+ * List governed = GovernedTool.governTools(tools, axonflow);
+ * }
+ */
+public final class GovernedTool implements Tool {
+
+ private final Tool wrapped;
+ private final AxonFlow client;
+ private final String connectorType;
+ private final String operation;
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+
+ private GovernedTool(Tool tool, AxonFlow client, String connectorType, String operation) {
+ this.wrapped = Objects.requireNonNull(tool, "tool cannot be null");
+ this.client = Objects.requireNonNull(client, "client cannot be null");
+ this.connectorType = Objects.requireNonNull(connectorType, "connectorType cannot be null");
+ this.operation = Objects.requireNonNull(operation, "operation cannot be null");
+ }
+
+ /**
+ * Wraps a tool with AxonFlow governance using default settings.
+ *
+ * Uses the tool's {@link Tool#name()} as the connector type and {@code "execute"} as the
+ * operation.
+ *
+ * @param tool the tool to wrap
+ * @param client the AxonFlow client
+ * @return a governed tool
+ */
+ public static GovernedTool wrap(Tool tool, AxonFlow client) {
+ Objects.requireNonNull(tool, "tool cannot be null");
+ return new GovernedTool(tool, client, tool.name(), "execute");
+ }
+
+ /**
+ * Creates a builder for a GovernedTool with custom options.
+ *
+ * @param tool the tool to wrap
+ * @param client the AxonFlow client
+ * @return a new builder
+ */
+ public static Builder builder(Tool tool, AxonFlow client) {
+ return new Builder(tool, client);
+ }
+
+ /**
+ * Wraps a list of tools with AxonFlow governance using default settings.
+ *
+ * @param tools the tools to wrap
+ * @param client the AxonFlow client
+ * @return list of governed tools
+ */
+ public static List governTools(List tools, AxonFlow client) {
+ return governTools(tools, client, null, "execute");
+ }
+
+ /**
+ * Wraps a list of tools with AxonFlow governance using custom settings.
+ *
+ * @param tools the tools to wrap
+ * @param client the AxonFlow client
+ * @param connectorTypeFn optional function mapping tool name to connector type (null for default)
+ * @param operation the operation type ({@code "execute"} or {@code "query"})
+ * @return list of governed tools
+ */
+ public static List governTools(
+ List tools,
+ AxonFlow client,
+ Function connectorTypeFn,
+ String operation) {
+ Objects.requireNonNull(tools, "tools cannot be null");
+ Objects.requireNonNull(client, "client cannot be null");
+
+ List result = new ArrayList<>(tools.size());
+ for (Tool tool : tools) {
+ String ct =
+ connectorTypeFn != null ? connectorTypeFn.apply(tool.name()) : tool.name();
+ result.add(new GovernedTool(tool, client, ct, operation != null ? operation : "execute"));
+ }
+ return result;
+ }
+
+ @Override
+ public String name() {
+ return wrapped.name();
+ }
+
+ @Override
+ public String description() {
+ return wrapped.description();
+ }
+
+ /**
+ * Invokes the wrapped tool with AxonFlow input/output governance.
+ *
+ *
+ * - Serializes the input (strings pass through, objects are JSON-serialized)
+ *
- Checks input against policies via {@code mcpCheckInput}
+ *
- If blocked, throws {@link PolicyViolationException} without running the tool
+ *
- Invokes the wrapped tool
+ *
- Serializes the output and checks it against policies via {@code mcpCheckOutput}
+ *
- If blocked, throws {@link PolicyViolationException}
+ *
- If redacted, returns the redacted data
+ *
- Otherwise, returns the original result
+ *
+ *
+ * @param input the tool input
+ * @return the tool result (possibly redacted)
+ * @throws PolicyViolationException if the input or output violates a policy
+ * @throws Exception if the tool invocation fails
+ */
+ @Override
+ public Object invoke(Object input) throws Exception {
+ // 1. Serialize input
+ String serialized = serializeContent(input);
+
+ // 2. Check input policy
+ MCPCheckInputResponse inputCheck = client.mcpCheckInput(connectorType, serialized);
+
+ // 3. If blocked, throw
+ if (!inputCheck.isAllowed()) {
+ throw new PolicyViolationException(
+ inputCheck.getBlockReason() != null
+ ? inputCheck.getBlockReason()
+ : "Tool call blocked by input policy");
+ }
+
+ // 4. Invoke wrapped tool
+ Object result = wrapped.invoke(input);
+
+ // 5. Serialize output and check
+ String serializedResult = serializeContent(result);
+ Map options = new HashMap<>();
+ options.put("message", serializedResult);
+ MCPCheckOutputResponse outputCheck = client.mcpCheckOutput(connectorType, null, options);
+
+ // 6. If blocked, throw
+ if (!outputCheck.isAllowed()) {
+ throw new PolicyViolationException(
+ outputCheck.getBlockReason() != null
+ ? outputCheck.getBlockReason()
+ : "Tool output blocked by policy");
+ }
+
+ // 7. If redacted, return redacted data
+ if (outputCheck.getRedactedData() != null) {
+ return outputCheck.getRedactedData();
+ }
+
+ // 8. Return original
+ return result;
+ }
+
+ @Override
+ public String toString() {
+ return String.format("GovernedTool(name=%s, connectorType=%s)", wrapped.name(), connectorType);
+ }
+
+ /**
+ * Serializes content to a string for policy evaluation. Strings pass through directly; other
+ * objects are JSON-serialized.
+ */
+ private static String serializeContent(Object content) {
+ if (content instanceof String) {
+ return (String) content;
+ }
+ try {
+ return MAPPER.writeValueAsString(content);
+ } catch (Exception e) {
+ return String.valueOf(content);
+ }
+ }
+
+ /** Builder for {@link GovernedTool}. */
+ public static final class Builder {
+
+ private final Tool tool;
+ private final AxonFlow client;
+ private Function connectorTypeFn;
+ private String operation = "execute";
+
+ Builder(Tool tool, AxonFlow client) {
+ this.tool = Objects.requireNonNull(tool, "tool cannot be null");
+ this.client = Objects.requireNonNull(client, "client cannot be null");
+ }
+
+ /**
+ * Sets a function to derive the connector type from the tool name.
+ *
+ * If not set, the tool's {@link Tool#name()} is used directly.
+ *
+ * @param fn the connector type function
+ * @return this builder
+ */
+ public Builder connectorTypeFn(Function fn) {
+ this.connectorTypeFn = fn;
+ return this;
+ }
+
+ /**
+ * Sets the operation type for input checks.
+ *
+ * Defaults to {@code "execute"}. Use {@code "query"} for read-only tools.
+ *
+ * @param operation the operation type
+ * @return this builder
+ */
+ public Builder operation(String operation) {
+ this.operation = Objects.requireNonNull(operation, "operation cannot be null");
+ return this;
+ }
+
+ /**
+ * Builds the governed tool.
+ *
+ * @return a new GovernedTool
+ */
+ public GovernedTool build() {
+ String ct = connectorTypeFn != null ? connectorTypeFn.apply(tool.name()) : tool.name();
+ return new GovernedTool(tool, client, ct, operation);
+ }
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/Tool.java b/src/main/java/com/getaxonflow/sdk/adapters/Tool.java
new file mode 100644
index 0000000..e4979f9
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/Tool.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright 2026 AxonFlow
+ *
+ * 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.getaxonflow.sdk.adapters;
+
+/**
+ * Framework-agnostic interface for any callable tool.
+ *
+ *
Implement this interface to wrap your tools with AxonFlow governance via {@link GovernedTool}.
+ *
+ *
Example usage:
+ *
+ *
{@code
+ * Tool searchTool = new Tool() {
+ * public String name() { return "web_search"; }
+ * public String description() { return "Search the web"; }
+ * public Object invoke(Object input) { return webSearch(input.toString()); }
+ * };
+ *
+ * GovernedTool governed = GovernedTool.wrap(searchTool, axonflow);
+ * Object result = governed.invoke("latest AI research");
+ * }
+ */
+public interface Tool {
+
+ /** Returns the tool name, used as the default connector type for policy checks. */
+ String name();
+
+ /** Returns a human-readable description of what this tool does. */
+ String description();
+
+ /**
+ * Invokes the tool with the given input.
+ *
+ * @param input the tool input (may be a String, Map, or any serializable object)
+ * @return the tool result
+ * @throws Exception if the tool invocation fails
+ */
+ Object invoke(Object input) throws Exception;
+}
diff --git a/src/main/java/com/getaxonflow/sdk/telemetry/TelemetryReporter.java b/src/main/java/com/getaxonflow/sdk/telemetry/TelemetryReporter.java
index 80b459f..bc0fd6b 100644
--- a/src/main/java/com/getaxonflow/sdk/telemetry/TelemetryReporter.java
+++ b/src/main/java/com/getaxonflow/sdk/telemetry/TelemetryReporter.java
@@ -98,14 +98,6 @@ static void sendPing(
return;
}
- // Suppress telemetry for localhost endpoints unless explicitly enabled.
- if (!Boolean.TRUE.equals(telemetryEnabled) && isLocalhostEndpoint(sdkEndpoint)) {
- if (debug) {
- logger.debug("Telemetry suppressed for localhost endpoint");
- }
- return;
- }
-
logger.info(
"AxonFlow: anonymous telemetry enabled. Opt out: AXONFLOW_TELEMETRY=off | https://docs.getaxonflow.com/docs/telemetry");
@@ -272,15 +264,6 @@ static String normalizeArch(String arch) {
return arch;
}
- /** Check whether the endpoint is a localhost address. */
- static boolean isLocalhostEndpoint(String endpoint) {
- if (endpoint == null || endpoint.isEmpty()) {
- return false;
- }
- String lower = endpoint.toLowerCase();
- return lower.contains("localhost") || lower.contains("127.0.0.1") || lower.contains("[::1]");
- }
-
private TelemetryReporter() {
// Utility class
}
diff --git a/src/main/java/com/getaxonflow/sdk/types/SDKCompatibility.java b/src/main/java/com/getaxonflow/sdk/types/SDKCompatibility.java
index 22cd27a..c844b05 100644
--- a/src/main/java/com/getaxonflow/sdk/types/SDKCompatibility.java
+++ b/src/main/java/com/getaxonflow/sdk/types/SDKCompatibility.java
@@ -17,6 +17,7 @@
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.Map;
import java.util.Objects;
/** SDK compatibility information returned by the AxonFlow platform health endpoint. */
@@ -24,26 +25,38 @@
public final class SDKCompatibility {
@JsonProperty("min_sdk_version")
- private final String minSdkVersion;
+ private final Map minSdkVersion;
@JsonProperty("recommended_sdk_version")
- private final String recommendedSdkVersion;
+ private final Map recommendedSdkVersion;
public SDKCompatibility(
- @JsonProperty("min_sdk_version") String minSdkVersion,
- @JsonProperty("recommended_sdk_version") String recommendedSdkVersion) {
+ @JsonProperty("min_sdk_version") Map minSdkVersion,
+ @JsonProperty("recommended_sdk_version") Map recommendedSdkVersion) {
this.minSdkVersion = minSdkVersion;
this.recommendedSdkVersion = recommendedSdkVersion;
}
- public String getMinSdkVersion() {
+ /** Returns the per-language minimum SDK version map (e.g. {"java":"5.0.0","python":"6.0.0"}). */
+ public Map getMinSdkVersion() {
return minSdkVersion;
}
- public String getRecommendedSdkVersion() {
+ /** Returns the minimum SDK version for a specific language, or null if not specified. */
+ public String getMinSdkVersionFor(String language) {
+ return minSdkVersion != null ? minSdkVersion.get(language) : null;
+ }
+
+ /** Returns the per-language recommended SDK version map. */
+ public Map getRecommendedSdkVersion() {
return recommendedSdkVersion;
}
+ /** Returns the recommended SDK version for a specific language, or null if not specified. */
+ public String getRecommendedSdkVersionFor(String language) {
+ return recommendedSdkVersion != null ? recommendedSdkVersion.get(language) : null;
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) return true;
@@ -60,10 +73,10 @@ public int hashCode() {
@Override
public String toString() {
- return "SDKCompatibility{minSdkVersion='"
+ return "SDKCompatibility{minSdkVersion="
+ minSdkVersion
- + "', recommendedSdkVersion='"
+ + ", recommendedSdkVersion="
+ recommendedSdkVersion
- + "'}";
+ + "}";
}
}
diff --git a/src/test/java/com/getaxonflow/sdk/adapters/GovernedToolTest.java b/src/test/java/com/getaxonflow/sdk/adapters/GovernedToolTest.java
new file mode 100644
index 0000000..05264f0
--- /dev/null
+++ b/src/test/java/com/getaxonflow/sdk/adapters/GovernedToolTest.java
@@ -0,0 +1,412 @@
+/*
+ * Copyright 2026 AxonFlow
+ *
+ * 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.getaxonflow.sdk.adapters;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+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 com.getaxonflow.sdk.AxonFlow;
+import com.getaxonflow.sdk.exceptions.PolicyViolationException;
+import com.getaxonflow.sdk.types.MCPCheckInputResponse;
+import com.getaxonflow.sdk.types.MCPCheckOutputResponse;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+@ExtendWith(MockitoExtension.class)
+class GovernedToolTest {
+
+ @Mock private AxonFlow client;
+
+ private MCPCheckInputResponse allowedInput;
+ private MCPCheckInputResponse blockedInput;
+ private MCPCheckOutputResponse allowedOutput;
+ private MCPCheckOutputResponse blockedOutput;
+ private MCPCheckOutputResponse redactedOutput;
+
+ @BeforeEach
+ void setUp() {
+ allowedInput = new MCPCheckInputResponse(true, null, 3, null);
+ blockedInput =
+ new MCPCheckInputResponse(false, "Dangerous SQL detected", 3, null);
+ allowedOutput = new MCPCheckOutputResponse(true, null, null, 2, null, null);
+ blockedOutput =
+ new MCPCheckOutputResponse(false, "PII detected in output", null, 2, null, null);
+ redactedOutput =
+ new MCPCheckOutputResponse(true, null, "[REDACTED:ssn] data", 2, null, null);
+ }
+
+ /** Creates a simple mock tool that returns the given result. */
+ private Tool mockTool(String name, Object result) {
+ return new Tool() {
+ boolean invoked = false;
+
+ @Override
+ public String name() {
+ return name;
+ }
+
+ @Override
+ public String description() {
+ return "Mock " + name + " tool";
+ }
+
+ @Override
+ public Object invoke(Object input) {
+ invoked = true;
+ return result;
+ }
+ };
+ }
+
+ /** Creates a mock tool that tracks whether it was invoked. */
+ private TrackableTool trackableTool(String name, Object result) {
+ return new TrackableTool(name, result);
+ }
+
+ @Test
+ void cleanCallAllowed() throws Exception {
+ when(client.mcpCheckInput(eq("web_search"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("web_search"), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool = mockTool("web_search", "search results");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ Object result = governed.invoke("latest AI research");
+
+ assertThat(result).isEqualTo("search results");
+ verify(client).mcpCheckInput(eq("web_search"), eq("latest AI research"));
+ verify(client).mcpCheckOutput(eq("web_search"), isNull(), any());
+ }
+
+ @Test
+ void inputBlocked() {
+ when(client.mcpCheckInput(eq("db_query"), any())).thenReturn(blockedInput);
+
+ TrackableTool tool = trackableTool("db_query", "should not see this");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ assertThatThrownBy(() -> governed.invoke("DROP TABLE users"))
+ .isInstanceOf(PolicyViolationException.class)
+ .hasMessageContaining("Dangerous SQL detected");
+
+ assertThat(tool.wasInvoked()).isFalse();
+ verify(client, never()).mcpCheckOutput(any(), any(), any());
+ }
+
+ @Test
+ void outputBlocked() {
+ when(client.mcpCheckInput(eq("db_query"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("db_query"), isNull(), any())).thenReturn(blockedOutput);
+
+ TrackableTool tool = trackableTool("db_query", "John Doe, SSN: 123-45-6789");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ assertThatThrownBy(() -> governed.invoke("SELECT * FROM users"))
+ .isInstanceOf(PolicyViolationException.class)
+ .hasMessageContaining("PII detected in output");
+
+ assertThat(tool.wasInvoked()).isTrue();
+ }
+
+ @Test
+ void outputRedacted() throws Exception {
+ when(client.mcpCheckInput(eq("db_query"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("db_query"), isNull(), any())).thenReturn(redactedOutput);
+
+ Tool tool = mockTool("db_query", "John Doe, SSN: 123-45-6789");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ Object result = governed.invoke("SELECT * FROM users");
+
+ assertThat(result).isEqualTo("[REDACTED:ssn] data");
+ }
+
+ @Test
+ void customConnectorTypeFn() throws Exception {
+ when(client.mcpCheckInput(eq("custom.web_search"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("custom.web_search"), isNull(), any()))
+ .thenReturn(allowedOutput);
+
+ Tool tool = mockTool("web_search", "results");
+ GovernedTool governed =
+ GovernedTool.builder(tool, client)
+ .connectorTypeFn(name -> "custom." + name)
+ .build();
+
+ Object result = governed.invoke("query");
+
+ assertThat(result).isEqualTo("results");
+ verify(client).mcpCheckInput(eq("custom.web_search"), any());
+ verify(client).mcpCheckOutput(eq("custom.web_search"), isNull(), any());
+ }
+
+ @Test
+ void customOperation() throws Exception {
+ when(client.mcpCheckInput(eq("db_query"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("db_query"), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool = mockTool("db_query", "rows");
+ GovernedTool governed =
+ GovernedTool.builder(tool, client).operation("query").build();
+
+ Object result = governed.invoke("SELECT 1");
+
+ assertThat(result).isEqualTo("rows");
+ }
+
+ @Test
+ void governToolsBatch() throws Exception {
+ when(client.mcpCheckInput(any(), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(any(), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool1 = mockTool("search", "result1");
+ Tool tool2 = mockTool("calculator", "result2");
+ Tool tool3 = mockTool("email", "result3");
+
+ List governed =
+ GovernedTool.governTools(Arrays.asList(tool1, tool2, tool3), client);
+
+ assertThat(governed).hasSize(3);
+ assertThat(governed.get(0).name()).isEqualTo("search");
+ assertThat(governed.get(1).name()).isEqualTo("calculator");
+ assertThat(governed.get(2).name()).isEqualTo("email");
+
+ // Invoke each to verify they work
+ assertThat(governed.get(0).invoke("q")).isEqualTo("result1");
+ assertThat(governed.get(1).invoke("1+1")).isEqualTo("result2");
+ assertThat(governed.get(2).invoke("send")).isEqualTo("result3");
+ }
+
+ @Test
+ void governToolsBatchWithCustomOptions() throws Exception {
+ when(client.mcpCheckInput(any(), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(any(), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool1 = mockTool("search", "result1");
+ Tool tool2 = mockTool("calculator", "result2");
+
+ List governed =
+ GovernedTool.governTools(
+ Arrays.asList(tool1, tool2), client, name -> "ns." + name, "query");
+
+ assertThat(governed).hasSize(2);
+
+ governed.get(0).invoke("q");
+ verify(client).mcpCheckInput(eq("ns.search"), any());
+
+ governed.get(1).invoke("1+1");
+ verify(client).mcpCheckInput(eq("ns.calculator"), any());
+ }
+
+ @Test
+ void stringInputPassthrough() throws Exception {
+ when(client.mcpCheckInput(eq("tool"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("tool"), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool = mockTool("tool", "ok");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ governed.invoke("plain string input");
+
+ // String should be passed through directly, not JSON-encoded
+ verify(client).mcpCheckInput(eq("tool"), eq("plain string input"));
+ }
+
+ @Test
+ void objectInputSerialized() throws Exception {
+ when(client.mcpCheckInput(eq("tool"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("tool"), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool = mockTool("tool", "ok");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ Map input = Map.of("query", "test", "limit", 10);
+ governed.invoke(input);
+
+ // Object should be JSON-serialized
+ ArgumentCaptor captor = ArgumentCaptor.forClass(String.class);
+ verify(client).mcpCheckInput(eq("tool"), captor.capture());
+ String serialized = captor.getValue();
+ assertThat(serialized).contains("\"query\"");
+ assertThat(serialized).contains("\"test\"");
+ assertThat(serialized).contains("\"limit\"");
+ assertThat(serialized).contains("10");
+ }
+
+ @Test
+ void builderPattern() throws Exception {
+ when(client.mcpCheckInput(eq("custom.myTool"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("custom.myTool"), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool = mockTool("myTool", "built result");
+ GovernedTool governed =
+ GovernedTool.builder(tool, client)
+ .connectorTypeFn(name -> "custom." + name)
+ .operation("query")
+ .build();
+
+ Object result = governed.invoke("input");
+
+ assertThat(result).isEqualTo("built result");
+ assertThat(governed.name()).isEqualTo("myTool");
+ assertThat(governed.description()).isEqualTo("Mock myTool tool");
+ }
+
+ @Test
+ void wrapFactory() throws Exception {
+ when(client.mcpCheckInput(eq("simpleTool"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("simpleTool"), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool = mockTool("simpleTool", "wrapped result");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ assertThat(governed.name()).isEqualTo("simpleTool");
+ assertThat(governed.description()).isEqualTo("Mock simpleTool tool");
+
+ Object result = governed.invoke("test");
+ assertThat(result).isEqualTo("wrapped result");
+ }
+
+ @Test
+ void toStringFormat() {
+ Tool tool = mockTool("myTool", null);
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ assertThat(governed.toString())
+ .isEqualTo("GovernedTool(name=myTool, connectorType=myTool)");
+ }
+
+ @Test
+ void toStringFormatWithCustomConnector() {
+ Tool tool = mockTool("myTool", null);
+ GovernedTool governed =
+ GovernedTool.builder(tool, client)
+ .connectorTypeFn(name -> "ns." + name)
+ .build();
+
+ assertThat(governed.toString())
+ .isEqualTo("GovernedTool(name=myTool, connectorType=ns.myTool)");
+ }
+
+ @Test
+ void outputCheckReceivesMessageInOptions() throws Exception {
+ when(client.mcpCheckInput(eq("tool"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("tool"), isNull(), any())).thenReturn(allowedOutput);
+
+ Tool tool = mockTool("tool", "tool output data");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ governed.invoke("input");
+
+ @SuppressWarnings("unchecked")
+ ArgumentCaptor> optionsCaptor = ArgumentCaptor.forClass(Map.class);
+ verify(client).mcpCheckOutput(eq("tool"), isNull(), optionsCaptor.capture());
+
+ Map capturedOptions = optionsCaptor.getValue();
+ assertThat(capturedOptions).containsKey("message");
+ assertThat(capturedOptions.get("message")).isEqualTo("tool output data");
+ }
+
+ @Test
+ void inputBlockedWithNullReason() {
+ MCPCheckInputResponse blockedNoReason =
+ new MCPCheckInputResponse(false, null, 1, null);
+ when(client.mcpCheckInput(eq("tool"), any())).thenReturn(blockedNoReason);
+
+ Tool tool = mockTool("tool", "should not run");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ assertThatThrownBy(() -> governed.invoke("input"))
+ .isInstanceOf(PolicyViolationException.class)
+ .hasMessageContaining("Tool call blocked by input policy");
+ }
+
+ @Test
+ void outputBlockedWithNullReason() {
+ MCPCheckOutputResponse blockedNoReason =
+ new MCPCheckOutputResponse(false, null, null, 1, null, null);
+ when(client.mcpCheckInput(eq("tool"), any())).thenReturn(allowedInput);
+ when(client.mcpCheckOutput(eq("tool"), isNull(), any())).thenReturn(blockedNoReason);
+
+ Tool tool = mockTool("tool", "result");
+ GovernedTool governed = GovernedTool.wrap(tool, client);
+
+ assertThatThrownBy(() -> governed.invoke("input"))
+ .isInstanceOf(PolicyViolationException.class)
+ .hasMessageContaining("Tool output blocked by policy");
+ }
+
+ @Test
+ void nullToolThrows() {
+ assertThatThrownBy(() -> GovernedTool.wrap(null, client))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining("tool cannot be null");
+ }
+
+ @Test
+ void nullClientThrows() {
+ Tool tool = mockTool("tool", null);
+ assertThatThrownBy(() -> GovernedTool.wrap(tool, null))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining("client cannot be null");
+ }
+
+ /** Helper class that tracks invocation state. */
+ private static class TrackableTool implements Tool {
+ private final String toolName;
+ private final Object result;
+ private boolean invoked;
+
+ TrackableTool(String name, Object result) {
+ this.toolName = name;
+ this.result = result;
+ }
+
+ @Override
+ public String name() {
+ return toolName;
+ }
+
+ @Override
+ public String description() {
+ return "Trackable " + toolName + " tool";
+ }
+
+ @Override
+ public Object invoke(Object input) {
+ invoked = true;
+ return result;
+ }
+
+ boolean wasInvoked() {
+ return invoked;
+ }
+ }
+}
diff --git a/src/test/java/com/getaxonflow/sdk/telemetry/TelemetryReporterTest.java b/src/test/java/com/getaxonflow/sdk/telemetry/TelemetryReporterTest.java
index 75177ea..a4f294e 100644
--- a/src/test/java/com/getaxonflow/sdk/telemetry/TelemetryReporterTest.java
+++ b/src/test/java/com/getaxonflow/sdk/telemetry/TelemetryReporterTest.java
@@ -145,7 +145,6 @@ void testCustomEndpoint(WireMockRuntimeInfo wmRuntimeInfo) throws Exception {
String customUrl = wmRuntimeInfo.getHttpBaseUrl() + "/v1/ping";
// Call sendPing with custom checkpoint URL, no env opt-outs, with credentials
- // telemetryEnabled=true overrides localhost guard (WireMock runs on localhost)
TelemetryReporter.sendPing(
"production",
"http://localhost:8080",
@@ -272,7 +271,7 @@ void testProductionModeWithoutCredentials(WireMockRuntimeInfo wmRuntimeInfo) thr
String customUrl = wmRuntimeInfo.getHttpBaseUrl() + "/v1/ping";
- // telemetryEnabled=true overrides localhost guard (WireMock runs on localhost)
+ // telemetryEnabled=true: explicit enable for this test
TelemetryReporter.sendPing(
"production",
"http://localhost:8080",
@@ -371,7 +370,7 @@ void testNon200ResponseNoCrash(WireMockRuntimeInfo wmRuntimeInfo) {
String customUrl = wmRuntimeInfo.getHttpBaseUrl() + "/v1/ping";
- // telemetryEnabled=true overrides localhost guard (WireMock runs on localhost)
+ // telemetryEnabled=true: explicit enable for this test
assertThatCode(
() -> {
TelemetryReporter.sendPing(
@@ -422,7 +421,7 @@ void testPayloadDeploymentModeEnterprise(WireMockRuntimeInfo wmRuntimeInfo) thro
String customUrl = wmRuntimeInfo.getHttpBaseUrl() + "/v1/ping";
- // telemetryEnabled=true overrides localhost guard (WireMock runs on localhost)
+ // telemetryEnabled=true: explicit enable for this test
// Use localhost:1 so detectPlatformVersion gets immediate connection-refused
// (localhost:8080 may have a running service that returns a version)
TelemetryReporter.sendPing(