From eebbd9d4c5ecc4399cb0e1606c0c7f0752bd1a46 Mon Sep 17 00:00:00 2001 From: chengliyao <2789191580@qq.com> Date: Thu, 3 Sep 2026 11:06:29 +0800 Subject: [PATCH] feat: add Spring AI MCP tool adapter --- sentinel-adapter/pom.xml | 1 + .../sentinel-spring-ai-mcp-adapter/README.md | 22 ++ .../sentinel-spring-ai-mcp-adapter/pom.xml | 65 ++++ ...SentinelMcpToolSpecificationDecorator.java | 251 ++++++++++++++ ...inelMcpToolSpecificationDecoratorTest.java | 323 ++++++++++++++++++ 5 files changed, 662 insertions(+) create mode 100644 sentinel-adapter/sentinel-spring-ai-mcp-adapter/README.md create mode 100644 sentinel-adapter/sentinel-spring-ai-mcp-adapter/pom.xml create mode 100644 sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/main/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecorator.java create mode 100644 sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecoratorTest.java diff --git a/sentinel-adapter/pom.xml b/sentinel-adapter/pom.xml index b1b87ac0a9..7e5ae4ac75 100755 --- a/sentinel-adapter/pom.xml +++ b/sentinel-adapter/pom.xml @@ -36,6 +36,7 @@ sentinel-zuul2-adapter sentinel-okhttp-adapter sentinel-spring-restclient-adapter + sentinel-spring-ai-mcp-adapter sentinel-jax-rs-adapter sentinel-quarkus-adapter sentinel-motan-adapter diff --git a/sentinel-adapter/sentinel-spring-ai-mcp-adapter/README.md b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/README.md new file mode 100644 index 0000000000..b0729fdabc --- /dev/null +++ b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/README.md @@ -0,0 +1,22 @@ +# Sentinel Spring AI MCP Adapter + +This module adds tool-level Sentinel protection to MCP tools exposed by Spring AI. +It decorates MCP tool specifications instead of intercepting an HTTP transport, so the +same resource works with stdio, SSE, Streamable HTTP, WebMVC, and WebFlux servers. + +The default resource name is `mcp:tool:`. Stateful and stateless tool +specifications are supported in both synchronous and asynchronous forms. Asynchronous +entries remain active until completion, error, or cancellation. + +```java +SyncToolSpecification protectedTool = + SentinelMcpToolSpecificationDecorator.decorate(toolSpecification); +``` + +```java +AsyncToolSpecification protectedTool = + SentinelMcpToolSpecificationDecorator.decorate(toolSpecification); +``` + +The build baseline is Spring AI 2.0.1 and MCP Java SDK 2.0.0. The module requires +Java 17. diff --git a/sentinel-adapter/sentinel-spring-ai-mcp-adapter/pom.xml b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/pom.xml new file mode 100644 index 0000000000..4ab330ca53 --- /dev/null +++ b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/pom.xml @@ -0,0 +1,65 @@ + + + 4.0.0 + + + com.alibaba.csp + sentinel-adapter + ${revision} + ../pom.xml + + + sentinel-spring-ai-mcp-adapter + ${project.groupId}:${project.artifactId} + + + 17 + 17 + 2.0.1 + 3.7.12 + + + + + com.alibaba.csp + sentinel-core + + + org.springframework.ai + spring-ai-mcp + ${spring-ai.version} + provided + + + + junit + junit + test + + + org.mockito + mockito-inline + test + + + io.projectreactor + reactor-test + ${reactor.version} + test + + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + 17 + + + + + diff --git a/sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/main/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecorator.java b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/main/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecorator.java new file mode 100644 index 0000000000..301818a610 --- /dev/null +++ b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/main/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecorator.java @@ -0,0 +1,251 @@ +/* + * Copyright 1999-2026 Alibaba Group Holding Ltd. + * + * 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.alibaba.csp.sentinel.adapter.spring.ai.mcp; + +import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; + +import com.alibaba.csp.sentinel.AsyncEntry; +import com.alibaba.csp.sentinel.Entry; +import com.alibaba.csp.sentinel.EntryType; +import com.alibaba.csp.sentinel.ResourceTypeConstants; +import com.alibaba.csp.sentinel.SphU; +import com.alibaba.csp.sentinel.Tracer; +import com.alibaba.csp.sentinel.slots.block.BlockException; + +import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolSpecification; +import io.modelcontextprotocol.server.McpServerFeatures.SyncToolSpecification; +import io.modelcontextprotocol.server.McpStatelessServerFeatures; +import io.modelcontextprotocol.spec.McpSchema.CallToolResult; +import reactor.core.publisher.Mono; + +/** + * Decorates Spring AI MCP tool specifications with tool-level Sentinel protection. + * + * @author chengliyao + */ +public final class SentinelMcpToolSpecificationDecorator { + + public static final String DEFAULT_RESOURCE_PREFIX = "mcp:tool:"; + + private static final InvocationFactory DEFAULT_INVOCATION_FACTORY = new SentinelInvocationFactory(); + + /** + * Decorates a synchronous MCP tool specification. + * + * @param specification original tool specification + * @return protected tool specification + */ + public static SyncToolSpecification decorate(SyncToolSpecification specification) { + return decorate(specification, DEFAULT_INVOCATION_FACTORY); + } + + /** + * Decorates an asynchronous MCP tool specification. + * + * @param specification original tool specification + * @return protected tool specification + */ + public static AsyncToolSpecification decorate(AsyncToolSpecification specification) { + return decorate(specification, DEFAULT_INVOCATION_FACTORY); + } + + /** + * Decorates a synchronous stateless MCP tool specification. + * + * @param specification original tool specification + * @return protected tool specification + */ + public static McpStatelessServerFeatures.SyncToolSpecification decorate( + McpStatelessServerFeatures.SyncToolSpecification specification) { + return decorate(specification, DEFAULT_INVOCATION_FACTORY); + } + + /** + * Decorates an asynchronous stateless MCP tool specification. + * + * @param specification original tool specification + * @return protected tool specification + */ + public static McpStatelessServerFeatures.AsyncToolSpecification decorate( + McpStatelessServerFeatures.AsyncToolSpecification specification) { + return decorate(specification, DEFAULT_INVOCATION_FACTORY); + } + + static SyncToolSpecification decorate(SyncToolSpecification specification, InvocationFactory factory) { + Objects.requireNonNull(specification, "specification"); + Objects.requireNonNull(factory, "factory"); + String resourceName = resourceName(specification.tool().name()); + return SyncToolSpecification.builder().tool(specification.tool()) + .callHandler((exchange, arguments) -> protectSync(resourceName, factory, + () -> specification.callHandler().apply(exchange, arguments))) + .build(); + } + + static AsyncToolSpecification decorate(AsyncToolSpecification specification, InvocationFactory factory) { + Objects.requireNonNull(specification, "specification"); + Objects.requireNonNull(factory, "factory"); + String resourceName = resourceName(specification.tool().name()); + return AsyncToolSpecification.builder().tool(specification.tool()) + .callHandler((exchange, arguments) -> protectAsync(resourceName, factory, + () -> specification.callHandler().apply(exchange, arguments))) + .build(); + } + + static McpStatelessServerFeatures.SyncToolSpecification decorate( + McpStatelessServerFeatures.SyncToolSpecification specification, InvocationFactory factory) { + Objects.requireNonNull(specification, "specification"); + Objects.requireNonNull(factory, "factory"); + String resourceName = resourceName(specification.tool().name()); + return McpStatelessServerFeatures.SyncToolSpecification.builder().tool(specification.tool()) + .callHandler((context, arguments) -> protectSync(resourceName, factory, + () -> specification.callHandler().apply(context, arguments))) + .build(); + } + + static McpStatelessServerFeatures.AsyncToolSpecification decorate( + McpStatelessServerFeatures.AsyncToolSpecification specification, InvocationFactory factory) { + Objects.requireNonNull(specification, "specification"); + Objects.requireNonNull(factory, "factory"); + String resourceName = resourceName(specification.tool().name()); + return McpStatelessServerFeatures.AsyncToolSpecification.builder().tool(specification.tool()) + .callHandler((context, arguments) -> protectAsync(resourceName, factory, + () -> specification.callHandler().apply(context, arguments))) + .build(); + } + + private static CallToolResult protectSync(String resourceName, InvocationFactory factory, + Supplier toolCall) { + Invocation invocation; + try { + invocation = factory.enter(resourceName, false); + } catch (BlockException ex) { + throw ex.toRuntimeException(); + } + try { + return toolCall.get(); + } catch (Throwable throwable) { + invocation.trace(throwable); + throw throwable; + } finally { + invocation.exit(); + } + } + + private static Mono protectAsync(String resourceName, InvocationFactory factory, + Supplier> toolCall) { + return Mono.defer(() -> { + Invocation invocation; + try { + invocation = factory.enter(resourceName, true); + } catch (BlockException ex) { + return Mono.error(ex); + } + + try { + Mono result = Objects.requireNonNull(toolCall.get(), + "MCP async tool returned null"); + return result.doOnError(invocation::trace).doFinally(signalType -> invocation.exit()); + } catch (Throwable throwable) { + invocation.trace(throwable); + invocation.exit(); + return Mono.error(throwable); + } + }); + } + + private static String resourceName(String toolName) { + return DEFAULT_RESOURCE_PREFIX + Objects.requireNonNull(toolName, "toolName"); + } + + interface InvocationFactory { + Invocation enter(String resourceName, boolean async) throws BlockException; + } + + interface Invocation { + void trace(Throwable throwable); + + void exit(); + } + + private static final class SentinelInvocationFactory implements InvocationFactory { + + @Override + public Invocation enter(String resourceName, boolean async) throws BlockException { + if (async) { + AsyncEntry entry = SphU.asyncEntry(resourceName, ResourceTypeConstants.COMMON_RPC, EntryType.IN); + return new AsyncSentinelInvocation(entry); + } + Entry entry = SphU.entry(resourceName, ResourceTypeConstants.COMMON_RPC, EntryType.IN); + return new SyncSentinelInvocation(entry); + } + } + + private abstract static class AbstractSentinelInvocation implements Invocation { + + private final AtomicBoolean exited = new AtomicBoolean(); + + @Override + public final void exit() { + if (exited.compareAndSet(false, true)) { + doExit(); + } + } + + abstract void doExit(); + } + + private static final class SyncSentinelInvocation extends AbstractSentinelInvocation { + + private final Entry entry; + + private SyncSentinelInvocation(Entry entry) { + this.entry = entry; + } + + @Override + public void trace(Throwable throwable) { + Tracer.traceEntry(throwable, entry); + } + + @Override + void doExit() { + entry.exit(); + } + } + + private static final class AsyncSentinelInvocation extends AbstractSentinelInvocation { + + private final AsyncEntry entry; + + private AsyncSentinelInvocation(AsyncEntry entry) { + this.entry = entry; + } + + @Override + public void trace(Throwable throwable) { + Tracer.traceContext(throwable, entry.getAsyncContext()); + } + + @Override + void doExit() { + entry.exit(); + } + } + + private SentinelMcpToolSpecificationDecorator() {} +} diff --git a/sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecoratorTest.java b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecoratorTest.java new file mode 100644 index 0000000000..97cd177ea5 --- /dev/null +++ b/sentinel-adapter/sentinel-spring-ai-mcp-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/ai/mcp/SentinelMcpToolSpecificationDecoratorTest.java @@ -0,0 +1,323 @@ +/* + * Copyright 1999-2026 Alibaba Group Holding Ltd. + * + * 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.alibaba.csp.sentinel.adapter.spring.ai.mcp; + +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; + +import com.alibaba.csp.sentinel.EntryType; +import com.alibaba.csp.sentinel.adapter.spring.ai.mcp.SentinelMcpToolSpecificationDecorator.Invocation; +import com.alibaba.csp.sentinel.adapter.spring.ai.mcp.SentinelMcpToolSpecificationDecorator.InvocationFactory; +import com.alibaba.csp.sentinel.node.ClusterNode; +import com.alibaba.csp.sentinel.slots.block.BlockException; +import com.alibaba.csp.sentinel.slots.block.RuleConstant; +import com.alibaba.csp.sentinel.slots.block.flow.FlowException; +import com.alibaba.csp.sentinel.slots.block.flow.FlowRule; +import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager; +import com.alibaba.csp.sentinel.slots.clusterbuilder.ClusterBuilderSlot; + +import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolSpecification; +import io.modelcontextprotocol.server.McpServerFeatures.SyncToolSpecification; +import io.modelcontextprotocol.server.McpStatelessServerFeatures; +import io.modelcontextprotocol.spec.McpSchema.CallToolRequest; +import io.modelcontextprotocol.spec.McpSchema.CallToolResult; +import io.modelcontextprotocol.spec.McpSchema.Tool; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class SentinelMcpToolSpecificationDecoratorTest { + + private static final CallToolResult RESULT = CallToolResult.builder() + .content(Collections.emptyList()).isError(false).build(); + private static final CallToolRequest REQUEST = CallToolRequest.builder("echo") + .arguments(Collections.emptyMap()).build(); + + private Tool tool; + + @Before + public void setUp() { + FlowRuleManager.loadRules(null); + ClusterBuilderSlot.getClusterNodeMap().clear(); + tool = mock(Tool.class); + when(tool.name()).thenReturn("echo"); + } + + @After + public void tearDown() { + FlowRuleManager.loadRules(null); + ClusterBuilderSlot.getClusterNodeMap().clear(); + } + + @Test + public void testSyncSuccess() { + RecordingFactory factory = new RecordingFactory(); + SyncToolSpecification original = syncTool((exchange, arguments) -> RESULT); + + CallToolResult actual = SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST); + + assertSame(RESULT, actual); + assertEquals("mcp:tool:echo", factory.resourceName); + assertEquals(1, factory.invocation.exitCount.get()); + assertEquals(0, factory.invocation.traceCount.get()); + } + + @Test + public void testSyncExceptionIsTracedAndEntryExits() { + RecordingFactory factory = new RecordingFactory(); + IllegalStateException failure = new IllegalStateException("boom"); + SyncToolSpecification original = syncTool((exchange, arguments) -> { + throw failure; + }); + + IllegalStateException actual = expectThrows(IllegalStateException.class, () -> + SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)); + + assertSame(failure, actual); + assertSame(failure, factory.invocation.lastFailure); + assertEquals(1, factory.invocation.traceCount.get()); + assertEquals(1, factory.invocation.exitCount.get()); + } + + @Test + public void testSyncBlockSkipsToolCall() { + RecordingFactory factory = new RecordingFactory(); + factory.blocked = true; + AtomicInteger calls = new AtomicInteger(); + SyncToolSpecification original = syncTool((exchange, arguments) -> { + calls.incrementAndGet(); + return RESULT; + }); + + RuntimeException actual = expectThrows(RuntimeException.class, () -> + SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)); + + assertTrue(BlockException.isBlockException(actual)); + assertEquals(0, calls.get()); + assertEquals(0, factory.invocation.exitCount.get()); + } + + @Test + public void testRealSentinelRuleBlocksSyncTool() { + String resourceName = "mcp:tool:echo"; + FlowRule rule = new FlowRule().setCount(0).setGrade(RuleConstant.FLOW_GRADE_QPS) + .setResource(resourceName).as(FlowRule.class); + FlowRuleManager.loadRules(Collections.singletonList(rule)); + SyncToolSpecification original = syncTool((exchange, arguments) -> RESULT); + + RuntimeException actual = expectThrows(RuntimeException.class, () -> + SentinelMcpToolSpecificationDecorator.decorate(original).callHandler().apply(null, REQUEST)); + + assertTrue(BlockException.isBlockException(actual)); + ClusterNode node = ClusterBuilderSlot.getClusterNode(resourceName, EntryType.IN); + assertEquals(1, node.blockRequest()); + } + + @Test + public void testAsyncCompletionExitsOnce() { + RecordingFactory factory = new RecordingFactory(); + AsyncToolSpecification original = asyncTool((exchange, arguments) -> Mono.just(RESULT)); + + StepVerifier.create(SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)) + .expectNext(RESULT) + .verifyComplete(); + + assertEquals(1, factory.invocation.exitCount.get()); + assertEquals(0, factory.invocation.traceCount.get()); + } + + @Test + public void testAsyncErrorIsTracedAndEntryExits() { + RecordingFactory factory = new RecordingFactory(); + IllegalStateException failure = new IllegalStateException("boom"); + AsyncToolSpecification original = asyncTool((exchange, arguments) -> Mono.error(failure)); + + StepVerifier.create(SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)) + .expectErrorMatches(error -> error == failure) + .verify(); + + assertSame(failure, factory.invocation.lastFailure); + assertEquals(1, factory.invocation.traceCount.get()); + assertEquals(1, factory.invocation.exitCount.get()); + } + + @Test + public void testAsyncCancellationExitsOnce() { + RecordingFactory factory = new RecordingFactory(); + AsyncToolSpecification original = asyncTool((exchange, arguments) -> Mono.never()); + + StepVerifier.create(SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)) + .thenCancel() + .verify(); + + assertEquals(1, factory.invocation.exitCount.get()); + assertEquals(0, factory.invocation.traceCount.get()); + } + + @Test + public void testAsyncBlockSkipsToolCall() { + RecordingFactory factory = new RecordingFactory(); + factory.blocked = true; + AtomicInteger calls = new AtomicInteger(); + AsyncToolSpecification original = asyncTool((exchange, arguments) -> { + calls.incrementAndGet(); + return Mono.just(RESULT); + }); + + StepVerifier.create(SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)) + .expectError(FlowException.class) + .verify(); + + assertEquals(0, calls.get()); + assertEquals(0, factory.invocation.exitCount.get()); + } + + @Test + public void testAsyncImmediateExceptionIsTracedAndEntryExits() { + RecordingFactory factory = new RecordingFactory(); + IllegalStateException failure = new IllegalStateException("boom"); + AsyncToolSpecification original = asyncTool((exchange, arguments) -> { + throw failure; + }); + + StepVerifier.create(SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)) + .expectErrorMatches(error -> error == failure) + .verify(); + + assertSame(failure, factory.invocation.lastFailure); + assertEquals(1, factory.invocation.traceCount.get()); + assertEquals(1, factory.invocation.exitCount.get()); + } + + @Test + public void testRealSentinelRuleBlocksAsyncTool() { + String resourceName = "mcp:tool:echo"; + FlowRule rule = new FlowRule().setCount(0).setGrade(RuleConstant.FLOW_GRADE_QPS) + .setResource(resourceName).as(FlowRule.class); + FlowRuleManager.loadRules(Collections.singletonList(rule)); + AsyncToolSpecification original = asyncTool((exchange, arguments) -> Mono.just(RESULT)); + + StepVerifier.create(SentinelMcpToolSpecificationDecorator.decorate(original) + .callHandler().apply(null, REQUEST)) + .expectErrorMatches(BlockException::isBlockException) + .verify(); + + } + + @Test + public void testStatelessSyncToolIsProtected() { + RecordingFactory factory = new RecordingFactory(); + McpStatelessServerFeatures.SyncToolSpecification original = + McpStatelessServerFeatures.SyncToolSpecification.builder().tool(tool) + .callHandler((context, arguments) -> RESULT).build(); + + CallToolResult actual = SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST); + + assertSame(RESULT, actual); + assertEquals("mcp:tool:echo", factory.resourceName); + assertEquals(1, factory.invocation.exitCount.get()); + } + + @Test + public void testStatelessAsyncToolIsProtected() { + RecordingFactory factory = new RecordingFactory(); + McpStatelessServerFeatures.AsyncToolSpecification original = + McpStatelessServerFeatures.AsyncToolSpecification.builder().tool(tool) + .callHandler((context, arguments) -> Mono.just(RESULT)).build(); + + StepVerifier.create(SentinelMcpToolSpecificationDecorator.decorate(original, factory) + .callHandler().apply(null, REQUEST)) + .expectNext(RESULT) + .verifyComplete(); + + assertEquals("mcp:tool:echo", factory.resourceName); + assertEquals(1, factory.invocation.exitCount.get()); + } + + private static T expectThrows(Class type, Runnable action) { + try { + action.run(); + } catch (Throwable throwable) { + assertTrue("Expected " + type.getName() + " but got " + throwable.getClass().getName(), + type.isInstance(throwable)); + return type.cast(throwable); + } + throw new AssertionError("Expected " + type.getName() + " to be thrown"); + } + + private SyncToolSpecification syncTool( + java.util.function.BiFunction handler) { + return SyncToolSpecification.builder().tool(tool).callHandler(handler).build(); + } + + private AsyncToolSpecification asyncTool( + java.util.function.BiFunction> handler) { + return AsyncToolSpecification.builder().tool(tool).callHandler(handler).build(); + } + + private static final class RecordingFactory implements InvocationFactory { + + private final RecordingInvocation invocation = new RecordingInvocation(); + private String resourceName; + private boolean blocked; + + @Override + public Invocation enter(String resourceName, boolean async) throws FlowException { + this.resourceName = resourceName; + if (blocked) { + throw new FlowException(resourceName); + } + return invocation; + } + } + + private static final class RecordingInvocation implements Invocation { + + private final AtomicInteger traceCount = new AtomicInteger(); + private final AtomicInteger exitCount = new AtomicInteger(); + private Throwable lastFailure; + + @Override + public void trace(Throwable throwable) { + lastFailure = throwable; + traceCount.incrementAndGet(); + } + + @Override + public void exit() { + exitCount.incrementAndGet(); + } + } +}