From 36ada1d5c3320a416eb5ef9368886a66cbe847e6 Mon Sep 17 00:00:00 2001 From: Saurabh Jain Date: Mon, 16 Mar 2026 14:52:05 +0530 Subject: [PATCH 1/4] feat: add LangGraph adapter with per-tool governance (#1243) Add LangGraphAdapter class with builder pattern and AutoCloseable for wrapping LangGraph workflows with AxonFlow governance gates. Provides per-tool governance via checkToolGate()/toolCompleted() and MCP policy enforcement via mcpToolInterceptor(). 12 source files, 47 tests. --- CHANGELOG.md | 17 + pom.xml | 1 + .../sdk/adapters/CheckGateOptions.java | 104 +++ .../sdk/adapters/CheckToolGateOptions.java | 102 ++ .../sdk/adapters/LangGraphAdapter.java | 564 +++++++++++ .../sdk/adapters/MCPInterceptorOptions.java | 94 ++ .../sdk/adapters/MCPToolHandler.java | 35 + .../sdk/adapters/MCPToolInterceptor.java | 134 +++ .../sdk/adapters/MCPToolRequest.java | 73 ++ .../sdk/adapters/StepCompletedOptions.java | 114 +++ .../sdk/adapters/ToolCompletedOptions.java | 114 +++ .../WorkflowApprovalRequiredError.java | 74 ++ .../sdk/adapters/WorkflowBlockedError.java | 76 ++ .../sdk/adapters/LangGraphAdapterTest.java | 884 ++++++++++++++++++ 14 files changed, 2386 insertions(+) create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/CheckGateOptions.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/CheckToolGateOptions.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/MCPInterceptorOptions.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/MCPToolHandler.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/MCPToolInterceptor.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/MCPToolRequest.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/StepCompletedOptions.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/ToolCompletedOptions.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/WorkflowApprovalRequiredError.java create mode 100644 src/main/java/com/getaxonflow/sdk/adapters/WorkflowBlockedError.java create mode 100644 src/test/java/com/getaxonflow/sdk/adapters/LangGraphAdapterTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 69fffdd..0d976a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,23 @@ All notable changes to the AxonFlow Java SDK will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- `LangGraphAdapter` class — wraps LangGraph workflows with AxonFlow governance gates and per-tool policy enforcement. Includes: + - `checkGate()` / `stepCompleted()` — step-level governance at LangGraph node boundaries + - `checkToolGate()` / `toolCompleted()` — per-tool governance within tool_call nodes (each tool gets its own gate check) + - `mcpToolInterceptor()` — factory returning an interceptor enforcing `mcpCheckInput → handler → mcpCheckOutput` around every MCP tool call + - `waitForApproval()` — poll until a step is approved or rejected + - `startWorkflow()` / `completeWorkflow()` / `abortWorkflow()` / `failWorkflow()` — workflow lifecycle management + - Builder pattern construction, implements `AutoCloseable` +- `WorkflowBlockedError` and `WorkflowApprovalRequiredError` exception classes +- Builder-based option classes: `CheckGateOptions`, `StepCompletedOptions`, `CheckToolGateOptions`, `ToolCompletedOptions` +- MCP interceptor types: `MCPInterceptorOptions`, `MCPToolRequest`, `MCPToolHandler`, `MCPToolInterceptor` + +--- + ## [4.1.0] - 2026-03-14 ### Added diff --git a/pom.xml b/pom.xml index 269ec0b..1241a76 100644 --- a/pom.xml +++ b/pom.xml @@ -178,6 +178,7 @@ maven-surefire-plugin ${maven-surefire-plugin.version} + @{argLine} -Dnet.bytebuddy.experimental=true **/*Test.java diff --git a/src/main/java/com/getaxonflow/sdk/adapters/CheckGateOptions.java b/src/main/java/com/getaxonflow/sdk/adapters/CheckGateOptions.java new file mode 100644 index 0000000..ea3352f --- /dev/null +++ b/src/main/java/com/getaxonflow/sdk/adapters/CheckGateOptions.java @@ -0,0 +1,104 @@ +/* + * 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.types.workflow.WorkflowTypes.ToolContext; + +import java.util.Map; + +/** + * Options for {@link LangGraphAdapter#checkGate}. + */ +public final class CheckGateOptions { + + private final String stepId; + private final Map stepInput; + private final String model; + private final String provider; + private final ToolContext toolContext; + + private CheckGateOptions(Builder builder) { + this.stepId = builder.stepId; + this.stepInput = builder.stepInput; + this.model = builder.model; + this.provider = builder.provider; + this.toolContext = builder.toolContext; + } + + public String getStepId() { + return stepId; + } + + public Map getStepInput() { + return stepInput; + } + + public String getModel() { + return model; + } + + public String getProvider() { + return provider; + } + + public ToolContext getToolContext() { + return toolContext; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private String stepId; + private Map stepInput; + private String model; + private String provider; + private ToolContext toolContext; + + private Builder() { + } + + public Builder stepId(String stepId) { + this.stepId = stepId; + return this; + } + + public Builder stepInput(Map stepInput) { + this.stepInput = stepInput; + return this; + } + + public Builder model(String model) { + this.model = model; + return this; + } + + public Builder provider(String provider) { + this.provider = provider; + return this; + } + + public Builder toolContext(ToolContext toolContext) { + this.toolContext = toolContext; + return this; + } + + public CheckGateOptions build() { + return new CheckGateOptions(this); + } + } +} diff --git a/src/main/java/com/getaxonflow/sdk/adapters/CheckToolGateOptions.java b/src/main/java/com/getaxonflow/sdk/adapters/CheckToolGateOptions.java new file mode 100644 index 0000000..82ad9a8 --- /dev/null +++ b/src/main/java/com/getaxonflow/sdk/adapters/CheckToolGateOptions.java @@ -0,0 +1,102 @@ +/* + * 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#checkToolGate}. + */ +public final class CheckToolGateOptions { + + private final String stepName; + private final String stepId; + private final Map toolInput; + private final String model; + private final String provider; + + private CheckToolGateOptions(Builder builder) { + this.stepName = builder.stepName; + this.stepId = builder.stepId; + this.toolInput = builder.toolInput; + this.model = builder.model; + this.provider = builder.provider; + } + + public String getStepName() { + return stepName; + } + + public String getStepId() { + return stepId; + } + + public Map getToolInput() { + return toolInput; + } + + public String getModel() { + return model; + } + + public String getProvider() { + return provider; + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + private String stepName; + private String stepId; + private Map toolInput; + private String model; + private String provider; + + private Builder() { + } + + public Builder stepName(String stepName) { + this.stepName = stepName; + return this; + } + + public Builder stepId(String stepId) { + this.stepId = stepId; + return this; + } + + public Builder toolInput(Map 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..46faca6 --- /dev/null +++ b/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java @@ -0,0 +1,564 @@ +/* + * 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. + * + *

"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 elapsed = 0; + while (elapsed < timeoutMs) { + 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); + elapsed += 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 = source; + 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: + *

    + *
  1. Derive connector type from the request
    2. + *
  2. Build a statement string and call {@code mcpCheckInput}
    3. + *
  3. If blocked, throw {@link PolicyViolationException}
    4. + *
  4. Execute the handler
    5. + *
  5. Call {@code mcpCheckOutput} with the result
    6. + *
  6. If blocked, throw {@link PolicyViolationException}
    7. + *
  7. If redacted data is present, return it instead of the raw result
    8. + *
+ * + * @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); + } +} From 205a027b0d5dcd869c9f031593ccd87f0ec36f95 Mon Sep 17 00:00:00 2001 From: Saurabh Jain Date: Mon, 16 Mar 2026 15:26:00 +0530 Subject: [PATCH 2/4] fix: wall-clock timeout, null source validation, thread-safety doc - waitForApproval now uses System.nanoTime() for accurate wall-clock timeout instead of accumulating sleep durations - Builder.source() rejects null with Objects.requireNonNull - Add thread-safety Javadoc: adapter is not thread-safe --- .../com/getaxonflow/sdk/adapters/LangGraphAdapter.java | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java b/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java index 46faca6..b974410 100644 --- a/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java +++ b/src/main/java/com/getaxonflow/sdk/adapters/LangGraphAdapter.java @@ -43,6 +43,9 @@ * 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: @@ -391,8 +394,8 @@ public boolean waitForApproval(String stepId, long pollIntervalMs, long timeoutM throws InterruptedException, TimeoutException { requireStarted(); - long elapsed = 0; - while (elapsed < timeoutMs) { + 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()) { @@ -410,7 +413,6 @@ public boolean waitForApproval(String stepId, long pollIntervalMs, long timeoutM } Thread.sleep(pollIntervalMs); - elapsed += pollIntervalMs; } throw new TimeoutException("Approval timeout after " + timeoutMs + "ms for step " + stepId); @@ -532,7 +534,7 @@ private Builder(AxonFlow client, String workflowName) { * @return this builder */ public Builder source(WorkflowSource source) { - this.source = source; + this.source = Objects.requireNonNull(source, "source cannot be null"); return this; } From 96813d010563222004a7fd9d110b1610fba4408e Mon Sep 17 00:00:00 2001 From: Saurabh Jain Date: Mon, 16 Mar 2026 15:55:29 +0530 Subject: [PATCH 3/4] =?UTF-8?q?chore:=20set=20release=20date=20v4.2.0=20?= =?UTF-8?q?=E2=80=94=202026-03-17?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d976a2..b763500 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ All notable changes to the AxonFlow Java SDK will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] +## [4.2.0] - 2026-03-17 ### Added From c70a5904560f911d37e2224855698e24d1e7575f Mon Sep 17 00:00:00 2001 From: Saurabh Jain Date: Mon, 16 Mar 2026 16:14:46 +0530 Subject: [PATCH 4/4] chore: bump version to 4.2.0 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 1241a76..95b2f2d 100644 --- a/pom.xml +++ b/pom.xml @@ -6,7 +6,7 @@ com.getaxonflow axonflow-sdk - 4.1.0 + 4.2.0 jar AxonFlow Java SDK