review fixes

Author: kabir

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

@@ -115,9 +115,7 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                                     LOGGER.debug("Waited {} timeouts for final event but it hasn't arrived - proceeding with normal timeout logic (queue={})",
                                         pollTimeoutsWhileAwaitingFinal, System.identityHashCode(queue));
                                     // Clear the flag on the queue itself, not just the local variable
-                                    if (queue instanceof EventQueue.ChildQueue) {
-                                        ((EventQueue.ChildQueue) queue).clearAwaitingFinalEvent();
-                                    }
+                                    queue.clearAwaitingFinalEvent();
                                     awaitingFinal = false; // Also update local variable for this iteration
                                 }
                             } else {

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

@@ -323,6 +323,17 @@ public boolean isAwaitingFinalEvent() {
         return false;
     }
 
+    /**
+     * Clears the awaiting final event flag.
+     * <p>
+     * Default implementation is a no-op for queues that don't track this state.
+     * ChildQueue overrides this to actually clear the flag.
+     * </p>
+     */
+    public void clearAwaitingFinalEvent() {
+        // Default no-op implementation - overridden by ChildQueue
+    }
+
     /**
      * Closes this event queue gracefully, allowing pending events to be consumed.
      */
@@ -754,12 +765,6 @@ 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));
                 }
@@ -824,7 +829,8 @@ void expectFinalEvent() {
          * Called by EventConsumer when it has waited too long for the final event.
          * This allows normal timeout logic to proceed if the final event never arrives.
          */
-        void clearAwaitingFinalEvent() {
+        @Override
+        public void clearAwaitingFinalEvent() {
             awaitingFinalEvent = false;
             LOGGER.debug("ChildQueue {} cleared awaitingFinalEvent flag (timeout)", System.identityHashCode(this));
         }