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: #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:

• [X] I have added or updated unit tests for my change.
• [X] 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

• [X] I have read the CONTRIBUTING.md document.
• [X] I have performed a self-review of my own code.
• [X] I have commented my code, particularly in hard-to-understand areas.
• [X] I have added tests that prove my fix is effective or that my feature works.
• [X] New and existing unit tests pass locally with my changes.
• [X] I have manually tested my changes end-to-end.
• [X] 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[bot]]

Summary of Changes

Hello @dylan-apex, I'm Gemini Code Assist[^1]! 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 :thumbsup: and :thumbsdown: 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.

[^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.

[dylan-apex]

/gemini review

[dylan-apex]

/gemini summary

[gemini-code-assist[bot]]

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].

[ryanaiagent]

Hi @dylan-apex , Thanks for you contribution. Can you fix the failing Unit tests?

[dylan-apex]

Hi @dylan-apex , Thanks for you contribution. Can you fix the failing Unit tests?

Working on it

[dylan-apex]

@ryanaiagent 2 of the 3 failing tests are also failing tests in main right now, but I did fix them (https://github.com/google/adk-python/commit/d85a301039003210b61a572ca3d98183d717768a)

All tests passing locally again 3158 passed, 2406 warnings in 44.66s

[dylan-apex]

I forgot to do uv sync and get the up to date version of a2a - my bad.

Tests are passing locally after uv sync and merging main into the branch 3201 passed, 2414 warnings in 53.69s

[ryanaiagent]

Hi @dylan-apex , Your PR has been received by the team and is currently under review. We will provide feedback as soon as we have an update to share.

[ryanaiagent]

Hi @DeanChensj , can you please review this.

[masterchop]

+1

[moralesjor-mms]

+1