Skip to content

fix: Support multiple messages to same task after non-blocking sendMe… #331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kabir wants to merge 5 commits into from
Closed

fix: Support multiple messages to same task after non-blocking sendMe… #331

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
746a903
fix: Support multiple messages to same task after non-blocking sendMe…
kabir Oct 7, 2025
f860fbf
Gemini review
kabir Oct 9, 2025
9d02bea
Another review suggestion
kabir Oct 9, 2025
d30ff35
Gemini review
kabir Oct 9, 2025
664c99a
Gemini review
kabir Oct 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Binary file removed .serena/cache/java/document_symbols_cache_v23-06-25.pkl
View file
Open in desktop
Binary file not shown.
11 changes: 11 additions & 0 deletions server-common/src/main/java/io/a2a/server/events/EventConsumer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,11 @@
import mutiny.zero.BackpressureStrategy;
import mutiny.zero.TubeConfiguration;
import mutiny.zero.ZeroPublisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class EventConsumer {
private static final Logger LOGGER = LoggerFactory.getLogger(EventConsumer.class);
private final EventQueue queue;
private Throwable error;

Expand All @@ -21,6 +24,7 @@ public class EventConsumer {

public EventConsumer(EventQueue queue) {
this.queue = queue;
LOGGER.debug("EventConsumer created with queue {}", System.identityHashCode(queue));
}

public Event consumeOne() throws A2AServerException, EventQueueClosedException {
Expand Down Expand Up @@ -107,4 +111,11 @@ public EnhancedRunnable.DoneCallback createAgentRunnableDoneCallback() {
}
};
}

public void close() {
// Close the queue to stop the polling loop in consumeAll()
// This will cause EventQueueClosedException and exit the while(true) loop
LOGGER.debug("EventConsumer closing queue {}", System.identityHashCode(queue));
queue.close();
}
}
3 changes: 3 additions & 0 deletions server-common/src/main/java/io/a2a/server/events/EventQueue.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,11 +111,14 @@ public Event dequeueEvent(int waitMilliSeconds) throws EventQueueClosedException
return event;
}
try {
LOGGER.trace("Polling queue {} (wait={}ms)", System.identityHashCode(this), waitMilliSeconds);
Event event = queue.poll(waitMilliSeconds, TimeUnit.MILLISECONDS);
if (event != null) {
// Call toString() since for errors we don't really want the full stacktrace
LOGGER.debug("Dequeued event (waiting) {} {}", this, event instanceof Throwable ? event.toString() : event);
semaphore.release();
} else {
LOGGER.trace("Dequeue timeout (null) from queue {}", System.identityHashCode(this));
}
return event;
} catch (InterruptedException e) {
Expand Down
16 changes: 15 additions & 1 deletion server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,13 @@
import java.util.concurrent.ConcurrentMap;

import jakarta.enterprise.context.ApplicationScoped;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@ApplicationScoped
public class InMemoryQueueManager implements QueueManager {
private static final Logger LOGGER = LoggerFactory.getLogger(InMemoryQueueManager.class);

private final ConcurrentMap<String, EventQueue> queues = new ConcurrentHashMap<>();
private final EventQueueFactory factory;

Expand Down Expand Up @@ -43,6 +47,9 @@ public void close(String taskId) {
if (existing == null) {
throw new NoTaskQueueException();
}
// Close the queue to stop EventConsumer polling loop
LOGGER.debug("Closing queue {} for task {}", System.identityHashCode(existing), taskId);
existing.close();
}

@Override
Expand All @@ -57,7 +64,14 @@ public EventQueue createOrTap(String taskId) {
// Make sure an existing queue has not been added in the meantime
existing = queues.putIfAbsent(taskId, newQueue);
}
return existing == null ? newQueue : existing.tap();
EventQueue result = existing == null ? newQueue : existing.tap();
if (existing == null) {
Copy link
Collaborator

@fjuma fjuma Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to check, can this be null here? Let me know if I'm missing something.

LOGGER.debug("Created new queue {} for task {}", System.identityHashCode(result), taskId);
} else {
LOGGER.debug("Tapped existing queue {} -> child {} for task {}",
System.identityHashCode(existing), System.identityHashCode(result), taskId);
}
return result;
}

@Override
Expand Down
70 changes: 61 additions & 9 deletions server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import io.a2a.server.events.EnhancedRunnable;
import io.a2a.server.events.EventConsumer;
import io.a2a.server.events.EventQueue;
import io.a2a.server.events.NoTaskQueueException;
import io.a2a.server.events.QueueManager;
import io.a2a.server.events.TaskQueueExistsException;
import io.a2a.server.tasks.PushNotificationConfigStore;
Expand Down Expand Up @@ -201,7 +202,7 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
// any errors thrown by the producerRunnable are not picked up by the consumer
producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
etai = resultAggregator.consumeAndBreakOnInterrupt(consumer, blocking, pushNotificationCallback);

if (etai == null) {
LOGGER.debug("No result, throwing InternalError");
throw new InternalError("No result");
Expand All @@ -218,10 +219,10 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
pushNotificationCallback.run();
} finally {
if (interruptedOrNonBlocking) {
CompletableFuture<Void> cleanupTask = CompletableFuture.runAsync(() -> cleanupProducer(taskId), executor);
CompletableFuture<Void> cleanupTask = CompletableFuture.runAsync(() -> cleanupProducer(taskId, false), executor);
trackBackgroundTask(cleanupTask);
} else {
cleanupProducer(taskId);
cleanupProducer(taskId, false);
}
}

Expand Down Expand Up @@ -342,7 +343,7 @@ private void startBackgroundConsumption() {
}
});
} finally {
CompletableFuture<Void> cleanupTask = CompletableFuture.runAsync(() -> cleanupProducer(taskId.get()), executor);
CompletableFuture<Void> cleanupTask = CompletableFuture.runAsync(() -> cleanupProducer(taskId.get(), true), executor);
trackBackgroundTask(cleanupTask);
}
}
Expand Down Expand Up @@ -473,8 +474,17 @@ public void run() {
.whenComplete((v, err) -> {
if (err != null) {
runnable.setError(err);
// Close queue on error
queue.close();
} else {
// Only close queue if task is in a final state
Task task = taskStore.get(taskId);
if (task != null && task.getStatus().state().isFinal()) {
queue.close();
} else {
LOGGER.debug("Task {} not in final state or not yet created, keeping queue open", taskId);
}
}
queue.close();
runnable.invokeDoneCallbacks();
});
runningAgents.put(taskId, cf);
Expand All @@ -499,11 +509,48 @@ private void trackBackgroundTask(CompletableFuture<Void> task) {
});
}

private void cleanupProducer(String taskId) {
private void cleanupProducer(String taskId, boolean isStreaming) {
LOGGER.debug("Starting cleanup for task {} (streaming={})", taskId, isStreaming);
// TODO the Python implementation waits for the producerRunnable
runningAgents.get(taskId)
.whenComplete((v, t) -> {
queueManager.close(taskId);
CompletableFuture<Void> agentFuture = runningAgents.get(taskId);
if (agentFuture == null) {
LOGGER.debug("No running agent found for task {}", taskId);
return;
}
agentFuture.whenComplete((v, t) -> {
LOGGER.debug("Agent completed for task {}", taskId);

boolean closeQueue = false;
if (isStreaming) {
// For streaming calls, always close queue when agent completes
// Each streaming call is independent and doesn't support multi-message reuse
LOGGER.debug("Streaming call, closing queue for task {}", taskId);
closeQueue = true;
} else {
// For non-streaming calls, only close queue if task is in final state
// For non-final states, queue must stay open for potential future messages to same taskId
// so we can handle the "fire and forget' case used e.g. in the TCK
Task task = taskStore.get(taskId);
if (task != null && task.getStatus().state().isFinal()) {
LOGGER.debug("Task in final state, closing queue for task {}", taskId);
closeQueue = true;
} else {
LOGGER.debug("Task not in final state, keeping queue open for task {}", taskId);
}
Comment on lines 530 to 539
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Oct 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic to check if a task is in a final state is duplicated here and in the whenComplete block of registerAndExecuteAgentAsync (lines 481-483).

To improve maintainability and reduce code duplication, consider extracting this check into a private helper method.

For example:

private boolean isTaskFinal(String taskId) {
    Task task = taskStore.get(taskId);
    return task != null && task.getStatus() != null && task.getStatus().state().isFinal();
}

You could then simplify the logic in both places.

In cleanupProducer:

// ...
} else {
    // For non-streaming calls, only close queue if task is in final state
    if (isTaskFinal(taskId)) {
        LOGGER.debug("Task in final state, closing queue for task {}", taskId);
        closeQueue = true;
    } else {
        LOGGER.debug("Task not in final state, keeping queue open for task {}", taskId);
    }
}
// ...

In registerAndExecuteAgentAsync:

// ...
} else {
    // Only close queue if task is in a final state
    if (isTaskFinal(taskId)) {
        queue.close();
    } else {
        LOGGER.debug("Task {} not in final state or not yet created, keeping queue open", taskId);
    }
}
// ...

Copy link
Collaborator Author

@kabir kabir Oct 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is clearer the way I have it. I removed the task.getStatus() != null bit though since it isn't necessary

}

if (closeQueue) {
try {
queueManager.close(taskId);
} catch (NoTaskQueueException e) {
// This can happen if the queue was already closed and removed, which is not an error in this cleanup context.
LOGGER.debug("Queue for task {} was already closed or does not exist.", taskId);
} catch (Exception e) {
LOGGER.debug("Error closing queue for task {}: {}", taskId, e.getMessage());
}
}

// Always remove from running agents
runningAgents.remove(taskId);
});
}
Expand Down Expand Up @@ -545,5 +592,10 @@ private void sendPushNotification(String taskId, ResultAggregator resultAggregat
}
}

private boolean isTaskFinalInTaskStore(String taskId) {
Copy link
Collaborator

@ehsavoie ehsavoie Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is not used, should we remove it ?

Copy link
Collaborator

@fjuma fjuma Oct 10, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like this could be used on lines 481 and 533

Task task = taskStore.get(taskId);
return task != null && task.getStatus().state().isFinal();
}

private record MessageSendSetup(TaskManager taskManager, Task task, RequestContext requestContext) {}
}
Loading