Skip to content

feat: Implement agent handoff tracking to skip after_agent_callback #4255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jayy-77 wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

feat: Implement agent handoff tracking to skip after_agent_callback #4255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6ceb56c
feat: Implement agent handoff tracking to skip after_agent_callback
jayy-77 Jan 23, 2026
6a1be4e
Merge branch 'main' into agent_handoff_fix
jayy-77 Jan 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions src/google/adk/agents/base_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -297,8 +297,10 @@ async def run_async(
if ctx.end_invocation:
return

if event := await self._handle_after_agent_callback(ctx):
yield event
# Skip after_agent_callback if the agent has handed off to another agent
if not ctx.agent_handoff_occurred:
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There's a potential logic issue here. When agent1 hands off to agent2, ctx.agent_handoff_occurred is set to True. This correctly skips the after_agent_callback for agent1. However, agent2 inherits the context with this flag set to True, so its after_agent_callback will also be skipped, which is likely unintended. The callback should only be skipped for the agent that initiates the handoff.

A possible solution is to reset this flag at the beginning of each agent's run. You could add ctx.agent_handoff_occurred = False at the start of the run_async method in this file, after ctx is created. Since this change is outside the current diff, I cannot provide a direct code suggestion for it.

if event := await self._handle_after_agent_callback(ctx):
yield event
Comment on lines +301 to +303
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For better readability, you can combine the nested if statements into a single condition.

      if not ctx.agent_handoff_occurred and (event := await self._handle_after_agent_callback(ctx)):
        yield event


@final
async def run_live(
Expand Down Expand Up @@ -327,8 +329,10 @@ async def run_live(
async for event in agen:
yield event

if event := await self._handle_after_agent_callback(ctx):
yield event
# Skip after_agent_callback if the agent has handed off to another agent
if not ctx.agent_handoff_occurred:
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Similar to my comment in run_async, there's a potential logic issue here that would cause the after_agent_callback to be skipped for the receiving agent as well as the transferring agent.

A possible solution is to reset this flag at the beginning of each agent's run. You could add ctx.agent_handoff_occurred = False at the start of the run_live method in this file, after ctx is created. Since this change is outside the current diff, I cannot provide a direct code suggestion for it.

if event := await self._handle_after_agent_callback(ctx):
yield event
Comment on lines +333 to +335
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For better readability, you can combine the nested if statements into a single condition.

      if not ctx.agent_handoff_occurred and (event := await self._handle_after_agent_callback(ctx)):
        yield event


async def _run_async_impl(
self, ctx: InvocationContext
Expand Down
7 changes: 7 additions & 0 deletions src/google/adk/agents/invocation_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,13 @@ class InvocationContext(BaseModel):

Set to True in callbacks or tools to terminate this invocation."""

agent_handoff_occurred: bool = False
"""Whether the current agent has handed off control to another agent.

Set to True when an agent transfers control to another agent via the
transfer_to_agent tool. Used to prevent after_agent_callback from running
on the transferring agent after handoff has occurred."""

live_request_queue: Optional[LiveRequestQueue] = None
"""The queue to receive live requests."""

Expand Down
4 changes: 4 additions & 0 deletions src/google/adk/flows/llm_flows/base_llm_flow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,8 @@ async def run_live(
transfer_to_agent = event.actions.transfer_to_agent
if transfer_to_agent:
logger.debug('Transferring to agent: %s', transfer_to_agent)
# Mark that handoff has occurred so after_agent_callback is skipped
invocation_context.agent_handoff_occurred = True
agent_to_run = self._get_agent_to_run(
invocation_context, transfer_to_agent
)
Expand Down Expand Up @@ -765,6 +767,8 @@ async def _postprocess_handle_function_calls_async(
yield final_event
transfer_to_agent = function_response_event.actions.transfer_to_agent
if transfer_to_agent:
# Mark that handoff has occurred so after_agent_callback is skipped
invocation_context.agent_handoff_occurred = True
agent_to_run = self._get_agent_to_run(
invocation_context, transfer_to_agent
)
Expand Down