fix(server): ensure DefaultRequestHandler cancels producer on blocking failure

[Wondr-design]

Fixes #609.

This PR addresses an issue where DefaultRequestHandler.on_message_send would hang indefinitely if the request failed or was cancelled while waiting for the agent execution in blocking mode. The previous implementation waited for the producer task to finish in the finally block without cancelling it, leading to a deadlock if the agent was still running (e.g., stuck or long-running).

Changes:

• Modified DefaultRequestHandler._cleanup_producer to accept an optional cancel argument.
• Updated DefaultRequestHandler.on_message_send to cancel the producer task if the request fails (exception raised) or is cancelled (e.g., client timeout) while in blocking mode.
• Ensured background tasks (non-blocking) continue to run as expected.

Verification:

• Reproduction Script: Verified with a script simulating a hanging agent and client timeout. Before fix: hung for 10s. After fix: exited immediately (~1s) upon timeout.
• Regression Tests: Ran all existing tests in a2a/server/request_handlers and they passed.

Type of Change:

• Bug fix (non-breaking change which fixes an issue)

[gemini-code-assist]

Summary of Changes

Hello @Wondr-design, 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 critical issue in the DefaultRequestHandler where blocking requests could lead to an indefinite hang if the underlying agent execution failed or was cancelled. The fix ensures proper resource management by explicitly cancelling the producer task in such scenarios, thereby preventing deadlocks and improving the robustness and responsiveness of the server when handling agent interactions.

Highlights

• Producer Task Cancellation: The DefaultRequestHandler.on_message_send method now ensures that the producer task is explicitly cancelled if a blocking request fails or is cancelled, preventing indefinite hangs.
• Flexible Producer Cleanup: The _cleanup_producer method has been enhanced to accept an optional cancel argument, allowing for conditional cancellation of the producer task during cleanup.
• Success Tracking: A success flag has been introduced within on_message_send to accurately determine if the agent execution completed without errors, guiding the decision to cancel the producer during cleanup.

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 effectively addresses a potential deadlock in on_message_send during blocking failures by ensuring the producer task is cancelled. The introduction of the success flag and the cancel parameter in _cleanup_producer is a clean approach. My review includes one suggestion to improve the exception handling within _cleanup_producer to prevent a CancelledError from masking the original failure exception, making the cleanup process more robust.

[gemini-code-assist]

When cancel=True, producer_task.cancel() is called. Subsequently, await producer_task will raise an asyncio.CancelledError. This exception is currently unhandled. If _cleanup_producer is called from a finally block (as it is in on_message_send), this unhandled CancelledError can mask the original exception that led to the finally block executing. To prevent this and handle the expected cancellation gracefully, it's best to wrap the await in a try...except block.

Suggested change

	await producer_task
	try:
	await producer_task
	except asyncio.CancelledError:
	# This is expected if the task was cancelled, so we can ignore.
	pass