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: 818822267

Author: google-genai-bot

src/google/adk/agents/readonly_context.py

@@ -55,21 +55,6 @@ def state(self) -> MappingProxyType[str, Any]:
     return MappingProxyType(self._invocation_context.session.state)
 
   @property
-  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
+  def session(self) -> Session:
+    """The current session for this invocation."""
+    return self._invocation_context.session

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

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

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

@@ -38,6 +38,7 @@
 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
 
@@ -80,10 +81,10 @@ def __init__(self, *, name: str = "adk_replay") -> None:
 
   @override
   async def before_run_callback(
-      self, *, callback_context: CallbackContext
+      self, *, invocation_context: InvocationContext
   ) -> Optional[types.Content]:
     """Load replay recordings when enabled."""
-    ctx = callback_context
+    ctx = CallbackContext(invocation_context)
     if self._is_replay_mode_on(ctx):
       # Load the replay state for this invocation
       self._load_invocation_state(ctx)
@@ -155,10 +156,10 @@ async def before_tool_callback(
 
   @override
   async def after_run_callback(
-      self, *, callback_context: CallbackContext
+      self, *, invocation_context: InvocationContext
   ) -> None:
     """Clean up replay state after invocation completes."""
-    ctx = callback_context
+    ctx = CallbackContext(invocation_context)
     if not self._is_replay_mode_on(ctx):
       return None
 

src/google/adk/plugins/base_plugin.py

@@ -111,125 +111,81 @@ 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, **kwargs: Any
+      self,
+      *,
+      invocation_context: InvocationContext,
+      user_message: types.Content,
   ) -> 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:
-      callback_context: The context for the callback execution.
+      invocation_context: The context for the entire invocation.
       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:
-      The modified user message or None if no modification is needed.
+      An optional `types.Content` to be returned to the ADK. Returning a
+      value to replace the user message. Returning `None` to proceed
+      normally.
     """
-    return None
+    pass
 
-  async def before_run_callback(self, **kwargs: Any) -> Optional[types.Content]:
+  async def before_run_callback(
+      self, *, invocation_context: InvocationContext
+  ) -> Optional[types.Content]:
     """Callback executed before the ADK runner runs.
 
-    This is the first lifecycle hook and is ideal for global setup, logging,
-    or checks that may stop the invocation from running.
+    This is the first callback to be called in the lifecycle, ideal for global
+    setup or initialization tasks.
 
     Args:
-      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.
+      invocation_context: The context for the entire invocation, containing
+        session information, the root agent, etc.
 
     Returns:
-      Optional `types.Content` to halt execution and return the value to the
-      caller. Return `None` to proceed normally.
+      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.
     """
-    return None
+    pass
 
-  async def on_event_callback(self, **kwargs: Any) -> Optional[Event]:
+  async def on_event_callback(
+      self, *, invocation_context: InvocationContext, event: Event
+  ) -> 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:
-      callback_context: The context for the callback execution.
+      invocation_context: The context for the entire invocation.
       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:
-      The modified event or None if no modification is needed.
+      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.
     """
-    return None
+    pass
 
-  async def after_run_callback(self, **kwargs: Any) -> None:
+  async def after_run_callback(
+      self, *, invocation_context: InvocationContext
+  ) -> 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:
-      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.
+      invocation_context: The context for the entire invocation.
 
     Returns:
       None
     """
-    return None
+    pass
 
   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 = callback_context
+    readonly_context = ReadonlyContext(callback_context.invocation_context)
     final_global_instruction = await self._resolve_global_instruction(
         readonly_context
     )

src/google/adk/plugins/logging_plugin.py

@@ -69,32 +69,38 @@ def __init__(self, name: str = "logging_plugin"):
   async def on_user_message_callback(
       self,
       *,
-      callback_context: CallbackContext,
+      invocation_context: InvocationContext,
       user_message: types.Content,
   ) -> Optional[types.Content]:
     """Log user message and invocation start."""
     self._log(f"🚀 USER MESSAGE RECEIVED")
-    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"   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"   User Content: {self._format_content(user_message)}")
-    if callback_context.branch:
-      self._log(f"   Branch: {callback_context.branch}")
+    if invocation_context.branch:
+      self._log(f"   Branch: {invocation_context.branch}")
     return None
 
   async def before_run_callback(
-      self, *, callback_context: CallbackContext
+      self, *, invocation_context: InvocationContext
   ) -> Optional[types.Content]:
     """Log invocation start."""
     self._log(f"🏃 INVOCATION STARTING")
-    self._log(f"   Invocation ID: {callback_context.invocation_id}")
-    self._log(f"   Starting Agent: {callback_context.agent_name}")
+    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'}"
+    )
     return None
 
   async def on_event_callback(
-      self, *, callback_context: CallbackContext, event: Event
+      self, *, invocation_context: InvocationContext, event: Event
   ) -> Optional[Event]:
     """Log events yielded from the runner."""
     self._log(f"📢 EVENT YIELDED")
@@ -117,12 +123,15 @@ async def on_event_callback(
     return None
 
   async def after_run_callback(
-      self, *, callback_context: CallbackContext
+      self, *, invocation_context: InvocationContext
   ) -> Optional[None]:
     """Log invocation completion."""
     self._log(f"✅ INVOCATION COMPLETED")
-    self._log(f"   Invocation ID: {callback_context.invocation_id}")
-    self._log(f"   Final Agent: {callback_context.agent_name}")
+    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'}"
+    )
     return None
 
   async def before_agent_callback(
@@ -132,8 +141,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.branch:
-      self._log(f"   Branch: {callback_context.branch}")
+    if callback_context._invocation_context.branch:
+      self._log(f"   Branch: {callback_context._invocation_context.branch}")
     return None
 
   async def after_agent_callback(