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

Author: GWeale

src/google/adk/agents/readonly_context.py

@@ -55,6 +55,21 @@ def state(self) -> MappingProxyType[str, Any]:
     return MappingProxyType(self._invocation_context.session.state)
 
   @property
-  def session(self) -> Session:
-    """The current session for this invocation."""
-    return self._invocation_context.session
+  def user_id(self) -> str:
+    """The user ID for the current invocation."""
+    return self._invocation_context.user_id
+
+  @property
+  def app_name(self) -> str:
+    """The application name for the current invocation."""
+    return self._invocation_context.app_name
+
+  @property
+  def session_id(self) -> str:
+    """The session ID for the current invocation."""
+    return self._invocation_context.session.id
+
+  @property
+  def branch(self) -> Optional[str]:
+    """The branch name for the current invocation, if any."""
+    return self._invocation_context.branch

src/google/adk/cli/plugins/recordings_plugin.py

@@ -39,7 +39,6 @@
 from .recordings_schema import ToolRecording
 
 if TYPE_CHECKING:
-  from ...agents.invocation_context import InvocationContext
   from ...tools.base_tool import BaseTool
   from ...tools.tool_context import ToolContext
 
@@ -75,10 +74,10 @@ def __init__(self, *, name: str = "adk_recordings") -> None:
 
   @override
   async def before_run_callback(
-      self, *, invocation_context: InvocationContext
+      self, *, callback_context: CallbackContext
   ) -> Optional[types.Content]:
     """Always create fresh per-invocation recording state when enabled."""
-    ctx = CallbackContext(invocation_context)
+    ctx = callback_context
     if self._is_record_mode_on(ctx):
       # Always create/overwrite the state for this invocation
       self._create_invocation_state(ctx)
@@ -280,10 +279,10 @@ async def on_tool_error_callback(
 
   @override
   async def after_run_callback(
-      self, *, invocation_context: InvocationContext
+      self, *, callback_context: CallbackContext
   ) -> None:
     """Finalize and persist recordings, then clean per-invocation state."""
-    ctx = CallbackContext(invocation_context)
+    ctx = callback_context
     if not self._is_record_mode_on(ctx):
       return None
 

src/google/adk/cli/plugins/replay_plugin.py

@@ -38,7 +38,6 @@
 from .recordings_schema import ToolRecording
 
 if TYPE_CHECKING:
-  from ...agents.invocation_context import InvocationContext
   from ...tools.base_tool import BaseTool
   from ...tools.tool_context import ToolContext
 
@@ -81,10 +80,10 @@ def __init__(self, *, name: str = "adk_replay") -> None:
 
   @override
   async def before_run_callback(
-      self, *, invocation_context: InvocationContext
+      self, *, callback_context: CallbackContext
   ) -> Optional[types.Content]:
     """Load replay recordings when enabled."""
-    ctx = CallbackContext(invocation_context)
+    ctx = callback_context
     if self._is_replay_mode_on(ctx):
       # Load the replay state for this invocation
       self._load_invocation_state(ctx)
@@ -156,10 +155,10 @@ async def before_tool_callback(
 
   @override
   async def after_run_callback(
-      self, *, invocation_context: InvocationContext
+      self, *, callback_context: CallbackContext
   ) -> None:
     """Clean up replay state after invocation completes."""
-    ctx = CallbackContext(invocation_context)
+    ctx = callback_context
     if not self._is_replay_mode_on(ctx):
       return None
 

src/google/adk/plugins/base_plugin.py

@@ -111,81 +111,125 @@ def __init__(self, name: str):
     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.
 
     This callback helps logging and modifying the user message before the
     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 @@ async def before_model_callback(
       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 @@ def __init__(self, name: str = "logging_plugin"):
   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 @@ async def on_event_callback(
     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 @@ async def before_agent_callback(
     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(