11package org .a2aproject .sdk .server .events ;
22
3+ import java .util .ArrayList ;
34import java .util .List ;
45import java .util .Objects ;
56import java .util .concurrent .BlockingQueue ;
910import java .util .concurrent .Semaphore ;
1011import java .util .concurrent .TimeUnit ;
1112import java .util .concurrent .atomic .AtomicBoolean ;
13+ import java .util .concurrent .atomic .AtomicInteger ;
1214
1315import org .a2aproject .sdk .server .tasks .TaskStateProvider ;
1416import org .a2aproject .sdk .spec .A2AError ;
@@ -513,20 +515,31 @@ public void enqueueItem(EventQueueItem item) {
513515
514516 boolean finalEvent = isFinalEvent (event );
515517
518+ // Notify current children before waiting for capacity. A consumer may already be
519+ // marked complete (for example, while waiting for a replicated final event) and
520+ // otherwise close its queue before this producer is able to submit the event.
521+ List <ChildQueue > awaitingChildren = new ArrayList <>();
522+ if (finalEvent ) {
523+ awaitingChildren .addAll (children );
524+ awaitingChildren .forEach (ChildQueue ::expectFinalEvent );
525+ }
526+
516527 // Acquire semaphore for backpressure
517528 try {
518529 semaphore .acquire ();
519530 } catch (InterruptedException e ) {
531+ awaitingChildren .forEach (ChildQueue ::cancelExpectedFinalEvent );
520532 Thread .currentThread ().interrupt ();
521533 throw new RuntimeException ("Unable to acquire the semaphore to enqueue the event" , e );
522534 }
523535
524- // Notify children only after acquiring capacity. An interrupted acquire must not
525- // leave them waiting for an event that was never submitted.
536+ // Include children that subscribed while this producer was waiting for capacity.
526537 if (finalEvent ) {
527- LOGGER .debug ("Final event detected, notifying {} children to expect it" , children .size ());
528538 for (ChildQueue child : children ) {
529- child .expectFinalEvent ();
539+ if (!awaitingChildren .contains (child )) {
540+ child .expectFinalEvent ();
541+ awaitingChildren .add (child );
542+ }
530543 }
531544 }
532545
@@ -544,11 +557,7 @@ public void enqueueItem(EventQueueItem item) {
544557 // Release the permit here to avoid leaking it and eventually blocking
545558 // all event processing for this task.
546559 semaphore .release ();
547- if (finalEvent ) {
548- for (ChildQueue child : children ) {
549- child .clearAwaitingFinalEvent ();
550- }
551- }
560+ awaitingChildren .forEach (ChildQueue ::cancelExpectedFinalEvent );
552561 throw e ;
553562 }
554563 }
@@ -786,7 +795,7 @@ static class ChildQueue extends EventQueue {
786795 private final MainQueue parent ;
787796 private final BlockingQueue <EventQueueItem > queue ;
788797 private volatile boolean immediateClose = false ;
789- private volatile boolean awaitingFinalEvent = false ;
798+ private final AtomicInteger expectedFinalEvents = new AtomicInteger () ;
790799
791800 public ChildQueue (MainQueue parent ) {
792801 this .parent = parent ;
@@ -824,8 +833,8 @@ private void internalEnqueueItem(EventQueueItem item) {
824833 LOGGER .debug ("Enqueued event {} {}" , event instanceof Throwable ? event .toString () : event , this );
825834
826835 // If we were awaiting a final event and this is it, clear the flag
827- if (awaitingFinalEvent && isFinalEvent (event )) {
828- awaitingFinalEvent = false ;
836+ if (isAwaitingFinalEvent () && isFinalEvent (event )) {
837+ expectedFinalEvents . set ( 0 ) ;
829838 LOGGER .debug ("ChildQueue {} received awaited final event" , System .identityHashCode (this ));
830839 }
831840 }
@@ -861,7 +870,7 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl
861870 // For immediate close: exit immediately even if queue is not empty (race with MainEventBusProcessor)
862871 // For graceful close: only exit when queue is empty (wait for all events to be consumed)
863872 // BUT: if awaiting final event, keep polling even if closed and empty
864- if (isClosed () && (queue .isEmpty () || immediateClose ) && !awaitingFinalEvent ) {
873+ if (isClosed () && (queue .isEmpty () || immediateClose ) && !isAwaitingFinalEvent () ) {
865874 LOGGER .debug ("ChildQueue is closed{}, sending termination message. {} (queueSize={})" ,
866875 immediateClose ? " (immediate)" : " and empty" ,
867876 this ,
@@ -906,7 +915,7 @@ public int size() {
906915
907916 @ Override
908917 public boolean isAwaitingFinalEvent () {
909- return awaitingFinalEvent ;
918+ return expectedFinalEvents . get () > 0 ;
910919 }
911920
912921 @ Override
@@ -938,17 +947,24 @@ protected void doClose(boolean immediate) {
938947 * This ensures the ChildQueue keeps polling until the final event arrives (after MainEventBusProcessor).
939948 */
940949 void expectFinalEvent () {
941- awaitingFinalEvent = true ;
950+ expectedFinalEvents . incrementAndGet () ;
942951 LOGGER .debug ("ChildQueue {} now awaiting final event" , System .identityHashCode (this ));
943952 }
944953
954+ /**
955+ * Removes one expectation when the corresponding final event could not be submitted.
956+ */
957+ void cancelExpectedFinalEvent () {
958+ expectedFinalEvents .updateAndGet (count -> Math .max (0 , count - 1 ));
959+ }
960+
945961 /**
946962 * Called by EventConsumer when it has waited too long for the final event.
947963 * This allows normal timeout logic to proceed if the final event never arrives.
948964 */
949965 @ Override
950966 public void clearAwaitingFinalEvent () {
951- awaitingFinalEvent = false ;
967+ expectedFinalEvents . set ( 0 ) ;
952968 LOGGER .debug ("ChildQueue {} cleared awaitingFinalEvent flag (timeout)" , System .identityHashCode (this ));
953969 }
954970
0 commit comments