chore: fixing flaky test

[ehsavoie]

The Problem
The second sendMessage call doesn't wait for its streaming subscription to be established. This creates a race condition:

• The agent receives message2 and immediately sends events (artifact-2 + completion)
• But the streamConsumer subscription might not be fully wired up yet
• The events get delivered to the resubConsumer (which is already established) but missed by the streamConsumer

The Fix
Add the awaitStreamingSubscription() pattern before the second sendMessage call:

Fixes #627 🦕

[gemini-code-assist]

Summary of Changes

Hello @ehsavoie, 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 a known flaky test (testNonBlockingWithMultipleMessages) that intermittently failed due to a race condition. The issue stemmed from the test sending a second message before the corresponding streaming subscription was fully active, causing events to be misdirected or lost. The fix introduces a synchronization mechanism to explicitly wait for the streaming subscription to be ready, thereby stabilizing the test and ensuring reliable event processing.

Highlights

• Flaky Test Fix: Addressed a race condition in testNonBlockingWithMultipleMessages by ensuring the streaming subscription is fully established before sending the second message, preventing events from being missed by the streamConsumer.

Changelog

• tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java
  • Introduced a CountDownLatch named streamSubscriptionLatch to manage synchronization.
  • Added a call to awaitStreamingSubscription() which decrements the streamSubscriptionLatch upon completion.
  • Implemented a wait (streamSubscriptionLatch.await) before sending the second message (message2) to ensure the streaming subscription is established.
  • Added an assertion to verify that the stream subscription is established within a 15-second timeout.

Activity

• No human activity (comments, reviews) has been recorded on 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 aims to fix a flaky test by introducing a wait for a streaming subscription to be established. While the intention is correct, the implementation contains a logical flaw where the wait happens after the action that triggers events, which doesn't fully eliminate the race condition. I've pointed out this issue and suggested a different approach, while also noting that a more fundamental change to the test API might be necessary for a robust solution.

[gemini-code-assist]

The await call is placed after getClient().sendMessage(...), which doesn't fully resolve the race condition. The sendMessage call triggers the agent to send events immediately. These events could be sent and missed by the client before this await completes, especially if the client-side consumer is not fully wired up.

To reliably fix this, you should wait for the subscription to be established before triggering the agent. The await should be before sendMessage. I've suggested this change below.

However, this will likely lead to a deadlock, as sendMessage is what triggers the subscription creation that awaitStreamingSubscription is waiting for. This indicates a deeper issue with the test's design or the sendMessage API, which might need a more comprehensive fix to allow waiting for the subscription to be ready before events are generated.

Suggested change

	getClient().sendMessage(message2, List.of(streamConsumer), null);
	// Ensure subscription is established before agent sends events
	assertTrue(streamSubscriptionLatch.await(15, TimeUnit.SECONDS),
	"Stream subscription should be established");
	// Ensure subscription is established before agent sends events
	assertTrue(streamSubscriptionLatch.await(15, TimeUnit.SECONDS),
	"Stream subscription should be established");
	getClient().sendMessage(message2, List.of(streamConsumer), null);