fix: Prevent client disconnect from stopping task execution

[ovidiutaralesca]

Issue

• Client disconnect triggered synchronous cleanup.
• That awaited the producer task, effectively tying producer lifetime to the client connection.
• Reconnecting with tasks/resubscribe would not receive further events because the producer had already been forced to finish.

This behaviour no longer raises a asyncio.exceptions.CancelledError like claimed in #296 due to this fix: #383, but tasks/resubscribe still didn't behave as expected.

How it's reproduced

In any streaming agent: Simply sending a (longer-running) message/stream, disconnecting, and then reconnecting to the task using tasks/resubscribe will no longer yield events, even though the task should have been still running.

Fix

Code

The fix is an one-liner. Now:

• Client disconnect schedules cleanup in the background and returns immediately.
• Producer continues; resubscribe taps the existing queue and receives subsequent events.
• Cleanup still runs once the producer completes.

Tests

Existing tests:

• Changed existing tests that asserted on AgentExecutor.execute by adding an asyncio.Event latch to wait until the background producer hits execute.

New tests:

• test_stream_disconnect_then_resubscribe_receives_future_events -- start streaming, disconnect, resubscribe, and confirm future events are received.
• test_on_message_send_stream_client_disconnect_triggers_background_cleanup_and_producer_continues -- to validate that disconnecting is non-blocking, producer continues, and cleanup completes afterward.

Fixes #296

[gemini-code-assist]

Summary of Changes

Hello @ovidiutaralesca, 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 an issue where client disconnections would prematurely terminate ongoing task execution by synchronously awaiting cleanup. The fix modifies the cleanup process to run as a background task, ensuring that producer tasks continue unhindered by client disconnects. This allows clients to successfully resubscribe to tasks and receive subsequent events, improving the robustness of streaming agent interactions.

Highlights

• Asynchronous Cleanup: The _cleanup_producer method is now scheduled as a non-blocking background task using asyncio.create_task, preventing client disconnects from blocking the main producer task.
• Task Continuity: Producer tasks can now continue execution even after a client disconnects, ensuring that long-running operations are not interrupted.
• Improved Resubscription: Clients can now resubscribe to an ongoing task via tasks/resubscribe and receive all subsequent events, resolving a critical issue where events were not delivered after reconnection.
• Comprehensive Testing: New dedicated tests (test_stream_disconnect_then_resubscribe_receives_future_events and test_on_message_send_stream_client_disconnect_triggers_background_cleanup_and_producer_continues) have been added to validate the fix's behavior regarding resubscription and background cleanup. Existing tests were also updated to correctly handle the asynchronous nature of the execute calls.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 correctly addresses an issue where a client disconnect would prematurely terminate a streaming task's producer. The fix, which defers cleanup to a background task, is appropriate and well-reasoned. The accompanying tests are thorough, especially the new tests that simulate disconnect and resubscription, which gives high confidence in the fix. I have one high-severity comment regarding potential unhandled exceptions in the new background task, and one medium-severity suggestion to improve the robustness of a new test.

[holtskinner]

/gemini review

[gemini-code-assist]

Code Review

This pull request effectively resolves a critical issue where a client disconnection would prematurely terminate a running task. The solution to decouple the producer's lifetime from the client connection by moving the cleanup logic into a tracked background task is well-implemented. The introduction of the _track_background_task method is a particularly strong addition, as it provides a robust way to manage these asynchronous operations, ensuring that exceptions are logged and tasks are not orphaned. The accompanying tests are comprehensive, covering the disconnect/resubscribe flow, the non-blocking nature of the new cleanup process, and the background task tracking mechanism itself. The changes are clean, well-reasoned, and significantly improve the reliability of streaming operations.