toolInput) {
+ this.toolInput = toolInput;
+ return this;
+ }
+
+ public Builder model(String model) {
+ this.model = model;
+ return this;
+ }
+
+ public Builder provider(String provider) {
+ this.provider = provider;
+ return this;
+ }
+
+ public CheckToolGateOptions build() {
+ return new CheckToolGateOptions(this);
+ }
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java b/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java
new file mode 100644
index 0000000..b974410
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java
@@ -0,0 +1,566 @@
+/*
+ * 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.getaxonflow.sdk.AxonFlow;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.ApprovalStatus;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.CreateWorkflowRequest;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.CreateWorkflowResponse;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.GateDecision;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.MarkStepCompletedRequest;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.StepGateRequest;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.StepGateResponse;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.StepType;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.ToolContext;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.WorkflowSource;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.WorkflowStatusResponse;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.WorkflowStepInfo;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
+
+/**
+ * Wraps LangGraph workflows with AxonFlow governance gates.
+ *
+ * This adapter provides a simple interface for integrating AxonFlow's
+ * Workflow Control Plane with LangGraph workflows. It handles workflow
+ * registration, step gate checks, per-tool governance, MCP tool interception,
+ * and workflow lifecycle management.
+ *
+ *
Thread safety: This class is not thread-safe. Each workflow
+ * execution should use its own adapter instance from a single thread.
+ *
+ *
"LangGraph runs the workflow. AxonFlow decides when it's allowed to move forward."
+ *
+ *
Example usage:
+ *
{@code
+ * try (LangGraphAdapter adapter = LangGraphAdapter.builder(client, "code-review")
+ * .autoBlock(true)
+ * .build()) {
+ * adapter.startWorkflow();
+ *
+ * if (adapter.checkGate("analyze", "llm_call")) {
+ * analyzeCode();
+ * adapter.stepCompleted("analyze");
+ * }
+ *
+ * adapter.completeWorkflow();
+ * }
+ * }
+ *
+ * @see CheckGateOptions
+ * @see MCPToolInterceptor
+ */
+public final class LangGraphAdapter implements AutoCloseable {
+
+ private final AxonFlow client;
+ private final String workflowName;
+ private final WorkflowSource source;
+ private final boolean autoBlock;
+ private String workflowId;
+ private int stepCounter;
+ private boolean closedNormally;
+
+ private LangGraphAdapter(Builder builder) {
+ this.client = builder.client;
+ this.workflowName = builder.workflowName;
+ this.source = builder.source;
+ this.autoBlock = builder.autoBlock;
+ }
+
+ /**
+ * Creates a new builder for a LangGraphAdapter.
+ *
+ * @param client the AxonFlow client instance
+ * @param workflowName human-readable name for the workflow
+ * @return a new builder
+ */
+ public static Builder builder(AxonFlow client, String workflowName) {
+ return new Builder(client, workflowName);
+ }
+
+ // ========================================================================
+ // Workflow Lifecycle
+ // ========================================================================
+
+ /**
+ * Registers the workflow with AxonFlow.
+ *
+ * Call this at the start of your LangGraph workflow execution.
+ *
+ * @param metadata additional workflow metadata (may be null)
+ * @param traceId external trace ID for correlation (may be null)
+ * @return the assigned workflow ID
+ */
+ public String startWorkflow(Map metadata, String traceId) {
+ CreateWorkflowRequest request = new CreateWorkflowRequest(
+ workflowName,
+ source,
+ metadata != null ? metadata : Collections.emptyMap(),
+ traceId
+ );
+
+ CreateWorkflowResponse response = client.createWorkflow(request);
+ this.workflowId = response.getWorkflowId();
+ return this.workflowId;
+ }
+
+ /**
+ * Registers the workflow with AxonFlow using default metadata.
+ *
+ * @return the assigned workflow ID
+ */
+ public String startWorkflow() {
+ return startWorkflow(null, null);
+ }
+
+ /**
+ * Marks the workflow as completed successfully.
+ *
+ * Call this when your LangGraph workflow finishes all steps.
+ *
+ * @throws IllegalStateException if workflow not started
+ */
+ public void completeWorkflow() {
+ requireStarted();
+ client.completeWorkflow(workflowId);
+ closedNormally = true;
+ }
+
+ /**
+ * Aborts the workflow.
+ *
+ * @param reason the reason for aborting (may be null)
+ * @throws IllegalStateException if workflow not started
+ */
+ public void abortWorkflow(String reason) {
+ requireStarted();
+ client.abortWorkflow(workflowId, reason);
+ closedNormally = true;
+ }
+
+ /**
+ * Fails the workflow.
+ *
+ * @param reason the reason for failing (may be null)
+ * @throws IllegalStateException if workflow not started
+ */
+ public void failWorkflow(String reason) {
+ requireStarted();
+ client.failWorkflow(workflowId, reason);
+ closedNormally = true;
+ }
+
+ // ========================================================================
+ // Step Governance
+ // ========================================================================
+
+ /**
+ * Checks if a step is allowed to proceed.
+ *
+ *
Call this before executing each LangGraph node.
+ *
+ * @param stepName human-readable step name
+ * @param stepType type of step (llm_call, tool_call, connector_call, human_task)
+ * @return true if allowed, false if blocked (when autoBlock is false)
+ * @throws WorkflowBlockedError if blocked and autoBlock is true
+ * @throws WorkflowApprovalRequiredError if approval is required
+ * @throws IllegalStateException if workflow not started
+ */
+ public boolean checkGate(String stepName, String stepType) {
+ return checkGate(stepName, stepType, null);
+ }
+
+ /**
+ * Checks if a step is allowed to proceed, with options.
+ *
+ * @param stepName human-readable step name
+ * @param stepType type of step (llm_call, tool_call, connector_call, human_task)
+ * @param options additional options (may be null)
+ * @return true if allowed, false if blocked (when autoBlock is false)
+ * @throws WorkflowBlockedError if blocked and autoBlock is true
+ * @throws WorkflowApprovalRequiredError if approval is required
+ * @throws IllegalStateException if workflow not started
+ */
+ public boolean checkGate(String stepName, String stepType, CheckGateOptions options) {
+ requireStarted();
+
+ StepType resolvedType = StepType.fromValue(stepType);
+
+ // Generate or use provided step ID
+ String stepId;
+ if (options != null && options.getStepId() != null) {
+ stepId = options.getStepId();
+ } else {
+ stepCounter++;
+ String safeName = stepName.toLowerCase().replace(" ", "-").replace("/", "-");
+ stepId = "step-" + stepCounter + "-" + safeName;
+ }
+
+ StepGateRequest request = new StepGateRequest(
+ stepName,
+ resolvedType,
+ options != null && options.getStepInput() != null ? options.getStepInput() : Collections.emptyMap(),
+ options != null ? options.getModel() : null,
+ options != null ? options.getProvider() : null,
+ options != null ? options.getToolContext() : null
+ );
+
+ StepGateResponse response = client.stepGate(workflowId, stepId, request);
+
+ if (response.getDecision() == GateDecision.BLOCK) {
+ if (autoBlock) {
+ throw new WorkflowBlockedError(
+ "Step '" + stepName + "' blocked: " + response.getReason(),
+ response.getStepId(),
+ response.getReason(),
+ response.getPolicyIds()
+ );
+ }
+ return false;
+ }
+
+ if (response.getDecision() == GateDecision.REQUIRE_APPROVAL) {
+ throw new WorkflowApprovalRequiredError(
+ "Step '" + stepName + "' requires approval",
+ response.getStepId(),
+ response.getApprovalUrl(),
+ response.getReason()
+ );
+ }
+
+ return true;
+ }
+
+ /**
+ * Marks a step as completed.
+ *
+ *
Call this after successfully executing a LangGraph node.
+ *
+ * @param stepName the step name (used to derive step ID if not provided)
+ * @throws IllegalStateException if workflow not started
+ */
+ public void stepCompleted(String stepName) {
+ stepCompleted(stepName, null);
+ }
+
+ /**
+ * Marks a step as completed, with options.
+ *
+ * @param stepName the step name
+ * @param options additional options (may be null)
+ * @throws IllegalStateException if workflow not started
+ */
+ public void stepCompleted(String stepName, StepCompletedOptions options) {
+ requireStarted();
+
+ // Generate step ID matching the current counter (same as last checkGate)
+ String stepId;
+ if (options != null && options.getStepId() != null) {
+ stepId = options.getStepId();
+ } else {
+ String safeName = stepName.toLowerCase().replace(" ", "-").replace("/", "-");
+ stepId = "step-" + stepCounter + "-" + safeName;
+ }
+
+ MarkStepCompletedRequest request = new MarkStepCompletedRequest(
+ options != null ? options.getOutput() : null,
+ options != null ? options.getMetadata() : null,
+ options != null ? options.getTokensIn() : null,
+ options != null ? options.getTokensOut() : null,
+ options != null ? options.getCostUsd() : null
+ );
+
+ client.markStepCompleted(workflowId, stepId, request);
+ }
+
+ // ========================================================================
+ // Per-Tool Governance
+ // ========================================================================
+
+ /**
+ * Checks if a specific tool invocation is allowed.
+ *
+ *
Convenience wrapper around {@link #checkGate} that creates a
+ * {@link ToolContext} and uses step type {@code tool_call}.
+ *
+ * @param toolName the tool name
+ * @param toolType the tool type (function, mcp, api)
+ * @return true if allowed, false if blocked (when autoBlock is false)
+ * @throws WorkflowBlockedError if blocked and autoBlock is true
+ * @throws WorkflowApprovalRequiredError if approval is required
+ * @throws IllegalStateException if workflow not started
+ */
+ public boolean checkToolGate(String toolName, String toolType) {
+ return checkToolGate(toolName, toolType, null);
+ }
+
+ /**
+ * Checks if a specific tool invocation is allowed, with options.
+ *
+ * @param toolName the tool name
+ * @param toolType the tool type (function, mcp, api)
+ * @param options additional options (may be null)
+ * @return true if allowed, false if blocked (when autoBlock is false)
+ */
+ public boolean checkToolGate(String toolName, String toolType, CheckToolGateOptions options) {
+ String stepName = (options != null && options.getStepName() != null)
+ ? options.getStepName()
+ : "tools/" + toolName;
+
+ ToolContext toolContext = ToolContext.builder(toolName)
+ .toolType(toolType)
+ .toolInput(options != null ? options.getToolInput() : null)
+ .build();
+
+ CheckGateOptions gateOptions = CheckGateOptions.builder()
+ .stepId(options != null ? options.getStepId() : null)
+ .model(options != null ? options.getModel() : null)
+ .provider(options != null ? options.getProvider() : null)
+ .toolContext(toolContext)
+ .build();
+
+ return checkGate(stepName, StepType.TOOL_CALL.getValue(), gateOptions);
+ }
+
+ /**
+ * Marks a tool invocation as completed.
+ *
+ * @param toolName the tool name
+ * @throws IllegalStateException if workflow not started
+ */
+ public void toolCompleted(String toolName) {
+ toolCompleted(toolName, null);
+ }
+
+ /**
+ * Marks a tool invocation as completed, with options.
+ *
+ * @param toolName the tool name
+ * @param options additional options (may be null)
+ * @throws IllegalStateException if workflow not started
+ */
+ public void toolCompleted(String toolName, ToolCompletedOptions options) {
+ String stepName = (options != null && options.getStepName() != null)
+ ? options.getStepName()
+ : "tools/" + toolName;
+
+ StepCompletedOptions stepOptions = null;
+ if (options != null) {
+ stepOptions = StepCompletedOptions.builder()
+ .stepId(options.getStepId())
+ .output(options.getOutput())
+ .tokensIn(options.getTokensIn())
+ .tokensOut(options.getTokensOut())
+ .costUsd(options.getCostUsd())
+ .build();
+ }
+
+ stepCompleted(stepName, stepOptions);
+ }
+
+ // ========================================================================
+ // Approval
+ // ========================================================================
+
+ /**
+ * Waits for a step to be approved by polling the workflow status.
+ *
+ * @param stepId the step ID to wait for
+ * @param pollIntervalMs milliseconds between polls
+ * @param timeoutMs maximum milliseconds to wait
+ * @return true if approved, false if rejected
+ * @throws InterruptedException if the thread is interrupted while waiting
+ * @throws TimeoutException if approval not received within timeout
+ * @throws IllegalStateException if workflow not started
+ */
+ public boolean waitForApproval(String stepId, long pollIntervalMs, long timeoutMs)
+ throws InterruptedException, TimeoutException {
+ requireStarted();
+
+ long deadlineNanos = System.nanoTime() + java.util.concurrent.TimeUnit.MILLISECONDS.toNanos(timeoutMs);
+ while (System.nanoTime() < deadlineNanos) {
+ WorkflowStatusResponse status = client.getWorkflow(workflowId);
+
+ for (WorkflowStepInfo step : status.getSteps()) {
+ if (stepId.equals(step.getStepId())) {
+ if (step.getApprovalStatus() != null) {
+ if (step.getApprovalStatus() == ApprovalStatus.APPROVED) {
+ return true;
+ }
+ if (step.getApprovalStatus() == ApprovalStatus.REJECTED) {
+ return false;
+ }
+ }
+ break;
+ }
+ }
+
+ Thread.sleep(pollIntervalMs);
+ }
+
+ throw new TimeoutException("Approval timeout after " + timeoutMs + "ms for step " + stepId);
+ }
+
+ // ========================================================================
+ // MCP Interceptor
+ // ========================================================================
+
+ /**
+ * Creates an MCP tool interceptor with default options.
+ *
+ *
The interceptor enforces AxonFlow input and output policies around
+ * every MCP tool call.
+ *
+ * @return a new MCP tool interceptor
+ */
+ public MCPToolInterceptor mcpToolInterceptor() {
+ return mcpToolInterceptor(null);
+ }
+
+ /**
+ * Creates an MCP tool interceptor with custom options.
+ *
+ * @param options interceptor options (may be null for defaults)
+ * @return a new MCP tool interceptor
+ */
+ public MCPToolInterceptor mcpToolInterceptor(MCPInterceptorOptions options) {
+ Function connectorTypeFn;
+ String operation;
+
+ if (options != null) {
+ connectorTypeFn = options.getConnectorTypeFn() != null
+ ? options.getConnectorTypeFn()
+ : req -> req.getServerName() + "." + req.getName();
+ operation = options.getOperation();
+ } else {
+ connectorTypeFn = req -> req.getServerName() + "." + req.getName();
+ operation = "execute";
+ }
+
+ return new MCPToolInterceptor(client, connectorTypeFn, operation);
+ }
+
+ // ========================================================================
+ // Accessors
+ // ========================================================================
+
+ /**
+ * Returns the workflow ID assigned after {@link #startWorkflow()}.
+ *
+ * @return the workflow ID, or null if not yet started
+ */
+ public String getWorkflowId() {
+ return workflowId;
+ }
+
+ /**
+ * Returns the current step counter value.
+ *
+ * @return the step counter
+ */
+ int getStepCounter() {
+ return stepCounter;
+ }
+
+ // ========================================================================
+ // AutoCloseable
+ // ========================================================================
+
+ /**
+ * Closes the adapter. If the workflow was started but not explicitly
+ * completed, aborted, or failed, it will be aborted automatically.
+ */
+ @Override
+ public void close() {
+ if (workflowId != null && !closedNormally) {
+ try {
+ abortWorkflow("Adapter closed without explicit completion");
+ } catch (Exception ignored) {
+ // Best-effort cleanup
+ }
+ }
+ }
+
+ // ========================================================================
+ // Internals
+ // ========================================================================
+
+ private void requireStarted() {
+ if (workflowId == null) {
+ throw new IllegalStateException("Workflow not started. Call startWorkflow() first.");
+ }
+ }
+
+ // ========================================================================
+ // Builder
+ // ========================================================================
+
+ /**
+ * Builder for {@link LangGraphAdapter}.
+ */
+ public static final class Builder {
+
+ private final AxonFlow client;
+ private final String workflowName;
+ private WorkflowSource source = WorkflowSource.LANGGRAPH;
+ private boolean autoBlock = true;
+
+ private Builder(AxonFlow client, String workflowName) {
+ this.client = Objects.requireNonNull(client, "client cannot be null");
+ this.workflowName = Objects.requireNonNull(workflowName, "workflowName cannot be null");
+ }
+
+ /**
+ * Sets the workflow source. Defaults to {@link WorkflowSource#LANGGRAPH}.
+ *
+ * @param source the workflow source
+ * @return this builder
+ */
+ public Builder source(WorkflowSource source) {
+ this.source = Objects.requireNonNull(source, "source cannot be null");
+ return this;
+ }
+
+ /**
+ * Sets whether to automatically throw on blocked steps.
+ * Defaults to {@code true}.
+ *
+ * When {@code true}, {@link LangGraphAdapter#checkGate} throws
+ * {@link WorkflowBlockedError} on block decisions.
+ * When {@code false}, it returns {@code false}.
+ *
+ * @param autoBlock whether to auto-block
+ * @return this builder
+ */
+ public Builder autoBlock(boolean autoBlock) {
+ this.autoBlock = autoBlock;
+ return this;
+ }
+
+ /**
+ * Builds the adapter.
+ *
+ * @return a new LangGraphAdapter
+ */
+ public LangGraphAdapter build() {
+ return new LangGraphAdapter(this);
+ }
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/MCPInterceptorOptions.java b/src/main/java/com/getaxonflow/sdk/adapters/MCPInterceptorOptions.java
new file mode 100644
index 0000000..cc3f584
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/MCPInterceptorOptions.java
@@ -0,0 +1,94 @@
+/*
+ * 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 java.util.function.Function;
+
+/**
+ * Options for {@link LangGraphAdapter#mcpToolInterceptor}.
+ *
+ *
Controls how MCP tool requests are mapped to connector types and what
+ * operation type is used for policy checks.
+ */
+public final class MCPInterceptorOptions {
+
+ private final Function connectorTypeFn;
+ private final String operation;
+
+ private MCPInterceptorOptions(Builder builder) {
+ this.connectorTypeFn = builder.connectorTypeFn;
+ this.operation = builder.operation;
+ }
+
+ /**
+ * Returns the function that maps an MCP request to a connector type string.
+ * May be null, in which case the default "{serverName}.{toolName}" is used.
+ *
+ * @return the connector type function, or null
+ */
+ public Function getConnectorTypeFn() {
+ return connectorTypeFn;
+ }
+
+ /**
+ * Returns the operation type passed to {@code mcpCheckInput}.
+ * Defaults to "execute".
+ *
+ * @return the operation type
+ */
+ public String getOperation() {
+ return operation;
+ }
+
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ public static final class Builder {
+ private Function connectorTypeFn;
+ private String operation = "execute";
+
+ private Builder() {
+ }
+
+ /**
+ * Sets a custom function to derive the connector type from an MCP request.
+ *
+ * @param connectorTypeFn mapping function
+ * @return this builder
+ */
+ public Builder connectorTypeFn(Function connectorTypeFn) {
+ this.connectorTypeFn = connectorTypeFn;
+ return this;
+ }
+
+ /**
+ * Sets the operation type. Defaults to "execute".
+ * Use "query" for known read-only tool calls.
+ *
+ * @param operation the operation type
+ * @return this builder
+ */
+ public Builder operation(String operation) {
+ this.operation = operation;
+ return this;
+ }
+
+ public MCPInterceptorOptions build() {
+ return new MCPInterceptorOptions(this);
+ }
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/MCPToolHandler.java b/src/main/java/com/getaxonflow/sdk/adapters/MCPToolHandler.java
new file mode 100644
index 0000000..a519680
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/MCPToolHandler.java
@@ -0,0 +1,35 @@
+/*
+ * 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;
+
+/**
+ * Functional interface for handling an MCP tool request.
+ *
+ * Implementations execute the actual tool call and return the result.
+ * Used by {@link MCPToolInterceptor} as the downstream handler.
+ */
+@FunctionalInterface
+public interface MCPToolHandler {
+
+ /**
+ * Handles an MCP tool request.
+ *
+ * @param request the tool request to handle
+ * @return the result of the tool invocation
+ * @throws Exception if the tool call fails
+ */
+ Object handle(MCPToolRequest request) throws Exception;
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/MCPToolInterceptor.java b/src/main/java/com/getaxonflow/sdk/adapters/MCPToolInterceptor.java
new file mode 100644
index 0000000..57d53af
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/MCPToolInterceptor.java
@@ -0,0 +1,134 @@
+/*
+ * 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.core.JsonProcessingException;
+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.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+
+/**
+ * Intercepts MCP tool calls with AxonFlow policy enforcement.
+ *
+ *
Wraps each MCP tool invocation with pre-call input validation and
+ * post-call output validation. If a policy blocks the input or output,
+ * a {@link PolicyViolationException} is thrown. If redaction is active,
+ * the redacted data is returned instead of the raw result.
+ *
+ *
Obtain an instance via {@link LangGraphAdapter#mcpToolInterceptor()}.
+ *
+ *
Example usage:
+ *
{@code
+ * MCPToolInterceptor interceptor = adapter.mcpToolInterceptor();
+ * Object result = interceptor.intercept(request, req -> callMCPServer(req));
+ * }
+ */
+public final class MCPToolInterceptor {
+
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ private final AxonFlow client;
+ private final Function connectorTypeFn;
+ private final String operation;
+
+ MCPToolInterceptor(AxonFlow client, Function connectorTypeFn, String operation) {
+ this.client = client;
+ this.connectorTypeFn = connectorTypeFn;
+ this.operation = operation;
+ }
+
+ /**
+ * Intercepts an MCP tool call with policy enforcement.
+ *
+ * Flow:
+ *
+ * - Derive connector type from the request
+ * - Build a statement string and call {@code mcpCheckInput}
+ * - If blocked, throw {@link PolicyViolationException}
+ * - Execute the handler
+ * - Call {@code mcpCheckOutput} with the result
+ * - If blocked, throw {@link PolicyViolationException}
+ * - If redacted data is present, return it instead of the raw result
+ *
+ *
+ * @param request the MCP tool request
+ * @param handler the downstream handler that executes the actual tool call
+ * @return the tool result, or redacted data if redaction policies are active
+ * @throws PolicyViolationException if input or output is blocked by policy
+ * @throws Exception if the handler throws
+ */
+ public Object intercept(MCPToolRequest request, MCPToolHandler handler) throws Exception {
+ String connectorType = connectorTypeFn.apply(request);
+ String argsStr = serializeArgs(request.getArgs());
+ String statement = connectorType + "(" + argsStr + ")";
+
+ // Pre-check: validate input
+ Map inputOptions = new HashMap<>();
+ inputOptions.put("operation", operation);
+ if (request.getArgs() != null && !request.getArgs().isEmpty()) {
+ inputOptions.put("parameters", request.getArgs());
+ }
+
+ MCPCheckInputResponse preCheck = client.mcpCheckInput(connectorType, statement, inputOptions);
+ if (!preCheck.isAllowed()) {
+ String reason = preCheck.getBlockReason() != null ? preCheck.getBlockReason() : "Tool call blocked by policy";
+ throw new PolicyViolationException(reason);
+ }
+
+ // Execute the tool
+ Object result = handler.handle(request);
+
+ // Post-check: validate output
+ String resultStr;
+ try {
+ resultStr = OBJECT_MAPPER.writeValueAsString(result);
+ } catch (JsonProcessingException e) {
+ resultStr = String.valueOf(result);
+ }
+
+ Map outputOptions = new HashMap<>();
+ outputOptions.put("message", resultStr);
+
+ MCPCheckOutputResponse outputCheck = client.mcpCheckOutput(connectorType, null, outputOptions);
+ if (!outputCheck.isAllowed()) {
+ String reason = outputCheck.getBlockReason() != null ? outputCheck.getBlockReason() : "Tool result blocked by policy";
+ throw new PolicyViolationException(reason);
+ }
+
+ if (outputCheck.getRedactedData() != null) {
+ return outputCheck.getRedactedData();
+ }
+
+ return result;
+ }
+
+ private static String serializeArgs(Map args) {
+ if (args == null || args.isEmpty()) {
+ return "{}";
+ }
+ try {
+ return OBJECT_MAPPER.writeValueAsString(args);
+ } catch (JsonProcessingException e) {
+ return args.toString();
+ }
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/MCPToolRequest.java b/src/main/java/com/getaxonflow/sdk/adapters/MCPToolRequest.java
new file mode 100644
index 0000000..b411ad2
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/MCPToolRequest.java
@@ -0,0 +1,73 @@
+/*
+ * 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 java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Represents an MCP tool invocation request.
+ *
+ * Used by {@link MCPToolInterceptor} to pass tool call information
+ * through the interceptor chain.
+ */
+public final class MCPToolRequest {
+
+ private final String serverName;
+ private final String name;
+ private final Map args;
+
+ /**
+ * Creates a new MCPToolRequest.
+ *
+ * @param serverName the MCP server name
+ * @param name the tool name
+ * @param args the tool arguments
+ */
+ public MCPToolRequest(String serverName, String name, Map args) {
+ this.serverName = Objects.requireNonNull(serverName, "serverName cannot be null");
+ this.name = Objects.requireNonNull(name, "name cannot be null");
+ this.args = args != null ? Collections.unmodifiableMap(args) : Collections.emptyMap();
+ }
+
+ /**
+ * Returns the MCP server name.
+ *
+ * @return the server name
+ */
+ public String getServerName() {
+ return serverName;
+ }
+
+ /**
+ * Returns the tool name.
+ *
+ * @return the tool name
+ */
+ public String getName() {
+ return name;
+ }
+
+ /**
+ * Returns the tool arguments.
+ *
+ * @return immutable map of arguments
+ */
+ public Map getArgs() {
+ return args;
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/StepCompletedOptions.java b/src/main/java/com/getaxonflow/sdk/adapters/StepCompletedOptions.java
new file mode 100644
index 0000000..3f06133
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/StepCompletedOptions.java
@@ -0,0 +1,114 @@
+/*
+ * 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 java.util.Map;
+
+/**
+ * Options for {@link LangGraphAdapter#stepCompleted}.
+ */
+public final class StepCompletedOptions {
+
+ private final String stepId;
+ private final Map output;
+ private final Map metadata;
+ private final Integer tokensIn;
+ private final Integer tokensOut;
+ private final Double costUsd;
+
+ private StepCompletedOptions(Builder builder) {
+ this.stepId = builder.stepId;
+ this.output = builder.output;
+ this.metadata = builder.metadata;
+ this.tokensIn = builder.tokensIn;
+ this.tokensOut = builder.tokensOut;
+ this.costUsd = builder.costUsd;
+ }
+
+ public String getStepId() {
+ return stepId;
+ }
+
+ public Map getOutput() {
+ return output;
+ }
+
+ public Map getMetadata() {
+ return metadata;
+ }
+
+ public Integer getTokensIn() {
+ return tokensIn;
+ }
+
+ public Integer getTokensOut() {
+ return tokensOut;
+ }
+
+ public Double getCostUsd() {
+ return costUsd;
+ }
+
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ public static final class Builder {
+ private String stepId;
+ private Map output;
+ private Map metadata;
+ private Integer tokensIn;
+ private Integer tokensOut;
+ private Double costUsd;
+
+ private Builder() {
+ }
+
+ public Builder stepId(String stepId) {
+ this.stepId = stepId;
+ return this;
+ }
+
+ public Builder output(Map output) {
+ this.output = output;
+ return this;
+ }
+
+ public Builder metadata(Map metadata) {
+ this.metadata = metadata;
+ return this;
+ }
+
+ public Builder tokensIn(Integer tokensIn) {
+ this.tokensIn = tokensIn;
+ return this;
+ }
+
+ public Builder tokensOut(Integer tokensOut) {
+ this.tokensOut = tokensOut;
+ return this;
+ }
+
+ public Builder costUsd(Double costUsd) {
+ this.costUsd = costUsd;
+ return this;
+ }
+
+ public StepCompletedOptions build() {
+ return new StepCompletedOptions(this);
+ }
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/ToolCompletedOptions.java b/src/main/java/com/getaxonflow/sdk/adapters/ToolCompletedOptions.java
new file mode 100644
index 0000000..6eab259
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/ToolCompletedOptions.java
@@ -0,0 +1,114 @@
+/*
+ * 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 java.util.Map;
+
+/**
+ * Options for {@link LangGraphAdapter#toolCompleted}.
+ */
+public final class ToolCompletedOptions {
+
+ private final String stepName;
+ private final String stepId;
+ private final Map output;
+ private final Integer tokensIn;
+ private final Integer tokensOut;
+ private final Double costUsd;
+
+ private ToolCompletedOptions(Builder builder) {
+ this.stepName = builder.stepName;
+ this.stepId = builder.stepId;
+ this.output = builder.output;
+ this.tokensIn = builder.tokensIn;
+ this.tokensOut = builder.tokensOut;
+ this.costUsd = builder.costUsd;
+ }
+
+ public String getStepName() {
+ return stepName;
+ }
+
+ public String getStepId() {
+ return stepId;
+ }
+
+ public Map getOutput() {
+ return output;
+ }
+
+ public Integer getTokensIn() {
+ return tokensIn;
+ }
+
+ public Integer getTokensOut() {
+ return tokensOut;
+ }
+
+ public Double getCostUsd() {
+ return costUsd;
+ }
+
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ public static final class Builder {
+ private String stepName;
+ private String stepId;
+ private Map output;
+ private Integer tokensIn;
+ private Integer tokensOut;
+ private Double costUsd;
+
+ private Builder() {
+ }
+
+ public Builder stepName(String stepName) {
+ this.stepName = stepName;
+ return this;
+ }
+
+ public Builder stepId(String stepId) {
+ this.stepId = stepId;
+ return this;
+ }
+
+ public Builder output(Map output) {
+ this.output = output;
+ return this;
+ }
+
+ public Builder tokensIn(Integer tokensIn) {
+ this.tokensIn = tokensIn;
+ return this;
+ }
+
+ public Builder tokensOut(Integer tokensOut) {
+ this.tokensOut = tokensOut;
+ return this;
+ }
+
+ public Builder costUsd(Double costUsd) {
+ this.costUsd = costUsd;
+ return this;
+ }
+
+ public ToolCompletedOptions build() {
+ return new ToolCompletedOptions(this);
+ }
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/WorkflowApprovalRequiredError.java b/src/main/java/com/getaxonflow/sdk/adapters/WorkflowApprovalRequiredError.java
new file mode 100644
index 0000000..0b11732
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/WorkflowApprovalRequiredError.java
@@ -0,0 +1,74 @@
+/*
+ * 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;
+
+/**
+ * Raised when a workflow step requires human approval before proceeding.
+ *
+ * This exception is thrown by {@link LangGraphAdapter#checkGate} when the
+ * gate decision is {@code REQUIRE_APPROVAL}. The caller should use
+ * {@link LangGraphAdapter#waitForApproval} to poll for approval.
+ */
+public class WorkflowApprovalRequiredError extends RuntimeException {
+
+ private static final long serialVersionUID = 1L;
+
+ private final String stepId;
+ private final String approvalUrl;
+ private final String reason;
+
+ /**
+ * Creates a new WorkflowApprovalRequiredError.
+ *
+ * @param message the error message
+ * @param stepId the step ID that requires approval
+ * @param approvalUrl the URL where approval can be granted
+ * @param reason the reason approval is required
+ */
+ public WorkflowApprovalRequiredError(String message, String stepId, String approvalUrl, String reason) {
+ super(message);
+ this.stepId = stepId;
+ this.approvalUrl = approvalUrl;
+ this.reason = reason;
+ }
+
+ /**
+ * Returns the step ID that requires approval.
+ *
+ * @return the step ID
+ */
+ public String getStepId() {
+ return stepId;
+ }
+
+ /**
+ * Returns the URL where approval can be granted.
+ *
+ * @return the approval URL
+ */
+ public String getApprovalUrl() {
+ return approvalUrl;
+ }
+
+ /**
+ * Returns the reason approval is required.
+ *
+ * @return the reason
+ */
+ public String getReason() {
+ return reason;
+ }
+}
diff --git a/src/main/java/com/getaxonflow/sdk/adapters/WorkflowBlockedError.java b/src/main/java/com/getaxonflow/sdk/adapters/WorkflowBlockedError.java
new file mode 100644
index 0000000..2186777
--- /dev/null
+++ b/src/main/java/com/getaxonflow/sdk/adapters/WorkflowBlockedError.java
@@ -0,0 +1,76 @@
+/*
+ * 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 java.util.Collections;
+import java.util.List;
+
+/**
+ * Raised when a workflow step is blocked by policy.
+ *
+ *
This exception is thrown by {@link LangGraphAdapter#checkGate} when
+ * {@code autoBlock} is {@code true} and the gate decision is {@code BLOCK}.
+ */
+public class WorkflowBlockedError extends RuntimeException {
+
+ private static final long serialVersionUID = 1L;
+
+ private final String stepId;
+ private final String reason;
+ private final List policyIds;
+
+ /**
+ * Creates a new WorkflowBlockedError.
+ *
+ * @param message the error message
+ * @param stepId the step ID that was blocked
+ * @param reason the reason the step was blocked
+ * @param policyIds the policy IDs that caused the block
+ */
+ public WorkflowBlockedError(String message, String stepId, String reason, List policyIds) {
+ super(message);
+ this.stepId = stepId;
+ this.reason = reason;
+ this.policyIds = policyIds != null ? Collections.unmodifiableList(policyIds) : Collections.emptyList();
+ }
+
+ /**
+ * Returns the step ID that was blocked.
+ *
+ * @return the step ID
+ */
+ public String getStepId() {
+ return stepId;
+ }
+
+ /**
+ * Returns the reason the step was blocked.
+ *
+ * @return the block reason
+ */
+ public String getReason() {
+ return reason;
+ }
+
+ /**
+ * Returns the policy IDs that caused the block.
+ *
+ * @return immutable list of policy IDs
+ */
+ public List getPolicyIds() {
+ return policyIds;
+ }
+}
diff --git a/src/test/java/com/getaxonflow/sdk/adapters/LangGraphAdapterTest.java b/src/test/java/com/getaxonflow/sdk/adapters/LangGraphAdapterTest.java
new file mode 100644
index 0000000..316a050
--- /dev/null
+++ b/src/test/java/com/getaxonflow/sdk/adapters/LangGraphAdapterTest.java
@@ -0,0 +1,884 @@
+/*
+ * 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.getaxonflow.sdk.AxonFlow;
+import com.getaxonflow.sdk.exceptions.PolicyViolationException;
+import com.getaxonflow.sdk.types.MCPCheckInputResponse;
+import com.getaxonflow.sdk.types.MCPCheckOutputResponse;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.ApprovalStatus;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.CreateWorkflowRequest;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.CreateWorkflowResponse;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.GateDecision;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.MarkStepCompletedRequest;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.StepGateRequest;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.StepGateResponse;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.StepType;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.WorkflowSource;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.WorkflowStatus;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.WorkflowStatusResponse;
+import com.getaxonflow.sdk.types.workflow.WorkflowTypes.WorkflowStepInfo;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Nested;
+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;
+
+import java.time.Instant;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeoutException;
+
+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.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@DisplayName("LangGraphAdapter")
+@ExtendWith(MockitoExtension.class)
+class LangGraphAdapterTest {
+
+ @Mock
+ private AxonFlow client;
+
+ private LangGraphAdapter adapter;
+
+ @BeforeEach
+ void setUp() {
+ adapter = LangGraphAdapter.builder(client, "test-workflow").build();
+ }
+
+ // ========================================================================
+ // Builder
+ // ========================================================================
+
+ @Nested
+ @DisplayName("Builder")
+ class BuilderTests {
+
+ @Test
+ @DisplayName("should use default source and autoBlock")
+ void shouldUseDefaults() {
+ LangGraphAdapter a = LangGraphAdapter.builder(client, "my-wf").build();
+ assertThat(a.getWorkflowId()).isNull();
+ }
+
+ @Test
+ @DisplayName("should accept custom source")
+ void shouldAcceptCustomSource() {
+ // Verify the adapter can be built with custom source without error
+ LangGraphAdapter a = LangGraphAdapter.builder(client, "wf")
+ .source(WorkflowSource.LANGCHAIN)
+ .build();
+ assertThat(a).isNotNull();
+ }
+
+ @Test
+ @DisplayName("should reject null client")
+ void shouldRejectNullClient() {
+ assertThatThrownBy(() -> LangGraphAdapter.builder(null, "wf"))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining("client");
+ }
+
+ @Test
+ @DisplayName("should reject null workflowName")
+ void shouldRejectNullWorkflowName() {
+ assertThatThrownBy(() -> LangGraphAdapter.builder(client, null))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining("workflowName");
+ }
+ }
+
+ // ========================================================================
+ // startWorkflow
+ // ========================================================================
+
+ @Nested
+ @DisplayName("startWorkflow")
+ class StartWorkflowTests {
+
+ @Test
+ @DisplayName("should create workflow and store ID")
+ void shouldCreateWorkflowAndStoreId() {
+ mockCreateWorkflow("wf-123");
+
+ String id = adapter.startWorkflow();
+
+ assertThat(id).isEqualTo("wf-123");
+ assertThat(adapter.getWorkflowId()).isEqualTo("wf-123");
+ }
+
+ @Test
+ @DisplayName("should pass metadata and traceId")
+ void shouldPassMetadataAndTraceId() {
+ mockCreateWorkflow("wf-456");
+
+ Map meta = Map.of("customer", "cust-1");
+ adapter.startWorkflow(meta, "trace-abc");
+
+ ArgumentCaptor captor = ArgumentCaptor.forClass(CreateWorkflowRequest.class);
+ verify(client).createWorkflow(captor.capture());
+
+ CreateWorkflowRequest req = captor.getValue();
+ assertThat(req.getWorkflowName()).isEqualTo("test-workflow");
+ assertThat(req.getSource()).isEqualTo(WorkflowSource.LANGGRAPH);
+ assertThat(req.getMetadata()).containsEntry("customer", "cust-1");
+ assertThat(req.getTraceId()).isEqualTo("trace-abc");
+ }
+ }
+
+ // ========================================================================
+ // checkGate
+ // ========================================================================
+
+ @Nested
+ @DisplayName("checkGate")
+ class CheckGateTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-1");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("should return true when allowed")
+ void shouldReturnTrueWhenAllowed() {
+ mockStepGate(GateDecision.ALLOW);
+
+ boolean result = adapter.checkGate("generate", "llm_call");
+
+ assertThat(result).isTrue();
+ }
+
+ @Test
+ @DisplayName("should throw WorkflowBlockedError when blocked and autoBlock=true")
+ void shouldThrowOnBlockWithAutoBlock() {
+ mockStepGate(GateDecision.BLOCK, "step-x", "Policy violation", List.of("pol-1"));
+
+ assertThatThrownBy(() -> adapter.checkGate("generate", "llm_call"))
+ .isInstanceOf(WorkflowBlockedError.class)
+ .hasMessageContaining("generate")
+ .hasMessageContaining("blocked");
+ }
+
+ @Test
+ @DisplayName("should return false when blocked and autoBlock=false")
+ void shouldReturnFalseOnBlockWithoutAutoBlock() {
+ LangGraphAdapter noBlock = LangGraphAdapter.builder(client, "wf")
+ .autoBlock(false)
+ .build();
+ mockCreateWorkflow("wf-nb");
+ noBlock.startWorkflow();
+
+ mockStepGate(GateDecision.BLOCK);
+
+ boolean result = noBlock.checkGate("generate", "llm_call");
+ assertThat(result).isFalse();
+ }
+
+ @Test
+ @DisplayName("should throw WorkflowApprovalRequiredError on require_approval")
+ void shouldThrowOnApprovalRequired() {
+ StepGateResponse resp = new StepGateResponse(
+ GateDecision.REQUIRE_APPROVAL, "step-approval", "Needs review",
+ Collections.emptyList(), "https://approve.me", null, null);
+ when(client.stepGate(anyString(), anyString(), any(StepGateRequest.class))).thenReturn(resp);
+
+ assertThatThrownBy(() -> adapter.checkGate("deploy", "human_task"))
+ .isInstanceOf(WorkflowApprovalRequiredError.class)
+ .satisfies(ex -> {
+ WorkflowApprovalRequiredError err = (WorkflowApprovalRequiredError) ex;
+ assertThat(err.getStepId()).isEqualTo("step-approval");
+ assertThat(err.getApprovalUrl()).isEqualTo("https://approve.me");
+ assertThat(err.getReason()).isEqualTo("Needs review");
+ });
+ }
+
+ @Test
+ @DisplayName("should throw IllegalStateException when workflow not started")
+ void shouldThrowWhenNotStarted() {
+ LangGraphAdapter fresh = LangGraphAdapter.builder(client, "wf").build();
+
+ assertThatThrownBy(() -> fresh.checkGate("step", "llm_call"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("not started");
+ }
+
+ @Test
+ @DisplayName("should pass model and provider from options")
+ void shouldPassModelAndProvider() {
+ mockStepGate(GateDecision.ALLOW);
+
+ CheckGateOptions opts = CheckGateOptions.builder()
+ .model("gpt-4")
+ .provider("openai")
+ .stepInput(Map.of("prompt", "hello"))
+ .build();
+
+ adapter.checkGate("generate", "llm_call", opts);
+
+ ArgumentCaptor captor = ArgumentCaptor.forClass(StepGateRequest.class);
+ verify(client).stepGate(eq("wf-1"), anyString(), captor.capture());
+
+ StepGateRequest req = captor.getValue();
+ assertThat(req.getModel()).isEqualTo("gpt-4");
+ assertThat(req.getProvider()).isEqualTo("openai");
+ assertThat(req.getStepInput()).containsEntry("prompt", "hello");
+ }
+
+ @Test
+ @DisplayName("should use custom stepId from options")
+ void shouldUseCustomStepId() {
+ mockStepGate(GateDecision.ALLOW);
+
+ CheckGateOptions opts = CheckGateOptions.builder()
+ .stepId("my-custom-step")
+ .build();
+
+ adapter.checkGate("generate", "llm_call", opts);
+
+ verify(client).stepGate(eq("wf-1"), eq("my-custom-step"), any(StepGateRequest.class));
+ }
+
+ @Test
+ @DisplayName("WorkflowBlockedError should contain policy details")
+ void blockedErrorShouldContainDetails() {
+ mockStepGate(GateDecision.BLOCK, "step-blocked-1", "Cost limit exceeded", List.of("cost-policy", "budget-policy"));
+
+ assertThatThrownBy(() -> adapter.checkGate("expensive", "llm_call"))
+ .isInstanceOf(WorkflowBlockedError.class)
+ .satisfies(ex -> {
+ WorkflowBlockedError err = (WorkflowBlockedError) ex;
+ assertThat(err.getStepId()).isEqualTo("step-blocked-1");
+ assertThat(err.getReason()).isEqualTo("Cost limit exceeded");
+ assertThat(err.getPolicyIds()).containsExactly("cost-policy", "budget-policy");
+ });
+ }
+ }
+
+ // ========================================================================
+ // stepCompleted
+ // ========================================================================
+
+ @Nested
+ @DisplayName("stepCompleted")
+ class StepCompletedTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-sc");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("should call markStepCompleted with matching step ID")
+ void shouldCallMarkStepCompleted() {
+ mockStepGate(GateDecision.ALLOW);
+ adapter.checkGate("analyze", "llm_call");
+
+ adapter.stepCompleted("analyze");
+
+ // Step counter is 1 after first checkGate, so step ID should be step-1-analyze
+ verify(client).markStepCompleted(eq("wf-sc"), eq("step-1-analyze"), any(MarkStepCompletedRequest.class));
+ }
+
+ @Test
+ @DisplayName("should pass output and metadata from options")
+ void shouldPassOptions() {
+ mockStepGate(GateDecision.ALLOW);
+ adapter.checkGate("generate", "llm_call");
+
+ StepCompletedOptions opts = StepCompletedOptions.builder()
+ .output(Map.of("code", "result"))
+ .metadata(Map.of("key", "val"))
+ .tokensIn(100)
+ .tokensOut(200)
+ .costUsd(0.05)
+ .build();
+
+ adapter.stepCompleted("generate", opts);
+
+ ArgumentCaptor captor = ArgumentCaptor.forClass(MarkStepCompletedRequest.class);
+ verify(client).markStepCompleted(eq("wf-sc"), anyString(), captor.capture());
+
+ MarkStepCompletedRequest req = captor.getValue();
+ assertThat(req.getOutput()).containsEntry("code", "result");
+ assertThat(req.getMetadata()).containsEntry("key", "val");
+ assertThat(req.getTokensIn()).isEqualTo(100);
+ assertThat(req.getTokensOut()).isEqualTo(200);
+ assertThat(req.getCostUsd()).isEqualTo(0.05);
+ }
+
+ @Test
+ @DisplayName("should throw IllegalStateException when workflow not started")
+ void shouldThrowWhenNotStarted() {
+ LangGraphAdapter fresh = LangGraphAdapter.builder(client, "wf").build();
+
+ assertThatThrownBy(() -> fresh.stepCompleted("step"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("not started");
+ }
+
+ @Test
+ @DisplayName("should use custom stepId from options")
+ void shouldUseCustomStepId() {
+ mockStepGate(GateDecision.ALLOW);
+ adapter.checkGate("generate", "llm_call");
+
+ StepCompletedOptions opts = StepCompletedOptions.builder()
+ .stepId("custom-id")
+ .build();
+
+ adapter.stepCompleted("generate", opts);
+
+ verify(client).markStepCompleted(eq("wf-sc"), eq("custom-id"), any(MarkStepCompletedRequest.class));
+ }
+ }
+
+ // ========================================================================
+ // checkToolGate
+ // ========================================================================
+
+ @Nested
+ @DisplayName("checkToolGate")
+ class CheckToolGateTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-tg");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("should use default step name tools/{toolName}")
+ void shouldUseDefaultStepName() {
+ mockStepGate(GateDecision.ALLOW);
+
+ adapter.checkToolGate("web_search", "function");
+
+ ArgumentCaptor captor = ArgumentCaptor.forClass(StepGateRequest.class);
+ verify(client).stepGate(eq("wf-tg"), anyString(), captor.capture());
+
+ StepGateRequest req = captor.getValue();
+ assertThat(req.getStepName()).isEqualTo("tools/web_search");
+ assertThat(req.getStepType()).isEqualTo(StepType.TOOL_CALL);
+ assertThat(req.getToolContext()).isNotNull();
+ assertThat(req.getToolContext().getToolName()).isEqualTo("web_search");
+ assertThat(req.getToolContext().getToolType()).isEqualTo("function");
+ }
+
+ @Test
+ @DisplayName("should use custom step name from options")
+ void shouldUseCustomStepName() {
+ mockStepGate(GateDecision.ALLOW);
+
+ CheckToolGateOptions opts = CheckToolGateOptions.builder()
+ .stepName("custom-tools/search")
+ .toolInput(Map.of("query", "test"))
+ .build();
+
+ adapter.checkToolGate("web_search", "function", opts);
+
+ ArgumentCaptor captor = ArgumentCaptor.forClass(StepGateRequest.class);
+ verify(client).stepGate(eq("wf-tg"), anyString(), captor.capture());
+
+ StepGateRequest req = captor.getValue();
+ assertThat(req.getStepName()).isEqualTo("custom-tools/search");
+ assertThat(req.getToolContext().getToolInput()).containsEntry("query", "test");
+ }
+ }
+
+ // ========================================================================
+ // toolCompleted
+ // ========================================================================
+
+ @Nested
+ @DisplayName("toolCompleted")
+ class ToolCompletedTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-tc");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("should delegate to stepCompleted with default step name")
+ void shouldDelegateWithDefaultName() {
+ mockStepGate(GateDecision.ALLOW);
+ adapter.checkToolGate("calculator", "function");
+
+ adapter.toolCompleted("calculator");
+
+ // Step ID should match the one generated in checkToolGate (tools/calculator -> tools-calculator)
+ verify(client).markStepCompleted(eq("wf-tc"), eq("step-1-tools-calculator"), any(MarkStepCompletedRequest.class));
+ }
+
+ @Test
+ @DisplayName("should pass options through to stepCompleted")
+ void shouldPassOptions() {
+ mockStepGate(GateDecision.ALLOW);
+ adapter.checkToolGate("calc", "function");
+
+ ToolCompletedOptions opts = ToolCompletedOptions.builder()
+ .output(Map.of("result", 42))
+ .tokensIn(10)
+ .tokensOut(20)
+ .costUsd(0.001)
+ .build();
+
+ adapter.toolCompleted("calc", opts);
+
+ ArgumentCaptor captor = ArgumentCaptor.forClass(MarkStepCompletedRequest.class);
+ verify(client).markStepCompleted(eq("wf-tc"), anyString(), captor.capture());
+
+ MarkStepCompletedRequest req = captor.getValue();
+ assertThat(req.getTokensIn()).isEqualTo(10);
+ assertThat(req.getTokensOut()).isEqualTo(20);
+ assertThat(req.getCostUsd()).isEqualTo(0.001);
+ }
+ }
+
+ // ========================================================================
+ // Workflow Lifecycle (complete/abort/fail)
+ // ========================================================================
+
+ @Nested
+ @DisplayName("Workflow Lifecycle")
+ class LifecycleTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-lc");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("completeWorkflow should delegate to client")
+ void completeWorkflowShouldDelegate() {
+ adapter.completeWorkflow();
+ verify(client).completeWorkflow("wf-lc");
+ }
+
+ @Test
+ @DisplayName("abortWorkflow should delegate with reason")
+ void abortWorkflowShouldDelegateWithReason() {
+ adapter.abortWorkflow("user cancelled");
+ verify(client).abortWorkflow("wf-lc", "user cancelled");
+ }
+
+ @Test
+ @DisplayName("failWorkflow should delegate with reason")
+ void failWorkflowShouldDelegateWithReason() {
+ adapter.failWorkflow("pipeline error");
+ verify(client).failWorkflow("wf-lc", "pipeline error");
+ }
+
+ @Test
+ @DisplayName("completeWorkflow should throw when not started")
+ void completeShouldThrowWhenNotStarted() {
+ LangGraphAdapter fresh = LangGraphAdapter.builder(client, "wf").build();
+ assertThatThrownBy(fresh::completeWorkflow)
+ .isInstanceOf(IllegalStateException.class);
+ }
+
+ @Test
+ @DisplayName("abortWorkflow should throw when not started")
+ void abortShouldThrowWhenNotStarted() {
+ LangGraphAdapter fresh = LangGraphAdapter.builder(client, "wf").build();
+ assertThatThrownBy(() -> fresh.abortWorkflow("reason"))
+ .isInstanceOf(IllegalStateException.class);
+ }
+
+ @Test
+ @DisplayName("failWorkflow should throw when not started")
+ void failShouldThrowWhenNotStarted() {
+ LangGraphAdapter fresh = LangGraphAdapter.builder(client, "wf").build();
+ assertThatThrownBy(() -> fresh.failWorkflow("reason"))
+ .isInstanceOf(IllegalStateException.class);
+ }
+ }
+
+ // ========================================================================
+ // Step Counter
+ // ========================================================================
+
+ @Nested
+ @DisplayName("Step Counter")
+ class StepCounterTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-cnt");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("should increment step counter on each checkGate")
+ void shouldIncrementCounter() {
+ mockStepGate(GateDecision.ALLOW);
+
+ adapter.checkGate("step-a", "llm_call");
+ assertThat(adapter.getStepCounter()).isEqualTo(1);
+
+ adapter.checkGate("step-b", "tool_call");
+ assertThat(adapter.getStepCounter()).isEqualTo(2);
+
+ adapter.checkGate("step-c", "llm_call");
+ assertThat(adapter.getStepCounter()).isEqualTo(3);
+ }
+
+ @Test
+ @DisplayName("should generate step IDs with counter and safe name")
+ void shouldGenerateStepIds() {
+ mockStepGate(GateDecision.ALLOW);
+
+ adapter.checkGate("My Step", "llm_call");
+ verify(client).stepGate(eq("wf-cnt"), eq("step-1-my-step"), any(StepGateRequest.class));
+
+ adapter.checkGate("tools/search", "tool_call");
+ verify(client).stepGate(eq("wf-cnt"), eq("step-2-tools-search"), any(StepGateRequest.class));
+ }
+ }
+
+ // ========================================================================
+ // waitForApproval
+ // ========================================================================
+
+ @Nested
+ @DisplayName("waitForApproval")
+ class WaitForApprovalTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-wa");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("should return true when step is approved")
+ void shouldReturnTrueOnApproval() throws Exception {
+ WorkflowStepInfo approvedStep = new WorkflowStepInfo(
+ "step-1", 1, "deploy", StepType.HUMAN_TASK,
+ GateDecision.REQUIRE_APPROVAL, null, ApprovalStatus.APPROVED,
+ "admin", Instant.now(), null);
+
+ WorkflowStatusResponse status = new WorkflowStatusResponse(
+ "wf-wa", "wf", WorkflowSource.LANGGRAPH, WorkflowStatus.IN_PROGRESS,
+ 1, null, Instant.now(), null, List.of(approvedStep));
+
+ when(client.getWorkflow("wf-wa")).thenReturn(status);
+
+ boolean result = adapter.waitForApproval("step-1", 50, 5000);
+ assertThat(result).isTrue();
+ }
+
+ @Test
+ @DisplayName("should return false when step is rejected")
+ void shouldReturnFalseOnRejection() throws Exception {
+ WorkflowStepInfo rejectedStep = new WorkflowStepInfo(
+ "step-1", 1, "deploy", StepType.HUMAN_TASK,
+ GateDecision.REQUIRE_APPROVAL, null, ApprovalStatus.REJECTED,
+ null, Instant.now(), null);
+
+ WorkflowStatusResponse status = new WorkflowStatusResponse(
+ "wf-wa", "wf", WorkflowSource.LANGGRAPH, WorkflowStatus.IN_PROGRESS,
+ 1, null, Instant.now(), null, List.of(rejectedStep));
+
+ when(client.getWorkflow("wf-wa")).thenReturn(status);
+
+ boolean result = adapter.waitForApproval("step-1", 50, 5000);
+ assertThat(result).isFalse();
+ }
+
+ @Test
+ @DisplayName("should throw TimeoutException when approval not received")
+ void shouldThrowOnTimeout() {
+ WorkflowStepInfo pendingStep = new WorkflowStepInfo(
+ "step-1", 1, "deploy", StepType.HUMAN_TASK,
+ GateDecision.REQUIRE_APPROVAL, null, ApprovalStatus.PENDING,
+ null, Instant.now(), null);
+
+ WorkflowStatusResponse status = new WorkflowStatusResponse(
+ "wf-wa", "wf", WorkflowSource.LANGGRAPH, WorkflowStatus.IN_PROGRESS,
+ 1, null, Instant.now(), null, List.of(pendingStep));
+
+ when(client.getWorkflow("wf-wa")).thenReturn(status);
+
+ assertThatThrownBy(() -> adapter.waitForApproval("step-1", 50, 120))
+ .isInstanceOf(TimeoutException.class)
+ .hasMessageContaining("timeout");
+ }
+
+ @Test
+ @DisplayName("should throw IllegalStateException when not started")
+ void shouldThrowWhenNotStarted() {
+ LangGraphAdapter fresh = LangGraphAdapter.builder(client, "wf").build();
+ assertThatThrownBy(() -> fresh.waitForApproval("step-1", 50, 1000))
+ .isInstanceOf(IllegalStateException.class);
+ }
+ }
+
+ // ========================================================================
+ // close()
+ // ========================================================================
+
+ @Nested
+ @DisplayName("close")
+ class CloseTests {
+
+ @Test
+ @DisplayName("should abort workflow if not closed normally")
+ void shouldAbortIfNotClosedNormally() {
+ mockCreateWorkflow("wf-close");
+ adapter.startWorkflow();
+
+ adapter.close();
+
+ verify(client).abortWorkflow(eq("wf-close"), eq("Adapter closed without explicit completion"));
+ }
+
+ @Test
+ @DisplayName("should not abort if completeWorkflow was called")
+ void shouldNotAbortIfCompleted() {
+ mockCreateWorkflow("wf-close");
+ adapter.startWorkflow();
+ adapter.completeWorkflow();
+
+ adapter.close();
+
+ verify(client, never()).abortWorkflow(eq("wf-close"), eq("Adapter closed without explicit completion"));
+ }
+
+ @Test
+ @DisplayName("should not abort if abortWorkflow was called")
+ void shouldNotAbortIfAborted() {
+ mockCreateWorkflow("wf-close");
+ adapter.startWorkflow();
+ adapter.abortWorkflow("manual");
+
+ adapter.close();
+
+ // abortWorkflow was called once (manual), not twice
+ verify(client, times(1)).abortWorkflow(anyString(), anyString());
+ }
+
+ @Test
+ @DisplayName("should not abort if failWorkflow was called")
+ void shouldNotAbortIfFailed() {
+ mockCreateWorkflow("wf-close");
+ adapter.startWorkflow();
+ adapter.failWorkflow("error");
+
+ adapter.close();
+
+ verify(client, never()).abortWorkflow(eq("wf-close"), eq("Adapter closed without explicit completion"));
+ }
+
+ @Test
+ @DisplayName("should do nothing if workflow was never started")
+ void shouldDoNothingIfNotStarted() {
+ adapter.close();
+ verify(client, never()).abortWorkflow(anyString(), anyString());
+ }
+
+ @Test
+ @DisplayName("should swallow exceptions during close abort")
+ void shouldSwallowExceptionsDuringCloseAbort() {
+ mockCreateWorkflow("wf-close-err");
+ adapter.startWorkflow();
+
+ doThrow(new RuntimeException("network error"))
+ .when(client).abortWorkflow(anyString(), anyString());
+
+ // Should not throw
+ adapter.close();
+ }
+ }
+
+ // ========================================================================
+ // MCPToolInterceptor
+ // ========================================================================
+
+ @Nested
+ @DisplayName("MCPToolInterceptor")
+ class MCPToolInterceptorTests {
+
+ @BeforeEach
+ void startWorkflow() {
+ mockCreateWorkflow("wf-mcp");
+ adapter.startWorkflow();
+ }
+
+ @Test
+ @DisplayName("should pass through when input and output are allowed")
+ void shouldPassThroughWhenAllowed() throws Exception {
+ MCPCheckInputResponse inputOk = new MCPCheckInputResponse(true, null, 1, null);
+ MCPCheckOutputResponse outputOk = new MCPCheckOutputResponse(true, null, null, 1, null, null);
+
+ when(client.mcpCheckInput(anyString(), anyString(), any())).thenReturn(inputOk);
+ when(client.mcpCheckOutput(anyString(), isNull(), any())).thenReturn(outputOk);
+
+ MCPToolInterceptor interceptor = adapter.mcpToolInterceptor();
+ MCPToolRequest request = new MCPToolRequest("postgres", "query", Map.of("sql", "SELECT 1"));
+
+ Object result = interceptor.intercept(request, req -> "result-data");
+
+ assertThat(result).isEqualTo("result-data");
+ }
+
+ @Test
+ @DisplayName("should throw PolicyViolationException when input is blocked")
+ void shouldThrowOnBlockedInput() {
+ MCPCheckInputResponse blocked = new MCPCheckInputResponse(false, "DROP not allowed", 1, null);
+ when(client.mcpCheckInput(anyString(), anyString(), any())).thenReturn(blocked);
+
+ MCPToolInterceptor interceptor = adapter.mcpToolInterceptor();
+ MCPToolRequest request = new MCPToolRequest("postgres", "execute", Map.of("sql", "DROP TABLE"));
+
+ assertThatThrownBy(() -> interceptor.intercept(request, req -> "ignored"))
+ .isInstanceOf(PolicyViolationException.class)
+ .hasMessageContaining("DROP not allowed");
+ }
+
+ @Test
+ @DisplayName("should throw PolicyViolationException when output is blocked")
+ void shouldThrowOnBlockedOutput() throws Exception {
+ MCPCheckInputResponse inputOk = new MCPCheckInputResponse(true, null, 1, null);
+ MCPCheckOutputResponse outputBlocked = new MCPCheckOutputResponse(false, "PII detected", null, 1, null, null);
+
+ when(client.mcpCheckInput(anyString(), anyString(), any())).thenReturn(inputOk);
+ when(client.mcpCheckOutput(anyString(), isNull(), any())).thenReturn(outputBlocked);
+
+ MCPToolInterceptor interceptor = adapter.mcpToolInterceptor();
+ MCPToolRequest request = new MCPToolRequest("postgres", "query", null);
+
+ assertThatThrownBy(() -> interceptor.intercept(request, req -> "sensitive-data"))
+ .isInstanceOf(PolicyViolationException.class)
+ .hasMessageContaining("PII detected");
+ }
+
+ @Test
+ @DisplayName("should return redacted data when available")
+ void shouldReturnRedactedData() throws Exception {
+ MCPCheckInputResponse inputOk = new MCPCheckInputResponse(true, null, 1, null);
+ MCPCheckOutputResponse redacted = new MCPCheckOutputResponse(true, null, "***REDACTED***", 1, null, null);
+
+ when(client.mcpCheckInput(anyString(), anyString(), any())).thenReturn(inputOk);
+ when(client.mcpCheckOutput(anyString(), isNull(), any())).thenReturn(redacted);
+
+ MCPToolInterceptor interceptor = adapter.mcpToolInterceptor();
+ MCPToolRequest request = new MCPToolRequest("db", "query", Map.of("q", "SELECT *"));
+
+ Object result = interceptor.intercept(request, req -> "raw-data-with-pii");
+
+ assertThat(result).isEqualTo("***REDACTED***");
+ }
+
+ @Test
+ @DisplayName("should use default connector type serverName.toolName")
+ void shouldUseDefaultConnectorType() throws Exception {
+ MCPCheckInputResponse inputOk = new MCPCheckInputResponse(true, null, 1, null);
+ MCPCheckOutputResponse outputOk = new MCPCheckOutputResponse(true, null, null, 1, null, null);
+
+ when(client.mcpCheckInput(anyString(), anyString(), any())).thenReturn(inputOk);
+ when(client.mcpCheckOutput(anyString(), isNull(), any())).thenReturn(outputOk);
+
+ MCPToolInterceptor interceptor = adapter.mcpToolInterceptor();
+ MCPToolRequest request = new MCPToolRequest("myserver", "mytool", Collections.emptyMap());
+
+ interceptor.intercept(request, req -> "ok");
+
+ verify(client).mcpCheckInput(eq("myserver.mytool"), anyString(), any());
+ verify(client).mcpCheckOutput(eq("myserver.mytool"), isNull(), any());
+ }
+
+ @Test
+ @DisplayName("should use custom connector type function")
+ void shouldUseCustomConnectorTypeFn() throws Exception {
+ MCPCheckInputResponse inputOk = new MCPCheckInputResponse(true, null, 1, null);
+ MCPCheckOutputResponse outputOk = new MCPCheckOutputResponse(true, null, null, 1, null, null);
+
+ when(client.mcpCheckInput(anyString(), anyString(), any())).thenReturn(inputOk);
+ when(client.mcpCheckOutput(anyString(), isNull(), any())).thenReturn(outputOk);
+
+ MCPInterceptorOptions opts = MCPInterceptorOptions.builder()
+ .connectorTypeFn(req -> "custom-" + req.getServerName())
+ .operation("query")
+ .build();
+
+ MCPToolInterceptor interceptor = adapter.mcpToolInterceptor(opts);
+ MCPToolRequest request = new MCPToolRequest("pg", "read", null);
+
+ interceptor.intercept(request, req -> "data");
+
+ ArgumentCaptor> inputOptsCaptor = ArgumentCaptor.forClass(Map.class);
+ verify(client).mcpCheckInput(eq("custom-pg"), anyString(), inputOptsCaptor.capture());
+ assertThat(inputOptsCaptor.getValue()).containsEntry("operation", "query");
+ }
+
+ @Test
+ @DisplayName("should not call handler when input is blocked")
+ void shouldNotCallHandlerWhenInputBlocked() {
+ MCPCheckInputResponse blocked = new MCPCheckInputResponse(false, "Blocked", 1, null);
+ when(client.mcpCheckInput(anyString(), anyString(), any())).thenReturn(blocked);
+
+ MCPToolInterceptor interceptor = adapter.mcpToolInterceptor();
+ MCPToolRequest request = new MCPToolRequest("srv", "tool", null);
+
+ MCPToolHandler handler = req -> {
+ throw new AssertionError("Handler should not be called");
+ };
+
+ assertThatThrownBy(() -> interceptor.intercept(request, handler))
+ .isInstanceOf(PolicyViolationException.class);
+ }
+ }
+
+ // ========================================================================
+ // Helpers
+ // ========================================================================
+
+ private void mockCreateWorkflow(String workflowId) {
+ CreateWorkflowResponse resp = new CreateWorkflowResponse(
+ workflowId, "test-workflow", WorkflowSource.LANGGRAPH,
+ WorkflowStatus.IN_PROGRESS, Instant.now());
+ when(client.createWorkflow(any(CreateWorkflowRequest.class))).thenReturn(resp);
+ }
+
+ private void mockStepGate(GateDecision decision) {
+ mockStepGate(decision, null, null, Collections.emptyList());
+ }
+
+ private void mockStepGate(GateDecision decision, String stepId, String reason, List policyIds) {
+ StepGateResponse resp = new StepGateResponse(
+ decision, stepId, reason, policyIds, null, null, null);
+ when(client.stepGate(anyString(), anyString(), any(StepGateRequest.class))).thenReturn(resp);
+ }
+}