feat: add stop_event to cancel SSE streaming mid-generation

[ferponse]

Closes #4823

Summary

• Adds an optional stop_event: asyncio.Event parameter to runner.run_async() for cancelling SSE streaming mid-generation
• When stop_event.set() is called, the flow stops yielding new chunks and returns cleanly
• Enables "stop generating" buttons in chat UIs

Problem

When using runner.run_async() with StreamingMode.SSE, there is no clean way for consumers to stop generation mid-stream. Task cancellation is a hard stop that may not clean up properly. Breaking out of the async for loop triggers aclose() but doesn't give the flow a chance to return gracefully.

Solution

An asyncio.Event-based cancellation mechanism:

1. Consumer creates an asyncio.Event and passes it to runner.run_async(stop_event=...)
2. The event is stored in InvocationContext.stop_event
3. The SSE flow checks stop_event.is_set() at two points:
   • Before each LLM call in run_async (prevents starting a new turn)
   • Before yielding each streaming chunk in _run_one_step_async (stops mid-stream)
4. When set, the flow returns cleanly without yielding further chunks

Usage

stop_event = asyncio.Event()

async for event in runner.run_async(
    user_id=user_id,
    session_id=session_id,
    new_message=content,
    run_config=RunConfig(streaming_mode=StreamingMode.SSE),
    stop_event=stop_event,
):
    send_to_client(event)

# When user clicks "stop generating":
stop_event.set()

Changes

File	Change
src/google/adk/agents/invocation_context.py	Added stop_event: Optional[asyncio.Event] field
src/google/adk/runners.py	Added stop_event parameter to run_async(), propagated to InvocationContext
src/google/adk/flows/llm_flows/base_llm_flow.py	Check stop_event before each LLM call and before yielding each chunk
tests/unittests/flows/llm_flows/test_stop_event.py	4 new tests

Test plan

• test_stop_event_stops_streaming_mid_chunks — setting stop_event mid-stream prevents further chunks
• test_stop_event_not_set_yields_all_chunks — providing stop_event without setting it yields all chunks normally
• test_stop_event_prevents_next_llm_call — setting stop_event between LLM calls prevents the next call
• test_no_stop_event_works_normally — not providing stop_event works exactly as before
• All existing + new tests pass

[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 need for controlled termination of streaming operations in the system. By integrating an asyncio.Event-based cancellation mechanism, it allows for a clean and graceful halt of ongoing generation processes, particularly for SSE streams. This enhancement improves the stability of the system by avoiding abrupt task cancellations and provides a foundation for more interactive user experiences in applications that consume these streams.

Highlights

• Graceful Streaming Cancellation: Introduced an optional stop_event: asyncio.Event parameter to runner.run_async() to allow consumers to gracefully cancel Server-Sent Events (SSE) streaming mid-generation.
• Clean Termination: When stop_event.set() is called, the streaming flow will stop yielding new chunks and return cleanly, providing a more robust alternative to hard task cancellation.
• Integration Points: The stop_event is propagated to InvocationContext and checked at two critical points within the SSE flow: before each LLM call (preventing new turns) and before yielding each streaming chunk (stopping mid-stream).
• Enhanced User Experience: This mechanism directly enables the implementation of 'stop generating' buttons in chat user interfaces, improving control for end-users.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/google/adk/agents/invocation_context.py
  • Added an optional stop_event: asyncio.Event field to the InvocationContext model.
• src/google/adk/flows/llm_flows/base_llm_flow.py
  • Implemented checks for the stop_event within the run_async loop to prevent further LLM calls.
  • Added checks for the stop_event within the _run_one_step_async method to stop yielding streaming chunks.
• src/google/adk/runners.py
  • Modified the run_async() method to accept an optional stop_event parameter.
  • Propagated the provided stop_event to the InvocationContext.
• tests/unittests/flows/llm_flows/test_stop_event.py
  • Added a new test file to cover the functionality of the stop_event cancellation mechanism.
  • Included tests for stopping mid-chunks, yielding all chunks when the event is not set, preventing subsequent LLM calls, and normal operation without the event.

Activity

• No human activity 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 introduces a valuable feature for cancelling SSE streaming via an asyncio.Event. The implementation is clean and logically sound, with checks correctly placed to stop generation both between and during LLM calls. The addition of comprehensive unit tests is also a great inclusion. I have a couple of minor suggestions to improve code readability and consistency, and one for improving the robustness of a test case.

[rohityan]

Hi @ferponse , Thank you for your contribution! We appreciate you taking the time to submit this pull request. Can you please fix the failing unit tests and formatting errors before we can proceed with a review.