TCK fixes

Author: kabir

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

@@ -65,7 +65,10 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                         }
                         event = item.getEvent();
 
+                        // Defensive logging for error handling
                         if (event instanceof Throwable thr) {
+                            LOGGER.info("EventConsumer detected Throwable event: {} - triggering tube.fail()",
+                                    thr.getClass().getSimpleName());
                             tube.fail(thr);
                             return;
                         }

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

@@ -167,6 +167,11 @@ private void processEvent(MainEventBusContext context) {
             }
 
             // Step 3: Then distribute to ChildQueues (clients see either event or error AFTER persistence attempt)
+            if (eventToDistribute == null) {
+                LOGGER.error("MainEventBusProcessor: eventToDistribute is NULL for task {} - this should never happen!", taskId);
+                eventToDistribute = new InternalError("Internal error: event processing failed");
+            }
+
             if (eventQueue instanceof EventQueue.MainQueue mainQueue) {
                 int childCount = mainQueue.getChildCount();
                 LOGGER.debug("MainEventBusProcessor: Distributing {} to {} children for task {}",

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

@@ -8,7 +8,6 @@
  * This interface is primarily intended for testing, allowing tests to synchronize
  * with the asynchronous MainEventBusProcessor. Production code should not rely on this.
  * </p>
- * <p>
  * Usage in tests:
  * <pre>
  * {@code
@@ -31,7 +30,6 @@
  * }
  * }
  * </pre>
- * </p>
  */
 public interface MainEventBusProcessorCallback {
 

server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java

@@ -16,6 +16,7 @@
 import io.a2a.spec.A2AError;
 import io.a2a.spec.Event;
 import io.a2a.spec.EventKind;
+import io.a2a.spec.InternalError;
 import io.a2a.spec.Message;
 import io.a2a.spec.Task;
 import io.a2a.spec.TaskState;
@@ -96,6 +97,7 @@ public EventKind consumeAll(EventConsumer consumer) throws A2AError {
     public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, boolean blocking) throws A2AError {
         Flow.Publisher<EventQueueItem> allItems = consumer.consumeAll();
         AtomicReference<Message> message = new AtomicReference<>();
+        AtomicReference<Task> capturedTask = new AtomicReference<>();  // Capture Task events
         AtomicBoolean interrupted = new AtomicBoolean(false);
         AtomicReference<Throwable> errorRef = new AtomicReference<>();
         CompletableFuture<Void> completionFuture = new CompletableFuture<>();
@@ -129,6 +131,11 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer,
                         return false;
                     }
 
+                    // Capture Task events (especially for new tasks where taskManager.getTask() would return null)
+                    if (event instanceof Task t) {
+                        capturedTask.set(t);
+                    }
+
                     // TaskStore update moved to MainEventBusProcessor
 
                     // Determine interrupt behavior
@@ -229,20 +236,20 @@ else if (blocking) {
             Utils.rethrow(error);
         }
 
-        EventKind eventType;
-        Message msg = message.get();
-        if (msg != null) {
-            eventType = msg;
-        } else {
-            Task task = taskManager.getTask();
-            if (task == null) {
-                throw new io.a2a.spec.InternalError("No task or message available after consuming events");
-            }
-            eventType = task;
+        // Return Message if captured, otherwise Task if captured, otherwise fetch from TaskStore
+        EventKind eventKind = message.get();
+        if (eventKind == null) {
+            eventKind = capturedTask.get();
+        }
+        if (eventKind == null) {
+            eventKind = taskManager.getTask();
+        }
+        if (eventKind == null) {
+            throw new InternalError("Could not find a Task/Message for " + taskManager.getTaskId());
         }
 
         return new EventTypeAndInterrupt(
-                eventType,
+                eventKind,
                 interrupted.get(),
                 consumptionCompletionFuture);
     }