fix: Add callback INSIDE registerAndExecuteAgentAsync before starting CompletableFuture

The previous fix moved EventConsumer creation before registerAndExecuteAgentAsync,
but there was STILL a race condition:

Previous (broken) timeline:
1. EventConsumer consumer = new EventConsumer(queue)
2. EnhancedRunnable r = registerAndExecuteAgentAsync(...)  ← Starts agent IMMEDIATELY
3. r.addDoneCallback(callback)  ← TOO LATE if agent already completed

The registerAndExecuteAgentAsync method calls CompletableFuture.runAsync() which
starts the agent task immediately. If the agent completes very fast (validation
errors complete in microseconds), the whenComplete handler fires and calls
invokeDoneCallbacks() BEFORE line 3 can register the callback.

Fixed timeline:
1. EventConsumer consumer = new EventConsumer(queue)
2. registerAndExecuteAgentAsync(..., callback)
   - Creates EnhancedRunnable
   - Adds callback to runnable  ← BEFORE starting!
   - Starts CompletableFuture.runAsync()
   - Returns runnable
3. Agent can complete any time, callback already registered

Changes:
- Add doneCallback parameter to registerAndExecuteAgentAsync()
- Register callback before CompletableFuture.runAsync()
- Update both callers (onMessageSend, onMessageSendStream)

This completely eliminates the race condition - callbacks are always registered
before the agent starts, regardless of execution speed.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: kabir

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

@@ -442,13 +442,10 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
 
         boolean interruptedOrNonBlocking = false;
 
-        // Create consumer and register callback BEFORE starting agent to avoid race condition
-        // If agent completes very fast (e.g., validation error), callback must already be registered
-        // Otherwise the agent's whenComplete handler invokes callbacks before they're registered
+        // Create consumer BEFORE starting agent - callback is registered inside registerAndExecuteAgentAsync
         EventConsumer consumer = new EventConsumer(queue);
 
-        EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue);
-        producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
+        EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue, consumer.createAgentRunnableDoneCallback());
 
         ResultAggregator.EventTypeAndInterrupt etai = null;
         EventKind kind = null;  // Declare outside try block so it's in scope for return
@@ -626,12 +623,10 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
 
         ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor, eventConsumerExecutor);
 
-        // Create consumer and register callback BEFORE starting agent to avoid race condition
-        // If agent completes very fast (e.g., validation error), callback must already be registered
+        // Create consumer BEFORE starting agent - callback is registered inside registerAndExecuteAgentAsync
         EventConsumer consumer = new EventConsumer(queue);
 
-        EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue);
-        producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
+        EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue, consumer.createAgentRunnableDoneCallback());
 
         // Store cancel callback in context for closeHandler to access
         // When client disconnects, closeHandler can call this to stop EventConsumer polling loop
@@ -869,8 +864,10 @@ private boolean shouldAddPushInfo(MessageSendParams params) {
      *
      * This design avoids blocking agent-executor threads waiting for consumer polling to start,
      * eliminating cascading delays when Vert.x worker threads are busy.
+     *
+     * @param doneCallback Callback to invoke when agent completes - MUST be added before starting CompletableFuture
      */
-    private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestContext requestContext, EventQueue queue) {
+    private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestContext requestContext, EventQueue queue, EnhancedRunnable.DoneCallback doneCallback) {
         LOGGER.debug("Registering agent execution for task {}, runningAgents.size() before: {}", taskId, runningAgents.size());
         logThreadStats("AGENT START");
         EnhancedRunnable runnable = new EnhancedRunnable() {
@@ -890,6 +887,10 @@ public void run() {
             }
         };
 
+        // CRITICAL: Add callback BEFORE starting CompletableFuture to avoid race condition
+        // If agent completes very fast, whenComplete can fire before caller adds callbacks
+        runnable.addDoneCallback(doneCallback);
+
         CompletableFuture<Void> cf = CompletableFuture.runAsync(runnable, executor)
                 .whenComplete((v, err) -> {
                     if (err != null) {