From c8b0f48c8dd1505db17e785a8ea8bb58ca9c3116 Mon Sep 17 00:00:00 2001 From: hanyu0504 Date: Wed, 23 Sep 2026 20:47:09 +0800 Subject: [PATCH] fix(engine): keep active notification streams alive and release them reliably Signed-off-by: hanyu0504 --- .../workflow/engine/client/A2ATransport.java | 1 + .../client/DefaultA2AJavaClientRuntime.java | 134 ++++++++++++++--- .../engine/client/A2ATransportHeaderTest.java | 39 +++++ .../client/NotificationStreamTimeoutTest.java | 138 ++++++++++++++++++ 4 files changed, 289 insertions(+), 23 deletions(-) create mode 100644 workflow-engine/src/test/java/dev/openan/workflow/engine/client/NotificationStreamTimeoutTest.java diff --git a/workflow-engine/src/main/java/dev/openan/workflow/engine/client/A2ATransport.java b/workflow-engine/src/main/java/dev/openan/workflow/engine/client/A2ATransport.java index f939274..c250777 100644 --- a/workflow-engine/src/main/java/dev/openan/workflow/engine/client/A2ATransport.java +++ b/workflow-engine/src/main/java/dev/openan/workflow/engine/client/A2ATransport.java @@ -661,6 +661,7 @@ public NotificationSubscription openNotificationStream( subscription.completeStream(); } } finally { + closeConversation(agentCard, contextId); notificationSubscriptions.remove(subscription); subscription.markStreamTerminated(); } diff --git a/workflow-engine/src/main/java/dev/openan/workflow/engine/client/DefaultA2AJavaClientRuntime.java b/workflow-engine/src/main/java/dev/openan/workflow/engine/client/DefaultA2AJavaClientRuntime.java index e8d5403..eb3feb9 100644 --- a/workflow-engine/src/main/java/dev/openan/workflow/engine/client/DefaultA2AJavaClientRuntime.java +++ b/workflow-engine/src/main/java/dev/openan/workflow/engine/client/DefaultA2AJavaClientRuntime.java @@ -35,6 +35,7 @@ import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import org.a2aproject.sdk.client.Client; @@ -185,8 +186,15 @@ private static void onEvent( List events, AtomicReference lastEventRef, Consumer eventSink, - CountDownLatch done) { - events.add(event); + CountDownLatch done, + AtomicLong lastActivityNanos, + AtomicLong eventCount) { + if (lastActivityNanos != null) { + lastActivityNanos.set(System.nanoTime()); + eventCount.incrementAndGet(); + } else { + events.add(event); + } lastEventRef.set(event); logEvent(agentName, event); ProtocolLogger.logResponseEvent(agentName, event); @@ -330,9 +338,21 @@ public Iterable sendMessage( ? null : callContext.getState().get(A2AJavaClientRuntime.TRANSPORT_ACTIVITY_STATE_KEY); Runnable activityListener = configuredActivity instanceof Runnable runnable ? runnable : null; + boolean notificationStream = isNotificationStream(callContext); + AtomicLong lastActivityNanos = notificationStream ? new AtomicLong(System.nanoTime()) : null; + AtomicLong eventCount = notificationStream ? new AtomicLong() : null; + Runnable observedActivity = + notificationStream + ? () -> { + lastActivityNanos.set(System.nanoTime()); + if (activityListener != null) activityListener.run(); + } + : activityListener; return TransportActivityMonitor.call( - activityListener, - () -> sendMessageObserved(agentCard, params, callContext, eventSink, logSink)); + observedActivity, + () -> + sendMessageObserved( + agentCard, params, callContext, eventSink, logSink, lastActivityNanos, eventCount)); } private Iterable sendMessageObserved( @@ -340,7 +360,9 @@ private Iterable sendMessageObserved( org.a2aproject.sdk.spec.MessageSendParams params, ClientCallContext callContext, Consumer eventSink, - Consumer logSink) { + Consumer logSink, + AtomicLong lastActivityNanos, + AtomicLong eventCount) { if (closed.get()) throw new IllegalStateException("A2A client runtime is closed"); String agentUrl = extractAgentUrl(agentCard); Client client = @@ -356,27 +378,55 @@ private Iterable sendMessageObserved( if (logSink != null) logSink.accept("[A2A] Sending message to " + agentCard.name()); try { - client.sendMessage( - params, - List.of( - (event, card) -> - onEvent(agentCard.name(), event, events, lastEventRef, eventSink, done)), - error -> onError(agentCard.name(), error, done, errorRef), - callContext); - } catch (A2AClientException e) { - throw new RuntimeException( - "A2A message:send failed for " + agentCard.name() + ": " + e.getMessage(), e); - } + try { + client.sendMessage( + params, + List.of( + (event, card) -> + onEvent( + agentCard.name(), + event, + events, + lastEventRef, + eventSink, + done, + lastActivityNanos, + eventCount)), + error -> onError(agentCard.name(), error, done, errorRef), + callContext); + } catch (A2AClientException e) { + throw new RuntimeException( + "A2A message:send failed for " + agentCard.name() + ": " + e.getMessage(), e); + } - awaitCompletion(agentCard.name(), done, events, lastEventRef); + if (lastActivityNanos != null) { + awaitNotificationActivity(agentCard.name(), done, lastActivityNanos::get); + } else { + awaitCompletion(agentCard.name(), done, events, lastEventRef); + } - if (errorRef.get() != null) { - throw new RuntimeException( - "A2A message:send failed for " + agentCard.name() + ": " + errorRef.get().getMessage(), - errorRef.get()); + if (errorRef.get() != null) { + throw new RuntimeException( + "A2A message:send failed for " + agentCard.name() + ": " + errorRef.get().getMessage(), + errorRef.get()); + } + log.info( + "[A2ARuntime] Completed for '{}': {} event(s)", + agentCard.name(), + eventCount != null ? eventCount.get() : events.size()); + return events; + } finally { + if (lastActivityNanos != null && params.message() != null) { + try { + closeConversation(agentCard, params.message().contextId()); + } catch (RuntimeException closeError) { + log.warn( + "[A2ARuntime] Failed to close Notification-T transport for '{}'", + agentCard.name(), + closeError); + } + } } - log.info("[A2ARuntime] Completed for '{}': {} event(s)", agentCard.name(), events.size()); - return events; } private Client getOrCreateClient(AgentCard agentCard, String agentUrl) { @@ -526,6 +576,44 @@ private void awaitCompletion( } } + /** + * Waits for a Notification-T stream to end without expecting task finality. + * + *

Subscription tasks stay {@code WORKING} for their whole lifetime and deliver events + * (heartbeats included) without ever completing, so a finality-based wait would tear down a + * healthy subscription once {@code sendTimeoutSeconds} elapsed. The stream instead ends through + * the error/terminal-event latch: a server-side close, a transport failure, or {@link + * dev.openan.workflow.engine.client.NotificationSubscription#close()}. The stream is only reaped + * when no transport activity or decoded event has arrived for {@code sendTimeoutSeconds}. + * + * @param agentName target agent display name, for diagnostics only + * @param done latch released on terminal event or stream error + * @param lastActivityNanos monotonic time of the most recent stream activity + */ + void awaitNotificationActivity( + String agentName, CountDownLatch done, java.util.function.LongSupplier lastActivityNanos) { + long idleBudgetNanos = TimeUnit.SECONDS.toNanos(sendTimeoutSeconds); + while (true) { + if (done.getCount() == 0) return; + long idleNanos = System.nanoTime() - lastActivityNanos.getAsLong(); + long remainingNanos = idleBudgetNanos - idleNanos; + if (remainingNanos <= 0) { + log.error( + "[A2ARuntime] TIMEOUT for '{}' after {}s without Notification-T activity", + agentName, + sendTimeoutSeconds); + throw new RuntimeException( + "A2A notification stream idle for " + sendTimeoutSeconds + "s for " + agentName); + } + try { + if (done.await(remainingNanos, TimeUnit.NANOSECONDS)) return; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("A2A notification stream interrupted for " + agentName, e); + } + } + } + @Override public Task getTask(AgentCard agentCard, String taskId, ClientCallContext callContext) { if (closed.get()) throw new IllegalStateException("A2A client runtime is closed"); diff --git a/workflow-engine/src/test/java/dev/openan/workflow/engine/client/A2ATransportHeaderTest.java b/workflow-engine/src/test/java/dev/openan/workflow/engine/client/A2ATransportHeaderTest.java index 1f46f10..0ac593e 100644 --- a/workflow-engine/src/test/java/dev/openan/workflow/engine/client/A2ATransportHeaderTest.java +++ b/workflow-engine/src/test/java/dev/openan/workflow/engine/client/A2ATransportHeaderTest.java @@ -345,6 +345,45 @@ public void close() {} assertTrue(exited.get(), "close must wait for the notification call to terminate"); } + @Test + void failedNotificationCallReleasesConversation() throws Exception { + AgentCard card = agentCard(); + AtomicInteger closes = new AtomicInteger(); + class FailingRuntime implements A2AJavaClientRuntime, ConversationScopedA2AJavaClientRuntime { + @Override + public Iterable sendMessage( + AgentCard agentCard, + MessageSendParams params, + ClientCallContext callContext, + Consumer eventSink, + Consumer logSink) { + throw new IllegalStateException("stream failed"); + } + + @Override + public void closeConversation(AgentCard agentCard, String contextId) { + assertEquals("context-failed", contextId); + closes.incrementAndGet(); + } + + @Override + public void close() {} + } + A2ATransport transport = + new A2ATransport( + List.of(card), new FailingRuntime(), WorkflowEngineClientConfig.builder().build()); + try { + NotificationSubscription subscription = + transport.openNotificationStream( + card, card.name(), content("subscribe", Map.of()), "context-failed", null); + + assertThrows(CompletionException.class, () -> subscription.completion().join()); + assertEquals(1, closes.get()); + } finally { + transport.close(); + } + } + @Test void borrowedClientDoesNotCloseCallerOwnedTransport() throws Exception { AtomicInteger closeCalls = new AtomicInteger(); diff --git a/workflow-engine/src/test/java/dev/openan/workflow/engine/client/NotificationStreamTimeoutTest.java b/workflow-engine/src/test/java/dev/openan/workflow/engine/client/NotificationStreamTimeoutTest.java new file mode 100644 index 0000000..a697266 --- /dev/null +++ b/workflow-engine/src/test/java/dev/openan/workflow/engine/client/NotificationStreamTimeoutTest.java @@ -0,0 +1,138 @@ +/* + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * All Rights Reserved. + * + * SPDX-License-Identifier: Apache-2.0 + * + * 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 dev.openan.workflow.engine.client; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import org.junit.jupiter.api.Test; + +/** + * Notification-T streams must not be reaped by the finality-based send timeout: subscription tasks + * stay WORKING forever, and periodic heartbeats are what keep the stream alive. + */ +class NotificationStreamTimeoutTest { + + private DefaultA2AJavaClientRuntime runtimeWithOneSecondIdleBudget() { + return new DefaultA2AJavaClientRuntime(true, null, 1L, null); + } + + @Test + void heartbeatsKeepNotificationStreamAliveBeyondSendTimeout() throws Exception { + DefaultA2AJavaClientRuntime runtime = runtimeWithOneSecondIdleBudget(); + CountDownLatch done = new CountDownLatch(1); + AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime()); + + // The stream outlives two 1s idle budgets while activity arrives every 300ms. + Thread feeder = + new Thread( + () -> { + try { + for (int i = 0; i < 9; i++) { + lastActivityNanos.set(System.nanoTime()); + Thread.sleep(300); + } + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } finally { + done.countDown(); + } + }); + feeder.start(); + runtime.awaitNotificationActivity("agent", done, lastActivityNanos::get); + feeder.join(TimeUnit.SECONDS.toMillis(4)); + + assertTrue(done.await(0, TimeUnit.MILLISECONDS), "stream should have ended"); + } + + @Test + void silentNotificationStreamIsReapedWithIdleTimeout() { + DefaultA2AJavaClientRuntime runtime = runtimeWithOneSecondIdleBudget(); + CountDownLatch done = new CountDownLatch(1); + AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime()); + + RuntimeException error = + assertThrows( + RuntimeException.class, + () -> runtime.awaitNotificationActivity("agent", done, lastActivityNanos::get)); + assertTrue(error.getMessage().contains("idle for 1s"), error.getMessage()); + } + + @Test + void idleDeadlineIsMeasuredFromLastActivity() throws Exception { + DefaultA2AJavaClientRuntime runtime = runtimeWithOneSecondIdleBudget(); + CountDownLatch done = new CountDownLatch(1); + AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime()); + Thread heartbeat = + new Thread( + () -> { + try { + Thread.sleep(500); + lastActivityNanos.set(System.nanoTime()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + heartbeat.start(); + long started = System.nanoTime(); + + assertThrows( + RuntimeException.class, + () -> runtime.awaitNotificationActivity("agent", done, lastActivityNanos::get)); + long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started); + heartbeat.join(); + + assertTrue(elapsedMillis >= 1_350, "must not expire at the initial 1s deadline"); + assertTrue(elapsedMillis < 2_000, "must not wait for a second fixed 1s window"); + } + + @Test + void alreadyIdleStreamTimesOutWithoutAnotherFullWait() { + DefaultA2AJavaClientRuntime runtime = runtimeWithOneSecondIdleBudget(); + long staleActivity = System.nanoTime() - TimeUnit.SECONDS.toNanos(2); + long started = System.nanoTime(); + + assertThrows( + RuntimeException.class, + () -> + runtime.awaitNotificationActivity("agent", new CountDownLatch(1), () -> staleActivity)); + + assertTrue( + TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started) < 500, + "an already idle stream should not receive a fresh timeout budget"); + } + + @Test + void terminalEventEndsNotificationStreamImmediately() throws Exception { + DefaultA2AJavaClientRuntime runtime = runtimeWithOneSecondIdleBudget(); + CountDownLatch done = new CountDownLatch(1); + done.countDown(); + long started = System.nanoTime(); + + runtime.awaitNotificationActivity("agent", done, () -> System.nanoTime()); + + assertTrue( + TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started) < 500, + "terminal stream should return immediately"); + } +}