fix: return background task from consume_and_break_on_interrupt to prevent GC

[kevinlu310]

Description

ResultAggregator.consume_and_break_on_interrupt creates a background asyncio.Task to continue consuming events after an interruption (non-blocking or auth_required), but discards the task reference. On Python 3.12+ the event loop only holds weak references to tasks, so the garbage collector can silently collect the task before it completes — dropping remaining events (completed/failed status) and push notification callbacks.

Return the background task as a third tuple element so callers can hold a strong reference. DefaultRequestHandler.on_message_send now tracks it via _track_background_task(), the same mechanism already used for other background work.

• Follow the CONTRIBUTING Guide.
• Make your Pull Request title in the https://www.conventionalcommits.org/ specification.
  • Important Prefixes for release-please:
    • fix: which represents bug fixes, and correlates to a SemVer patch.
    • feat: represents a new feature, and correlates to a SemVer minor.
    • feat!:, or fix!:, refactor!:, etc., which represent a breaking change (indicated by the !) and will result in a SemVer major.
• Ensure the tests and linter pass (Run bash scripts/format.sh from the repository root to format)

Fixes #774 🦕

[gemini-code-assist]

Summary of Changes

Hello, 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 addresses a critical issue in Python 3.12+ where background asyncio.Task objects, created for continuing event consumption after an interruption, could be prematurely garbage collected due to the event loop holding only weak references. The solution involves modifying the consume_and_break_on_interrupt method to return a strong reference to this background task, allowing the calling DefaultRequestHandler to explicitly track it. This ensures that all events are processed and push notification callbacks are executed reliably, preventing data loss and unexpected behavior.

Highlights

• Prevented premature garbage collection: The consume_and_break_on_interrupt method now returns the asyncio.Task it creates for background consumption, ensuring a strong reference is held and preventing its premature garbage collection in Python 3.12+.
• Integrated task tracking: DefaultRequestHandler.on_message_send was updated to receive and track this background task using _track_background_task, preventing its premature termination.
• Updated method signature and documentation: The return type hint and docstring for consume_and_break_on_interrupt were adjusted to reflect the new asyncio.Task return value and explain the necessity for callers to hold a strong reference.
• Comprehensive test updates: Existing unit tests were modified across both request_handlers and result_aggregator to accommodate the new return value from consume_and_break_on_interrupt, ensuring correctness and robustness.

Changelog

• src/a2a/server/request_handlers/default_request_handler.py
  • Updated on_message_send to capture the background task returned by consume_and_break_on_interrupt.
  • Added logic to set the name of the captured background task and track it using _track_background_task.
• src/a2a/server/tasks/result_aggregator.py
  • Modified the return type hint of consume_and_break_on_interrupt to include asyncio.Task | None as a third element.
  • Updated the docstring for consume_and_break_on_interrupt to describe the new background task return value and the requirement for callers to hold a strong reference.
  • Assigned the asyncio.create_task result to a bg_task variable and included it in the return tuple.
• tests/server/request_handlers/test_default_request_handler.py
  • Modified mock return values for consume_and_break_on_interrupt in various test cases to include None or a MagicMock(spec=asyncio.Task) for the new background task.
• tests/server/tasks/test_result_aggregator.py
  • Updated test assertions to check for the new bg_task return value, asserting None or assertIsNotNone as appropriate.
  • Adjusted mock setup for asyncio.create_task to return a sentinel_task for testing purposes.
  • Modified calls to consume_and_break_on_interrupt to unpack the now three-element return tuple.

Activity

• The pull request addresses and fixes issue [Bug]: ResultAggregator._continue_consuming background task is not referenced and may be garbage collected #774, which likely details the problem of background tasks being garbage collected.
• The author provided a clear description outlining the problem with asyncio.Task garbage collection in Python 3.12+ and the proposed solution.
• The author included a standard checklist for PR submission, indicating adherence to contribution guidelines, conventional commits, passing tests/linters, and documentation updates.

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 addresses a critical garbage collection issue with asyncio.Task on Python 3.12+ by ensuring a strong reference is held for background tasks. The approach of returning the background task from consume_and_break_on_interrupt and tracking it in DefaultRequestHandler is correct and well-implemented. The changes are clean, include necessary type hint updates and excellent docstring explanations. The tests have been thoroughly updated to reflect the new return signature. I've found one minor redundancy in a test mock setup that could be cleaned up.

[guglielmo-san]

I believe that this comment should be also removed now