diff --git a/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AAutoConfiguration.java b/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AAutoConfiguration.java index 8041024..ebba756 100644 --- a/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AAutoConfiguration.java +++ b/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AAutoConfiguration.java @@ -234,8 +234,12 @@ public RestHandler restHandler(AgentCard card, RequestHandler handler, ExecutorS @Bean @ConditionalOnMissingBean public A2AController a2aController( - RestHandler restHandler, RequestHandler requestHandler, AgentCard agentCard) { - return new A2AController(restHandler, requestHandler, agentCard); + RestHandler restHandler, + RequestHandler requestHandler, + AgentCard agentCard, + A2AProperties properties) { + return new A2AController( + restHandler, requestHandler, agentCard, properties.getHeartbeatIntervalSeconds() * 1000L); } /** diff --git a/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AController.java b/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AController.java index fa89296..3c61ec1 100644 --- a/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AController.java +++ b/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AController.java @@ -26,7 +26,10 @@ import java.util.Locale; import java.util.Map; import java.util.Set; +import java.util.concurrent.Executors; import java.util.concurrent.Flow; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -95,12 +98,40 @@ public class A2AController { private final AgentCard agentCard; private final AtomicInteger activeStreams = new AtomicInteger(); private final Object streamCompletionMonitor = new Object(); + private final long heartbeatIntervalMillis; + private final ScheduledExecutorService heartbeatScheduler; public A2AController( RestHandler restHandler, RequestHandler requestHandler, AgentCard agentCard) { + this(restHandler, requestHandler, agentCard, 0L); + } + + public A2AController( + RestHandler restHandler, + RequestHandler requestHandler, + AgentCard agentCard, + long heartbeatIntervalMillis) { this.restHandler = restHandler; this.requestHandler = requestHandler; this.agentCard = agentCard; + this.heartbeatIntervalMillis = heartbeatIntervalMillis; + this.heartbeatScheduler = + heartbeatIntervalMillis > 0 + ? Executors.newSingleThreadScheduledExecutor( + runnable -> { + Thread thread = new Thread(runnable, "a2a-sse-heartbeat"); + thread.setDaemon(true); + return thread; + }) + : null; + } + + /** Releases the heartbeat scheduler; called by Spring on context shutdown. */ + @jakarta.annotation.PreDestroy + public void shutdown() { + if (heartbeatScheduler != null) { + heartbeatScheduler.shutdownNow(); + } } @PostMapping("${a2at.server.path-prefix}/message:send") @@ -158,9 +189,14 @@ private SseEmitter subscribeToEmitter( Flow.Publisher publisher, ServerCallContext context) { AtomicReference subscription = new AtomicReference<>(); AtomicBoolean released = new AtomicBoolean(); + AtomicReference> heartbeat = new AtomicReference<>(); Runnable release = () -> { if (released.compareAndSet(false, true)) { + ScheduledFuture activeHeartbeat = heartbeat.get(); + if (activeHeartbeat != null) { + activeHeartbeat.cancel(false); + } Flow.Subscription active = subscription.get(); if (active != null) { try { @@ -179,6 +215,20 @@ private SseEmitter subscribeToEmitter( }; SseEmitter emitter = new SseEmitter(0L); Runnable finish = track(emitter, release); + if (heartbeatIntervalMillis > 0 && heartbeatScheduler != null) { + heartbeat.set( + heartbeatScheduler.scheduleAtFixedRate( + () -> { + try { + emitter.send(SseEmitter.event().comment("heartbeat")); + } catch (Exception e) { + release.run(); + } + }, + heartbeatIntervalMillis, + heartbeatIntervalMillis, + TimeUnit.MILLISECONDS)); + } final AtomicLong seq = new AtomicLong(0); Flow.Subscriber subscriber = new Flow.Subscriber<>() { diff --git a/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AProperties.java b/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AProperties.java index df91261..5a5c2ce 100644 --- a/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AProperties.java +++ b/spring-boot-starter/src/main/java/dev/openan/workflow/engine/spring/A2AProperties.java @@ -66,11 +66,23 @@ public class A2AProperties { /** Idle timeout for server executor threads above the core size. */ private int executorKeepAliveSeconds = 60; + /** + * Interval in seconds for SSE heartbeat comments on streaming endpoints. Keeps intermediaries + * (gateways, load balancers) from idle-timing-out long-running SSE connections during agent + * execution. Set to 0 to disable. Default: 15 seconds. + */ + private int heartbeatIntervalSeconds = 15; + private static int positive(int value, String name) { if (value <= 0) throw new IllegalArgumentException(name + " must be positive"); return value; } + private static int nonNegative(int value, String name) { + if (value < 0) throw new IllegalArgumentException(name + " must be non-negative"); + return value; + } + public boolean isEnabled() { return enabled; } @@ -160,4 +172,13 @@ public int getExecutorKeepAliveSeconds() { public void setExecutorKeepAliveSeconds(int executorKeepAliveSeconds) { this.executorKeepAliveSeconds = positive(executorKeepAliveSeconds, "executorKeepAliveSeconds"); } + + public int getHeartbeatIntervalSeconds() { + return heartbeatIntervalSeconds; + } + + public void setHeartbeatIntervalSeconds(int heartbeatIntervalSeconds) { + this.heartbeatIntervalSeconds = + nonNegative(heartbeatIntervalSeconds, "heartbeatIntervalSeconds"); + } } diff --git a/spring-boot-starter/src/test/java/dev/openan/workflow/engine/spring/A2AAutoConfigurationTest.java b/spring-boot-starter/src/test/java/dev/openan/workflow/engine/spring/A2AAutoConfigurationTest.java index d5d9dc2..5200403 100644 --- a/spring-boot-starter/src/test/java/dev/openan/workflow/engine/spring/A2AAutoConfigurationTest.java +++ b/spring-boot-starter/src/test/java/dev/openan/workflow/engine/spring/A2AAutoConfigurationTest.java @@ -622,4 +622,81 @@ void requestHandlerIsCreatedWithoutAnAuthorizationProvider() { .withPropertyValues("a2at.server.enabled=true") .run(context -> assertTrue(context.containsBean("requestHandler"))); } + + @Test + void heartbeatIntervalDefaultsToFifteenSeconds() { + A2AProperties properties = new A2AProperties(); + assertEquals(15, properties.getHeartbeatIntervalSeconds()); + } + + @Test + void heartbeatIntervalZeroDisablesHeartbeat() { + A2AProperties properties = new A2AProperties(); + properties.setHeartbeatIntervalSeconds(0); + assertEquals(0, properties.getHeartbeatIntervalSeconds()); + } + + @Test + void heartbeatIntervalRejectsNegative() { + A2AProperties properties = new A2AProperties(); + assertThrows(IllegalArgumentException.class, () -> properties.setHeartbeatIntervalSeconds(-1)); + } + + @Test + void heartbeatCommentsSentDuringLongRunningStream() throws Exception { + RestHandler restHandler = mock(RestHandler.class); + RequestHandler requestHandler = mock(RequestHandler.class); + AgentCard agentCard = mock(AgentCard.class); + when(agentCard.capabilities()).thenReturn(AgentCapabilities.builder().streaming(true).build()); + var agentInterface = mock(org.a2aproject.sdk.spec.AgentInterface.class); + when(agentInterface.protocolVersion()).thenReturn("1.0"); + when(agentCard.supportedInterfaces()).thenReturn(java.util.List.of(agentInterface)); + + AtomicReference> subscriberRef = + new AtomicReference<>(); + Flow.Publisher publisher = + s -> { + subscriberRef.set(s); + s.onSubscribe( + new Flow.Subscription() { + @Override + public void request(long n) {} + + @Override + public void cancel() {} + }); + }; + when(requestHandler.onMessageSendStream(any(), any())).thenReturn(publisher); + + A2AController controller = new A2AController(restHandler, requestHandler, agentCard, 200L); + try { + MockMvc mockMvc = + MockMvcBuilders.standaloneSetup(controller) + .addPlaceholderValue("a2at.server.path-prefix", "/a2a/json") + .build(); + + MvcResult result = + mockMvc + .perform( + post("/a2a/json/message:stream") + .header("A2A-Version", "1.0") + .contentType("application/json") + .accept("text/event-stream") + .content( + "{\"message\":{\"messageId\":\"m1\",\"role\":\"ROLE_USER\"," + + "\"parts\":[{\"text\":\"test\"}]}}")) + .andExpect(request().asyncStarted()) + .andReturn(); + + Thread.sleep(350); + subscriberRef.get().onComplete(); + mockMvc.perform(asyncDispatch(result)).andExpect(status().isOk()); + + // MockMvc's mock response may not capture cross-thread SSE comment writes; the + // wire-level heartbeat verification is covered by E2E tests. Here we verify the + // stream stays alive through multiple heartbeat intervals and completes cleanly. + } finally { + controller.shutdown(); + } + } }