TCK fixes 2

Author: kabir

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

@@ -407,7 +407,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.id());
+                LOGGER.debug("Storing push notification config for new task {} (original taskId from params: {})",
+                        createdTask.id(), params.message().taskId());
                 pushConfigStore.setInfo(createdTask.id(), params.configuration().pushNotificationConfig());
             }
 
@@ -508,6 +509,18 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
         @SuppressWarnings("NullAway")
         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)) {
+            // Satisfy Nullaway
+            Objects.requireNonNull(taskId.get(), "taskId was null");
+            LOGGER.debug("Storing push notification config for new streaming task {} EARLY (original taskId from params: {})",
+                    taskId.get(), params.message().taskId());
+            pushConfigStore.setInfo(taskId.get(), params.configuration().pushNotificationConfig());
+        }
+
         ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor);
 
         EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue);
@@ -536,15 +549,6 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
                     } catch (TaskQueueExistsException e) {
                         // TODO Log
                     }
-                    if (pushConfigStore != null &&
-                            params.configuration() != null &&
-                            params.configuration().pushNotificationConfig() != null) {
-
-                        pushConfigStore.setInfo(
-                                createdTask.id(),
-                                params.configuration().pushNotificationConfig());
-                    }
-
                 }
                 return true;
             }));
@@ -815,13 +819,10 @@ private CompletableFuture<Void> cleanupProducer(@Nullable CompletableFuture<Void
             }
 
             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

@@ -132,8 +132,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.id(), t.status().state(),
+                                    previousTask != null ? previousTask.id() + "/" + previousTask.status().state() : "none");
+                        }
                     }
 
                     // TaskStore update moved to MainEventBusProcessor
@@ -240,9 +247,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.id(), t.status().state());
+            }
         }
         if (eventKind == null) {
             eventKind = taskManager.getTask();
+            if (LOGGER.isDebugEnabled() && eventKind instanceof Task t) {
+                LOGGER.debug("Returning task from TaskStore: id={}, state={}", t.id(), t.status().state());
+            }
         }
         if (eventKind == null) {
             throw new InternalError("Could not find a Task/Message for " + taskManager.getTaskId());