fix: prevent premature stream closure in EventConsumer grace period

The EventConsumer grace period logic could close streams prematurely when
final events were still in-transit through MainEventBusProcessor. This
manifested as intermittent test failures where the stream would close
before all events were delivered.

Root Cause:
- When agent execution completes, EventConsumer enters a grace period
- It polls the ChildQueue with 500ms timeout, allowing 3 consecutive
  timeouts (1.5s total) before closing the stream
- The original logic only checked queue.size() == 0
- However, final events can be in-transit: MainQueue → MainEventBus →
  MainEventBusProcessor → ChildQueue
- This timing window (typically <500ms) could result in premature closure
  when the local queue was empty but the final event hadn't arrived yet

Solution:
- Added EventQueue.isAwaitingFinalEvent() method
- MainQueue calls child.expectFinalEvent() when enqueueing final events
- EventConsumer checks awaitingFinalEvent flag before starting timeout
  counter: agentCompleted && queueSize == 0 && !awaitingFinal
- ChildQueue clears the flag only when a FINAL event is dequeued (not on
  any event, to avoid clearing it too early when non-final events arrive)
- This ensures the grace period doesn't start counting down while a final
  event is still being distributed

The fix handles both local execution (events available immediately) and
replicated scenarios (events may arrive via Kafka with delays).

Author: kabir

server-common/src/main/java/io/a2a/server/events/EventConsumer.java

@@ -82,8 +82,9 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                         item = queue.dequeueEventItem(QUEUE_WAIT_MILLISECONDS);
                         if (item == null) {
                             int queueSize = queue.size();
-                            LOGGER.debug("EventConsumer poll timeout (null item), agentCompleted={}, queue.size()={}, timeoutCount={}",
-                                agentCompleted, queueSize, pollTimeoutsAfterAgentCompleted);
+                            boolean awaitingFinal = queue.isAwaitingFinalEvent();
+                            LOGGER.debug("EventConsumer poll timeout (null item), agentCompleted={}, queue.size()={}, awaitingFinalEvent={}, timeoutCount={}",
+                                agentCompleted, queueSize, awaitingFinal, pollTimeoutsAfterAgentCompleted);
                             // If agent completed, a poll timeout means no more events are coming
                             // MainEventBusProcessor has 500ms to distribute events from MainEventBus
                             // If we timeout with agentCompleted=true, all events have been distributed
@@ -94,8 +95,12 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                             // CRITICAL: Do NOT close if task is in interrupted state (INPUT_REQUIRED, AUTH_REQUIRED)
                             // Per A2A spec, interrupted states are NOT terminal - the stream must stay open
                             // for future state updates even after agent completes (agent will be re-invoked later).
+                            //
+                            // CRITICAL: Don't start timeout counter if we're awaiting a final event.
+                            // The awaitingFinalEvent flag is set when MainQueue enqueues a final event
+                            // but it hasn't been distributed to this ChildQueue yet.
                             boolean isInterruptedState = lastSeenTaskState != null && lastSeenTaskState.isInterrupted();
-                            if (agentCompleted && queueSize == 0 && !isInterruptedState) {
+                            if (agentCompleted && queueSize == 0 && !isInterruptedState && !awaitingFinal) {
                                 pollTimeoutsAfterAgentCompleted++;
                                 if (pollTimeoutsAfterAgentCompleted >= MAX_POLL_TIMEOUTS_AFTER_AGENT_COMPLETED) {
                                     LOGGER.debug("Agent completed with {} consecutive poll timeouts and empty queue, closing for graceful completion (queue={})",
@@ -112,10 +117,10 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                                 LOGGER.debug("Agent completed but task is in interrupted state ({}), stream must remain open (queue={})",
                                     lastSeenTaskState, System.identityHashCode(queue));
                                 pollTimeoutsAfterAgentCompleted = 0; // Reset counter
-                            } else if (agentCompleted && queueSize > 0) {
-                                LOGGER.debug("Agent completed but queue has {} pending events, resetting timeout counter and continuing to poll (queue={})",
-                                    queueSize, System.identityHashCode(queue));
-                                pollTimeoutsAfterAgentCompleted = 0; // Reset counter when events arrive
+                            } else if (agentCompleted && (queueSize > 0 || awaitingFinal)) {
+                                LOGGER.debug("Agent completed but queue has {} pending events or awaitingFinalEvent={}, resetting timeout counter and continuing to poll (queue={})",
+                                    queueSize, awaitingFinal, System.identityHashCode(queue));
+                                pollTimeoutsAfterAgentCompleted = 0; // Reset counter when events arrive or awaiting final
                             }
                             continue;
                         }

server-common/src/main/java/io/a2a/server/events/EventQueue.java

@@ -304,6 +304,25 @@ public void taskDone() {
      */
     public abstract int size();
 
+    /**
+     * Returns whether this queue is awaiting a final event to be delivered.
+     * <p>
+     * This is used by EventConsumer to determine if it should keep polling even when
+     * the queue is empty. A final event may still be in-transit through MainEventBusProcessor.
+     * </p>
+     * <p>
+     * For MainQueue: always returns false (MainQueue cannot be consumed).
+     * For ChildQueue: returns true if {@link ChildQueue#expectFinalEvent()} was called
+     * but the final event hasn't been received yet.
+     * </p>
+     *
+     * @return true if awaiting a final event, false otherwise
+     */
+    public boolean isAwaitingFinalEvent() {
+        // Default implementation - overridden by ChildQueue
+        return false;
+    }
+
     /**
      * Closes this event queue gracefully, allowing pending events to be consumed.
      */
@@ -735,6 +754,12 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl
                 if (item != null) {
                     Event event = item.getEvent();
                     LOGGER.debug("Dequeued event item (waiting) {} {}", this, event instanceof Throwable ? event.toString() : event);
+                    // Clear the awaiting flag only if this is a final event
+                    // This allows EventConsumer grace period logic to proceed correctly
+                    if (awaitingFinalEvent && isFinalEvent(event)) {
+                        awaitingFinalEvent = false;
+                        LOGGER.debug("ChildQueue {} received final event while awaiting - flag cleared", System.identityHashCode(this));
+                    }
                 } else {
                     LOGGER.trace("Dequeue timeout (null) from ChildQueue {}", System.identityHashCode(this));
                 }
@@ -757,6 +782,11 @@ public int size() {
             return queue.size();
         }
 
+        @Override
+        public boolean isAwaitingFinalEvent() {
+            return awaitingFinalEvent;
+        }
+
         @Override
         public void awaitQueuePollerStart() throws InterruptedException {
             parent.awaitQueuePollerStart();