fix(callbacks): make after_agent_callback replace instead of append

[dylan-apex]

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

• Closes: after_agent_callback is not replacing the agent output but add the callback output to the original agent output #3502

Problem:
(after_agent_callback)[https://google.github.io/adk-docs/callbacks/types-of-callbacks/#after-agent-callback] clearly states that when you use this, it will replace the original event with the new event, however, in practice it simply appends the event and both events are marked as is_final_response. In multi-agent workflows, this means that not only can you get multiple final responses from the individual agents, but you can get more than 1 final response per agent, therefore making it more difficult to tell which responses matter.

Solution:
When after_agent_callback returns a Content object, it now replaces the agent's original output instead of appending as an additional event.

The implementation marks the original final response event as partial when a callback exists, then yields the callback's content as the true final response. This ensures only one final response event is emitted while preserving session history for sequential agents and tool calling.

It also updates unit tests that relied on the old behavior of sending 2 events when after_agent_callback is used to only send and expect 1 event.is_final_response

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

====================== 3135 passed, 2396 warnings in 47.72s =======================

Manual End-to-End (E2E) Tests:

Working E2E Test

#!/usr/bin/env python3 import asyncio from google.genai import types from google.adk.agents import LlmAgent from google.adk.agents.callback_context import CallbackContext from google.adk.runners import InMemoryRunner from typing import Optional import dotenv dotenv.load_dotenv() async def replacement_callback(callback_context: CallbackContext) -> Optional[types.Content]: """According to docs, this should REPLACE the agent's output.""" print("\n[CALLBACK INVOKED]") print(f"State keys: {list(callback_context.state.to_dict().keys())}") return types.Content( parts=[types.Part(text="REPLACEMENT: This should be the ONLY final response (according to docs)")], role="model" ) async def main(): agent = LlmAgent( model='gemini-2.0-flash-exp', name='test_agent', instruction="Return a simple greeting message.", output_key='greeting', after_agent_callback=replacement_callback ) app_name = 'test_app' user_id = 'test_user' runner = InMemoryRunner(agent=agent, app_name=app_name) session = await runner.session_service.create_session( app_name=app_name, user_id=user_id ) content = types.Content(role='user', parts=[types.Part.from_text(text='Say hello')]) events = [] async for event in runner.run_async( user_id=user_id, session_id=session.id, new_message=content, ): events.append(event) print(f"\n[EVENT {len(events)}]") print(f" Author: {event.author}") print(f" is_final_response(): {event.is_final_response()}") print(f" partial: {event.partial}") if event.content and event.content.parts: for i, part in enumerate(event.content.parts): if hasattr(part, 'text') and part.text: print(f" Text: {part.text[:200]}") break final_events = [e for e in events if e.is_final_response()] if len(final_events) == 1: print("Only one final response") print(f"\nFinal response text: {final_events[0].content.parts[0].text[:100]}") else: print(f"Got {len(final_events)} final responses") if __name__ == "__main__": asyncio.run(main()) 

Output

[EVENT 1] Author: test_agent is_final_response(): False partial: True Text: Hello! [CALLBACK INVOKED] State keys: [] [EVENT 2] Author: test_agent is_final_response(): True partial: None Text: REPLACEMENT: This should be the ONLY final response (according to docs) Only one final response

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules.

Additional context

Add any other context or screenshots about the feature request here.

[gemini-code-assist]

Summary of Changes

Hello @dylan-apex, 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 bug in the after_agent_callback mechanism where it was incorrectly appending agent responses instead of replacing them, leading to an undesirable state of multiple final responses. The changes ensure that the callback now functions as intended, providing a single, definitive final response from the agent, which is crucial for reliable multi-agent workflows and consistent session history.

Highlights

• Callback Behavior Fix: The after_agent_callback now correctly replaces the agent's original output instead of appending it, resolving an issue where multiple final responses could be emitted.
• Single Final Response Guarantee: When an after_agent_callback returns content, the original final response event is marked as partial, and the callback's content is yielded as the sole true final response, ensuring only one final response event per agent.
• Updated Documentation: The docstring for after_agent_callback has been updated to accurately reflect its new replacement behavior.
• Comprehensive Testing: Unit tests have been updated to assert the single final response behavior, and a manual end-to-end test confirms the fix in practice.

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 correctly addresses the issue of after_agent_callback appending instead of replacing the agent's final response. The approach of marking the original response as partial and yielding a new final response from the callback is sound. However, I've identified a critical bug in the implementation in base_agent.py where the agent's final response can be dropped if a callback modifies the session state but doesn't return any new content. I have left detailed comments with suggested fixes for both run_async and run_live methods to address this data loss scenario. The updates to the unit tests look good and align with the intended new behavior.

[dylan-apex]

/gemini review

[gemini-code-assist]

Code Review

This pull request correctly adjusts the after_agent_callback to replace the original agent's output, aligning its behavior with the documentation. The implementation is sound and covers various callback scenarios, and the accompanying test updates validate the new logic. My review focuses on improving maintainability by addressing significant code duplication. I've suggested refactoring a large, repeated block of logic in base_agent.py into a helper method and consolidating a recurring filtering pattern in the test suite.

[dylan-apex]

/gemini summary

[gemini-code-assist]

Summary of Changes

This pull request addresses a critical bug in the after_agent_callback mechanism within the BaseAgent class, where it was incorrectly appending agent responses instead of replacing them. This led to scenarios with multiple final responses, complicating multi-agent workflows. The changes ensure that when an after_agent_callback returns content, it now precisely replaces the agent's original output, guaranteeing a single, definitive final response. The core callback handling logic has been centralized into a new _run_with_callbacks method, and comprehensive unit tests have been updated and expanded to validate this corrected behavior, including cases where callbacks only modify state.

Highlights

• Corrected after_agent_callback behavior: The after_agent_callback now correctly replaces the agent's original output with the callback's content, rather than appending it, resolving an issue where multiple final responses could be emitted.
• Ensured Single Final Response: The implementation now marks the original final response event as partial when a callback exists and returns content, yielding the callback's content as the sole true final response. If the callback returns None, the original final response is preserved.
• Centralized Callback Logic: A new private asynchronous method, _run_with_callbacks, has been introduced in BaseAgent to encapsulate and reuse the complex logic for handling after_agent_callback across both run_async and run_live methods, improving maintainability.
• Updated Documentation: The docstring for after_agent_callback in BaseAgent has been updated to accurately reflect its new replacement behavior.
• Enhanced Unit Testing: New helper functions (_after_agent_callback_state_only, _get_final_events) and a new unit test (test_run_async_after_agent_callback_state_only) have been added to thoroughly validate the corrected callback behavior, especially for state-only modifications, and existing tests were updated to reflect the single final response expectation.

Changelog

• contributing/samples/gepa/experiment.py
  • Removed an unnecessary blank line in imports.
• contributing/samples/gepa/run_experiment.py
  • Removed an unnecessary blank line in imports.
• src/google/adk/agents/base_agent.py
  • Updated the docstring for after_agent_callback to explicitly state that it replaces the agent's original output.
  • Introduced a new private async generator method _run_with_callbacks to wrap agent implementations and manage after_agent_callback logic. This method handles marking original final responses as partial and yielding the callback's content as the new final response, or preserving the original if the callback returns None.
  • Refactored run_async and run_live methods to utilize the new _run_with_callbacks method, centralizing the callback handling.
• tests/unittests/agents/test_base_agent.py
  • Added a new helper function _after_agent_callback_state_only to simulate callbacks that only update state without returning content.
  • Introduced a new helper function _get_final_events to simplify filtering for final response events in tests.
  • Modified mock_sync_agent_cb_side_effect test data to expect a single final response.
  • Updated test_before_agent_callbacks_chain and test_after_agent_callbacks_chain to filter for and assert against final events using _get_final_events.
  • Adjusted assertions in test_run_async_after_agent_callback_use_plugin, test_run_async_after_agent_callback_noop, test_run_async_with_async_after_agent_callback_noop, test_run_async_after_agent_callback_append_reply, and test_run_async_with_async_after_agent_callback_append_reply to expect and verify a single final response event.
  • Added a new test test_run_async_after_agent_callback_state_only to specifically test scenarios where the after_agent_callback modifies state but returns no content, ensuring the original final response is still emitted alongside the state update.

Activity

• gemini-code-assist[bot] provided an initial summary of the PR.
• gemini-code-assist[bot] identified a critical bug in the callback logic where state-only callbacks could cause the original final response to be lost, providing a suggested fix.
• gemini-code-assist[bot] recommended extracting duplicated callback handling logic into a shared private method for improved maintainability.
• gemini-code-assist[bot] suggested creating a helper function for filtering final events in tests to reduce code duplication.
• dylan-apex explicitly requested a review and a summary from gemini-code-assist[bot].