fix(agents): trigger after_agent_callback on cancellation

When agent execution was cancelled via asyncio.CancelledError (e.g. client
disconnect or RPC timeout), _handle_after_agent_callback was bypassed,
causing execution metrics and timers in plugins to leak. This adds an explicit
exception handler for CancelledError to ensure after_agent_callback runs for
side-effects during cancellation, while preserving the invariant that
end_invocation short-circuits and unhandled errors do not trigger
after_agent_callback.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 980466771

Author: DeanChensj

src/google/adk/agents/base_agent.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import abc
+import asyncio
 import inspect
 import logging
 from typing import Any
@@ -341,8 +342,12 @@ async def run_async(
     async def _run() -> AsyncGenerator[Event, None]:
       ctx = self._create_invocation_context(parent_context)
       async with _instrumentation.record_agent_invocation(ctx, self):
+        before_callback_completed = False
+        after_callback_called = False
         try:
-          if event := await self._handle_before_agent_callback(ctx):
+          event = await self._handle_before_agent_callback(ctx)
+          before_callback_completed = True
+          if event:
             yield event
           if ctx.end_invocation:
             return
@@ -354,8 +359,25 @@ async def _run() -> AsyncGenerator[Event, None]:
           if ctx.end_invocation:
             return
 
+          after_callback_called = True
           if event := await self._handle_after_agent_callback(ctx):
             yield event
+        except asyncio.CancelledError:
+          if (
+              before_callback_completed
+              and not after_callback_called
+              and not ctx.end_invocation
+          ):
+            try:
+              await self._handle_after_agent_callback(ctx)
+            except asyncio.CancelledError:
+              raise
+            except Exception:  # pylint: disable=broad-except
+              logger.exception(
+                  'after_agent_callback raised on cancellation;'
+                  ' suppressing so original cancellation propagates.'
+              )
+          raise
         except Exception as e:
           await self._handle_agent_error_callback(ctx, e)
           raise
@@ -403,8 +425,12 @@ async def run_live(
     async def _run() -> AsyncGenerator[Event, None]:
       ctx = self._create_invocation_context(parent_context)
       async with _instrumentation.record_agent_invocation(ctx, self):
+        before_callback_completed = False
+        after_callback_called = False
         try:
-          if event := await self._handle_before_agent_callback(ctx):
+          event = await self._handle_before_agent_callback(ctx)
+          before_callback_completed = True
+          if event:
             yield event
           if ctx.end_invocation:
             return
@@ -413,8 +439,28 @@ async def _run() -> AsyncGenerator[Event, None]:
             async for event in agen:
               yield event
 
+          if ctx.end_invocation:
+            return
+
+          after_callback_called = True
           if event := await self._handle_after_agent_callback(ctx):
             yield event
+        except asyncio.CancelledError:
+          if (
+              before_callback_completed
+              and not after_callback_called
+              and not ctx.end_invocation
+          ):
+            try:
+              await self._handle_after_agent_callback(ctx)
+            except asyncio.CancelledError:
+              raise
+            except Exception:  # pylint: disable=broad-except
+              logger.exception(
+                  'after_agent_callback raised on cancellation;'
+                  ' suppressing so original cancellation propagates.'
+              )
+          raise
         except Exception as e:
           await self._handle_agent_error_callback(ctx, e)
           raise