From 6ceb56cde0369e861e60fa5a5940ef7b488abebb Mon Sep 17 00:00:00 2001
From: jayy-77 <1427jay@gmail.com>
Date: Sat, 24 Jan 2026 03:37:30 +0530
Subject: [PATCH] feat: Implement agent handoff tracking to skip
 after_agent_callback

---
 src/google/adk/agents/base_agent.py             | 12 ++++++++----
 src/google/adk/agents/invocation_context.py     |  7 +++++++
 src/google/adk/flows/llm_flows/base_llm_flow.py |  4 ++++
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py
index e15f9af981..082ee35560 100644
--- a/src/google/adk/agents/base_agent.py
+++ b/src/google/adk/agents/base_agent.py
@@ -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:
+        if event := await self._handle_after_agent_callback(ctx):
+          yield event
 
   @final
   async def run_live(
@@ -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:
+        if event := await self._handle_after_agent_callback(ctx):
+          yield event
 
   async def _run_async_impl(
       self, ctx: InvocationContext
diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py
index 24fdce9d59..0318529de5 100644
--- a/src/google/adk/agents/invocation_context.py
+++ b/src/google/adk/agents/invocation_context.py
@@ -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."""
 
diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py
index 91b57cb873..2073adf320 100644
--- a/src/google/adk/flows/llm_flows/base_llm_flow.py
+++ b/src/google/adk/flows/llm_flows/base_llm_flow.py
@@ -194,6 +194,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
                     )
@@ -714,6 +716,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
         )