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

kabir · 2025-10-09T13:26:45Z

…ssage

Add test for non-blocking multi-message pattern where multiple sendMessage calls to same taskId continue processing. Test uncovered EventConsumer polling loops never terminating, causing resource leaks. Fix: Close EventQueue at appropriate times to stop infinite polling:

EventConsumer.close() explicitly closes queue
InMemoryQueueManager.close() fails fast on missing queue
DefaultRequestHandler cleanup: streaming (always close) vs non-blocking (close only if task final, for multi-message support)
Refactor ResultAggregator to use AsyncUtils.consumer() consistently

…ssage Add test for non-blocking multi-message pattern where multiple sendMessage calls to same taskId continue processing. Test uncovered EventConsumer polling loops never terminating, causing resource leaks. Fix: Close EventQueue at appropriate times to stop infinite polling: - EventConsumer.close() explicitly closes queue - InMemoryQueueManager.close() fails fast on missing queue - DefaultRequestHandler cleanup: streaming (always close) vs non-blocking (close only if task final, for multi-message support) - Refactor ResultAggregator to use AsyncUtils.consumer() consistently

gemini-code-assist · 2025-10-09T13:27:06Z

Summary of Changes

Hello @kabir, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves critical resource leak issues arising from infinite polling loops in the event consumption mechanism, especially when handling non-blocking messages. By introducing precise queue closure logic and refining event aggregation, it ensures that multiple messages can be sent to the same task without exhausting resources, thereby supporting a more robust 'fire-and-forget' pattern. The changes improve the stability and efficiency of task processing by preventing unmanaged event queues.

Highlights

Resource Leak Fix: Addressed and fixed an issue where EventConsumer polling loops would never terminate, leading to resource leaks, particularly in non-blocking multi-message scenarios.
Multi-Message Support: Enabled proper handling of multiple non-blocking sendMessage calls to the same task ID, ensuring continuous processing without resource exhaustion.
Queue Management Enhancements: Implemented explicit queue closing mechanisms in EventConsumer and InMemoryQueueManager, and refined cleanup logic in DefaultRequestHandler to close queues appropriately based on streaming vs. non-blocking calls and task finality.
ResultAggregator Refactoring: Refactored ResultAggregator to consistently use AsyncUtils.consumer() and manage event consumption and interruption more robustly, including running subscriptions on background threads to prevent blocking.
New Test Cases: Introduced new test cases (testBlockingMessageContinueTask and testNonBlockingWithMultipleMessages) to validate the multi-message pattern and ensure correct queue behavior and task lifecycle management.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request provides a comprehensive fix for resource leaks and event loop blocking issues related to event queue consumption. The changes correctly introduce logic to close event queues at appropriate times, supporting multi-message patterns for non-blocking calls by keeping queues open for non-final tasks. The refactoring of ResultAggregator to run polling on a background thread is a critical improvement. The addition of thorough unit and integration tests to validate the new behavior is excellent.

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

kabir · 2025-10-09T13:35:53Z

/gemini review

gemini-code-assist

Code Review

This pull request effectively addresses a resource leak issue with non-terminating EventConsumer loops by introducing explicit queue closing mechanisms. The refactoring in ResultAggregator to run consumption on a background thread is a crucial improvement to prevent blocking worker threads. The new integration tests are comprehensive and provide good coverage for the new multi-message and non-blocking scenarios.

I've identified a potential logic issue in ResultAggregator for blocking calls that could prevent them from returning as early as expected, and a minor improvement for error handling in DefaultRequestHandler. Overall, this is a solid fix for a complex problem.

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

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

kabir · 2025-10-09T13:54:23Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces important fixes for handling multiple messages to a single task in non-blocking scenarios and resolves a resource leak by ensuring event queues are closed properly. The changes are well-structured, particularly the refactoring in ResultAggregator to handle consumption in a background thread, which improves robustness. The addition of specific tests for these scenarios is also a great improvement.

I have a couple of suggestions to improve code maintainability by reducing duplication and simplifying conditional logic. Overall, this is a solid contribution that addresses a critical part of the server's behavior.

gemini-code-assist · 2025-10-09T13:56:55Z

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

+                        // 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() != 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);
+                        }


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); } } // ...

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

ehsavoie · 2025-10-10T13:34:36Z

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

        }
    }

+    private boolean isTaskFinalInTaskStore(String taskId) {


It is not used, should we remove it ?

Looks like this could be used on lines 481 and 533

fjuma · 2025-10-10T13:49:28Z

server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java

        }
-        return existing == null ? newQueue : existing.tap();
+        EventQueue result = existing == null ? newQueue : existing.tap();
+        if (existing == null) {


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

fjuma · 2025-10-10T13:53:26Z

transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java

 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;


Just minor, looks like this isn't needed right?

kabir · 2025-10-10T14:10:17Z

Replaced by #333

@amithsourya ^^ contains this and the second half of what I was mentioning. It undos some of the changes in this PR

amithsourya · 2025-10-11T16:40:30Z

@kabir It took me a while to wrap my head around everything that's happening in #333 & #331, but I can really appreciate the elegance of the solution. This is exactly the kind of sophisticated code I hope to write someday 😃

gemini-code-assist bot reviewed Oct 9, 2025

View reviewed changes

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

Gemini review

f860fbf

gemini-code-assist bot reviewed Oct 9, 2025

View reviewed changes

server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java Outdated Show resolved Hide resolved

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

kabir added 2 commits October 9, 2025 14:42

Another review suggestion

9d02bea

Gemini review

d30ff35

gemini-code-assist bot reviewed Oct 9, 2025

View reviewed changes

Gemini review

664c99a

kabir mentioned this pull request Oct 10, 2025

Add a test to JSONRPCHandlerTest for sending a streaming message where multiple events are received from the server #52

Closed

ehsavoie reviewed Oct 10, 2025

View reviewed changes

fjuma reviewed Oct 10, 2025

View reviewed changes

kabir closed this Oct 13, 2025

kabir deleted the event-queues-take2 branch November 3, 2025 13:02

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

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

Uh oh!

Conversation

kabir commented Oct 9, 2025

Uh oh!

gemini-code-assist bot commented Oct 9, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

kabir commented Oct 9, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

kabir commented Oct 9, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Oct 9, 2025

Choose a reason for hiding this comment

Uh oh!

kabir Oct 9, 2025

Choose a reason for hiding this comment

Uh oh!

ehsavoie Oct 10, 2025

Choose a reason for hiding this comment

Uh oh!

fjuma Oct 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fjuma Oct 10, 2025

Choose a reason for hiding this comment

Uh oh!

fjuma Oct 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kabir commented Oct 10, 2025

Uh oh!

amithsourya commented Oct 11, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

fjuma Oct 10, 2025 •

edited

Loading

fjuma Oct 10, 2025 •

edited

Loading