fix(server): reconcile blocking result with TaskStore (a2aproject#991)

* fix(server): reconcile blocking result with TaskStore

- Make reconciliation timeout configurable, add logging
- Keep the existing four-argument constructor delegating to the original one-second reconciliation timeout, and cover the new configuration wiring and default value.

Fixes a2aproject/A2A#2049

Author: jstar0

docs/content/configuration.md

@@ -40,6 +40,9 @@ a2a.blocking.agent.timeout.seconds=30
 
 # Timeout for event consumption in blocking calls (default: 5 seconds)
 a2a.blocking.consumption.timeout.seconds=5
+
+# Timeout for TaskStore reconciliation polling in blocking calls (default: 1 second)
+a2a.blocking.reconciliation.timeout.seconds=1
 ```
 
 ### Tuning Guidelines
@@ -48,6 +51,7 @@ a2a.blocking.consumption.timeout.seconds=5
 - **Resource Management**: The dedicated executor prevents streaming operations from competing with the ForkJoinPool.
 - **Concurrency**: In production with high concurrent streaming, increase pool sizes accordingly.
 - **Agent Timeouts**: LLM-based agents may need longer timeouts (60-120s) compared to simple agents.
+- **Reconciliation Timeout**: Increase if blocking calls fail with "Could not find a Task/Message" under heavy load or with slow TaskStore implementations.
 
 ## MicroProfile Config Integration
 

integrations/microprofile-config/README.md

@@ -37,6 +37,7 @@ a2a.executor.max-pool-size=100
 # Timeout configuration
 a2a.blocking.agent.timeout.seconds=60
 a2a.blocking.consumption.timeout.seconds=10
+a2a.blocking.reconciliation.timeout.seconds=2
 ```
 
 **Environment variables:**

integrations/microprofile-config/src/test/java/org/a2aproject/sdk/integrations/microprofile/MicroProfileConfigProviderTest.java

@@ -24,6 +24,8 @@ public class MicroProfileConfigProviderTest {
     private static final String A2A_EXECUTOR_CORE_POOL_SIZE = "a2a.executor.core-pool-size";
     private static final String A2A_EXECUTOR_MAX_POOL_SIZE = "a2a.executor.max-pool-size";
     private static final String A2A_EXECUTOR_KEEP_ALIVE_SECONDS = "a2a.executor.keep-alive-seconds";
+    private static final String A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS =
+            "a2a.blocking.reconciliation.timeout.seconds";
 
     @Inject
     A2AConfigProvider configProvider;
@@ -59,6 +61,12 @@ public void testGetValueAnotherDefault() {
         assertEquals("60", value, "Should fall back to default value");
     }
 
+    @Test
+    public void testGetReconciliationTimeoutDefault() {
+        String value = configProvider.getValue(A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS);
+        assertEquals("1", value, "Should fall back to the reconciliation timeout default");
+    }
+
     @Test
     public void testGetOptionalValueFromMicroProfileConfig() {
         // Test optional value that exists in application.properties

server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java

@@ -142,8 +142,9 @@
  *   <li><b>Blocking (configuration.blocking=true):</b> Client waits for first event or final task state</li>
  *   <li><b>Streaming:</b> Client receives events as they arrive via reactive streams</li>
  *   <li>Both modes support fire-and-forget (agent continues after client disconnect)</li>
- *   <li>Configurable timeouts via {@code a2a.blocking.agent.timeout.seconds} and
- *       {@code a2a.blocking.consumption.timeout.seconds}</li>
+ *   <li>Configurable timeouts via {@code a2a.blocking.agent.timeout.seconds},
+ *       {@code a2a.blocking.consumption.timeout.seconds}, and
+ *       {@code a2a.blocking.reconciliation.timeout.seconds}</li>
  * </ul>
  *
  * <h2>CDI Dependencies</h2>
@@ -189,6 +190,7 @@ public class DefaultRequestHandler implements RequestHandler {
 
     private static final String A2A_BLOCKING_AGENT_TIMEOUT_SECONDS = "a2a.blocking.agent.timeout.seconds";
     private static final String A2A_BLOCKING_CONSUMPTION_TIMEOUT_SECONDS = "a2a.blocking.consumption.timeout.seconds";
+    private static final String A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS = "a2a.blocking.reconciliation.timeout.seconds";
 
     @Inject
     A2AConfigProvider configProvider;
@@ -215,6 +217,19 @@ public class DefaultRequestHandler implements RequestHandler {
      */
     int consumptionCompletionTimeoutSeconds;
 
+    /**
+     * Timeout in seconds for TaskStore reconciliation polling in blocking calls.
+     * When the in-memory event capture is empty, the aggregator polls TaskStore
+     * with a bounded timeout to handle the race where MainEventBusProcessor
+     * has not yet persisted the task.
+     * <p>
+     * Property: {@code a2a.blocking.reconciliation.timeout.seconds}<br>
+     * Default: 1 second<br>
+     * Note: Property override requires a configurable {@link A2AConfigProvider} on the classpath
+     * (e.g., MicroProfileConfigProvider in reference implementations).
+     */
+    int reconciliationTimeoutSeconds;
+
     // Fields set by constructor injection cannot be final. We need a noargs constructor for
     // Jakarta compatibility, and it seems that making fields set by constructor injection
     // final, is not proxyable in all runtimes
@@ -276,6 +291,8 @@ void initConfig() {
                 configProvider.getValue(A2A_BLOCKING_AGENT_TIMEOUT_SECONDS));
         consumptionCompletionTimeoutSeconds = Integer.parseInt(
                 configProvider.getValue(A2A_BLOCKING_CONSUMPTION_TIMEOUT_SECONDS));
+        reconciliationTimeoutSeconds = Integer.parseInt(
+                configProvider.getValue(A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS));
     }
 
 
@@ -291,6 +308,7 @@ public static DefaultRequestHandler create(AgentExecutor agentExecutor, TaskStor
                         mainEventBusProcessor, executor, eventConsumerExecutor);
         handler.agentCompletionTimeoutSeconds = 5;
         handler.consumptionCompletionTimeoutSeconds = 2;
+        handler.reconciliationTimeoutSeconds = 1;
 
         return handler;
     }
@@ -377,7 +395,8 @@ public Task onCancelTask(CancelTaskParams params, ServerCallContext context) thr
                 taskStore,
                 null);
 
-        ResultAggregator resultAggregator = new ResultAggregator(taskManager, null, executor, eventConsumerExecutor);
+        ResultAggregator resultAggregator = new ResultAggregator(taskManager, null, executor, eventConsumerExecutor,
+                        SECONDS.toNanos(reconciliationTimeoutSeconds));
 
         EventQueue queue = queueManager.createOrTap(task.id());
         EventConsumer consumer = new EventConsumer(queue, eventConsumerExecutor);
@@ -450,7 +469,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
         // Create queue with real taskId (no tempId parameter needed)
         EventQueue queue = queueManager.createOrTap(queueTaskId);
         final java.util.concurrent.atomic.AtomicReference<@NonNull String> taskId = new java.util.concurrent.atomic.AtomicReference<>(queueTaskId);
-        ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor, eventConsumerExecutor);
+        ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor, eventConsumerExecutor,
+                        SECONDS.toNanos(reconciliationTimeoutSeconds));
 
         // Default to blocking per A2A spec (returnImmediately defaults to false, meaning wait for completion)
         boolean returnImmediately = params.configuration() != null && Boolean.TRUE.equals(params.configuration().returnImmediately());
@@ -668,7 +688,8 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
                     .taskId(taskId.get()).build(), version);
         }
 
-        ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor, eventConsumerExecutor);
+        ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor, eventConsumerExecutor,
+                        SECONDS.toNanos(reconciliationTimeoutSeconds));
 
         // Create consumer BEFORE starting agent - callback is registered inside registerAndExecuteAgentAsync
         EventConsumer consumer = new EventConsumer(queue, eventConsumerExecutor);
@@ -857,7 +878,8 @@ public Flow.Publisher<StreamingEventKind> onSubscribeToTask(TaskIdParams params,
         }
 
         TaskManager taskManager = new TaskManager(task.id(), task.contextId(), taskStore, null);
-        ResultAggregator resultAggregator = new ResultAggregator(taskManager, null, executor, eventConsumerExecutor);
+        ResultAggregator resultAggregator = new ResultAggregator(taskManager, null, executor, eventConsumerExecutor,
+                        SECONDS.toNanos(reconciliationTimeoutSeconds));
         EventQueue queue = queueManager.tap(task.id());
         LOGGER.debug("onSubscribeToTask - tapped queue: {}", queue != null ? System.identityHashCode(queue) : "null");
 

server-common/src/main/java/org/a2aproject/sdk/server/tasks/ResultAggregator.java

@@ -8,6 +8,7 @@
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Flow;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -28,17 +29,21 @@
 
 public class ResultAggregator {
     private static final Logger LOGGER = LoggerFactory.getLogger(ResultAggregator.class);
+    private static final long TASK_STORE_RECONCILIATION_POLL_MILLIS = 10;
 
     private final TaskManager taskManager;
     private final Executor executor;
     private final Executor eventConsumerExecutor;
+    private final long reconciliationTimeoutNanos;
     private volatile @Nullable Message message;
 
-    public ResultAggregator(TaskManager taskManager, @Nullable Message message, Executor executor, Executor eventConsumerExecutor) {
+    public ResultAggregator(TaskManager taskManager, @Nullable Message message, Executor executor,
+            Executor eventConsumerExecutor, long reconciliationTimeoutNanos) {
         this.taskManager = taskManager;
         this.message = message;
         this.executor = executor;
         this.eventConsumerExecutor = eventConsumerExecutor;
+        this.reconciliationTimeoutNanos = reconciliationTimeoutNanos;
     }
 
     public @Nullable EventKind getCurrentResult() {
@@ -235,7 +240,8 @@ else if (blocking) {
             Utils.rethrow(error);
         }
 
-        // Return Message if captured, otherwise Task if captured, otherwise fetch from TaskStore
+        // Return Message if captured, otherwise Task if captured,
+        // otherwise poll TaskStore with bounded timeout (blocking) or single read (non-blocking).
         EventKind eventKind = message.get();
         if (eventKind == null) {
             eventKind = capturedTask.get();
@@ -244,7 +250,7 @@ else if (blocking) {
             }
         }
         if (eventKind == null) {
-            eventKind = taskManager.getTask();
+            eventKind = reconcileTaskStore(blocking);
             if (LOGGER.isDebugEnabled() && eventKind instanceof Task t) {
                 LOGGER.debug("Returning task from TaskStore: id={}, state={}", t.id(), t.status().state());
             }
@@ -259,6 +265,33 @@ else if (blocking) {
                 consumptionCompletionFuture);
     }
 
+    private @Nullable Task reconcileTaskStore(boolean blocking) throws A2AError {
+        Task task = taskManager.getTask();
+        if (task != null || !blocking) {
+            return task;
+        }
+
+        LOGGER.debug("TaskStore reconciliation: task not found on first read, polling for up to {}ms",
+                TimeUnit.NANOSECONDS.toMillis(reconciliationTimeoutNanos));
+
+        long deadline = System.nanoTime() + reconciliationTimeoutNanos;
+        while (System.nanoTime() < deadline) {
+            try {
+                Thread.sleep(TASK_STORE_RECONCILIATION_POLL_MILLIS);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new InternalError("Interrupted while reconciling TaskStore for " + taskManager.getTaskId());
+            }
+
+            task = taskManager.getTask();
+            if (task != null) {
+                LOGGER.debug("TaskStore reconciliation: task {} found after polling", task.id());
+                return task;
+            }
+        }
+        return null;
+    }
+
     private String taskIdForLogging() {
         Task task = taskManager.getTask();
         return task != null ? task.id() : "unknown";

server-common/src/main/resources/META-INF/a2a-defaults.properties

@@ -10,6 +10,10 @@ a2a.blocking.agent.timeout.seconds=30
 # Ensures TaskStore is fully updated before returning to client
 a2a.blocking.consumption.timeout.seconds=5
 
+# Timeout for TaskStore reconciliation polling in blocking calls (seconds)
+# When in-memory event capture is empty, polls TaskStore with this bounded timeout
+a2a.blocking.reconciliation.timeout.seconds=1
+
 # AsyncExecutorProducer - Thread pool configuration
 # Core pool size for async agent execution
 a2a.executor.core-pool-size=5

server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java

@@ -6,6 +6,8 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.util.List;
 import java.util.Map;
@@ -20,6 +22,7 @@
 import org.a2aproject.sdk.server.ServerCallContext;
 import org.a2aproject.sdk.server.agentexecution.AgentExecutor;
 import org.a2aproject.sdk.server.agentexecution.RequestContext;
+import org.a2aproject.sdk.server.config.A2AConfigProvider;
 import org.a2aproject.sdk.server.events.EventQueue;
 import org.a2aproject.sdk.server.events.EventQueueItem;
 import org.a2aproject.sdk.server.events.EventQueueUtil;
@@ -146,6 +149,22 @@ protected interface AgentExecutorMethod {
         void invoke(RequestContext context, AgentEmitter agentEmitter) throws A2AError;
     }
 
+    @Test
+    void testInitConfigReadsBlockingTimeouts() {
+        A2AConfigProvider configProvider = mock(A2AConfigProvider.class);
+        when(configProvider.getValue("a2a.blocking.agent.timeout.seconds")).thenReturn("30");
+        when(configProvider.getValue("a2a.blocking.consumption.timeout.seconds")).thenReturn("5");
+        when(configProvider.getValue("a2a.blocking.reconciliation.timeout.seconds")).thenReturn("7");
+
+        DefaultRequestHandler handler = new DefaultRequestHandler();
+        handler.configProvider = configProvider;
+        handler.initConfig();
+
+        assertEquals(30, handler.agentCompletionTimeoutSeconds);
+        assertEquals(5, handler.consumptionCompletionTimeoutSeconds);
+        assertEquals(7, handler.reconciliationTimeoutSeconds);
+    }
+
     /**
      * Test 1: Non-streaming AUTH_REQUIRED returns immediately while agent continues.
      * Verifies: