feat: Integrating Plugin with ADK

This change integrates the plugin system with ADK. PluginManager is attached to the invocation context similar to session/artifact/memory.

It includes integrations with following ADK internal callbacks:

* App callbacks: Integrated in the BaseRunner class, in run_async and run_live
* On Message callbacks: Integrated in the BaseRunner class, triggers on run_async.
* Agent callbacks: Integrated in the BaseAgent class. Leveraging the existing *callback functions
* Model callbacks: Integrating in the base_llm_flow.
* Tool callbacks: Integrated in functions.py, wrapped around the code for agent tool_callbacks

Sample code to use plugins:

```python
# Add plugins to Runner

runner = Runner(
      app_name="my-app",
      agent=root_agent,
      artifact_service=artifact_service,
      session_service=session_service,
      memory_service=memory_service,
      plugins=[
        MySamplePlugin(),
        LoggingPlugin(),
      ],
  )

```

PiperOrigin-RevId: 781746586

src/google/adk/agents/base_agent.py

@@ -227,11 +227,18 @@ class BaseAgent(BaseModel):
     """
     with tracer.start_as_current_span(f'agent_run [{self.name}]'):
       ctx = self._create_invocation_context(parent_context)
-      # TODO(hangfei): support before/after_agent_callback
+
+      if event := await self.__handle_before_agent_callback(ctx):
+        yield event
+      if ctx.end_invocation:
+        return
 
       async for event in self._run_live_impl(ctx):
         yield event
 
+      if event := await self.__handle_after_agent_callback(ctx):
+        yield event
+
   async def _run_async_impl(
       self, ctx: InvocationContext
   ) -> AsyncGenerator[Event, None]:
@@ -335,73 +342,99 @@ class BaseAgent(BaseModel):
   ) -> Optional[Event]:
     """Runs the before_agent_callback if it exists.
 
+    Args:
+      ctx: InvocationContext, the invocation context for this agent.
+
     Returns:
       Optional[Event]: an event if callback provides content or changed state.
     """
-    ret_event = None
-
-    if not self.canonical_before_agent_callbacks:
-      return ret_event
-
     callback_context = CallbackContext(ctx)
 
-    for callback in self.canonical_before_agent_callbacks:
-      before_agent_callback_content = callback(
-          callback_context=callback_context
-      )
-      if inspect.isawaitable(before_agent_callback_content):
-        before_agent_callback_content = await before_agent_callback_content
-      if before_agent_callback_content:
-        ret_event = Event(
-            invocation_id=ctx.invocation_id,
-            author=self.name,
-            branch=ctx.branch,
-            content=before_agent_callback_content,
-            actions=callback_context._event_actions,
+    # Run callbacks from the plugins.
+    before_agent_callback_content = (
+        await ctx.plugin_manager.run_before_agent_callback(
+            agent=self, callback_context=callback_context
         )
-        ctx.end_invocation = True
-        return ret_event
+    )
+
+    # If no overrides are provided from the plugins, further run the canonical
+    # callbacks.
+    if (
+        not before_agent_callback_content
+        and self.canonical_before_agent_callbacks
+    ):
+      for callback in self.canonical_before_agent_callbacks:
+        before_agent_callback_content = callback(
+            callback_context=callback_context
+        )
+        if inspect.isawaitable(before_agent_callback_content):
+          before_agent_callback_content = await before_agent_callback_content
+        if before_agent_callback_content:
+          break
+
+    # Process the override content if exists, and further process the state
+    # change if exists.
+    if before_agent_callback_content:
+      ret_event = Event(
+          invocation_id=ctx.invocation_id,
+          author=self.name,
+          branch=ctx.branch,
+          content=before_agent_callback_content,
+          actions=callback_context._event_actions,
+      )
+      ctx.end_invocation = True
+      return ret_event
 
     if callback_context.state.has_delta():
-      ret_event = Event(
+      return Event(
           invocation_id=ctx.invocation_id,
           author=self.name,
           branch=ctx.branch,
           actions=callback_context._event_actions,
       )
 
-    return ret_event
+    return None
 
   async def __handle_after_agent_callback(
       self, invocation_context: InvocationContext
   ) -> Optional[Event]:
     """Runs the after_agent_callback if it exists.
 
+    Args:
+      invocation_context: InvocationContext, the invocation context for this
+        agent.
+
     Returns:
       Optional[Event]: an event if callback provides content or changed state.
     """
-    ret_event = None
-
-    if not self.canonical_after_agent_callbacks:
-      return ret_event
 
     callback_context = CallbackContext(invocation_context)
 
-    for callback in self.canonical_after_agent_callbacks:
-      after_agent_callback_content = callback(callback_context=callback_context)
-      if inspect.isawaitable(after_agent_callback_content):
-        after_agent_callback_content = await after_agent_callback_content
-      if after_agent_callback_content:
-        ret_event = Event(
-            invocation_id=invocation_context.invocation_id,
-            author=self.name,
-            branch=invocation_context.branch,
-            content=after_agent_callback_content,
-            actions=callback_context._event_actions,
+    # Run callbacks from the plugins.
+    after_agent_callback_content = (
+        await invocation_context.plugin_manager.run_after_agent_callback(
+            agent=self, callback_context=callback_context
         )
-        return ret_event
+    )
 
-    if callback_context.state.has_delta():
+    # If no overrides are provided from the plugins, further run the canonical
+    # callbacks.
+    if (
+        not after_agent_callback_content
+        and self.canonical_after_agent_callbacks
+    ):
+      for callback in self.canonical_after_agent_callbacks:
+        after_agent_callback_content = callback(
+            callback_context=callback_context
+        )
+        if inspect.isawaitable(after_agent_callback_content):
+          after_agent_callback_content = await after_agent_callback_content
+        if after_agent_callback_content:
+          break
+
+    # Process the override content if exists, and further process the state
+    # change if exists.
+    if after_agent_callback_content:
       ret_event = Event(
           invocation_id=invocation_context.invocation_id,
           author=self.name,
@@ -409,8 +442,17 @@ class BaseAgent(BaseModel):
           content=after_agent_callback_content,
           actions=callback_context._event_actions,
       )
+      return ret_event
 
-    return ret_event
+    if callback_context.state.has_delta():
+      return Event(
+          invocation_id=invocation_context.invocation_id,
+          author=self.name,
+          branch=invocation_context.branch,
+          content=after_agent_callback_content,
+          actions=callback_context._event_actions,
+      )
+    return None
 
   @override
   def model_post_init(self, __context: Any) -> None:

src/google/adk/agents/invocation_context.py

@@ -24,6 +24,7 @@ from pydantic import ConfigDict
 from ..artifacts.base_artifact_service import BaseArtifactService
 from ..auth.credential_service.base_credential_service import BaseCredentialService
 from ..memory.base_memory_service import BaseMemoryService
+from ..plugins.plugin_manager import PluginManager
 from ..sessions.base_session_service import BaseSessionService
 from ..sessions.session import Session
 from .active_streaming_tool import ActiveStreamingTool
@@ -153,6 +154,9 @@ class InvocationContext(BaseModel):
   run_config: Optional[RunConfig] = None
   """Configurations for live agents under this invocation."""
 
+  plugin_manager: PluginManager = PluginManager()
+  """The manager for keeping track of plugins in this invocation."""
+
   _invocation_cost_manager: _InvocationCostManager = _InvocationCostManager()
   """A container to keep track of different kinds of costs incurred as a part
   of this invocation.

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -565,21 +565,32 @@ class BaseLlmFlow(ABC):
     if not isinstance(agent, LlmAgent):
       return
 
-    if not agent.canonical_before_model_callbacks:
-      return
-
     callback_context = CallbackContext(
         invocation_context, event_actions=model_response_event.actions
     )
 
+    # First run callbacks from the plugins.
+    callback_response = (
+        await invocation_context.plugin_manager.run_before_model_callback(
+            callback_context=callback_context,
+            llm_request=llm_request,
+        )
+    )
+    if callback_response:
+      return callback_response
+
+    # If no overrides are provided from the plugins, further run the canonical
+    # callbacks.
+    if not agent.canonical_before_model_callbacks:
+      return
     for callback in agent.canonical_before_model_callbacks:
-      before_model_callback_content = callback(
+      callback_response = callback(
           callback_context=callback_context, llm_request=llm_request
       )
-      if inspect.isawaitable(before_model_callback_content):
-        before_model_callback_content = await before_model_callback_content
-      if before_model_callback_content:
-        return before_model_callback_content
+      if inspect.isawaitable(callback_response):
+        callback_response = await callback_response
+      if callback_response:
+        return callback_response
 
   async def _handle_after_model_callback(
       self,
@@ -593,21 +604,32 @@ class BaseLlmFlow(ABC):
     if not isinstance(agent, LlmAgent):
       return
 
-    if not agent.canonical_after_model_callbacks:
-      return
-
     callback_context = CallbackContext(
         invocation_context, event_actions=model_response_event.actions
     )
 
+    # First run callbacks from the plugins.
+    callback_response = (
+        await invocation_context.plugin_manager.run_after_model_callback(
+            callback_context=CallbackContext(invocation_context),
+            llm_response=llm_response,
+        )
+    )
+    if callback_response:
+      return callback_response
+
+    # If no overrides are provided from the plugins, further run the canonical
+    # callbacks.
+    if not agent.canonical_after_model_callbacks:
+      return
     for callback in agent.canonical_after_model_callbacks:
-      after_model_callback_content = callback(
+      callback_response = callback(
           callback_context=callback_context, llm_response=llm_response
       )
-      if inspect.isawaitable(after_model_callback_content):
-        after_model_callback_content = await after_model_callback_content
-      if after_model_callback_content:
-        return after_model_callback_content
+      if inspect.isawaitable(callback_response):
+        callback_response = await callback_response
+      if callback_response:
+        return callback_response
 
   def _finalize_model_response_event(
       self,

src/google/adk/flows/llm_flows/functions.py

@@ -153,37 +153,67 @@ async def handle_function_calls_async(
       # do not use "args" as the variable name, because it is a reserved keyword
       # in python debugger.
       function_args = function_call.args or {}
-      function_response: Optional[dict] = None
 
-      for callback in agent.canonical_before_tool_callbacks:
-        function_response = callback(
-            tool=tool, args=function_args, tool_context=tool_context
-        )
-        if inspect.isawaitable(function_response):
-          function_response = await function_response
-        if function_response:
-          break
+      # Step 1: Check if plugin before_tool_callback overrides the function
+      # response.
+      function_response = (
+          await invocation_context.plugin_manager.run_before_tool_callback(
+              tool=tool, tool_args=function_args, tool_context=tool_context
+          )
+      )
 
-      if not function_response:
+      # Step 2: If no overrides are provided from the plugins, further run the
+      # canonical callback.
+      if function_response is None:
+        for callback in agent.canonical_before_tool_callbacks:
+          function_response = callback(
+              tool=tool, args=function_args, tool_context=tool_context
+          )
+          if inspect.isawaitable(function_response):
+            function_response = await function_response
+          if function_response:
+            break
+
+      # Step 3: Otherwise, proceed calling the tool normally.
+      if function_response is None:
         function_response = await __call_tool_async(
             tool, args=function_args, tool_context=tool_context
         )
 
-      for callback in agent.canonical_after_tool_callbacks:
-        altered_function_response = callback(
-            tool=tool,
-            args=function_args,
-            tool_context=tool_context,
-            tool_response=function_response,
-        )
-        if inspect.isawaitable(altered_function_response):
-          altered_function_response = await altered_function_response
-        if altered_function_response is not None:
-          function_response = altered_function_response
-          break
+      # Step 4: Check if plugin after_tool_callback overrides the function
+      # response.
+      altered_function_response = (
+          await invocation_context.plugin_manager.run_after_tool_callback(
+              tool=tool,
+              tool_args=function_args,
+              tool_context=tool_context,
+              result=function_response,
+          )
+      )
+
+      # Step 5: If no overrides are provided from the plugins, further run the
+      # canonical after_tool_callbacks.
+      if altered_function_response is None:
+        for callback in agent.canonical_after_tool_callbacks:
+          altered_function_response = callback(
+              tool=tool,
+              args=function_args,
+              tool_context=tool_context,
+              tool_response=function_response,
+          )
+          if inspect.isawaitable(altered_function_response):
+            altered_function_response = await altered_function_response
+          if altered_function_response:
+            break
+
+      # Step 6: If alternative response exists from after_tool_callback, use it
+      # instead of the original function response.
+      if altered_function_response is not None:
+        function_response = altered_function_response
 
       if tool.is_long_running:
-        # Allow long running function to return None to not provide function response.
+        # Allow long running function to return None to not provide function
+        # response.
         if not function_response:
           continue
 
@@ -264,6 +294,7 @@ async def handle_function_calls_live(
       #   )
       #   if new_response:
       #     function_response = new_response
+      altered_function_response = None
       if agent.after_tool_callback:
         altered_function_response = agent.after_tool_callback(
             tool=tool,
@@ -273,8 +304,8 @@ async def handle_function_calls_live(
         )
         if inspect.isawaitable(altered_function_response):
           altered_function_response = await altered_function_response
-        if altered_function_response is not None:
-          function_response = altered_function_response
+      if altered_function_response is not None:
+        function_response = altered_function_response
 
       if tool.is_long_running:
         # Allow async function to return None to not provide function response.

src/google/adk/runners.py

@@ -18,7 +18,9 @@ import asyncio
 import logging
 import queue
 from typing import AsyncGenerator
+from typing import Callable
 from typing import Generator
+from typing import List
 from typing import Optional
 import warnings
 
@@ -40,6 +42,8 @@ from .flows.llm_flows.functions import find_matching_function_call
 from .memory.base_memory_service import BaseMemoryService
 from .memory.in_memory_memory_service import InMemoryMemoryService
 from .platform.thread import create_thread
+from .plugins.base_plugin import BasePlugin
+from .plugins.plugin_manager import PluginManager
 from .sessions.base_session_service import BaseSessionService
 from .sessions.in_memory_session_service import InMemorySessionService
 from .sessions.session import Session
@@ -60,6 +64,7 @@ class Runner:
       app_name: The application name of the runner.
       agent: The root agent to run.
       artifact_service: The artifact service for the runner.
+      plugin_manager: The plugin manager for the runner.
       session_service: The session service for the runner.
       memory_service: The memory service for the runner.
   """
@@ -70,6 +75,8 @@ class Runner:
   """The root agent to run."""
   artifact_service: Optional[BaseArtifactService] = None
   """The artifact service for the runner."""
+  plugin_manager: PluginManager
+  """The plugin manager for the runner."""
   session_service: BaseSessionService
   """The session service for the runner."""
   memory_service: Optional[BaseMemoryService] = None
@@ -82,6 +89,7 @@ class Runner:
       *,
       app_name: str,
       agent: BaseAgent,
+      plugins: Optional[List[BasePlugin]] = None,
       artifact_service: Optional[BaseArtifactService] = None,
       session_service: BaseSessionService,
       memory_service: Optional[BaseMemoryService] = None,
@@ -102,6 +110,7 @@ class Runner:
     self.session_service = session_service
     self.memory_service = memory_service
     self.credential_service = credential_service
+    self.plugin_manager = PluginManager(plugins=plugins)
 
   def run(
       self,
@@ -191,6 +200,15 @@ class Runner:
       )
       root_agent = self.agent
 
+      # Modify user message before execution.
+      modified_user_message = (
+          await invocation_context.plugin_manager.run_on_user_message_callback(
+              invocation_context=invocation_context, user_message=new_message
+          )
+      )
+      if modified_user_message is not None:
+        new_message = modified_user_message
+
       if new_message:
         await self._append_new_message_to_session(
             session,
@@ -200,10 +218,64 @@ class Runner:
         )
 
       invocation_context.agent = self._find_agent_to_run(session, root_agent)
-      async for event in invocation_context.agent.run_async(invocation_context):
+
+      async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
+        async for event in ctx.agent.run_async(ctx):
+          yield event
+
+      async for event in self._exec_with_plugin(
+          invocation_context, session, execute
+      ):
+        yield event
+
+  async def _exec_with_plugin(
+      self,
+      invocation_context: InvocationContext,
+      session: Session,
+      execute_fn: Callable[[InvocationContext], AsyncGenerator[Event, None]],
+  ) -> AsyncGenerator[Event, None]:
+    """Wraps execution with plugin callbacks.
+
+    Args:
+      invocation_context: The invocation context
+      session: The current session
+      execute_fn: A callable that returns an AsyncGenerator of Events
+
+    Yields:
+      Events from the execution, including any generated by plugins
+    """
+
+    plugin_manager = invocation_context.plugin_manager
+
+    # Step 1: Run the before_run callbacks to see if we should early exit.
+    early_exit_result = await plugin_manager.run_before_run_callback(
+        invocation_context=invocation_context
+    )
+    if isinstance(early_exit_result, Event):
+      await self.session_service.append_event(
+          session=session,
+          event=Event(
+              invocation_id=invocation_context.invocation_id,
+              author='model',
+              content=early_exit_result,
+          ),
+      )
+      yield early_exit_result
+    else:
+      # Step 2: Otherwise continue with normal execution
+      async for event in execute_fn(invocation_context):
         if not event.partial:
           await self.session_service.append_event(session=session, event=event)
-        yield event
+        # Step 3: Run the on_event callbacks to optionally modify the event.
+        modified_event = await plugin_manager.run_on_event_callback(
+            invocation_context=invocation_context, event=event
+        )
+        yield (modified_event if modified_event else event)
+
+    # Step 4: Run the after_run callbacks to optionally modify the context.
+    await plugin_manager.run_after_run_callback(
+        invocation_context=invocation_context
+    )
 
   async def _append_new_message_to_session(
       self,
@@ -345,8 +417,14 @@ class Runner:
             invocation_context.active_streaming_tools[tool.__name__] = (
                 active_streaming_tool
             )
-    async for event in invocation_context.agent.run_live(invocation_context):
-      await self.session_service.append_event(session=session, event=event)
+
+    async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
+      async for event in ctx.agent.run_live(ctx):
+        yield event
+
+    async for event in self._exec_with_plugin(
+        invocation_context, session, execute
+    ):
       yield event
 
   def _find_agent_to_run(
@@ -450,6 +528,7 @@ class Runner:
         session_service=self.session_service,
         memory_service=self.memory_service,
         credential_service=self.credential_service,
+        plugin_manager=self.plugin_manager,
         invocation_id=invocation_id,
         agent=self.agent,
         session=session,
@@ -538,7 +617,13 @@ class InMemoryRunner(Runner):
         session service for the runner.
   """
 
-  def __init__(self, agent: BaseAgent, *, app_name: str = 'InMemoryRunner'):
+  def __init__(
+      self,
+      agent: BaseAgent,
+      *,
+      app_name: str = 'InMemoryRunner',
+      plugins: Optional[list[BasePlugin]] = None,
+  ):
     """Initializes the InMemoryRunner.
 
     Args:
@@ -551,6 +636,7 @@ class InMemoryRunner(Runner):
         app_name=app_name,
         agent=agent,
         artifact_service=InMemoryArtifactService(),
+        plugins=plugins,
         session_service=self._in_memory_session_service,
         memory_service=InMemoryMemoryService(),
     )

tests/unittests/agents/test_base_agent.py

@@ -26,6 +26,8 @@ from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.events import Event
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.plugins.plugin_manager import PluginManager
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.genai import types
 import pytest
@@ -83,6 +85,35 @@ async def _async_after_agent_callback_append_agent_reply(
   )
 
 
+class MockPlugin(BasePlugin):
+  before_agent_text = 'before_agent_text from MockPlugin'
+  after_agent_text = 'after_agent_text from MockPlugin'
+
+  def __init__(self, name='mock_plugin'):
+    self.name = name
+    self.enable_before_agent_callback = False
+    self.enable_after_agent_callback = False
+
+  async def before_agent_callback(
+      self, *, agent: BaseAgent, callback_context: CallbackContext
+  ) -> Optional[types.Content]:
+    if not self.enable_before_agent_callback:
+      return None
+    return types.Content(parts=[types.Part(text=self.before_agent_text)])
+
+  async def after_agent_callback(
+      self, *, agent: BaseAgent, callback_context: CallbackContext
+  ) -> Optional[types.Content]:
+    if not self.enable_after_agent_callback:
+      return None
+    return types.Content(parts=[types.Part(text=self.after_agent_text)])
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
 class _IncompleteAgent(BaseAgent):
   pass
 
@@ -113,7 +144,10 @@ class _TestingAgent(BaseAgent):
 
 
 async def _create_parent_invocation_context(
-    test_name: str, agent: BaseAgent, branch: Optional[str] = None
+    test_name: str,
+    agent: BaseAgent,
+    branch: Optional[str] = None,
+    plugins: list[BasePlugin] = [],
 ) -> InvocationContext:
   session_service = InMemorySessionService()
   session = await session_service.create_session(
@@ -125,6 +159,7 @@ async def _create_parent_invocation_context(
       agent=agent,
       session=session,
       session_service=session_service,
+      plugin_manager=PluginManager(plugins=plugins),
   )
 
 
@@ -190,6 +225,36 @@ async def test_run_async_before_agent_callback_noop(
   spy_run_async_impl.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_run_async_before_agent_callback_use_plugin(
+    request: pytest.FixtureRequest,
+    mocker: pytest_mock.MockerFixture,
+    mock_plugin: MockPlugin,
+):
+  """Test that the before agent callback uses the plugin response if both plugin callback and canonical agent callbacks are present."""
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      before_agent_callback=_before_agent_callback_bypass_agent,
+  )
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, agent, plugins=[mock_plugin]
+  )
+  mock_plugin.enable_before_agent_callback = True
+  spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
+  spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
+
+  # Act
+  events = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  spy_before_agent_callback.assert_not_called()
+  spy_run_async_impl.assert_not_called()
+
+  assert len(events) == 1
+  assert events[0].content.parts[0].text == MockPlugin.before_agent_text
+
+
 @pytest.mark.asyncio
 async def test_run_async_with_async_before_agent_callback_noop(
     request: pytest.FixtureRequest,
@@ -486,6 +551,34 @@ async def test_after_agent_callbacks_chain(
         mock_cb.assert_called(expected_calls_count)
 
 
+@pytest.mark.asyncio
+async def test_run_async_after_agent_callback_use_plugin(
+    request: pytest.FixtureRequest,
+    mocker: pytest_mock.MockerFixture,
+    mock_plugin: MockPlugin,
+):
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      after_agent_callback=_after_agent_callback_noop,
+  )
+  mock_plugin.enable_after_agent_callback = True
+  parent_ctx = await _create_parent_invocation_context(
+      request.function.__name__, agent, plugins=[mock_plugin]
+  )
+  spy_after_agent_callback = mocker.spy(agent, 'after_agent_callback')
+
+  # Act
+  events = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  spy_after_agent_callback.assert_not_called()
+  # The first event is regular model response, the second event is
+  # after_agent_callback response.
+  assert len(events) == 2
+  assert events[1].content.parts[0].text == mock_plugin.after_agent_text
+
+
 @pytest.mark.asyncio
 async def test_run_async_after_agent_callback_noop(
     request: pytest.FixtureRequest,
@@ -757,3 +850,7 @@ def test_set_parent_agent_for_sub_agent_twice(
         name=f'{request.function.__name__}_parent_2',
         sub_agents=[sub_agent],
     )
+
+
+if __name__ == '__main__':
+  pytest.main([__file__])

tests/unittests/agents/test_langgraph_agent.py

@@ -17,6 +17,7 @@ from unittest.mock import MagicMock
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.langgraph_agent import LangGraphAgent
 from google.adk.events import Event
+from google.adk.plugins.plugin_manager import PluginManager
 from google.genai import types
 from langchain_core.messages import AIMessage
 from langchain_core.messages import HumanMessage
@@ -169,6 +170,7 @@ async def test_langgraph_agent(
   mock_session.events = events_list
   mock_parent_context.invocation_id = "test_invocation_id"
   mock_parent_context.model_copy.return_value = mock_parent_context
+  mock_parent_context.plugin_manager = PluginManager(plugins=[])
 
   weather_agent = LangGraphAgent(
       name="weather_agent",

tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py

@@ -0,0 +1,128 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from google.adk.agents import Agent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.models import LlmRequest
+from google.adk.models import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class MockPlugin(BasePlugin):
+  before_model_text = 'before_model_text from MockPlugin'
+  after_model_text = 'after_model_text from MockPlugin'
+
+  def __init__(self, name='mock_plugin'):
+    self.name = name
+    self.enable_before_model_callback = False
+    self.enable_after_model_callback = False
+    self.before_model_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.before_model_text)]
+        )
+    )
+    self.after_model_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.after_model_text)]
+        )
+    )
+
+  async def before_model_callback(
+      self, *, callback_context: CallbackContext, llm_request: LlmRequest
+  ) -> Optional[LlmResponse]:
+    if not self.enable_before_model_callback:
+      return None
+    return self.before_model_response
+
+  async def after_model_callback(
+      self, *, callback_context: CallbackContext, llm_response: LlmResponse
+  ) -> Optional[LlmResponse]:
+    if not self.enable_after_model_callback:
+      return None
+    return self.after_model_response
+
+
+CANONICAL_MODEL_CALLBACK_CONTENT = 'canonical_model_callback_content'
+
+
+def canonical_agent_model_callback(**kwargs) -> Optional[LlmResponse]:
+  return LlmResponse(
+      content=testing_utils.ModelContent(
+          [types.Part.from_text(text=CANONICAL_MODEL_CALLBACK_CONTENT)]
+      )
+  )
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
+def test_before_model_callback_with_plugin(mock_plugin):
+  """Tests that the model response is overridden by before_model_callback from the plugin."""
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  mock_plugin.enable_before_model_callback = True
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', mock_plugin.before_model_text),
+  ]
+
+
+def test_before_model_fallback_canonical_callback(mock_plugin):
+  """Tests that when plugin returns empty response, the model response is overridden by the canonical agent model callback."""
+  responses = ['model_response']
+  mock_plugin.enable_before_model_callback = False
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=canonical_agent_model_callback,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', CANONICAL_MODEL_CALLBACK_CONTENT),
+  ]
+
+
+def test_before_model_callback_fallback_model(mock_plugin):
+  """Tests that the model response is executed normally when both plugin and canonical agent model callback return empty response."""
+  responses = ['model_response']
+  mock_plugin.enable_before_model_callback = False
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
+if __name__ == '__main__':
+  pytest.main([__file__])

tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py

@@ -0,0 +1,128 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from google.adk.agents import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import handle_function_calls_async
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class MockPlugin(BasePlugin):
+  before_tool_response = {"MockPlugin": "before_tool_response from MockPlugin"}
+  after_tool_response = {"MockPlugin": "after_tool_response from MockPlugin"}
+
+  def __init__(self, name="mock_plugin"):
+    self.name = name
+    self.enable_before_tool_callback = False
+    self.enable_after_tool_callback = False
+
+  async def before_tool_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+  ) -> Optional[dict]:
+    if not self.enable_before_tool_callback:
+      return None
+    return self.before_tool_response
+
+  async def after_tool_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+      result: dict,
+  ) -> Optional[dict]:
+    if not self.enable_after_tool_callback:
+      return None
+    return self.after_tool_response
+
+
+@pytest.fixture
+def mock_tool():
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  return FunctionTool(simple_fn)
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
+async def invoke_tool_with_plugin(mock_tool, mock_plugin) -> Optional[Event]:
+  """Invokes a tool with a plugin."""
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[mock_tool],
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content="", plugins=[mock_plugin]
+  )
+  # Build function call event
+  function_call = types.FunctionCall(name=mock_tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {mock_tool.name: mock_tool}
+  return await handle_function_calls_async(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+
+@pytest.mark.asyncio
+async def test_async_before_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_before_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.before_tool_response
+
+
+@pytest.mark.asyncio
+async def test_async_after_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_after_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.after_tool_response
+
+
+if __name__ == "__main__":
+  pytest.main([__file__])

tests/unittests/test_runners.py

@@ -15,13 +15,20 @@
 from typing import Optional
 
 from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
 from google.adk.events.event import Event
+from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.sessions.session import Session
 from google.genai import types
+import pytest
+
+TEST_APP_ID = "test_app"
+TEST_USER_ID = "test_user"
+TEST_SESSION_ID = "test_session"
 
 
 class MockAgent(BaseAgent):
@@ -72,6 +79,51 @@ class MockLlmAgent(LlmAgent):
     )
 
 
+class MockPlugin(BasePlugin):
+  """Mock plugin for unit testing."""
+
+  ON_USER_CALLBACK_MSG = (
+      "Modified user message ON_USER_CALLBACK_MSG from MockPlugin"
+  )
+  ON_EVENT_CALLBACK_MSG = "Modified event ON_EVENT_CALLBACK_MSG from MockPlugin"
+
+  def __init__(self):
+    super().__init__(name="mock_plugin")
+    self.enable_user_message_callback = False
+    self.enable_event_callback = False
+
+  async def on_user_message_callback(
+      self,
+      *,
+      invocation_context: InvocationContext,
+      user_message: types.Content,
+  ) -> Optional[types.Content]:
+    if not self.enable_user_message_callback:
+      return None
+    return types.Content(
+        role="model",
+        parts=[types.Part(text=self.ON_USER_CALLBACK_MSG)],
+    )
+
+  async def on_event_callback(
+      self, *, invocation_context: InvocationContext, event: Event
+  ) -> Optional[Event]:
+    if not self.enable_event_callback:
+      return None
+    return Event(
+        invocation_id="",
+        author="",
+        content=types.Content(
+            parts=[
+                types.Part(
+                    text=self.ON_EVENT_CALLBACK_MSG,
+                )
+            ],
+            role=event.content.role,
+        ),
+    )
+
+
 class TestRunnerFindAgentToRun:
   """Tests for Runner._find_agent_to_run method."""
 
@@ -308,3 +360,73 @@ class TestRunnerFindAgentToRun:
     # MockAgent inherits from BaseAgent, not LlmAgent, so it should return False
     result = self.runner._is_transferable_across_agent_tree(non_llm_agent)
     assert result is False
+
+
+class TestRunnerWithPlugins:
+  """Tests for Runner with plugins."""
+
+  def setup_method(self):
+    self.plugin = MockPlugin()
+    self.session_service = InMemorySessionService()
+    self.artifact_service = InMemoryArtifactService()
+    self.root_agent = MockLlmAgent("root_agent")
+    self.runner = Runner(
+        app_name="test_app",
+        agent=MockLlmAgent("test_agent"),
+        session_service=self.session_service,
+        artifact_service=self.artifact_service,
+        plugins=[self.plugin],
+    )
+
+  async def run_test(self, original_user_input="Hello") -> list[Event]:
+    """Prepares the test by creating a session and running the runner."""
+    await self.session_service.create_session(
+        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+    )
+    events = []
+    async for event in self.runner.run_async(
+        user_id=TEST_USER_ID,
+        session_id=TEST_SESSION_ID,
+        new_message=types.Content(
+            role="user", parts=[types.Part(text=original_user_input)]
+        ),
+    ):
+      events.append(event)
+    return events
+
+  @pytest.mark.asyncio
+  async def test_runner_is_initialized_with_plugins(self):
+    """Test that the runner is initialized with plugins."""
+    await self.run_test()
+
+    assert self.runner.plugin_manager is not None
+
+  @pytest.mark.asyncio
+  async def test_runner_modifies_user_message_before_execution(self):
+    """Test that the runner modifies the user message before execution."""
+    original_user_input = "original_input"
+    self.plugin.enable_user_message_callback = True
+
+    await self.run_test(original_user_input=original_user_input)
+    session = await self.session_service.get_session(
+        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
+    )
+    generated_event = session.events[0]
+    modified_user_message = generated_event.content.parts[0].text
+
+    assert modified_user_message == MockPlugin.ON_USER_CALLBACK_MSG
+
+  @pytest.mark.asyncio
+  async def test_runner_modifies_event_after_execution(self):
+    """Test that the runner modifies the event after execution."""
+    self.plugin.enable_event_callback = True
+
+    events = await self.run_test()
+    generated_event = events[0]
+    modified_event_message = generated_event.content.parts[0].text
+
+    assert modified_event_message == MockPlugin.ON_EVENT_CALLBACK_MSG
+
+
+if __name__ == "__main__":
+  pytest.main([__file__])

tests/unittests/testing_utils.py

@@ -30,6 +30,8 @@ from google.adk.models.base_llm import BaseLlm
 from google.adk.models.base_llm_connection import BaseLlmConnection
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.plugins.plugin_manager import PluginManager
 from google.adk.runners import InMemoryRunner as AfInMemoryRunner
 from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
@@ -57,7 +59,10 @@ class ModelContent(types.Content):
 
 
 async def create_invocation_context(
-    agent: Agent, user_content: str = '', run_config: RunConfig = None
+    agent: Agent,
+    user_content: str = '',
+    run_config: RunConfig = None,
+    plugins: list[BasePlugin] = [],
 ):
   invocation_id = 'test_id'
   artifact_service = InMemoryArtifactService()
@@ -67,6 +72,7 @@ async def create_invocation_context(
       artifact_service=artifact_service,
       session_service=session_service,
       memory_service=memory_service,
+      plugin_manager=PluginManager(plugins=plugins),
       invocation_id=invocation_id,
       agent=agent,
       session=await session_service.create_session(
@@ -165,6 +171,7 @@ class InMemoryRunner:
       self,
       root_agent: Union[Agent, LlmAgent],
       response_modalities: list[str] = None,
+      plugins: list[BasePlugin] = [],
   ):
     self.root_agent = root_agent
     self.runner = Runner(
@@ -173,6 +180,7 @@ class InMemoryRunner:
         artifact_service=InMemoryArtifactService(),
         session_service=InMemorySessionService(),
         memory_service=InMemoryMemoryService(),
+        plugins=plugins,
     )
     self.session_id = None