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

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