TCK fixes 2

Author: kabir

server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java

@@ -304,7 +304,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
             // Store push notification config for newly created tasks (mirrors streaming logic)
             // Only for NEW tasks - existing tasks are handled by initMessageSend()
             if (mss.task() == null && kind instanceof Task createdTask && shouldAddPushInfo(params)) {
-                LOGGER.debug("Storing push notification config for new task {}", createdTask.getId());
+                LOGGER.debug("Storing push notification config for new task {} (original taskId from params: {})",
+                        createdTask.getId(), params.message().getTaskId());
                 pushConfigStore.setInfo(createdTask.getId(), params.configuration().pushNotificationConfig());
             }
 
@@ -397,6 +398,16 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
         AtomicReference<String> taskId = new AtomicReference<>(mss.requestContext.getTaskId());
         EventQueue queue = queueManager.createOrTap(taskId.get());
         LOGGER.debug("Created/tapped queue for task {}: {}", taskId.get(), queue);
+
+        // Store push notification config SYNCHRONOUSLY for new tasks before agent starts
+        // This ensures config is available when MainEventBusProcessor sends push notifications
+        // For existing tasks, config was already stored in initMessageSend()
+        if (mss.task() == null && shouldAddPushInfo(params)) {
+            LOGGER.debug("Storing push notification config for new streaming task {} EARLY (original taskId from params: {})",
+                    taskId.get(), params.message().getTaskId());
+            pushConfigStore.setInfo(taskId.get(), params.configuration().pushNotificationConfig());
+        }
+
         ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor);
 
         EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(taskId.get(), mss.requestContext, queue);
@@ -425,14 +436,9 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
                     } catch (TaskQueueExistsException e) {
                         // TODO Log
                     }
-                    if (pushConfigStore != null &&
-                            params.configuration() != null &&
-                            params.configuration().pushNotificationConfig() != null) {
 
-                        pushConfigStore.setInfo(
-                                createdTask.getId(),
-                                params.configuration().pushNotificationConfig());
-                    }
+                    // Push notification config already stored synchronously at start of onMessageSendStream
+                    // for new tasks, or in initMessageSend for existing tasks. No need to store again here.
 
                 }
                 // Push notifications now sent by MainEventBusProcessor after persistence
@@ -710,13 +716,10 @@ private CompletableFuture<Void> cleanupProducer(CompletableFuture<Void> agentFut
             }
 
             if (isStreaming) {
-                // For streaming: DON'T close the ChildQueue here
-                // The ChildQueue must stay open so MainEventBusProcessor can distribute events to it
-                // and the subscriber can consume them. EventConsumer will close the queue when:
-                // 1. A final event is detected, or
-                // 2. The subscriber cancels, or
-                // 3. EventConsumer completes naturally
-                LOGGER.debug("Streaming call, NOT closing ChildQueue in cleanup for task {} (EventConsumer will close it)", taskId);
+                // For streaming: Queue lifecycle managed by EventConsumer
+                // EventConsumer closes queue when it detects final event (or QueueClosedEvent from replication)
+                // For fire-and-forget tasks, MainQueue stays open per architectural principle
+                LOGGER.debug("Streaming call for task {} - queue lifecycle managed by EventConsumer", taskId);
             } else {
                 // For non-streaming: close the ChildQueue and notify the parent MainQueue
                 // The parent will close itself when all children are closed (childClosing logic)

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

@@ -125,8 +125,15 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer,
                     }
 
                     // Capture Task events (especially for new tasks where taskManager.getTask() would return null)
+                    // We capture the LATEST task to ensure we get the most up-to-date state
                     if (event instanceof Task t) {
+                        Task previousTask = capturedTask.get();
                         capturedTask.set(t);
+                        if (LOGGER.isDebugEnabled()) {
+                            LOGGER.debug("Captured Task event: id={}, state={} (previous: {})",
+                                    t.getId(), t.getStatus().state(),
+                                    previousTask != null ? previousTask.getId() + "/" + previousTask.getStatus().state() : "none");
+                        }
                     }
 
                     // TaskStore update moved to MainEventBusProcessor
@@ -233,9 +240,15 @@ else if (blocking) {
         EventKind eventKind = message.get();
         if (eventKind == null) {
             eventKind = capturedTask.get();
+            if (LOGGER.isDebugEnabled() && eventKind instanceof Task t) {
+                LOGGER.debug("Returning capturedTask: id={}, state={}", t.getId(), t.getStatus().state());
+            }
         }
         if (eventKind == null) {
             eventKind = taskManager.getTask();
+            if (LOGGER.isDebugEnabled() && eventKind instanceof Task t) {
+                LOGGER.debug("Returning task from TaskStore: id={}, state={}", t.getId(), t.getStatus().state());
+            }
         }
 
         return new EventTypeAndInterrupt(