fix: Wait for agent completion and ensure all events processed in blo… · anFatum/a2a-java@465a5e4 · GitHub
Skip to content

Commit 465a5e4

Browse files
authored
fix: Wait for agent completion and ensure all events processed in blo… (a2aproject#431)
…cking calls Fixes race condition where DefaultRequestHandler.onMessageSend() returns before all task events are fully processed and persisted to TaskStore, resulting in incomplete Task objects being returned to clients (missing artifacts, incorrect state). Root Cause: - Blocking calls interrupted immediately after first event and returned to client before background event consumption completed - Agent execution and event processing happened asynchronously in background - No synchronization to ensure all events were consumed and persisted before returning Task to client Solution (4-step process): 1. Wait for agent to finish enqueueing events (configurable timeout) 2. Close the queue to signal consumption can complete (breaks dependency) 3. Wait for consumption to finish processing events (configurable timeout) 4. Fetch final task state from TaskStore (has all artifacts and correct state) This ensures blocking calls return complete Task objects with all artifacts and correct state, including support for fire-and-forget tasks that never emit final state events. Added unit tests: - testBlockingFireAndForgetReturnsNonFinalTask: Validates fire-and-forget pattern - testBlockingCallReturnsCompleteTaskWithArtifacts: Ensures all artifacts included Fixes a2aproject#428 🦕
1 parent 5c5970d commit 465a5e4

6 files changed

Lines changed: 308 additions & 64 deletions

File tree

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

Lines changed: 106 additions & 5 deletions

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

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,20 +11,20 @@
1111
import java.util.concurrent.atomic.AtomicBoolean;
1212
import java.util.concurrent.atomic.AtomicReference;
1313

14-
import org.slf4j.Logger;
15-
import org.slf4j.LoggerFactory;
16-
1714
import io.a2a.server.events.EventConsumer;
1815
import io.a2a.server.events.EventQueueItem;
1916
import io.a2a.spec.A2AServerException;
2017
import io.a2a.spec.Event;
2118
import io.a2a.spec.EventKind;
19+
import io.a2a.spec.InternalError;
2220
import io.a2a.spec.JSONRPCError;
2321
import io.a2a.spec.Message;
2422
import io.a2a.spec.Task;
2523
import io.a2a.spec.TaskState;
2624
import io.a2a.spec.TaskStatusUpdateEvent;
2725
import io.a2a.util.Utils;
26+
import org.slf4j.Logger;
27+
import org.slf4j.LoggerFactory;
2828

2929
public class ResultAggregator {
3030
private static final Logger LOGGER = LoggerFactory.getLogger(ResultAggregator.class);
@@ -106,10 +106,6 @@ public EventKind consumeAll(EventConsumer consumer) throws JSONRPCError {
106106
}
107107

108108
public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, boolean blocking) throws JSONRPCError {
109-
return consumeAndBreakOnInterrupt(consumer, blocking, null);
110-
}
111-
112-
public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, boolean blocking, Runnable eventCallback) throws JSONRPCError {
113109
Flow.Publisher<EventQueueItem> allItems = consumer.consumeAll();
114110
AtomicReference<Message> message = new AtomicReference<>();
115111
AtomicBoolean interrupted = new AtomicBoolean(false);
@@ -180,11 +176,11 @@ else if (!blocking) {
180176
shouldInterrupt = true;
181177
continueInBackground = true;
182178
}
183-
else {
184-
// For ALL blocking calls (both final and non-final events), use background consumption
185-
// This ensures all events are processed and persisted to TaskStore in background
186-
// Queue lifecycle is now managed by DefaultRequestHandler.cleanupProducer()
187-
// which waits for BOTH agent and consumption futures before closing queues
179+
else if (blocking) {
180+
// For blocking calls: Interrupt to free Vert.x thread, but continue in background
181+
// Python's async consumption doesn't block threads, but Java's does
182+
// So we interrupt to return quickly, then rely on background consumption
183+
// DefaultRequestHandler will fetch the final state from TaskStore
188184
shouldInterrupt = true;
189185
continueInBackground = true;
190186
if (LOGGER.isDebugEnabled()) {
@@ -198,10 +194,17 @@ else if (!blocking) {
198194
interrupted.set(true);
199195
completionFuture.complete(null);
200196

201-
// Signal that cleanup can proceed while consumption continues in background.
202-
// This prevents infinite hangs for fire-and-forget agents that never emit final events.
203-
// Processing continues (return true below) and all events are still persisted to TaskStore.
204-
consumptionCompletionFuture.complete(null);
197+
// For blocking calls, DON'T complete consumptionCompletionFuture here.
198+
// Let it complete naturally when subscription finishes (onComplete callback below).
199+
// This ensures all events are processed and persisted to TaskStore before
200+
// DefaultRequestHandler.cleanupProducer() proceeds with cleanup.
201+
//
202+
// For non-blocking and auth-required calls, complete immediately to allow
203+
// cleanup to proceed while consumption continues in background.
204+
if (!blocking) {
205+
consumptionCompletionFuture.complete(null);
206+
}
207+
// else: blocking calls wait for actual consumption completion in onComplete
205208

206209
// Continue consuming in background - keep requesting events
207210
// Note: continueInBackground is always true when shouldInterrupt is true
@@ -244,8 +247,8 @@ else if (!blocking) {
244247
}
245248
}
246249

247-
// Background consumption continues automatically via the subscription
248-
// returning true in the consumer function keeps the subscription alive
250+
// Note: For blocking calls that were interrupted, the wait logic has been moved
251+
// to DefaultRequestHandler.onMessageSend() to avoid blocking Vert.x worker threads.
249252
// Queue lifecycle is managed by DefaultRequestHandler.cleanupProducer()
250253

251254
Throwable error = errorRef.get();

server-common/src/test/java/io/a2a/server/requesthandlers/AbstractA2ARequestHandlerTest.java

Lines changed: 2 additions & 1 deletion

0 commit comments

Comments
 (0)