Skip to content

Commit 46f7691

Browse files
kabirkabir
committed
feat: Implement MainEventBus architecture and resolve multi-instance replication race conditions
Introduce centralized event bus with single background processor to guarantee events persist to TaskStore before distribution to clients. This eliminates race conditions when multiple concurrent requests update the same task. **Key Components:** - MainEventBus: Shared BlockingQueue for all MainQueue events - MainEventBusProcessor: Single background thread for ordered processing - Processing sequence: TaskStore.save() → PushNotifications → distributeToChildren() - TaskStateProvider interface: Task state queries for queue lifecycle management **Event Flow:** ``` AgentExecutor → MainQueue → MainEventBus → MainEventBusProcessor → TaskStore (persist first) → Push Notifications → ChildQueues → Clients ``` **Benefits:** - Events persist before clients receive them (no stale data) - Serial processing prevents concurrent TaskStore updates - Platform-agnostic ChildQueue synchronization (works across gRPC/JSONRPC/REST) - Clean separation: MainQueue (no local queue) vs ChildQueue (local queue for clients) Implement two-level protection to keep MainQueues open for fire-and-forget tasks and late resubscriptions while cleaning up finalized tasks: **Level 1** - Cleanup Callback: Check TaskStateProvider.isTaskFinalized() before removing queue from QueueManager map **Level 2** - Auto-Close Prevention: MainQueue.childClosing() checks finality before closing when last ChildQueue disconnects **Result:** Non-final tasks keep queues open for resubscription; finalized tasks clean up immediately ReplicatedQueueManager.onTaskFinalized() sent full Task objects to remote instances via Kafka, while local instances sent TaskStatusUpdateEvent. Client auto-close logic only checked for TaskStatusUpdateEvent.isFinal(), causing connection leaks on remote instances. **ReplicatedQueueManager.onTaskFinalized():** Convert Task to TaskStatusUpdateEvent before sending to Kafka, ensuring consistent event types across all instances **EventConsumer:** Add 50ms delay before tube.complete() to allow SSE buffer flush in replicated scenarios where events arrive via Kafka with timing variations **SSEEventListener (JSONRPC):** Check both TaskStatusUpdateEvent.isFinal() and Task.status().state().isFinal() for auto-close **RestSSEEventListener (REST):** Add complete auto-close logic (was missing entirely) **Benefits:** - Handles late subscriptions to completed tasks gracefully - Prevents connection leaks in all scenarios - Consistent behavior across JSONRPC and REST transports - Defensive programming for edge cases **MultiInstanceReplicationTest:** - Add TaskEvent handling and container log dumping on failure - Verify both APP1 and APP2 receive all events including final states - Test late-arriving replicated events and poison pill ordering **Integration Tests:** - EventConsumerTest: Grace period mechanism for replicated scenarios - ReplicatedQueueManagerTest: Event type conversion validation - All existing tests updated for MainEventBus architecture **EventQueue.Builder:** Now requires MainEventBus parameter (validates non-null) **QueueManager Implementations:** Must handle TaskStateProvider for lifecycle checks Existing code continues to work - InMemoryQueueManager and ReplicatedQueueManager automatically inject MainEventBus via CDI. Custom QueueManager implementations should inject MainEventBus and pass to EventQueue.Builder. **Core Architecture:** - server-common/.../events/MainEventBus.java (new) - server-common/.../events/MainEventBusProcessor.java (new) - server-common/.../events/EventQueue.java (requires MainEventBus) - server-common/.../events/InMemoryQueueManager.java (queue lifecycle) **Replication:** - extras/queue-manager-replicated/core/.../ReplicatedQueueManager.java (event conversion) - extras/queue-manager-replicated/core/.../ReplicatedEventQueueItem.java (Task support) **Client Transports:** - client/transport/jsonrpc/.../SSEEventListener.java (enhanced auto-close) - client/transport/rest/.../RestSSEEventListener.java (add auto-close) **Event Processing:** - server-common/.../events/EventConsumer.java (grace period + buffer flush) - server-common/.../requesthandlers/DefaultRequestHandler.java (MainEventBus integration) **Task Management:** - server-common/.../tasks/TaskStore.java (TaskStateProvider interface) - extras/task-store-database-jpa/.../JpaDatabaseTaskStore.java (implement TaskStateProvider) ✅ All unit tests pass (150+ tests) ✅ MultiInstanceReplicationTest passes (both instances receive all events) ✅ TCK tests pass (no connection leaks) ✅ Integration tests pass (EventConsumer, QueueManager, TaskStore)
1 parent 659df81 commit 46f7691

56 files changed

Lines changed: 3236 additions & 2210 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎client/transport/jsonrpc/src/main/java/io/a2a/client/transport/jsonrpc/sse/SSEEventListener.java‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
import io.a2a.jsonrpc.common.json.JsonProcessingException;
1111
import io.a2a.spec.A2AError;
1212
import io.a2a.spec.StreamingEventKind;
13+
import io.a2a.spec.Task;
14+
import io.a2a.spec.TaskState;
1315
import io.a2a.spec.TaskStatusUpdateEvent;
1416
import org.jspecify.annotations.Nullable;
1517

@@ -64,11 +66,23 @@ private void handleMessage(String message, @Nullable Future<Void> future) {
6466

6567
StreamingEventKind event = ProtoUtils.FromProto.streamingEventKind(response);
6668
eventHandler.accept(event);
67-
if (event instanceof TaskStatusUpdateEvent && ((TaskStatusUpdateEvent) event).isFinal()) {
68-
if (future != null) {
69-
future.cancel(true); // close SSE channel
69+
70+
// Client-side auto-close on final events to prevent connection leaks
71+
// Handles both TaskStatusUpdateEvent and Task objects with final states
72+
// This covers late subscriptions to completed tasks and ensures no connection leaks
73+
boolean shouldClose = false;
74+
if (event instanceof TaskStatusUpdateEvent tue && tue.isFinal()) {
75+
shouldClose = true;
76+
} else if (event instanceof Task task) {
77+
TaskState state = task.status().state();
78+
if (state.isFinal()) {
79+
shouldClose = true;
7080
}
7181
}
82+
83+
if (shouldClose && future != null) {
84+
future.cancel(true); // close SSE channel
85+
}
7286
} catch (A2AError error) {
7387
if (errorHandler != null) {
7488
errorHandler.accept(error);

‎client/transport/rest/src/main/java/io/a2a/client/transport/rest/sse/RestSSEEventListener.java‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@
1010
import io.a2a.grpc.StreamResponse;
1111
import io.a2a.grpc.utils.ProtoUtils;
1212
import io.a2a.spec.StreamingEventKind;
13+
import io.a2a.spec.Task;
14+
import io.a2a.spec.TaskState;
15+
import io.a2a.spec.TaskStatusUpdateEvent;
1316
import org.jspecify.annotations.Nullable;
1417

1518
public class RestSSEEventListener {
@@ -29,7 +32,7 @@ public void onMessage(String message, @Nullable Future<Void> completableFuture)
2932
log.fine("Streaming message received: " + message);
3033
io.a2a.grpc.StreamResponse.Builder builder = io.a2a.grpc.StreamResponse.newBuilder();
3134
JsonFormat.parser().merge(message, builder);
32-
handleMessage(builder.build());
35+
handleMessage(builder.build(), completableFuture);
3336
} catch (InvalidProtocolBufferException e) {
3437
errorHandler.accept(RestErrorMapper.mapRestError(message, 500));
3538
}
@@ -44,7 +47,7 @@ public void onError(Throwable throwable, @Nullable Future<Void> future) {
4447
}
4548
}
4649

47-
private void handleMessage(StreamResponse response) {
50+
private void handleMessage(StreamResponse response, @Nullable Future<Void> future) {
4851
StreamingEventKind event;
4952
switch (response.getPayloadCase()) {
5053
case MESSAGE ->
@@ -62,6 +65,23 @@ private void handleMessage(StreamResponse response) {
6265
}
6366
}
6467
eventHandler.accept(event);
68+
69+
// Client-side auto-close on final events to prevent connection leaks
70+
// Handles both TaskStatusUpdateEvent and Task objects with final states
71+
// This covers late subscriptions to completed tasks and ensures no connection leaks
72+
boolean shouldClose = false;
73+
if (event instanceof TaskStatusUpdateEvent tue && tue.isFinal()) {
74+
shouldClose = true;
75+
} else if (event instanceof Task task) {
76+
TaskState state = task.status().state();
77+
if (state.isFinal()) {
78+
shouldClose = true;
79+
}
80+
}
81+
82+
if (shouldClose && future != null) {
83+
future.cancel(true); // close SSE channel
84+
}
6585
}
6686

6787
}

‎examples/cloud-deployment/scripts/deploy.sh‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,22 @@ echo ""
212212
echo "Deploying PostgreSQL..."
213213
kubectl apply -f ../k8s/01-postgres.yaml
214214
echo "Waiting for PostgreSQL to be ready..."
215+
216+
# Wait for pod to be created (StatefulSet takes time to create pod)
217+
for i in {1..30}; do
218+
if kubectl get pod -l app=postgres -n a2a-demo 2>/dev/null | grep -q postgres; then
219+
echo "PostgreSQL pod found, waiting for ready state..."
220+
break
221+
fi
222+
if [ $i -eq 30 ]; then
223+
echo -e "${RED}ERROR: PostgreSQL pod not created after 30 seconds${NC}"
224+
kubectl get statefulset -n a2a-demo
225+
exit 1
226+
fi
227+
sleep 1
228+
done
229+
230+
# Now wait for pod to be ready
215231
kubectl wait --for=condition=Ready pod -l app=postgres -n a2a-demo --timeout=120s
216232
echo -e "${GREEN}✓ PostgreSQL deployed${NC}"
217233

‎extras/common/src/main/java/io/a2a/extras/common/events/TaskFinalizedEvent.java‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,22 +5,28 @@
55
* This event is fired AFTER the database transaction commits, making it safe for downstream
66
* components to assume the task is durably stored.
77
*
8-
* <p>Used by the replicated queue manager to send poison pill events after ensuring
9-
* the final task state is committed to the database, eliminating race conditions.
8+
* <p>Used by the replicated queue manager to send the final task state before the poison pill,
9+
* ensuring correct event ordering across instances and eliminating race conditions.
1010
*/
1111
public class TaskFinalizedEvent {
1212
private final String taskId;
13+
private final Object task; // Task type from io.a2a.spec - using Object to avoid dependency
1314

14-
public TaskFinalizedEvent(String taskId) {
15+
public TaskFinalizedEvent(String taskId, Object task) {
1516
this.taskId = taskId;
17+
this.task = task;
1618
}
1719

1820
public String getTaskId() {
1921
return taskId;
2022
}
2123

24+
public Object getTask() {
25+
return task;
26+
}
27+
2228
@Override
2329
public String toString() {
24-
return "TaskFinalizedEvent{taskId='" + taskId + "'}";
30+
return "TaskFinalizedEvent{taskId='" + taskId + "', task=" + task + "}";
2531
}
2632
}

‎extras/push-notification-config-store-database-jpa/src/main/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/JpaDatabasePushNotificationConfigStore.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,4 +164,5 @@ public void deleteInfo(String taskId, String configId) {
164164
taskId, configId);
165165
}
166166
}
167+
167168
}

‎extras/queue-manager-replicated/core/src/main/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedEventQueueItem.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,16 @@ public void setClosedEvent(boolean closedEvent) {
149149
}
150150
}
151151

152+
/**
153+
* Check if this event is a Task event.
154+
* Task events should always be processed even for inactive tasks,
155+
* as they carry the final task state.
156+
* @return true if this is a Task event
157+
*/
158+
public boolean isTaskEvent() {
159+
return event instanceof io.a2a.spec.Task;
160+
}
161+
152162
@Override
153163
public String toString() {
154164
return "ReplicatedEventQueueItem{" +

‎extras/queue-manager-replicated/core/src/main/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManager.java‎

Lines changed: 69 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,10 @@
1313
import io.a2a.server.events.EventQueueFactory;
1414
import io.a2a.server.events.EventQueueItem;
1515
import io.a2a.server.events.InMemoryQueueManager;
16+
import io.a2a.server.events.MainEventBus;
1617
import io.a2a.server.events.QueueManager;
1718
import io.a2a.server.tasks.TaskStateProvider;
19+
import org.jspecify.annotations.Nullable;
1820
import org.slf4j.Logger;
1921
import org.slf4j.LoggerFactory;
2022

@@ -45,10 +47,12 @@ protected ReplicatedQueueManager() {
4547
}
4648

4749
@Inject
48-
public ReplicatedQueueManager(ReplicationStrategy replicationStrategy, TaskStateProvider taskStateProvider) {
50+
public ReplicatedQueueManager(ReplicationStrategy replicationStrategy,
51+
TaskStateProvider taskStateProvider,
52+
MainEventBus mainEventBus) {
4953
this.replicationStrategy = replicationStrategy;
5054
this.taskStateProvider = taskStateProvider;
51-
this.delegate = new InMemoryQueueManager(new ReplicatingEventQueueFactory(), taskStateProvider);
55+
this.delegate = new InMemoryQueueManager(new ReplicatingEventQueueFactory(), taskStateProvider, mainEventBus);
5256
}
5357

5458

@@ -77,8 +81,7 @@ public void close(String taskId) {
7781

7882
@Override
7983
public EventQueue createOrTap(String taskId) {
80-
EventQueue queue = delegate.createOrTap(taskId);
81-
return queue;
84+
return delegate.createOrTap(taskId);
8285
}
8386

8487
@Override
@@ -87,48 +90,81 @@ public void awaitQueuePollerStart(EventQueue eventQueue) throws InterruptedExcep
8790
}
8891

8992
public void onReplicatedEvent(@Observes ReplicatedEventQueueItem replicatedEvent) {
90-
// Check if task is still active before processing replicated event (unless it's a QueueClosedEvent)
91-
// QueueClosedEvent should always be processed to terminate streams, even for inactive tasks
93+
// Check if task is still active before processing replicated event
94+
// Always allow QueueClosedEvent and Task events (they carry final state)
95+
// Skip other event types for inactive tasks to prevent queue creation for expired tasks
9296
if (!replicatedEvent.isClosedEvent()
97+
&& !replicatedEvent.isTaskEvent()
9398
&& !taskStateProvider.isTaskActive(replicatedEvent.getTaskId())) {
9499
// Task is no longer active - skip processing this replicated event
95100
// This prevents creating queues for tasks that have been finalized beyond the grace period
96101
LOGGER.debug("Skipping replicated event for inactive task {}", replicatedEvent.getTaskId());
97102
return;
98103
}
99104

100-
// Get or create a ChildQueue for this task (creates MainQueue if it doesn't exist)
101-
EventQueue childQueue = delegate.createOrTap(replicatedEvent.getTaskId());
102-
103-
try {
104-
// Get the MainQueue to enqueue the replicated event item
105-
// We must use enqueueItem (not enqueueEvent) to preserve the isReplicated() flag
106-
// and avoid triggering the replication hook again (which would cause a replication loop)
107-
EventQueue mainQueue = delegate.get(replicatedEvent.getTaskId());
108-
if (mainQueue != null) {
109-
mainQueue.enqueueItem(replicatedEvent);
110-
} else {
111-
LOGGER.warn("MainQueue not found for task {}, cannot enqueue replicated event. This may happen if the queue was already cleaned up.",
112-
replicatedEvent.getTaskId());
113-
}
114-
} finally {
115-
// Close the temporary ChildQueue to prevent leaks
116-
// The MainQueue remains open for other consumers
117-
childQueue.close();
105+
// Get the MainQueue to enqueue the replicated event item
106+
// We must use enqueueItem (not enqueueEvent) to preserve the isReplicated() flag
107+
// and avoid triggering the replication hook again (which would cause a replication loop)
108+
//
109+
// IMPORTANT: We must NOT create a ChildQueue here! Creating and immediately closing
110+
// a ChildQueue means there are zero children when MainEventBusProcessor distributes
111+
// the event. Existing ChildQueues (from active client subscriptions) will receive
112+
// the event when MainEventBusProcessor distributes it to all children.
113+
//
114+
// If MainQueue doesn't exist, create it. This handles late-arriving replicated events
115+
// for tasks that were created on another instance.
116+
EventQueue mainQueue = delegate.get(replicatedEvent.getTaskId());
117+
if (mainQueue == null) {
118+
// MainQueue doesn't exist - create it by calling createOrTap and then get the MainQueue
119+
// Replicated events should always have real task IDs (not temp IDs) because
120+
// replication now happens AFTER TaskStore persistence in MainEventBusProcessor
121+
LOGGER.debug("Creating MainQueue for replicated event on task {}", replicatedEvent.getTaskId());
122+
delegate.createOrTap(replicatedEvent.getTaskId());
123+
mainQueue = delegate.get(replicatedEvent.getTaskId());
124+
}
125+
126+
if (mainQueue != null) {
127+
mainQueue.enqueueItem(replicatedEvent);
128+
} else {
129+
LOGGER.warn("MainQueue not found for task {}, cannot enqueue replicated event. This may happen if the queue was already cleaned up.",
130+
replicatedEvent.getTaskId());
118131
}
119132
}
120133

121134
/**
122135
* Observes task finalization events fired AFTER database transaction commits.
123-
* This guarantees the task's final state is durably stored before sending the poison pill.
136+
* This guarantees the task's final state is durably stored before replication.
137+
*
138+
* Sends TaskStatusUpdateEvent (not full Task) FIRST, then the poison pill (QueueClosedEvent),
139+
* ensuring correct event ordering across instances and eliminating race conditions where
140+
* the poison pill arrives before the final task state.
124141
*
125-
* @param event the task finalized event containing the task ID
142+
* IMPORTANT: We send TaskStatusUpdateEvent instead of full Task to maintain consistency
143+
* with local event distribution. Clients expect TaskStatusUpdateEvent for status changes,
144+
* and sending the full Task causes issues in remote instances where clients don't handle
145+
* bare Task objects the same way they handle TaskStatusUpdateEvent.
146+
*
147+
* @param event the task finalized event containing the task ID and final Task
126148
*/
127149
public void onTaskFinalized(@Observes(during = TransactionPhase.AFTER_SUCCESS) TaskFinalizedEvent event) {
128150
String taskId = event.getTaskId();
129-
LOGGER.debug("Task {} finalized - sending poison pill (QueueClosedEvent) after transaction commit", taskId);
151+
io.a2a.spec.Task finalTask = (io.a2a.spec.Task) event.getTask(); // Cast from Object
152+
153+
LOGGER.debug("Task {} finalized - sending TaskStatusUpdateEvent then poison pill (QueueClosedEvent) after transaction commit", taskId);
154+
155+
// Convert final Task to TaskStatusUpdateEvent to match local event distribution
156+
// This ensures remote instances receive the same event type as local instances
157+
io.a2a.spec.TaskStatusUpdateEvent finalStatusEvent = io.a2a.spec.TaskStatusUpdateEvent.builder()
158+
.taskId(taskId)
159+
.contextId(finalTask.contextId())
160+
.status(finalTask.status())
161+
.isFinal(true)
162+
.build();
163+
164+
// Send TaskStatusUpdateEvent FIRST to ensure it arrives before poison pill
165+
replicationStrategy.send(taskId, finalStatusEvent);
130166

131-
// Send poison pill directly via replication strategy
167+
// Then send poison pill
132168
// The transaction has committed, so the final state is guaranteed to be in the database
133169
io.a2a.server.events.QueueClosedEvent closedEvent = new io.a2a.server.events.QueueClosedEvent(taskId);
134170
replicationStrategy.send(taskId, closedEvent);
@@ -152,12 +188,11 @@ public EventQueue.EventQueueBuilder builder(String taskId) {
152188
// which sends the QueueClosedEvent after the database transaction commits.
153189
// This ensures proper ordering and transactional guarantees.
154190

155-
// Return the builder with callbacks
156-
return delegate.getEventQueueBuilder(taskId)
157-
.taskId(taskId)
158-
.hook(new ReplicationHook(taskId))
159-
.addOnCloseCallback(delegate.getCleanupCallback(taskId))
160-
.taskStateProvider(taskStateProvider);
191+
// Call createBaseEventQueueBuilder() directly to avoid infinite recursion
192+
// (getEventQueueBuilder() would delegate back to this factory, creating a loop)
193+
// The base builder already includes: taskId, cleanup callback, taskStateProvider, mainEventBus
194+
return delegate.createBaseEventQueueBuilder(taskId)
195+
.hook(new ReplicationHook(taskId));
161196
}
162197
}
163198

0 commit comments

Comments
 (0)