fix(engine): keep active notification streams alive and release them reliably

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();
                       }

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<ClientEvent> events,
       AtomicReference<ClientEvent> lastEventRef,
       Consumer<ClientEvent> 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,17 +338,31 @@ public Iterable<ClientEvent> 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<ClientEvent> sendMessageObserved(
       AgentCard agentCard,
       org.a2aproject.sdk.spec.MessageSendParams params,
       ClientCallContext callContext,
       Consumer<ClientEvent> eventSink,
-      Consumer<String> logSink) {
+      Consumer<String> 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<ClientEvent> 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.
+   *
+   * <p>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");

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<ClientEvent> sendMessage(
+          AgentCard agentCard,
+          MessageSendParams params,
+          ClientCallContext callContext,
+          Consumer<ClientEvent> eventSink,
+          Consumer<String> 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();

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");
+  }
+}