Skip to content

Commit 3a98eb4

Browse files
kabirkabir
committed
Gemini review
1 parent 2255492 commit 3a98eb4

5 files changed

Lines changed: 30 additions & 17 deletions

File tree

‎extras/queue-manager-replicated/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -290,7 +290,7 @@ public void save(Task task) {
290290
```java
291291
public void onTaskFinalized(@Observes(during = TransactionPhase.AFTER_SUCCESS) TaskFinalizedEvent event) {
292292
String taskId = event.getTaskId();
293-
LOGGER.info("Task {} finalized - sending poison pill after transaction commit", taskId);
293+
LOGGER.debug("Task {} finalized - sending poison pill after transaction commit", taskId);
294294

295295
// Send QueueClosedEvent to all nodes via replication
296296
QueueClosedEvent closedEvent = new QueueClosedEvent(taskId);

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,9 @@ public class ReplicatedQueueManager implements QueueManager {
3333

3434
@Inject
3535
public ReplicatedQueueManager(ReplicationStrategy replicationStrategy, TaskStateProvider taskStateProvider) {
36-
this.delegate = new InMemoryQueueManager(new ReplicatingEventQueueFactory(), taskStateProvider);
37-
this.taskStateProvider = taskStateProvider;
3836
this.replicationStrategy = replicationStrategy;
37+
this.taskStateProvider = taskStateProvider;
38+
this.delegate = new InMemoryQueueManager(new ReplicatingEventQueueFactory(), taskStateProvider);
3939
}
4040

4141

@@ -118,9 +118,7 @@ public void onTaskFinalized(@Observes(during = TransactionPhase.AFTER_SUCCESS) T
118118
// Send poison pill directly via replication strategy
119119
// The transaction has committed, so the final state is guaranteed to be in the database
120120
io.a2a.server.events.QueueClosedEvent closedEvent = new io.a2a.server.events.QueueClosedEvent(taskId);
121-
if (replicationStrategy != null) {
122-
replicationStrategy.send(taskId, closedEvent);
123-
}
121+
replicationStrategy.send(taskId, closedEvent);
124122
}
125123

126124
@Override
@@ -163,7 +161,7 @@ public void onEnqueue(EventQueueItem item) {
163161
if (!item.isReplicated()) {
164162
// Only replicate if this isn't already a replicated event
165163
// This prevents replication loops
166-
if (replicationStrategy != null && taskId != null) {
164+
if (taskId != null) {
167165
replicationStrategy.send(taskId, item.getEvent());
168166
}
169167
}

‎extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ class ReplicatedQueueManagerTest {
3838

3939
@BeforeEach
4040
void setUp() {
41-
queueManager = new ReplicatedQueueManager(null, new MockTaskStateProvider(true));
41+
queueManager = new ReplicatedQueueManager(new NoOpReplicationStrategy(), new MockTaskStateProvider(true));
4242
testEvent = new TaskStatusUpdateEvent.Builder()
4343
.taskId("test-task")
4444
.contextId("test-context")
@@ -485,6 +485,13 @@ void testReplicatedQueueClosedEventTerminatesConsumer() throws InterruptedExcept
485485
"Second event should be QueueClosedEvent");
486486
}
487487

488+
private static class NoOpReplicationStrategy implements ReplicationStrategy {
489+
@Override
490+
public void send(String taskId, Event event) {
491+
// No-op for tests that don't care about replication
492+
}
493+
}
494+
488495
private static class CountingReplicationStrategy implements ReplicationStrategy {
489496
private final AtomicInteger callCount = new AtomicInteger(0);
490497
private volatile String lastTaskId;

‎extras/task-store-database-jpa/src/main/java/io/a2a/extras/taskstore/database/jpa/JpaTask.java‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -87,20 +87,26 @@ public void setTask(Task task) throws JsonProcessingException {
8787
id = task.getId();
8888
}
8989
this.task = task;
90-
if (task.getStatus() != null && task.getStatus().state() != null) {
91-
// Set finalizedAt timestamp idempotently if task is in final state
92-
setFinalizedAt(Instant.now(), task.getStatus().state().isFinal());
93-
}
90+
updateFinalizedTimestamp(task);
9491
}
9592

9693
static JpaTask createFromTask(Task task) throws JsonProcessingException {
9794
String json = Utils.OBJECT_MAPPER.writeValueAsString(task);
9895
JpaTask jpaTask = new JpaTask(task.getId(), json);
9996
jpaTask.task = task;
97+
jpaTask.updateFinalizedTimestamp(task);
98+
return jpaTask;
99+
}
100+
101+
/**
102+
* Updates the finalizedAt timestamp if the task is in a final state.
103+
* This method is idempotent and only sets the timestamp on first finalization.
104+
*
105+
* @param task the task to check for finalization
106+
*/
107+
private void updateFinalizedTimestamp(Task task) {
100108
if (task.getStatus() != null && task.getStatus().state() != null) {
101-
// Set finalizedAt timestamp idempotently if task is in final state
102-
jpaTask.setFinalizedAt(Instant.now(), task.getStatus().state().isFinal());
109+
setFinalizedAt(Instant.now(), task.getStatus().state().isFinal());
103110
}
104-
return jpaTask;
105111
}
106112
}

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -197,8 +197,10 @@ else if (!blocking) {
197197
// Complete the future to unblock the main thread
198198
interrupted.set(true);
199199
completionFuture.complete(null);
200-
// For interrupted/background tasks, complete the consumption future immediately
201-
// so cleanup doesn't wait for the entire subscription to finish
200+
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.
202204
consumptionCompletionFuture.complete(null);
203205

204206
// Continue consuming in background - keep requesting events

0 commit comments

Comments
 (0)