fix: flaky Agent to Agent tests

[kabir]

No description provided.

[gemini-code-assist]

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 flakiness in Agent-to-Agent tests by overhauling the internal delegation mechanism. The previous asynchronous, callback-driven approach, which was prone to race conditions, has been replaced with a simplified synchronous blocking client call. This ensures that the test reliably receives the final state of delegated tasks, significantly improving test stability and providing better diagnostics for unexpected states.

Highlights

• Flaky Test Fix: Addressed flakiness in Agent-to-Agent tests by refactoring the delegation handling mechanism to use a blocking client call instead of an asynchronous, latch-based approach. This change aims to eliminate race conditions that caused test instability.
• Synchronous Delegation: Modified the handleDelegation method to perform a synchronous client call (streaming=false), ensuring that the final task state is returned directly without relying on CountDownLatch or asynchronous callbacks.
• Concurrency Simplification: Removed CountDownLatch, TimeUnit, and BiConsumer imports and their associated logic, simplifying the concurrency model within the test utility.
• Race Condition Diagnostics: Introduced a diagnostic check to verify that the blocking delegation call returns a task in a final state. If a non-final task is returned, it indicates a potential server-side race condition and logs a detailed error message.

Changelog

• tests/server-common/src/test/java/io/a2a/server/apps/common/AgentExecutorProducer.java
  • Removed imports for CountDownLatch, TimeUnit, and BiConsumer.
  • Updated the handleDelegation method to use a blocking client call for delegation, eliminating the need for explicit latch-based waiting.
  • Replaced asynchronous task result capture with direct assignment from the blocking call's event consumer.
  • Added a diagnostic check to log and fail if a blocking delegation call returns a non-final task state, indicating a potential race condition.
  • Removed InterruptedException handling as the asynchronous wait mechanism was removed.

Activity

• No specific activity (comments, reviews) has been recorded for this pull request yet.

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.

Footnotes

1. 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 refactors the agent-to-agent delegation test logic to use a blocking client call instead of an asynchronous one with a CountDownLatch. This change aims to eliminate flakiness in the tests by simplifying the synchronization logic. The new implementation also includes a helpful diagnostic check to ensure the blocking call returns a final task state as expected. The overall change is a good improvement for test stability. No issues were found that contradict the repository's rules.

[kabir]

/gemini review

[gemini-code-assist]

Code Review

This pull request effectively addresses flaky agent-to-agent tests by refactoring the delegation logic to use a synchronous, blocking client call instead of an asynchronous, latch-based approach. This simplification in AgentExecutorProducer makes the code more robust and easier to follow. The corresponding change to use a non-streaming client in AbstractA2AServerTest is consistent and correct. I particularly appreciate the addition of the diagnostic check to verify that the blocking call returns a final task state, which is an excellent safeguard against future regressions.

[gemini-code-assist]

The consumer lambda for sendMessage includes a check for TaskUpdateEvent. However, since this sendMessage call is on a non-streaming client (streaming=false), the client will only invoke the consumer with a final TaskEvent or MessageEvent after the blocking call completes. It will not receive intermediate TaskUpdateEvents. Therefore, the else if (event instanceof TaskUpdateEvent tue) branch is effectively dead code in this context. Removing it would make the code clearer and more accurately reflect the behavior of a non-streaming client call.

                        if (event instanceof TaskEvent te) {
                            taskRef.set(te.getTask());
                        }