feat: migrate invocation_context to callback_context

Update plugin manager and built-in plugins to prioritize CallbackContext. Keep InvocationContext access for legacy plugins with adapter. Change callback docs/tests to cover the new context.

PiperOrigin-RevId: 818798087

src/google/adk/plugins/base_plugin.py

@@ -111,11 +111,61 @@ class BasePlugin(ABC):
     super().__init__()
     self.name = name
 
+  if TYPE_CHECKING:
+
+    async def on_user_message_callback(
+        self,
+        *,
+        callback_context: Optional[CallbackContext] = None,
+        user_message: types.Content,
+        invocation_context: Optional[InvocationContext] = None,
+    ) -> Optional[types.Content]:
+      """Callback executed when a user message is received before an invocation starts.
+
+      Plugins can implement this with either callback_context (new) or
+      invocation_context (deprecated) or both.
+      """
+
+    async def before_run_callback(
+        self,
+        *,
+        callback_context: Optional[CallbackContext] = None,
+        invocation_context: Optional[InvocationContext] = None,
+    ) -> Optional[types.Content]:
+      """Callback executed before the ADK runner runs.
+
+      Plugins can implement this with either callback_context (new) or
+      invocation_context (deprecated) or both.
+      """
+
+    async def on_event_callback(
+        self,
+        *,
+        callback_context: Optional[CallbackContext] = None,
+        event: Event,
+        invocation_context: Optional[InvocationContext] = None,
+    ) -> Optional[Event]:
+      """Callback executed after an event is yielded from runner.
+
+      Plugins can implement this with either callback_context (new) or
+      invocation_context (deprecated) or both.
+      """
+
+    async def after_run_callback(
+        self,
+        *,
+        callback_context: Optional[CallbackContext] = None,
+        invocation_context: Optional[InvocationContext] = None,
+    ) -> None:
+      """Callback executed after an ADK runner run has completed.
+
+      Plugins can implement this with either callback_context (new) or
+      invocation_context (deprecated) or both.
+      """
+
+  # Runtime implementation accepts both via **kwargs
   async def on_user_message_callback(
-      self,
-      *,
-      invocation_context: InvocationContext,
-      user_message: types.Content,
+      self, **kwargs: Any
   ) -> Optional[types.Content]:
     """Callback executed when a user message is received before an invocation starts.
 
@@ -123,69 +173,63 @@ class BasePlugin(ABC):
     runner starts the invocation.
 
     Args:
-      invocation_context: The context for the entire invocation.
+      callback_context: The context for the callback execution.
       user_message: The message content input by user.
+      invocation_context: DEPRECATED. Use callback_context instead. The context
+        for the entire invocation. This parameter is maintained for backward
+        compatibility and will be removed in a future version.
 
     Returns:
-      An optional `types.Content` to be returned to the ADK. Returning a
-      value to replace the user message. Returning `None` to proceed
-      normally.
+      The modified user message or None if no modification is needed.
     """
-    pass
+    return None
 
-  async def before_run_callback(
-      self, *, invocation_context: InvocationContext
-  ) -> Optional[types.Content]:
+  async def before_run_callback(self, **kwargs: Any) -> Optional[types.Content]:
     """Callback executed before the ADK runner runs.
 
-    This is the first callback to be called in the lifecycle, ideal for global
-    setup or initialization tasks.
+    This is the first lifecycle hook and is ideal for global setup, logging,
+    or checks that may stop the invocation from running.
 
     Args:
-      invocation_context: The context for the entire invocation, containing
-        session information, the root agent, etc.
+      callback_context: The context for the callback execution.
+      invocation_context: DEPRECATED. Use callback_context instead. The context
+        for the entire invocation. This parameter is maintained for backward
+        compatibility and will be removed in a future version.
 
     Returns:
-      An optional `Event` to be returned to the ADK. Returning a value to
-      halt execution of the runner and ends the runner with that event. Return
-      `None` to proceed normally.
+      Optional `types.Content` to halt execution and return the value to the
+      caller. Return `None` to proceed normally.
     """
-    pass
+    return None
 
-  async def on_event_callback(
-      self, *, invocation_context: InvocationContext, event: Event
-  ) -> Optional[Event]:
+  async def on_event_callback(self, **kwargs: Any) -> Optional[Event]:
     """Callback executed after an event is yielded from runner.
 
-    This is the ideal place to make modification to the event before the event
-    is handled by the underlying agent app.
-
     Args:
-      invocation_context: The context for the entire invocation.
+      callback_context: The context for the callback execution.
       event: The event raised by the runner.
+      invocation_context: DEPRECATED. Use callback_context instead. The context
+        for the entire invocation. This parameter is maintained for backward
+        compatibility and will be removed in a future version.
 
     Returns:
-      An optional value. A non-`None` return may be used by the framework to
-      modify or replace the response. Returning `None` allows the original
-      response to be used.
+      The modified event or None if no modification is needed.
     """
-    pass
+    return None
 
-  async def after_run_callback(
-      self, *, invocation_context: InvocationContext
-  ) -> None:
+  async def after_run_callback(self, **kwargs: Any) -> None:
     """Callback executed after an ADK runner run has completed.
 
-    This is the final callback in the ADK lifecycle, suitable for cleanup, final
-    logging, or reporting tasks.
-
     Args:
-      invocation_context: The context for the entire invocation.
+      callback_context: The context for the callback execution.
+      invocation_context: DEPRECATED. Use callback_context instead. The context
+        for the entire invocation. This parameter is maintained for backward
+        compatibility and will be removed in a future version.
 
     Returns:
       None
     """
-    pass
+    return None
 
   async def before_agent_callback(
       self, *, agent: BaseAgent, callback_context: CallbackContext

src/google/adk/plugins/global_instruction_plugin.py

@@ -79,7 +79,7 @@ class GlobalInstructionPlugin(BasePlugin):
       return None
 
     # Resolve the global instruction (handle both string and InstructionProvider)
-    readonly_context = ReadonlyContext(callback_context.invocation_context)
+    readonly_context = callback_context
     final_global_instruction = await self._resolve_global_instruction(
         readonly_context
     )

src/google/adk/plugins/logging_plugin.py

@@ -69,38 +69,32 @@ class LoggingPlugin(BasePlugin):
   async def on_user_message_callback(
       self,
       *,
-      invocation_context: InvocationContext,
+      callback_context: CallbackContext,
       user_message: types.Content,
   ) -> Optional[types.Content]:
     """Log user message and invocation start."""
     self._log(f"🚀 USER MESSAGE RECEIVED")
-    self._log(f"   Invocation ID: {invocation_context.invocation_id}")
-    self._log(f"   Session ID: {invocation_context.session.id}")
-    self._log(f"   User ID: {invocation_context.user_id}")
-    self._log(f"   App Name: {invocation_context.app_name}")
-    self._log(
-        "   Root Agent:"
-        f" {invocation_context.agent.name if hasattr(invocation_context.agent, 'name') else 'Unknown'}"
-    )
+    self._log(f"   Invocation ID: {callback_context.invocation_id}")
+    self._log(f"   Session ID: {callback_context.session_id}")
+    self._log(f"   User ID: {callback_context.user_id}")
+    self._log(f"   App Name: {callback_context.app_name}")
+    self._log(f"   Root Agent: {callback_context.agent_name}")
     self._log(f"   User Content: {self._format_content(user_message)}")
-    if invocation_context.branch:
-      self._log(f"   Branch: {invocation_context.branch}")
+    if callback_context.branch:
+      self._log(f"   Branch: {callback_context.branch}")
     return None
 
   async def before_run_callback(
-      self, *, invocation_context: InvocationContext
+      self, *, callback_context: CallbackContext
   ) -> Optional[types.Content]:
     """Log invocation start."""
     self._log(f"🏃 INVOCATION STARTING")
-    self._log(f"   Invocation ID: {invocation_context.invocation_id}")
-    self._log(
-        "   Starting Agent:"
-        f" {invocation_context.agent.name if hasattr(invocation_context.agent, 'name') else 'Unknown'}"
-    )
+    self._log(f"   Invocation ID: {callback_context.invocation_id}")
+    self._log(f"   Starting Agent: {callback_context.agent_name}")
     return None
 
   async def on_event_callback(
-      self, *, invocation_context: InvocationContext, event: Event
+      self, *, callback_context: CallbackContext, event: Event
   ) -> Optional[Event]:
     """Log events yielded from the runner."""
     self._log(f"📢 EVENT YIELDED")
@@ -123,15 +117,12 @@ class LoggingPlugin(BasePlugin):
     return None
 
   async def after_run_callback(
-      self, *, invocation_context: InvocationContext
+      self, *, callback_context: CallbackContext
   ) -> Optional[None]:
     """Log invocation completion."""
     self._log(f"✅ INVOCATION COMPLETED")
-    self._log(f"   Invocation ID: {invocation_context.invocation_id}")
-    self._log(
-        "   Final Agent:"
-        f" {invocation_context.agent.name if hasattr(invocation_context.agent, 'name') else 'Unknown'}"
-    )
+    self._log(f"   Invocation ID: {callback_context.invocation_id}")
+    self._log(f"   Final Agent: {callback_context.agent_name}")
     return None
 
   async def before_agent_callback(
@@ -141,8 +132,8 @@ class LoggingPlugin(BasePlugin):
     self._log(f"🤖 AGENT STARTING")
     self._log(f"   Agent Name: {callback_context.agent_name}")
     self._log(f"   Invocation ID: {callback_context.invocation_id}")
-    if callback_context._invocation_context.branch:
-      self._log(f"   Branch: {callback_context._invocation_context.branch}")
+    if callback_context.branch:
+      self._log(f"   Branch: {callback_context.branch}")
     return None
 
   async def after_agent_callback(

src/google/adk/plugins/plugin_manager.py

@@ -14,20 +14,22 @@
 
 from __future__ import annotations
 
+import inspect
 import logging
 from typing import Any
 from typing import List
 from typing import Literal
 from typing import Optional
 from typing import TYPE_CHECKING
+import warnings
 
 from google.genai import types
 
+from ..agents.callback_context import CallbackContext
 from .base_plugin import BasePlugin
 
 if TYPE_CHECKING:
   from ..agents.base_agent import BaseAgent
-  from ..agents.callback_context import CallbackContext
   from ..agents.invocation_context import InvocationContext
   from ..events.event import Event
   from ..models.llm_request import LlmRequest
@@ -113,35 +115,39 @@ class PluginManager:
       invocation_context: InvocationContext,
   ) -> Optional[types.Content]:
     """Runs the `on_user_message_callback` for all plugins."""
+    callback_context = CallbackContext(invocation_context)
     return await self._run_callbacks(
         "on_user_message_callback",
         user_message=user_message,
-        invocation_context=invocation_context,
+        callback_context=callback_context,
     )
 
   async def run_before_run_callback(
       self, *, invocation_context: InvocationContext
   ) -> Optional[types.Content]:
     """Runs the `before_run_callback` for all plugins."""
+    callback_context = CallbackContext(invocation_context)
     return await self._run_callbacks(
-        "before_run_callback", invocation_context=invocation_context
+        "before_run_callback", callback_context=callback_context
     )
 
   async def run_after_run_callback(
       self, *, invocation_context: InvocationContext
   ) -> Optional[None]:
     """Runs the `after_run_callback` for all plugins."""
+    callback_context = CallbackContext(invocation_context)
     return await self._run_callbacks(
-        "after_run_callback", invocation_context=invocation_context
+        "after_run_callback", callback_context=callback_context
     )
 
   async def run_on_event_callback(
       self, *, invocation_context: InvocationContext, event: Event
   ) -> Optional[Event]:
     """Runs the `on_event_callback` for all plugins."""
+    callback_context = CallbackContext(invocation_context)
     return await self._run_callbacks(
         "on_event_callback",
-        invocation_context=invocation_context,
+        callback_context=callback_context,
         event=event,
     )
 
@@ -277,8 +283,14 @@ class PluginManager:
       # Each plugin might not implement all callbacks. The base class provides
       # default `pass` implementations, so `getattr` will always succeed.
       callback_method = getattr(plugin, callback_name)
+
+      # Backward compatibility: Support both callback_context and invocation_context
+      adapted_kwargs = self._adapt_kwargs_for_plugin(
+          plugin, callback_method, kwargs
+      )
+
       try:
-        result = await callback_method(**kwargs)
+        result = await callback_method(**adapted_kwargs)
         if result is not None:
           # Early exit: A plugin has returned a value. We stop
           # processing further plugins and return this value immediately.
@@ -297,3 +309,80 @@ class PluginManager:
         raise RuntimeError(error_message) from e
 
     return None
+
+  def _adapt_kwargs_for_plugin(
+      self, plugin: BasePlugin, callback_method: Any, kwargs: dict[str, Any]
+  ) -> dict[str, Any]:
+    """Adapts keyword arguments for backward compatibility with legacy plugins.
+
+    This method handles the migration from invocation_context to
+    callback_context
+    by inspecting the plugin's callback method signature and providing the
+    appropriate parameter name. For maximum compatibility, it may pass both
+    parameters when the signature is ambiguous.
+
+    Args:
+      plugin: The plugin instance.
+      callback_method: The callback method to be invoked.
+      kwargs: The original keyword arguments.
+
+    Returns:
+      Adapted keyword arguments that match the plugin's expected signature.
+    """
+    # If no callback_context in kwargs, no adaptation needed
+    if "callback_context" not in kwargs:
+      return kwargs.copy()
+
+    callback_context = kwargs["callback_context"]
+
+    try:
+      # Inspect the callback method signature
+      sig = inspect.signature(callback_method)
+      params = sig.parameters
+
+      # Case 1: Method explicitly wants only invocation_context
+      if "invocation_context" in params and "callback_context" not in params:
+        # Legacy plugin - pass only invocation_context
+        warnings.warn(
+            f"Plugin '{plugin.name}' uses deprecated 'invocation_context' "
+            "parameter in callback methods. Please update to use "
+            "'callback_context' instead. Support for 'invocation_context' "
+            "will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+        adapted_kwargs = kwargs.copy()
+        adapted_kwargs["invocation_context"] = (
+            callback_context._invocation_context
+        )
+        del adapted_kwargs["callback_context"]
+        return adapted_kwargs
+
+      # Case 2: Method explicitly wants only callback_context
+      elif "callback_context" in params and "invocation_context" not in params:
+        # Modern plugin - pass only callback_context
+        return kwargs.copy()
+
+      # Case 3: Method wants both, uses **kwargs, or signature is unclear
+      else:
+        # Pass both parameters for maximum compatibility
+        # This handles: **kwargs, both parameters explicitly, or unknown cases
+        adapted_kwargs = kwargs.copy()
+        adapted_kwargs["invocation_context"] = (
+            callback_context._invocation_context
+        )
+        return adapted_kwargs
+
+    except (ValueError, TypeError) as e:
+      # Fallback: Pass both parameters for safety
+      logger.debug(
+          "Failed to inspect plugin '%s' callback signature: %s. "
+          "Passing both callback_context and invocation_context for safety.",
+          plugin.name,
+          e,
+      )
+      adapted_kwargs = kwargs.copy()
+      adapted_kwargs["invocation_context"] = (
+          callback_context._invocation_context
+      )
+      return adapted_kwargs