chore: import upstream snapshot with attribution

tests/unittests/plugins/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 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.

tests/unittests/plugins/test_auto_tracing_plugin.py

@@ -0,0 +1,461 @@
+# Copyright 2026 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 __future__ import annotations
+
+import asyncio
+import sys
+import types
+from typing import Any
+from unittest import mock
+
+from google.adk.plugins import auto_tracing_helpers
+from google.adk.plugins import auto_tracing_plugin
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace import export as trace_export
+from opentelemetry.sdk.trace.export import in_memory_span_exporter
+import pytest
+
+_FIXTURE_MODULE_NAME = (
+    "google.adk.tests.unittests.plugins.synthetic_test_fixture"
+)
+
+
+def _sync_fn(x: int) -> int:
+  return x + 1
+
+
+async def _async_fn(x: int) -> int:
+  return x * 2
+
+
+def _build_fixture_module() -> types.ModuleType:
+  module = types.ModuleType(_FIXTURE_MODULE_NAME)
+  module.__name__ = _FIXTURE_MODULE_NAME
+  for fn in (_sync_fn, _async_fn):
+    fn.__module__ = _FIXTURE_MODULE_NAME
+
+  def _method(unused_self, x: int) -> int:
+    return x - 1
+
+  async def _async_method(unused_self, x: int) -> int:
+    return x + 10
+
+  _method.__module__ = _FIXTURE_MODULE_NAME
+  _async_method.__module__ = _FIXTURE_MODULE_NAME
+  cls = type("C", (), {"method": _method, "async_method": _async_method})
+  cls.__module__ = _FIXTURE_MODULE_NAME
+
+  module.sync_fn = _sync_fn
+  module.async_fn = _async_fn
+  module.C = cls
+  return module
+
+
+def _install_module(name: str, fn) -> types.ModuleType:
+  mod = types.ModuleType(name)
+  mod.__name__ = name
+  fn.__module__ = name
+  mod.fn = fn
+  sys.modules[name] = mod
+  return mod
+
+
+def _run_sync(module):
+  return module.sync_fn(7)
+
+
+def _run_async(module):
+  return asyncio.run(module.async_fn(4))
+
+
+def _run_class_method(module):
+  return module.C().method(5)
+
+
+def _run_class_async_method(module):
+  return asyncio.run(module.C().async_method(5))
+
+
+class _Ctx:
+
+  def __init__(self, agent):
+    self.agent = agent
+
+
+def _build_slot_agent(module: str, slots, attr: str, value):
+  cls = type("_Agent", (), {"__slots__": slots, "__module__": module})
+  obj = cls()
+  setattr(obj, attr, value)
+  return obj
+
+
+def _sub_helper():
+  return 7
+
+
+@pytest.fixture
+def fixture():
+  exporter = in_memory_span_exporter.InMemorySpanExporter()
+  provider = trace_sdk.TracerProvider()
+  provider.add_span_processor(trace_export.SimpleSpanProcessor(exporter))
+  tracer = provider.get_tracer("test")
+  module = _build_fixture_module()
+  sys.modules[_FIXTURE_MODULE_NAME] = module
+  yield types.SimpleNamespace(exporter=exporter, tracer=tracer, module=module)
+  sys.modules.pop(_FIXTURE_MODULE_NAME, None)
+
+
+def _span_names(exporter) -> list[str]:
+  return [s.name for s in exporter.get_finished_spans()]
+
+
+def _attrs_for(exporter, substr: str) -> dict[str, Any]:
+  matches = [
+      dict(s.attributes or {})
+      for s in exporter.get_finished_spans()
+      if substr in s.name
+  ]
+  assert matches, f"no span matched {substr!r} in {_span_names(exporter)}"
+  return matches[0]
+
+
+def _instrument(
+    tracer, scope_prefixes=(_FIXTURE_MODULE_NAME,)
+) -> auto_tracing_plugin.AutoTracingPlugin:
+  plugin = auto_tracing_plugin.AutoTracingPlugin(
+      tracer=tracer, extra_scope_prefixes=scope_prefixes
+  )
+  asyncio.run(plugin.before_run_callback(invocation_context=None))
+  return plugin
+
+
+@pytest.mark.parametrize(
+    "run_fn,expected_substr",
+    [
+        (_run_sync, "_sync_fn"),
+        (_run_async, "_async_fn"),
+        (_run_class_method, "._method"),
+        (_run_class_async_method, "._async_method"),
+    ],
+)
+def test_emits_span(fixture, run_fn, expected_substr):
+  _instrument(fixture.tracer)
+  run_fn(fixture.module)
+  assert any(
+      expected_substr in n for n in _span_names(fixture.exporter)
+  ), f"missing {expected_substr!r} in {_span_names(fixture.exporter)}"
+
+
+@pytest.mark.parametrize(
+    "run_fn,expected_substr,expected_attrs",
+    [
+        (
+            _run_sync,
+            "_sync_fn",
+            {"adk.fn.arg.x": "7", "adk.fn.return": "8"},
+        ),
+        (_run_async, "_async_fn", {"adk.fn.return": "8"}),
+    ],
+)
+def test_records_io(fixture, run_fn, expected_substr, expected_attrs):
+  _instrument(fixture.tracer)
+  run_fn(fixture.module)
+  attrs = _attrs_for(fixture.exporter, expected_substr)
+  assert {k: attrs.get(k) for k in expected_attrs} == expected_attrs
+
+
+@pytest.mark.parametrize("attr", ["sync_fn", "async_fn"])
+def test_repeat_instrument_is_idempotent(fixture, attr):
+  plugin = _instrument(fixture.tracer)
+  first = getattr(fixture.module, attr)
+  asyncio.run(plugin.before_run_callback(invocation_context=None))
+  assert getattr(fixture.module, attr) is first
+
+
+@pytest.mark.parametrize("attr", ["sync_fn", "async_fn"])
+def test_wrapper_marker_is_true(fixture, attr):
+  _instrument(fixture.tracer)
+  assert (
+      getattr(getattr(fixture.module, attr), auto_tracing_helpers.WRAPPED_ATTR)
+      is True
+  )
+
+
+def test_out_of_scope_module_is_not_instrumented(fixture):
+  name = "auto_tracing_plugin_test_not_in_scope"
+  mod = _install_module(name, lambda: 42)
+  try:
+    _instrument(fixture.tracer)
+    mod.fn()
+    assert f"{name}.fn" not in _span_names(fixture.exporter)
+  finally:
+    sys.modules.pop(name, None)
+
+
+def test_records_exception(fixture):
+  name = "auto_tracing_plugin_test_boom"
+
+  def boom():
+    raise ValueError("kaboom")
+
+  mod = _install_module(name, boom)
+  try:
+    _instrument(fixture.tracer, scope_prefixes=(name,))
+    with pytest.raises(ValueError, match="kaboom"):
+      mod.fn()
+    attrs = _attrs_for(fixture.exporter, "boom")
+    assert attrs.get("adk.fn.exc_type") == "ValueError"
+    assert "kaboom" in attrs.get("adk.fn.exc_repr", "")
+  finally:
+    sys.modules.pop(name, None)
+
+
+def test_walk_returns_quickly_on_none_agent(fixture):
+  plugin = auto_tracing_plugin.AutoTracingPlugin(tracer=fixture.tracer)
+  asyncio.run(plugin.before_run_callback(invocation_context=_Ctx(None)))
+  assert _span_names(fixture.exporter) == []
+
+
+def test_add_agent_scope_picks_up_agent_package(fixture):
+  pkg = "auto_tracing_plugin_test_agent_pkg"
+  mod_name = f"{pkg}.helpers"
+
+  def helper():
+    return 99
+
+  mod = _install_module(mod_name, helper)
+  try:
+
+    class _Agent:
+      __module__ = f"{pkg}.agent"
+
+    plugin = auto_tracing_plugin.AutoTracingPlugin(tracer=fixture.tracer)
+    asyncio.run(plugin.before_run_callback(invocation_context=_Ctx(_Agent())))
+    mod.fn()
+    assert any("helper" in n for n in _span_names(fixture.exporter)), (
+        f"agent pkg {pkg!r} was not absorbed;"
+        f" spans={_span_names(fixture.exporter)}"
+    )
+  finally:
+    sys.modules.pop(mod_name, None)
+
+
+@pytest.mark.parametrize(
+    "pkg,slots",
+    [
+        ("auto_tracing_plugin_test_slots_pkg", ("child",)),
+        ("auto_tracing_plugin_test_str_slot_pkg", "child"),
+    ],
+)
+def test_add_agent_scope_walks_slots_attrs(fixture, pkg, slots):
+  sub_mod_name = f"{pkg}.sub"
+  mod = _install_module(sub_mod_name, _sub_helper)
+  try:
+    sub = type("_Sub", (), {"__module__": sub_mod_name})()
+    agent = _build_slot_agent(f"{pkg}.agent", slots, "child", sub)
+    plugin = auto_tracing_plugin.AutoTracingPlugin(tracer=fixture.tracer)
+    asyncio.run(plugin.before_run_callback(invocation_context=_Ctx(agent)))
+    mod.fn()
+    assert any("_sub_helper" in n for n in _span_names(fixture.exporter)), (
+        f"slot-referenced pkg {pkg!r} was not absorbed;"
+        f" spans={_span_names(fixture.exporter)}"
+    )
+  finally:
+    sys.modules.pop(sub_mod_name, None)
+
+
+def test_add_agent_scope_does_not_fire_property_descriptors(fixture):
+  fired: list[str] = []
+
+  class _Agent:
+    __module__ = "auto_tracing_plugin_test_no_descriptor_pkg.agent"
+
+    @property
+    def expensive(self):
+      fired.append("expensive")
+      raise RuntimeError("should never be invoked during scope walk")
+
+  plugin = auto_tracing_plugin.AutoTracingPlugin(tracer=fixture.tracer)
+  asyncio.run(plugin.before_run_callback(invocation_context=_Ctx(_Agent())))
+  assert fired == [], f"@property fired during agent-scope walk: {fired!r}"
+
+
+def test_module_removed_mid_iteration_does_not_log_exception(fixture):
+  name = "auto_tracing_plugin_test_disappearing"
+  _install_module(name, lambda: 1)
+  try:
+    plugin = auto_tracing_plugin.AutoTracingPlugin(
+        tracer=fixture.tracer, extra_scope_prefixes=(name,)
+    )
+
+    class _DroppingModules(dict):
+
+      def get(self, key, default=None):
+        if key == name:
+          return None
+        return super().get(key, default)
+
+    dropping = _DroppingModules(sys.modules)
+    with (
+        mock.patch.object(
+            auto_tracing_plugin.logger, "exception", autospec=True
+        ) as log_exc,
+        mock.patch.object(auto_tracing_plugin.sys, "modules", new=dropping),
+    ):
+      asyncio.run(plugin.before_run_callback(invocation_context=None))
+    assert (
+        not log_exc.called
+    ), f"unexpected logger.exception calls: {log_exc.call_args_list}"
+    assert name not in plugin._wrapped_modules
+  finally:
+    sys.modules.pop(name, None)
+
+
+def test_repeat_instrument_does_not_rewrap(fixture):
+  plugin = _instrument(fixture.tracer)
+  assert getattr(
+      fixture.module.sync_fn, auto_tracing_helpers.WRAPPED_ATTR, False
+  )
+  assert _FIXTURE_MODULE_NAME in plugin._wrapped_modules
+  with mock.patch.object(plugin, "_wrap_module", autospec=True) as wrap_module:
+    asyncio.run(plugin.before_run_callback(invocation_context=None))
+  wrap_module.assert_not_called()
+
+
+class _Slotted:
+  __slots__ = ("a", "b")
+
+  def __init__(self):
+    self.a = 1
+    self.b = "x"
+
+
+class _Bare:
+  __slots__ = ()
+
+
+@pytest.mark.parametrize(
+    "instance,expected_substrings",
+    [
+        (_Slotted(), ("_Slotted", "a=1", "b='x'")),
+        (_Bare(), ("<_Bare>",)),
+    ],
+)
+def test_summarize_default(instance, expected_substrings):
+  rendered = auto_tracing_helpers.safe_repr(
+      instance, auto_tracing_helpers.Caps()
+  )
+  for s in expected_substrings:
+    assert s in rendered, rendered
+
+
+def test_add_agent_scope_picks_up_top_level_module(fixture):
+  top_mod_name = "auto_tracing_plugin_test_top_level_pkg"
+
+  def top_helper():
+    return 1
+
+  mod = _install_module(top_mod_name, top_helper)
+  try:
+
+    class _Agent:
+      __module__ = top_mod_name
+
+    plugin = auto_tracing_plugin.AutoTracingPlugin(tracer=fixture.tracer)
+    asyncio.run(plugin.before_run_callback(invocation_context=_Ctx(_Agent())))
+    mod.fn()
+    assert any("top_helper" in n for n in _span_names(fixture.exporter)), (
+        f"top-level module {top_mod_name!r} not absorbed;"
+        f" spans={_span_names(fixture.exporter)}"
+    )
+  finally:
+    sys.modules.pop(top_mod_name, None)
+
+
+def test_signature_introspection_happens_once_per_wrap(fixture):
+  with mock.patch.object(
+      auto_tracing_helpers.inspect, "signature", autospec=True
+  ) as sig:
+    sig.side_effect = auto_tracing_helpers.inspect.signature
+    _instrument(fixture.tracer)
+    wrap_calls = sig.call_count
+    for _ in range(5):
+      _run_sync(fixture.module)
+      _run_async(fixture.module)
+  assert (
+      sig.call_count == wrap_calls
+  ), f"inspect.signature called per-call: {wrap_calls} -> {sig.call_count}"
+
+
+def test_async_gen_caps_buffered_items(fixture):
+  cap = 3
+  total_yields = 100
+  name = "auto_tracing_plugin_test_async_gen_cap"
+
+  async def producer():
+    for i in range(total_yields):
+      yield i
+
+  mod = _install_module(name, producer)
+  try:
+    plugin = auto_tracing_plugin.AutoTracingPlugin(
+        tracer=fixture.tracer,
+        extra_scope_prefixes=(name,),
+        max_recorded_yields=cap,
+    )
+    asyncio.run(plugin.before_run_callback(invocation_context=None))
+
+    async def drive():
+      seen = []
+      async for x in mod.fn():
+        seen.append(x)
+      return seen
+
+    out = asyncio.run(drive())
+    assert out == list(range(total_yields))
+    attrs = _attrs_for(fixture.exporter, "producer")
+    rendered = attrs.get("adk.fn.return", "")
+    assert f"{total_yields} items yielded" in rendered, rendered
+    assert f"first {cap}:" in rendered, rendered
+    assert f"+ {total_yields - cap} more" in rendered, rendered
+  finally:
+    sys.modules.pop(name, None)
+
+
+def test_sync_gen_caps_buffered_items(fixture):
+  cap = 2
+  total_yields = 50
+  name = "auto_tracing_plugin_test_sync_gen_cap"
+
+  def producer():
+    for i in range(total_yields):
+      yield i
+
+  mod = _install_module(name, producer)
+  try:
+    plugin = auto_tracing_plugin.AutoTracingPlugin(
+        tracer=fixture.tracer,
+        extra_scope_prefixes=(name,),
+        max_recorded_yields=cap,
+    )
+    asyncio.run(plugin.before_run_callback(invocation_context=None))
+    out = list(mod.fn())
+    assert out == list(range(total_yields))
+    attrs = _attrs_for(fixture.exporter, "producer")
+    rendered = attrs.get("adk.fn.return", "")
+    assert f"{total_yields} items yielded" in rendered, rendered
+    assert f"first {cap}:" in rendered, rendered
+  finally:
+    sys.modules.pop(name, None)

tests/unittests/plugins/test_base_plugin.py

@@ -0,0 +1,280 @@
+# Copyright 2026 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 __future__ import annotations
+
+from unittest.mock import Mock
+
+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.event import Event
+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.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+
+class TestablePlugin(BasePlugin):
+  __test__ = False
+  """A concrete implementation of BasePlugin for testing purposes."""
+  pass
+
+
+class FullOverridePlugin(BasePlugin):
+  __test__ = False
+
+  """A plugin that overrides every single callback method for testing."""
+
+  def __init__(self, name: str = "full_override"):
+    super().__init__(name)
+
+  async def on_user_message_callback(self, **kwargs) -> str:
+    return "overridden_on_user_message"
+
+  async def before_run_callback(self, **kwargs) -> str:
+    return "overridden_before_run"
+
+  async def after_run_callback(self, **kwargs) -> str:
+    return "overridden_after_run"
+
+  async def on_event_callback(self, **kwargs) -> str:
+    return "overridden_on_event"
+
+  async def before_agent_callback(self, **kwargs) -> str:
+    return "overridden_before_agent"
+
+  async def after_agent_callback(self, **kwargs) -> str:
+    return "overridden_after_agent"
+
+  async def before_tool_callback(self, **kwargs) -> str:
+    return "overridden_before_tool"
+
+  async def after_tool_callback(self, **kwargs) -> str:
+    return "overridden_after_tool"
+
+  async def on_tool_error_callback(self, **kwargs) -> str:
+    return "overridden_on_tool_error"
+
+  async def before_model_callback(self, **kwargs) -> str:
+    return "overridden_before_model"
+
+  async def after_model_callback(self, **kwargs) -> str:
+    return "overridden_after_model"
+
+  async def on_model_error_callback(self, **kwargs) -> str:
+    return "overridden_on_model_error"
+
+
+def test_base_plugin_initialization():
+  """Tests that a plugin is initialized with the correct name."""
+  plugin_name = "my_test_plugin"
+  plugin = TestablePlugin(name=plugin_name)
+  assert plugin.name == plugin_name
+
+
+@pytest.mark.asyncio
+async def test_base_plugin_default_callbacks_return_none():
+  """Tests that the default (non-overridden) callbacks in BasePlugin exist
+
+  and return None as expected.
+  """
+  plugin = TestablePlugin(name="default_plugin")
+
+  # Mocking all necessary context objects
+  mock_context = Mock()
+  mock_user_message = Mock()
+
+  # The default implementations should do nothing and return None.
+  assert (
+      await plugin.on_user_message_callback(
+          user_message=mock_user_message,
+          invocation_context=mock_context,
+      )
+      is None
+  )
+  assert (
+      await plugin.before_run_callback(invocation_context=mock_context) is None
+  )
+  assert (
+      await plugin.after_run_callback(invocation_context=mock_context) is None
+  )
+  assert (
+      await plugin.on_event_callback(
+          invocation_context=mock_context, event=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.before_agent_callback(
+          agent=mock_context, callback_context=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.after_agent_callback(
+          agent=mock_context, callback_context=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.before_tool_callback(
+          tool=mock_context, tool_args={}, tool_context=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.after_tool_callback(
+          tool=mock_context, tool_args={}, tool_context=mock_context, result={}
+      )
+      is None
+  )
+  assert (
+      await plugin.on_tool_error_callback(
+          tool=mock_context,
+          tool_args={},
+          tool_context=mock_context,
+          error=Exception(),
+      )
+      is None
+  )
+  assert (
+      await plugin.before_model_callback(
+          callback_context=mock_context, llm_request=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.after_model_callback(
+          callback_context=mock_context, llm_response=mock_context
+      )
+      is None
+  )
+  assert (
+      await plugin.on_model_error_callback(
+          callback_context=mock_context,
+          llm_request=mock_context,
+          error=Exception(),
+      )
+      is None
+  )
+
+
+@pytest.mark.asyncio
+async def test_base_plugin_all_callbacks_can_be_overridden():
+  """Verifies that a user can create a subclass of BasePlugin and that all
+
+  overridden methods are correctly called.
+  """
+  plugin = FullOverridePlugin()
+
+  # Create mock objects for all required arguments. We don't need real
+  # objects, just placeholders to satisfy the method signatures.
+  mock_user_message = Mock(spec=types.Content)
+  mock_invocation_context = Mock(spec=InvocationContext)
+  mock_callback_context = Mock(spec=CallbackContext)
+  mock_agent = Mock(spec=BaseAgent)
+  mock_tool = Mock(spec=BaseTool)
+  mock_tool_context = Mock(spec=ToolContext)
+  mock_llm_request = Mock(spec=LlmRequest)
+  mock_llm_response = Mock(spec=LlmResponse)
+  mock_event = Mock(spec=Event)
+  mock_error = Mock(spec=Exception)
+
+  # Call each method and assert it returns the unique string from the override.
+  # This proves that the subclass's method was executed.
+  assert (
+      await plugin.on_user_message_callback(
+          user_message=mock_user_message,
+          invocation_context=mock_invocation_context,
+      )
+      == "overridden_on_user_message"
+  )
+  assert (
+      await plugin.before_run_callback(
+          invocation_context=mock_invocation_context
+      )
+      == "overridden_before_run"
+  )
+  assert (
+      await plugin.after_run_callback(
+          invocation_context=mock_invocation_context
+      )
+      == "overridden_after_run"
+  )
+  assert (
+      await plugin.on_event_callback(
+          invocation_context=mock_invocation_context, event=mock_event
+      )
+      == "overridden_on_event"
+  )
+  assert (
+      await plugin.before_agent_callback(
+          agent=mock_agent, callback_context=mock_callback_context
+      )
+      == "overridden_before_agent"
+  )
+  assert (
+      await plugin.after_agent_callback(
+          agent=mock_agent, callback_context=mock_callback_context
+      )
+      == "overridden_after_agent"
+  )
+  assert (
+      await plugin.before_model_callback(
+          callback_context=mock_callback_context, llm_request=mock_llm_request
+      )
+      == "overridden_before_model"
+  )
+  assert (
+      await plugin.after_model_callback(
+          callback_context=mock_callback_context, llm_response=mock_llm_response
+      )
+      == "overridden_after_model"
+  )
+  assert (
+      await plugin.before_tool_callback(
+          tool=mock_tool, tool_args={}, tool_context=mock_tool_context
+      )
+      == "overridden_before_tool"
+  )
+  assert (
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args={},
+          tool_context=mock_tool_context,
+          result={},
+      )
+      == "overridden_after_tool"
+  )
+  assert (
+      await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args={},
+          tool_context=mock_tool_context,
+          error=mock_error,
+      )
+      == "overridden_on_tool_error"
+  )
+  assert (
+      await plugin.on_model_error_callback(
+          callback_context=mock_callback_context,
+          llm_request=mock_llm_request,
+          error=mock_error,
+      )
+      == "overridden_on_model_error"
+  )

tests/unittests/plugins/test_context_filtering_plugin.py

@@ -0,0 +1,468 @@
+# Copyright 2026 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.
+
+"""Unit tests for the ContextFilteringPlugin."""
+
+from unittest import mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.models.llm_request import LlmRequest
+from google.adk.plugins.context_filter_plugin import ContextFilterPlugin
+from google.genai import types
+import pytest
+
+
+def _create_content(role: str, text: str) -> types.Content:
+  return types.Content(parts=[types.Part(text=text)], role=role)
+
+
+@pytest.mark.asyncio
+async def test_filter_last_n_invocations():
+  """Tests that the context is truncated to the last N invocations."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=1)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  assert len(llm_request.contents) == 2
+  assert llm_request.contents[0].parts[0].text == "user_prompt_2"
+  assert llm_request.contents[1].parts[0].text == "model_response_2"
+
+
+@pytest.mark.asyncio
+async def test_filter_with_function():
+  """Tests that a custom filter function is applied to the context."""
+
+  def remove_model_responses(contents):
+    return [c for c in contents if c.role != "model"]
+
+  plugin = ContextFilterPlugin(custom_filter=remove_model_responses)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  assert len(llm_request.contents) == 2
+  assert all(c.role == "user" for c in llm_request.contents)
+
+
+@pytest.mark.asyncio
+async def test_filter_with_function_and_last_n_invocations():
+  """Tests that both filtering methods are applied correctly."""
+
+  def remove_first_invocation(contents):
+    return contents[2:]
+
+  plugin = ContextFilterPlugin(
+      num_invocations_to_keep=1, custom_filter=remove_first_invocation
+  )
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  assert len(llm_request.contents) == 0
+
+
+@pytest.mark.asyncio
+async def test_no_filtering_when_no_options_provided():
+  """Tests that no filtering occurs when no options are provided."""
+  plugin = ContextFilterPlugin()
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+  original_contents = list(llm_request.contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  assert llm_request.contents == original_contents
+
+
+@pytest.mark.asyncio
+async def test_last_n_invocations_with_multiple_user_turns():
+  """Tests filtering with multiple user turns in a single invocation."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=1)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2a"),
+      _create_content("user", "user_prompt_2b"),
+      _create_content("model", "model_response_2"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  assert len(llm_request.contents) == 3
+  assert llm_request.contents[0].parts[0].text == "user_prompt_2a"
+  assert llm_request.contents[1].parts[0].text == "user_prompt_2b"
+  assert llm_request.contents[2].parts[0].text == "model_response_2"
+
+
+@pytest.mark.asyncio
+async def test_last_n_invocations_more_than_existing_invocations():
+  """Tests that no filtering occurs if last_n_invocations is greater than
+
+  the number of invocations.
+  """
+  plugin = ContextFilterPlugin(num_invocations_to_keep=3)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+  original_contents = list(llm_request.contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  assert llm_request.contents == original_contents
+
+
+@pytest.mark.asyncio
+async def test_filter_function_raises_exception():
+  """Tests that the plugin handles exceptions from the filter function."""
+
+  def faulty_filter(contents):
+    raise ValueError("Filter error")
+
+  plugin = ContextFilterPlugin(custom_filter=faulty_filter)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+  original_contents = list(llm_request.contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  assert llm_request.contents == original_contents
+
+
+def _create_function_call_content(name: str, call_id: str) -> types.Content:
+  """Creates a model content with a function call."""
+  return types.Content(
+      parts=[
+          types.Part(
+              function_call=types.FunctionCall(id=call_id, name=name, args={})
+          )
+      ],
+      role="model",
+  )
+
+
+def _create_function_response_content(name: str, call_id: str) -> types.Content:
+  """Creates a user content with a function response."""
+  return types.Content(
+      parts=[
+          types.Part(
+              function_response=types.FunctionResponse(
+                  id=call_id, name=name, response={"result": "ok"}
+              )
+          )
+      ],
+      role="user",
+  )
+
+
+@pytest.mark.asyncio
+async def test_filter_preserves_function_call_response_pairs():
+  """Tests that function_call and function_response pairs are kept together.
+
+  This tests the fix for issue #4027 where filtering could create orphaned
+  function_response messages without their corresponding function_call.
+  """
+  plugin = ContextFilterPlugin(num_invocations_to_keep=2)
+
+  # Simulate conversation from issue #4027:
+  # user -> model -> user -> model(function_call) -> user(function_response)
+  # -> model -> user -> model(function_call) -> user(function_response)
+  contents = [
+      _create_content("user", "Hello"),
+      _create_content("model", "Hi there!"),
+      _create_content("user", "I want to know about X"),
+      _create_function_call_content("knowledge_base", "call_1"),
+      _create_function_response_content("knowledge_base", "call_1"),
+      _create_content("model", "I found some information..."),
+      _create_content("user", "can you explain more about Y"),
+      _create_function_call_content("knowledge_base", "call_2"),
+      _create_function_response_content("knowledge_base", "call_2"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # Verify function_call for call_1 is included (not orphaned function_response)
+  call_ids_present = set()
+  response_ids_present = set()
+  for content in llm_request.contents:
+    if content.parts:
+      for part in content.parts:
+        if part.function_call and part.function_call.id:
+          call_ids_present.add(part.function_call.id)
+        if part.function_response and part.function_response.id:
+          response_ids_present.add(part.function_response.id)
+
+  # Every function_response should have a matching function_call
+  assert response_ids_present.issubset(call_ids_present), (
+      "Orphaned function_responses found. "
+      f"Responses: {response_ids_present}, Calls: {call_ids_present}"
+  )
+
+
+@pytest.mark.asyncio
+async def test_filter_with_nested_function_calls():
+  """Tests filtering with multiple nested function call sequences."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=1)
+
+  contents = [
+      _create_content("user", "Hello"),
+      _create_content("model", "Hi!"),
+      _create_content("user", "Do task"),
+      _create_function_call_content("tool_a", "call_a"),
+      _create_function_response_content("tool_a", "call_a"),
+      _create_function_call_content("tool_b", "call_b"),
+      _create_function_response_content("tool_b", "call_b"),
+      _create_content("model", "Done with tasks"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # Verify no orphaned function_responses
+  call_ids = set()
+  response_ids = set()
+  for content in llm_request.contents:
+    if content.parts:
+      for part in content.parts:
+        if part.function_call and part.function_call.id:
+          call_ids.add(part.function_call.id)
+        if part.function_response and part.function_response.id:
+          response_ids.add(part.function_response.id)
+
+  texts = []
+  for content in llm_request.contents:
+    if content.parts:
+      for part in content.parts:
+        if part.text:
+          texts.append(part.text)
+
+  assert "Do task" in texts
+  assert "Done with tasks" in texts
+  assert "Hello" not in texts
+  assert "Hi!" not in texts
+
+  assert response_ids.issubset(call_ids)
+
+
+@pytest.mark.asyncio
+async def test_last_invocation_with_tool_call_keeps_user_prompt():
+  """Tests that multi-model-turn invocations keep the initial user prompt."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=1)
+
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_function_call_content("get_weather", "call_1"),
+      _create_function_response_content("get_weather", "call_1"),
+      _create_content("model", "final_answer_2"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  texts = []
+  for content in llm_request.contents:
+    if content.parts:
+      for part in content.parts:
+        if part.text:
+          texts.append(part.text)
+
+  assert "user_prompt_2" in texts
+  assert "final_answer_2" in texts
+
+
+@pytest.mark.asyncio
+async def test_filter_with_remove_amount():
+  """Tests that remove_amount correctly removes additional invocations."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=2, remove_amount=1)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # With num_invocations_to_keep=2 and remove_amount=1, keeps last 2.
+  assert len(llm_request.contents) == 4
+  assert llm_request.contents[0].parts[0].text == "user_prompt_2"
+  assert llm_request.contents[1].parts[0].text == "model_response_2"
+  assert llm_request.contents[2].parts[0].text == "user_prompt_3"
+  assert llm_request.contents[3].parts[0].text == "model_response_3"
+
+
+@pytest.mark.asyncio
+async def test_filter_with_higher_remove_amount():
+  """Tests remove_amount with a higher value to remove more invocations."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=3, remove_amount=2)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+      _create_content("user", "user_prompt_4"),
+      _create_content("model", "model_response_4"),
+      _create_content("user", "user_prompt_5"),
+      _create_content("model", "model_response_5"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # With num_invocations_to_keep=3 and remove_amount=2, keeps last 3.
+  assert len(llm_request.contents) == 6
+  assert llm_request.contents[0].parts[0].text == "user_prompt_3"
+  assert llm_request.contents[1].parts[0].text == "model_response_3"
+  assert llm_request.contents[2].parts[0].text == "user_prompt_4"
+  assert llm_request.contents[3].parts[0].text == "model_response_4"
+  assert llm_request.contents[4].parts[0].text == "user_prompt_5"
+  assert llm_request.contents[5].parts[0].text == "model_response_5"
+
+
+def test_invalid_remove_amount():
+  """Tests that initializing with remove_amount < 1 raises ValueError."""
+  with pytest.raises(ValueError, match="remove_amount must be at least 1"):
+    ContextFilterPlugin(num_invocations_to_keep=1, remove_amount=0)
+
+  with pytest.raises(ValueError, match="remove_amount must be at least 1"):
+    ContextFilterPlugin(num_invocations_to_keep=1, remove_amount=-1)
+
+
+@pytest.mark.asyncio
+async def test_filter_remove_amount_with_multiple_user_turns():
+  """Tests remove_amount with multiple user turns in invocations."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=2, remove_amount=1)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2a"),
+      _create_content("user", "user_prompt_2b"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # Should keep last 2 invocations including multiple user turns
+  assert len(llm_request.contents) == 5
+  assert llm_request.contents[0].parts[0].text == "user_prompt_2a"
+  assert llm_request.contents[1].parts[0].text == "user_prompt_2b"
+  assert llm_request.contents[2].parts[0].text == "model_response_2"
+  assert llm_request.contents[3].parts[0].text == "user_prompt_3"
+  assert llm_request.contents[4].parts[0].text == "model_response_3"
+
+
+@pytest.mark.asyncio
+async def test_filter_bypass_when_under_remove_threshold():
+  """Tests that filtering is bypassed when total invocations are between keep limit and keep+remove limit."""
+  plugin = ContextFilterPlugin(num_invocations_to_keep=2, remove_amount=2)
+  contents = [
+      _create_content("user", "user_prompt_1"),
+      _create_content("model", "model_response_1"),
+      _create_content("user", "user_prompt_2"),
+      _create_content("model", "model_response_2"),
+      _create_content("user", "user_prompt_3"),
+      _create_content("model", "model_response_3"),
+  ]
+  llm_request = LlmRequest(contents=contents)
+  original_contents = list(llm_request.contents)
+
+  await plugin.before_model_callback(
+      callback_context=mock.create_autospec(CallbackContext, instance=True),
+      llm_request=llm_request,
+  )
+
+  # With num_invocations_to_keep=2 and remove_amount=2, threshold is 4.
+  # We have 3 invocations, so no filtering should occur.
+  assert llm_request.contents == original_contents

tests/unittests/plugins/test_debug_logging_plugin.py

@@ -0,0 +1,605 @@
+# Copyright 2026 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 __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import Mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events.event import Event
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.debug_logging_plugin import DebugLoggingPlugin
+from google.adk.sessions.session import Session
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+import yaml
+
+
+@pytest.fixture
+def debug_output_file(tmp_path):
+  """Fixture to provide a temporary file path for debug output."""
+  return tmp_path / "debug_output.yaml"
+
+
+@pytest.fixture
+def mock_session():
+  """Create a mock session."""
+  session = Mock(spec=Session)
+  session.id = "test-session-id"
+  session.app_name = "test-app"
+  session.user_id = "test-user"
+  session.state = {"key1": "value1", "key2": 123}
+  session.events = []
+  return session
+
+
+@pytest.fixture
+def mock_invocation_context(mock_session):
+  """Create a mock invocation context."""
+  ctx = Mock(spec=InvocationContext)
+  ctx.invocation_id = "test-invocation-id"
+  ctx.session = mock_session
+  ctx.user_id = "test-user"
+  ctx.app_name = "test-app"
+  ctx.branch = None
+  ctx.agent = Mock()
+  ctx.agent.name = "test-agent"
+  return ctx
+
+
+@pytest.fixture
+def mock_callback_context(mock_invocation_context):
+  """Create a mock callback context."""
+  ctx = Mock(spec=CallbackContext)
+  ctx.invocation_id = mock_invocation_context.invocation_id
+  ctx.agent_name = "test-agent"
+  ctx._invocation_context = mock_invocation_context
+  ctx.state = {}
+  return ctx
+
+
+@pytest.fixture
+def mock_tool_context(mock_invocation_context):
+  """Create a mock tool context."""
+  ctx = Mock(spec=ToolContext)
+  ctx.invocation_id = mock_invocation_context.invocation_id
+  ctx.agent_name = "test-agent"
+  ctx.function_call_id = "test-function-call-id"
+  return ctx
+
+
+class TestDebugLoggingPluginInitialization:
+  """Tests for DebugLoggingPlugin initialization."""
+
+  def test_default_initialization(self):
+    """Test plugin initialization with default values."""
+    plugin = DebugLoggingPlugin()
+    assert plugin.name == "debug_logging_plugin"
+    assert plugin._output_path == Path("adk_debug.yaml")
+    assert plugin._include_session_state is True
+    assert plugin._include_system_instruction is True
+
+  def test_custom_initialization(self, debug_output_file):
+    """Test plugin initialization with custom values."""
+    plugin = DebugLoggingPlugin(
+        name="custom_debug",
+        output_path=str(debug_output_file),
+        include_session_state=False,
+        include_system_instruction=False,
+    )
+    assert plugin.name == "custom_debug"
+    assert plugin._output_path == debug_output_file
+    assert plugin._include_session_state is False
+    assert plugin._include_system_instruction is False
+
+
+class TestDebugLoggingPluginCallbacks:
+  """Tests for DebugLoggingPlugin callback methods."""
+
+  async def test_before_run_callback_initializes_state(
+      self, debug_output_file, mock_invocation_context
+  ):
+    """Test that before_run_callback initializes debug state."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    result = await plugin.before_run_callback(
+        invocation_context=mock_invocation_context
+    )
+
+    assert result is None
+    assert mock_invocation_context.invocation_id in plugin._invocation_states
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    assert state.invocation_id == mock_invocation_context.invocation_id
+    assert state.session_id == mock_invocation_context.session.id
+    assert len(state.entries) == 1
+    assert state.entries[0].entry_type == "invocation_start"
+
+  async def test_on_user_message_callback_logs_message(
+      self, debug_output_file, mock_invocation_context
+  ):
+    """Test that on_user_message_callback logs user messages."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    user_message = types.Content(
+        role="user", parts=[types.Part.from_text(text="Hello, world!")]
+    )
+
+    result = await plugin.on_user_message_callback(
+        invocation_context=mock_invocation_context, user_message=user_message
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    user_message_entries = [
+        e for e in state.entries if e.entry_type == "user_message"
+    ]
+    assert len(user_message_entries) == 1
+    assert user_message_entries[0].data["content"]["role"] == "user"
+    assert user_message_entries[0].data["content"]["parts"][0]["text"] == (
+        "Hello, world!"
+    )
+
+  async def test_before_model_callback_logs_request(
+      self, debug_output_file, mock_invocation_context, mock_callback_context
+  ):
+    """Test that before_model_callback logs LLM requests."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user", parts=[types.Part.from_text(text="Test prompt")]
+            )
+        ],
+    )
+    llm_request.config.system_instruction = "You are a helpful assistant."
+
+    result = await plugin.before_model_callback(
+        callback_context=mock_callback_context, llm_request=llm_request
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    llm_entries = [e for e in state.entries if e.entry_type == "llm_request"]
+    assert len(llm_entries) == 1
+    assert llm_entries[0].data["model"] == "gemini-2.5-flash"
+    assert llm_entries[0].data["content_count"] == 1
+    assert "config" in llm_entries[0].data
+    assert (
+        llm_entries[0].data["config"]["system_instruction"]
+        == "You are a helpful assistant."
+    )
+
+  async def test_after_model_callback_logs_response(
+      self, debug_output_file, mock_invocation_context, mock_callback_context
+  ):
+    """Test that after_model_callback logs LLM responses."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    llm_response = LlmResponse(
+        content=types.Content(
+            role="model",
+            parts=[types.Part.from_text(text="Hello! How can I help?")],
+        ),
+        turn_complete=True,
+    )
+
+    result = await plugin.after_model_callback(
+        callback_context=mock_callback_context, llm_response=llm_response
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    llm_entries = [e for e in state.entries if e.entry_type == "llm_response"]
+    assert len(llm_entries) == 1
+    assert llm_entries[0].data["turn_complete"] is True
+    assert llm_entries[0].data["content"]["role"] == "model"
+
+  async def test_before_tool_callback_logs_tool_call(
+      self, debug_output_file, mock_invocation_context, mock_tool_context
+  ):
+    """Test that before_tool_callback logs tool calls."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    mock_tool = Mock(spec=BaseTool)
+    mock_tool.name = "test_tool"
+    tool_args = {"param1": "value1", "param2": 42}
+
+    result = await plugin.before_tool_callback(
+        tool=mock_tool, tool_args=tool_args, tool_context=mock_tool_context
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    tool_entries = [e for e in state.entries if e.entry_type == "tool_call"]
+    assert len(tool_entries) == 1
+    assert tool_entries[0].data["tool_name"] == "test_tool"
+    assert tool_entries[0].data["args"]["param1"] == "value1"
+    assert tool_entries[0].data["args"]["param2"] == 42
+
+  async def test_after_tool_callback_logs_tool_response(
+      self, debug_output_file, mock_invocation_context, mock_tool_context
+  ):
+    """Test that after_tool_callback logs tool responses."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    mock_tool = Mock(spec=BaseTool)
+    mock_tool.name = "test_tool"
+    tool_args = {"param1": "value1"}
+    result_data = {"output": "success", "data": [1, 2, 3]}
+
+    result = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=tool_args,
+        tool_context=mock_tool_context,
+        result=result_data,
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    tool_entries = [e for e in state.entries if e.entry_type == "tool_response"]
+    assert len(tool_entries) == 1
+    assert tool_entries[0].data["tool_name"] == "test_tool"
+    assert tool_entries[0].data["result"]["output"] == "success"
+
+  async def test_on_event_callback_logs_event(
+      self, debug_output_file, mock_invocation_context
+  ):
+    """Test that on_event_callback logs events."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    event = Event(
+        author="test-agent",
+        content=types.Content(
+            role="model",
+            parts=[types.Part.from_text(text="Response text")],
+        ),
+    )
+
+    result = await plugin.on_event_callback(
+        invocation_context=mock_invocation_context, event=event
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    event_entries = [e for e in state.entries if e.entry_type == "event"]
+    assert len(event_entries) == 1
+    assert event_entries[0].data["author"] == "test-agent"
+    assert event_entries[0].data["event_id"] == event.id
+
+  async def test_on_model_error_callback_logs_error(
+      self, debug_output_file, mock_invocation_context, mock_callback_context
+  ):
+    """Test that on_model_error_callback logs LLM errors."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    llm_request = LlmRequest(model="gemini-2.5-flash")
+    error = ValueError("Test error message")
+
+    result = await plugin.on_model_error_callback(
+        callback_context=mock_callback_context,
+        llm_request=llm_request,
+        error=error,
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    error_entries = [e for e in state.entries if e.entry_type == "llm_error"]
+    assert len(error_entries) == 1
+    assert error_entries[0].data["error_type"] == "ValueError"
+    assert error_entries[0].data["error_message"] == "Test error message"
+
+  async def test_on_tool_error_callback_logs_error(
+      self, debug_output_file, mock_invocation_context, mock_tool_context
+  ):
+    """Test that on_tool_error_callback logs tool errors."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state first
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    mock_tool = Mock(spec=BaseTool)
+    mock_tool.name = "test_tool"
+    tool_args = {"param1": "value1"}
+    error = RuntimeError("Tool execution failed")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    assert result is None
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    error_entries = [e for e in state.entries if e.entry_type == "tool_error"]
+    assert len(error_entries) == 1
+    assert error_entries[0].data["tool_name"] == "test_tool"
+    assert error_entries[0].data["error_type"] == "RuntimeError"
+
+
+class TestDebugLoggingPluginFileOutput:
+  """Tests for DebugLoggingPlugin file output."""
+
+  async def test_after_run_callback_writes_to_file(
+      self, debug_output_file, mock_invocation_context
+  ):
+    """Test that after_run_callback writes debug data to file."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # Initialize state
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    # Add some entries
+    user_message = types.Content(
+        role="user", parts=[types.Part.from_text(text="Test message")]
+    )
+    await plugin.on_user_message_callback(
+        invocation_context=mock_invocation_context, user_message=user_message
+    )
+
+    # Finalize
+    await plugin.after_run_callback(invocation_context=mock_invocation_context)
+
+    # Verify file was written
+    assert debug_output_file.exists()
+
+    # Parse and verify content (YAML format with --- separator)
+    with open(debug_output_file, "r") as f:
+      documents = list(yaml.safe_load_all(f))
+
+    assert len(documents) == 1
+    data = documents[0]
+    assert data["invocation_id"] == "test-invocation-id"
+    assert data["session_id"] == "test-session-id"
+    assert (
+        len(data["entries"]) >= 2
+    )  # At least invocation_start and user_message
+
+  async def test_after_run_callback_includes_session_state(
+      self, debug_output_file, mock_invocation_context
+  ):
+    """Test that session state is included when enabled."""
+    plugin = DebugLoggingPlugin(
+        output_path=str(debug_output_file), include_session_state=True
+    )
+
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+    await plugin.after_run_callback(invocation_context=mock_invocation_context)
+
+    with open(debug_output_file, "r") as f:
+      documents = list(yaml.safe_load_all(f))
+
+    data = documents[0]
+    session_state_entries = [
+        e
+        for e in data["entries"]
+        if e["entry_type"] == "session_state_snapshot"
+    ]
+    assert len(session_state_entries) == 1
+    assert session_state_entries[0]["data"]["state"]["key1"] == "value1"
+
+  async def test_after_run_callback_excludes_session_state_when_disabled(
+      self, debug_output_file, mock_invocation_context
+  ):
+    """Test that session state is excluded when disabled."""
+    plugin = DebugLoggingPlugin(
+        output_path=str(debug_output_file), include_session_state=False
+    )
+
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+    await plugin.after_run_callback(invocation_context=mock_invocation_context)
+
+    with open(debug_output_file, "r") as f:
+      documents = list(yaml.safe_load_all(f))
+
+    data = documents[0]
+    session_state_entries = [
+        e
+        for e in data["entries"]
+        if e["entry_type"] == "session_state_snapshot"
+    ]
+    assert not session_state_entries
+
+  async def test_multiple_invocations_append_to_file(
+      self, debug_output_file, mock_session
+  ):
+    """Test that multiple invocations append to the same file."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    # First invocation
+    ctx1 = Mock(spec=InvocationContext)
+    ctx1.invocation_id = "invocation-1"
+    ctx1.session = mock_session
+    ctx1.user_id = "test-user"
+    ctx1.branch = None
+    ctx1.agent = Mock()
+    ctx1.agent.name = "agent-1"
+
+    await plugin.before_run_callback(invocation_context=ctx1)
+    await plugin.after_run_callback(invocation_context=ctx1)
+
+    # Second invocation
+    ctx2 = Mock(spec=InvocationContext)
+    ctx2.invocation_id = "invocation-2"
+    ctx2.session = mock_session
+    ctx2.user_id = "test-user"
+    ctx2.branch = None
+    ctx2.agent = Mock()
+    ctx2.agent.name = "agent-2"
+
+    await plugin.before_run_callback(invocation_context=ctx2)
+    await plugin.after_run_callback(invocation_context=ctx2)
+
+    # Verify both invocations are in the file (as separate YAML documents)
+    with open(debug_output_file, "r") as f:
+      documents = list(yaml.safe_load_all(f))
+
+    assert len(documents) == 2
+    assert documents[0]["invocation_id"] == "invocation-1"
+    assert documents[1]["invocation_id"] == "invocation-2"
+
+  async def test_after_run_callback_cleans_up_state(
+      self, debug_output_file, mock_invocation_context
+  ):
+    """Test that invocation state is cleaned up after writing."""
+    plugin = DebugLoggingPlugin(output_path=str(debug_output_file))
+
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+    assert mock_invocation_context.invocation_id in plugin._invocation_states
+
+    await plugin.after_run_callback(invocation_context=mock_invocation_context)
+    assert (
+        mock_invocation_context.invocation_id not in plugin._invocation_states
+    )
+
+
+class TestDebugLoggingPluginSerialization:
+  """Tests for content serialization."""
+
+  def test_serialize_content_with_text(self):
+    """Test serialization of text content."""
+    plugin = DebugLoggingPlugin()
+    content = types.Content(
+        role="user", parts=[types.Part.from_text(text="Hello")]
+    )
+
+    result = plugin._serialize_content(content)
+
+    assert result["role"] == "user"
+    assert len(result["parts"]) == 1
+    assert result["parts"][0]["text"] == "Hello"
+
+  def test_serialize_content_with_function_call(self):
+    """Test serialization of function call content."""
+    plugin = DebugLoggingPlugin()
+    content = types.Content(
+        role="model",
+        parts=[
+            types.Part(
+                function_call=types.FunctionCall(
+                    id="fc-1", name="test_func", args={"arg1": "val1"}
+                )
+            )
+        ],
+    )
+
+    result = plugin._serialize_content(content)
+
+    assert result["parts"][0]["function_call"]["name"] == "test_func"
+    assert result["parts"][0]["function_call"]["args"]["arg1"] == "val1"
+
+  def test_serialize_content_with_none(self):
+    """Test serialization of None content."""
+    plugin = DebugLoggingPlugin()
+    result = plugin._serialize_content(None)
+    assert result is None
+
+  def test_safe_serialize_handles_bytes(self):
+    """Test that bytes are safely serialized."""
+    plugin = DebugLoggingPlugin()
+    result = plugin._safe_serialize(b"binary data")
+    assert result == "<bytes: 11 bytes>"
+
+  def test_safe_serialize_handles_nested_structures(self):
+    """Test that nested structures are serialized."""
+    plugin = DebugLoggingPlugin()
+    data = {
+        "list": [1, 2, {"nested": "value"}],
+        "tuple": (3, 4),
+        "string": "text",
+    }
+
+    result = plugin._safe_serialize(data)
+
+    assert result["list"] == [1, 2, {"nested": "value"}]
+    assert result["tuple"] == [3, 4]  # Tuple becomes list
+    assert result["string"] == "text"
+
+
+class TestDebugLoggingPluginSystemInstructionConfig:
+  """Tests for system instruction configuration."""
+
+  async def test_system_instruction_included_when_enabled(
+      self, debug_output_file, mock_invocation_context, mock_callback_context
+  ):
+    """Test that full system instruction is included when enabled."""
+    plugin = DebugLoggingPlugin(
+        output_path=str(debug_output_file), include_system_instruction=True
+    )
+
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    llm_request = LlmRequest(model="gemini-2.5-flash")
+    llm_request.config.system_instruction = "Full system instruction text"
+
+    await plugin.before_model_callback(
+        callback_context=mock_callback_context, llm_request=llm_request
+    )
+
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    llm_entries = [e for e in state.entries if e.entry_type == "llm_request"]
+    assert (
+        llm_entries[0].data["config"]["system_instruction"]
+        == "Full system instruction text"
+    )
+
+  async def test_system_instruction_length_only_when_disabled(
+      self, debug_output_file, mock_invocation_context, mock_callback_context
+  ):
+    """Test that only length is included when system instruction is disabled."""
+    plugin = DebugLoggingPlugin(
+        output_path=str(debug_output_file), include_system_instruction=False
+    )
+
+    await plugin.before_run_callback(invocation_context=mock_invocation_context)
+
+    llm_request = LlmRequest(model="gemini-2.5-flash")
+    llm_request.config.system_instruction = "Full system instruction text"
+
+    await plugin.before_model_callback(
+        callback_context=mock_callback_context, llm_request=llm_request
+    )
+
+    state = plugin._invocation_states[mock_invocation_context.invocation_id]
+    llm_entries = [e for e in state.entries if e.entry_type == "llm_request"]
+    assert "system_instruction" not in llm_entries[0].data.get("config", {})
+    assert llm_entries[0].data["config"]["system_instruction_length"] == 28

tests/unittests/plugins/test_global_instruction_plugin.py

@@ -0,0 +1,208 @@
+# Copyright 2026 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 unittest.mock import Mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.models.llm_request import LlmRequest
+from google.adk.plugins.global_instruction_plugin import GlobalInstructionPlugin
+from google.adk.sessions.session import Session
+from google.genai import types
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_global_instruction_plugin_with_string():
+  """Test GlobalInstructionPlugin with a string global instruction."""
+  plugin = GlobalInstructionPlugin(
+      global_instruction=(
+          "You are a helpful assistant with a friendly personality."
+      )
+  )
+
+  # Create mock objects
+  mock_session = Session(
+      app_name="test_app", user_id="test_user", id="test_session", state={}
+  )
+
+  mock_invocation_context = Mock(spec=InvocationContext)
+  mock_invocation_context.session = mock_session
+
+  mock_callback_context = Mock(spec=CallbackContext)
+  mock_callback_context._invocation_context = mock_invocation_context
+
+  llm_request = LlmRequest(
+      model="gemini-2.5-flash",
+      config=types.GenerateContentConfig(system_instruction=""),
+  )
+
+  # Execute the plugin's before_model_callback
+  result = await plugin.before_model_callback(
+      callback_context=mock_callback_context, llm_request=llm_request
+  )
+
+  # Plugin should return None to allow normal processing
+  assert result is None
+
+  # System instruction should now contain the global instruction
+  assert (
+      "You are a helpful assistant with a friendly personality."
+      in llm_request.config.system_instruction
+  )
+
+
+@pytest.mark.asyncio
+async def test_global_instruction_plugin_with_instruction_provider():
+  """Test GlobalInstructionPlugin with an InstructionProvider function."""
+
+  async def build_global_instruction(readonly_context: ReadonlyContext) -> str:
+    return f"You are assistant for user {readonly_context.session.user_id}."
+
+  plugin = GlobalInstructionPlugin(global_instruction=build_global_instruction)
+
+  # Create mock objects
+  mock_session = Session(
+      app_name="test_app", user_id="alice", id="test_session", state={}
+  )
+
+  mock_invocation_context = Mock(spec=InvocationContext)
+
+  mock_callback_context = Mock(spec=CallbackContext)
+  mock_callback_context._invocation_context = mock_invocation_context
+  mock_callback_context.session = mock_session
+
+  llm_request = LlmRequest(
+      model="gemini-2.5-flash",
+      config=types.GenerateContentConfig(system_instruction=""),
+  )
+
+  # Execute the plugin's before_model_callback
+  result = await plugin.before_model_callback(
+      callback_context=mock_callback_context, llm_request=llm_request
+  )
+
+  # Plugin should return None to allow normal processing
+  assert result is None
+
+  # System instruction should contain the dynamically generated instruction
+  assert (
+      "You are assistant for user alice."
+      in llm_request.config.system_instruction
+  )
+
+
+@pytest.mark.asyncio
+async def test_global_instruction_plugin_empty_instruction():
+  """Test GlobalInstructionPlugin with empty global instruction."""
+  plugin = GlobalInstructionPlugin(global_instruction="")
+
+  # Create mock objects
+  mock_session = Session(
+      app_name="test_app", user_id="test_user", id="test_session", state={}
+  )
+
+  mock_invocation_context = Mock(spec=InvocationContext)
+  mock_invocation_context.session = mock_session
+
+  mock_callback_context = Mock(spec=CallbackContext)
+  mock_callback_context._invocation_context = mock_invocation_context
+
+  llm_request = LlmRequest(
+      model="gemini-2.5-flash",
+      config=types.GenerateContentConfig(
+          system_instruction="Original instruction"
+      ),
+  )
+
+  # Execute the plugin's before_model_callback
+  result = await plugin.before_model_callback(
+      callback_context=mock_callback_context, llm_request=llm_request
+  )
+
+  # Plugin should return None to allow normal processing
+  assert result is None
+
+  # System instruction should remain unchanged
+  assert llm_request.config.system_instruction == "Original instruction"
+
+
+@pytest.mark.asyncio
+async def test_global_instruction_plugin_leads_existing():
+  """Test that GlobalInstructionPlugin prepends global instructions."""
+  plugin = GlobalInstructionPlugin(
+      global_instruction="You are a helpful assistant."
+  )
+
+  # Create mock objects
+  mock_session = Session(
+      app_name="test_app", user_id="test_user", id="test_session", state={}
+  )
+
+  mock_invocation_context = Mock(spec=InvocationContext)
+  mock_invocation_context.session = mock_session
+
+  mock_callback_context = Mock(spec=CallbackContext)
+  mock_callback_context._invocation_context = mock_invocation_context
+
+  llm_request = LlmRequest(
+      model="gemini-2.5-flash",
+      config=types.GenerateContentConfig(
+          system_instruction="Existing instructions."
+      ),
+  )
+
+  # Execute the plugin's before_model_callback
+  result = await plugin.before_model_callback(
+      callback_context=mock_callback_context, llm_request=llm_request
+  )
+
+  # Plugin should return None to allow normal processing
+  assert result is None
+
+  # System instruction should contain global instruction before existing ones
+  expected = "You are a helpful assistant.\n\nExisting instructions."
+  assert llm_request.config.system_instruction == expected
+
+
+@pytest.mark.asyncio
+async def test_global_instruction_plugin_prepends_to_list():
+  """Test GlobalInstructionPlugin prepends to a list of instructions."""
+  plugin = GlobalInstructionPlugin(global_instruction="Global instruction.")
+
+  mock_session = Session(
+      app_name="test_app", user_id="test_user", id="test_session", state={}
+  )
+
+  mock_invocation_context = Mock(spec=InvocationContext)
+  mock_invocation_context.session = mock_session
+
+  mock_callback_context = Mock(spec=CallbackContext)
+  mock_callback_context._invocation_context = mock_invocation_context
+
+  llm_request = LlmRequest(
+      model="gemini-2.5-flash",
+      config=types.GenerateContentConfig(
+          system_instruction=["Existing instruction."]
+      ),
+  )
+
+  await plugin.before_model_callback(
+      callback_context=mock_callback_context, llm_request=llm_request
+  )
+
+  expected = ["Global instruction.", "Existing instruction."]
+  assert llm_request.config.system_instruction == expected

tests/unittests/plugins/test_multimodal_tool_results_plugin.py

@@ -0,0 +1,154 @@
+# Copyright 2026 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 __future__ import annotations
+
+from typing import Any
+from unittest.mock import Mock
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.models.llm_request import LlmRequest
+from google.adk.plugins.multimodal_tool_results_plugin import MultimodalToolResultsPlugin
+from google.adk.plugins.multimodal_tool_results_plugin import PARTS_RETURNED_BY_TOOLS_ID
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from .. import testing_utils
+
+
+@pytest.fixture
+def plugin() -> MultimodalToolResultsPlugin:
+  """Create a default plugin instance for testing."""
+  return MultimodalToolResultsPlugin()
+
+
+@pytest.fixture
+def mock_tool() -> MockTool:
+  """Create a mock tool for testing."""
+  return Mock(spec=BaseTool)
+
+
+@pytest.fixture
+async def tool_context() -> ToolContext:
+  """Create a mock tool context."""
+  return ToolContext(
+      invocation_context=await testing_utils.create_invocation_context(
+          agent=Mock(spec=BaseAgent)
+      )
+  )
+
+
+@pytest.mark.asyncio
+async def test_tool_returning_parts_are_added_to_llm_request(
+    plugin: MultimodalToolResultsPlugin,
+    mock_tool: MockTool,
+    tool_context: ToolContext,
+):
+  """Test that parts returned by a tool are present in the llm_request later."""
+  parts = [types.Part(text="part1"), types.Part(text="part2")]
+
+  result = await plugin.after_tool_callback(
+      tool=mock_tool,
+      tool_args={},
+      tool_context=tool_context,
+      result=parts,
+  )
+
+  assert result == None
+  assert PARTS_RETURNED_BY_TOOLS_ID in tool_context.state
+  assert tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] == parts
+
+  callback_context = Mock(spec=CallbackContext)
+  callback_context.state = tool_context.state
+  llm_request = LlmRequest(contents=[types.Content(parts=[])])
+
+  await plugin.before_model_callback(
+      callback_context=callback_context, llm_request=llm_request
+  )
+
+  assert llm_request.contents[-1].parts == parts
+
+
+@pytest.mark.asyncio
+async def test_tool_returning_non_list_of_parts_is_unchanged(
+    plugin: MultimodalToolResultsPlugin,
+    mock_tool: MockTool,
+    tool_context: ToolContext,
+):
+  """Test where tool returning non list of parts, has this result unchanged."""
+  original_result = {"some": "data"}
+
+  result = await plugin.after_tool_callback(
+      tool=mock_tool,
+      tool_args={},
+      tool_context=tool_context,
+      result=original_result,
+  )
+
+  assert result == original_result
+  assert PARTS_RETURNED_BY_TOOLS_ID not in tool_context.state
+
+  callback_context = Mock(spec=CallbackContext)
+  callback_context.state = tool_context.state
+  llm_request = LlmRequest(
+      contents=[types.Content(parts=[types.Part(text="original")])]
+  )
+  original_parts = list(llm_request.contents[-1].parts)
+
+  await plugin.before_model_callback(
+      callback_context=callback_context, llm_request=llm_request
+  )
+
+  assert llm_request.contents[-1].parts == original_parts
+
+
+@pytest.mark.asyncio
+async def test_multiple_tools_returning_parts_are_accumulated(
+    plugin: ToolReturningGenAiPartsPlugin,
+    mock_tool: MockTool,
+    tool_context: ToolContext,
+):
+  """Test that parts from multiple tool calls are accumulated."""
+  parts1 = [types.Part(text="part1")]
+  parts2 = [types.Part(text="part2")]
+
+  await plugin.after_tool_callback(
+      tool=mock_tool,
+      tool_args={},
+      tool_context=tool_context,
+      result=parts1,
+  )
+
+  await plugin.after_tool_callback(
+      tool=mock_tool,
+      tool_args={},
+      tool_context=tool_context,
+      result=parts2,
+  )
+
+  assert PARTS_RETURNED_BY_TOOLS_ID in tool_context.state
+  assert tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] == parts1 + parts2
+
+  callback_context = Mock(spec=CallbackContext)
+  callback_context.state = tool_context.state
+  llm_request = LlmRequest(contents=[types.Content(parts=[])])
+
+  await plugin.before_model_callback(
+      callback_context=callback_context, llm_request=llm_request
+  )
+
+  assert llm_request.contents[-1].parts == parts1 + parts2

tests/unittests/plugins/test_notification_error_callbacks.py

@@ -0,0 +1,936 @@
+# Copyright 2026 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.
+
+"""Tests for on_agent_error_callback and on_run_error_callback.
+
+Validates RFC #5044: agent-level and runner-level error callbacks.
+"""
+
+import asyncio
+from typing import AsyncGenerator
+from typing import Optional
+from unittest.mock import Mock
+
+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.event 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.adk.workflow._base_node import BaseNode
+from google.genai import types
+import pytest
+from typing_extensions import override
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class _CrashingNode(BaseNode):
+  """A workflow node whose _run_impl always raises.
+
+  A root ``BaseNode`` (that is not a ``BaseAgent``) is executed through the
+  node runtime (``Runner._run_node_async``), so this exercises the
+  ``on_run_error_callback`` dispatch site added there.
+  """
+
+  __test__ = False
+
+  @override
+  async def _run_impl(self, *, ctx, node_input):
+    raise RuntimeError("node crashed")
+    yield  # pragma: no cover - makes this an async generator
+
+
+class _CrashingAgent(BaseAgent):
+  """Agent whose _run_async_impl always raises."""
+
+  crash_error: Exception = RuntimeError("agent crashed")
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    raise self.crash_error
+    yield  # make it an async generator
+
+  @override
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    raise self.crash_error
+    yield
+
+
+class _SuccessAgent(BaseAgent):
+  """Agent that completes successfully."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        branch=ctx.branch,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(parts=[types.Part(text="ok")]),
+    )
+
+  @override
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(
+        author=self.name,
+        branch=ctx.branch,
+        invocation_id=ctx.invocation_id,
+        content=types.Content(parts=[types.Part(text="ok live")]),
+    )
+
+
+class _ErrorTrackingPlugin(BasePlugin):
+  """Plugin that records which error callbacks were called."""
+
+  __test__ = False
+
+  def __init__(self, name: str = "error_tracker"):
+    super().__init__(name)
+    self.agent_errors: list[tuple[str, Exception]] = []
+    self.run_errors: list[Exception] = []
+    self.after_agent_called = False
+    self.after_run_called = False
+
+  async def on_agent_error_callback(
+      self,
+      *,
+      agent: BaseAgent,
+      callback_context: CallbackContext,
+      error: Exception,
+  ) -> None:
+    self.agent_errors.append((agent.name, error))
+
+  async def on_run_error_callback(
+      self,
+      *,
+      invocation_context: InvocationContext,
+      error: Exception,
+  ) -> None:
+    self.run_errors.append(error)
+
+  async def after_agent_callback(
+      self,
+      *,
+      agent: BaseAgent,
+      callback_context: CallbackContext,
+  ) -> Optional[types.Content]:
+    self.after_agent_called = True
+    return None
+
+  async def after_run_callback(
+      self,
+      *,
+      invocation_context: InvocationContext,
+  ) -> None:
+    self.after_run_called = True
+
+
+async def _create_ctx(
+    agent: BaseAgent,
+    plugins: list[BasePlugin] | None = None,
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name="test_app", user_id="test_user"
+  )
+  return InvocationContext(
+      invocation_id="test_invocation",
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      plugin_manager=PluginManager(plugins=plugins or []),
+  )
+
+
+# ---------------------------------------------------------------------------
+# Agent-level error callback tests
+# ---------------------------------------------------------------------------
+
+
+class TestAgentErrorCallback:
+  """Tests for on_agent_error_callback in base_agent.py."""
+
+  @pytest.mark.asyncio
+  async def test_agent_error_callback_fires_on_crash(self):
+    """Error callback fires when _run_async_impl raises."""
+    plugin = _ErrorTrackingPlugin()
+    agent = _CrashingAgent(name="crash_agent")
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    with pytest.raises(RuntimeError, match="agent crashed"):
+      _ = [e async for e in agent.run_async(ctx)]
+
+    assert len(plugin.agent_errors) == 1
+    assert plugin.agent_errors[0][0] == "crash_agent"
+    assert str(plugin.agent_errors[0][1]) == "agent crashed"
+
+  @pytest.mark.asyncio
+  async def test_agent_error_callback_fires_on_live_crash(self):
+    """Error callback fires when _run_live_impl raises."""
+    plugin = _ErrorTrackingPlugin()
+    agent = _CrashingAgent(name="crash_agent")
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    with pytest.raises(RuntimeError, match="agent crashed"):
+      _ = [e async for e in agent.run_live(ctx)]
+
+    assert len(plugin.agent_errors) == 1
+    assert plugin.agent_errors[0][0] == "crash_agent"
+
+  @pytest.mark.asyncio
+  async def test_after_agent_not_called_on_crash(self):
+    """after_agent_callback (success-only) is NOT called on failure."""
+    plugin = _ErrorTrackingPlugin()
+    agent = _CrashingAgent(name="crash_agent")
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    with pytest.raises(RuntimeError):
+      _ = [e async for e in agent.run_async(ctx)]
+
+    assert not plugin.after_agent_called
+
+  @pytest.mark.asyncio
+  async def test_agent_error_callback_fires_on_before_callback_failure(self):
+    """Error callback fires when before_agent_callback raises.
+
+    The error handler wraps the full agent lifecycle, so lifecycle-callback
+    failures (not just _run_async_impl) are surfaced to on_agent_error_callback.
+    """
+    plugin = _ErrorTrackingPlugin()
+
+    def _boom(callback_context):
+      raise RuntimeError("before boom")
+
+    agent = _SuccessAgent(name="good_agent", before_agent_callback=_boom)
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    with pytest.raises(RuntimeError, match="before boom"):
+      _ = [e async for e in agent.run_async(ctx)]
+
+    assert len(plugin.agent_errors) == 1
+    assert plugin.agent_errors[0][0] == "good_agent"
+    assert not plugin.after_agent_called
+
+  @pytest.mark.asyncio
+  async def test_agent_error_callback_fires_on_after_callback_failure(self):
+    """Error callback fires when after_agent_callback raises."""
+    plugin = _ErrorTrackingPlugin()
+
+    def _boom(callback_context):
+      raise RuntimeError("after boom")
+
+    agent = _SuccessAgent(name="good_agent", after_agent_callback=_boom)
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    with pytest.raises(RuntimeError, match="after boom"):
+      _ = [e async for e in agent.run_async(ctx)]
+
+    assert len(plugin.agent_errors) == 1
+    assert plugin.agent_errors[0][0] == "good_agent"
+
+  @pytest.mark.asyncio
+  async def test_exception_is_reraised_after_agent_error_callback(self):
+    """The original exception propagates after the error callback."""
+    plugin = _ErrorTrackingPlugin()
+    err = ValueError("specific error")
+    agent = _CrashingAgent(name="crash_agent", crash_error=err)
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    with pytest.raises(ValueError, match="specific error"):
+      _ = [e async for e in agent.run_async(ctx)]
+
+  @pytest.mark.asyncio
+  async def test_agent_error_callback_not_fired_on_success(self):
+    """Error callback does NOT fire when agent succeeds."""
+    plugin = _ErrorTrackingPlugin()
+    agent = _SuccessAgent(name="good_agent")
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    events = [e async for e in agent.run_async(ctx)]
+
+    assert len(events) > 0
+    assert len(plugin.agent_errors) == 0
+    # after_agent_callback should still fire on success
+    assert plugin.after_agent_called
+
+  @pytest.mark.asyncio
+  async def test_cancelled_error_does_not_trigger_agent_error_callback(
+      self,
+  ):
+    """asyncio.CancelledError (BaseException) does NOT trigger error callback."""
+
+    class _CancellingAgent(BaseAgent):
+
+      @override
+      async def _run_async_impl(self, ctx):
+        raise asyncio.CancelledError()
+        yield
+
+      @override
+      async def _run_live_impl(self, ctx):
+        raise asyncio.CancelledError()
+        yield
+
+    plugin = _ErrorTrackingPlugin()
+    agent = _CancellingAgent(name="cancel_agent")
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    with pytest.raises(asyncio.CancelledError):
+      _ = [e async for e in agent.run_async(ctx)]
+
+    assert len(plugin.agent_errors) == 0
+
+
+# ---------------------------------------------------------------------------
+# Runner-level error callback tests
+# ---------------------------------------------------------------------------
+
+
+class TestRunErrorCallback:
+  """Tests for on_run_error_callback in runners.py."""
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_fires_on_crash(self):
+    """on_run_error_callback fires when execute_fn raises."""
+    from google.adk.runners import Runner
+
+    plugin = _ErrorTrackingPlugin()
+    agent = _CrashingAgent(name="crash_agent")
+    runner = Runner(
+        agent=agent,
+        app_name="test_app",
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError, match="agent crashed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    assert len(plugin.run_errors) == 1
+    assert str(plugin.run_errors[0]) == "agent crashed"
+
+  @pytest.mark.asyncio
+  async def test_after_run_not_called_on_crash(self):
+    """after_run_callback (success-only) is NOT called on failure."""
+    from google.adk.runners import Runner
+
+    plugin = _ErrorTrackingPlugin()
+    agent = _CrashingAgent(name="crash_agent")
+    runner = Runner(
+        agent=agent,
+        app_name="test_app",
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    assert not plugin.after_run_called
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_not_fired_on_success(self):
+    """on_run_error_callback does NOT fire on success."""
+    from google.adk.runners import Runner
+
+    plugin = _ErrorTrackingPlugin()
+    agent = _SuccessAgent(name="good_agent")
+    runner = Runner(
+        agent=agent,
+        app_name="test_app",
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    events = [
+        e
+        async for e in runner.run_async(
+            user_id="test_user",
+            session_id=session.id,
+            new_message=types.Content(parts=[types.Part(text="hello")]),
+        )
+    ]
+
+    assert len(events) > 0
+    assert len(plugin.run_errors) == 0
+    assert plugin.after_run_called
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_fires_on_after_run_failure(self):
+    """An after_run_callback failure notifies on_run_error and re-raises.
+
+    after_run runs on the success path, after the main-execution catch. A
+    failing after_run plugin (which PluginManager surfaces as a RuntimeError)
+    is an unhandled runner error, so it must still notify on_run_error_callback
+    exactly once while the original error propagates.
+    """
+    from google.adk.runners import Runner
+
+    class _FailingAfterRunPlugin(BasePlugin):
+
+      async def after_run_callback(
+          self, *, invocation_context: InvocationContext
+      ) -> None:
+        raise RuntimeError("after_run failed")
+
+    tracker = _ErrorTrackingPlugin()
+    agent = _SuccessAgent(name="good_agent")
+    runner = Runner(
+        agent=agent,
+        app_name="test_app",
+        session_service=InMemorySessionService(),
+        plugins=[_FailingAfterRunPlugin(name="failing_after_run"), tracker],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError, match="after_run failed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    # Exactly one run-error notification for the after_run failure.
+    assert len(tracker.run_errors) == 1
+    # PluginManager wraps plugin exceptions with plugin/callback context.
+    assert "after_run failed" in str(tracker.run_errors[0])
+
+
+# ---------------------------------------------------------------------------
+# Exactly-once-per-layer tests
+# ---------------------------------------------------------------------------
+
+
+class TestExactlyOncePerLayer:
+  """Verify each error callback fires exactly once at its own layer."""
+
+  @pytest.mark.asyncio
+  async def test_agent_crash_fires_both_callbacks_once_each(self):
+    """A crashing agent fires on_agent_error_callback once AND
+    on_run_error_callback once (the re-raised exception propagates)."""
+    from google.adk.runners import Runner
+
+    plugin = _ErrorTrackingPlugin()
+    agent = _CrashingAgent(name="crash_agent")
+    runner = Runner(
+        agent=agent,
+        app_name="test_app",
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError, match="agent crashed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    # Agent error callback: exactly 1 call
+    assert len(plugin.agent_errors) == 1
+    assert plugin.agent_errors[0][0] == "crash_agent"
+
+    # Run error callback: exactly 1 call (same exception bubbled up)
+    assert len(plugin.run_errors) == 1
+    assert plugin.run_errors[0] is plugin.agent_errors[0][1]
+
+    # Neither after callback should fire
+    assert not plugin.after_agent_called
+    assert not plugin.after_run_called
+
+
+# ---------------------------------------------------------------------------
+# PluginManager dispatch tests
+# ---------------------------------------------------------------------------
+
+
+class TestPluginManagerErrorCallbackDispatch:
+  """Test PluginManager correctly dispatches the new error callbacks."""
+
+  @pytest.mark.asyncio
+  async def test_run_on_agent_error_callback_dispatches(self):
+    """run_on_agent_error_callback calls all plugins."""
+    plugin1 = _ErrorTrackingPlugin(name="p1")
+    plugin2 = _ErrorTrackingPlugin(name="p2")
+    pm = PluginManager(plugins=[plugin1, plugin2])
+
+    mock_agent = Mock(spec=BaseAgent)
+    mock_agent.name = "test_agent"
+    mock_ctx = Mock(spec=CallbackContext)
+    err = RuntimeError("boom")
+
+    await pm.run_on_agent_error_callback(
+        agent=mock_agent,
+        callback_context=mock_ctx,
+        error=err,
+    )
+
+    assert len(plugin1.agent_errors) == 1
+    assert len(plugin2.agent_errors) == 1
+
+  @pytest.mark.asyncio
+  async def test_run_on_run_error_callback_dispatches(self):
+    """run_on_run_error_callback calls all plugins."""
+    plugin1 = _ErrorTrackingPlugin(name="p1")
+    plugin2 = _ErrorTrackingPlugin(name="p2")
+    pm = PluginManager(plugins=[plugin1, plugin2])
+
+    mock_ctx = Mock(spec=InvocationContext)
+    err = RuntimeError("boom")
+
+    await pm.run_on_run_error_callback(
+        invocation_context=mock_ctx,
+        error=err,
+    )
+
+    assert len(plugin1.run_errors) == 1
+    assert len(plugin2.run_errors) == 1
+
+  @pytest.mark.asyncio
+  async def test_agent_error_callback_does_not_short_circuit(self):
+    """on_agent_error_callback is notification-only: a non-None return
+    from one plugin does NOT skip subsequent plugins."""
+
+    class _ReturningPlugin(BasePlugin):
+      __test__ = False
+
+      def __init__(self, name):
+        super().__init__(name)
+        self.agent_error_called = False
+
+      async def on_agent_error_callback(self, **kwargs):
+        self.agent_error_called = True
+        return "should be ignored"
+
+    p1 = _ReturningPlugin(name="p1")
+    p2 = _ReturningPlugin(name="p2")
+    pm = PluginManager(plugins=[p1, p2])
+
+    await pm.run_on_agent_error_callback(
+        agent=Mock(spec=BaseAgent),
+        callback_context=Mock(spec=CallbackContext),
+        error=RuntimeError("x"),
+    )
+
+    # Both plugins must be called even though p1 returns non-None.
+    assert p1.agent_error_called
+    assert p2.agent_error_called
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_does_not_short_circuit(self):
+    """on_run_error_callback is notification-only: a non-None return
+    from one plugin does NOT skip subsequent plugins."""
+
+    class _ReturningPlugin(BasePlugin):
+      __test__ = False
+
+      def __init__(self, name):
+        super().__init__(name)
+        self.run_error_called = False
+
+      async def on_run_error_callback(self, **kwargs):
+        self.run_error_called = True
+        return "should be ignored"
+
+    p1 = _ReturningPlugin(name="p1")
+    p2 = _ReturningPlugin(name="p2")
+    pm = PluginManager(plugins=[p1, p2])
+
+    await pm.run_on_run_error_callback(
+        invocation_context=Mock(spec=InvocationContext),
+        error=RuntimeError("x"),
+    )
+
+    # Both plugins must be called even though p1 returns non-None.
+    assert p1.run_error_called
+    assert p2.run_error_called
+
+  @pytest.mark.asyncio
+  async def test_plugin_callback_failure_does_not_mask_app_error(self):
+    """When a plugin's error callback raises, iteration continues
+    and the original application exception is what the caller sees."""
+
+    class _FailingPlugin(BasePlugin):
+      __test__ = False
+
+      def __init__(self, name):
+        super().__init__(name)
+        self.agent_error_called = False
+        self.run_error_called = False
+
+      async def on_agent_error_callback(self, **kwargs):
+        self.agent_error_called = True
+        raise ValueError("plugin boom")
+
+      async def on_run_error_callback(self, **kwargs):
+        self.run_error_called = True
+        raise ValueError("plugin boom")
+
+    p1 = _FailingPlugin(name="p1")
+    p2 = _ErrorTrackingPlugin(name="p2")
+    pm = PluginManager(plugins=[p1, p2])
+
+    # Agent error callback: p1 raises, p2 must still be notified.
+    mock_agent = Mock(spec=BaseAgent)
+    mock_agent.name = "test_agent"
+    await pm.run_on_agent_error_callback(
+        agent=mock_agent,
+        callback_context=Mock(spec=CallbackContext),
+        error=RuntimeError("app crash"),
+    )
+    assert p1.agent_error_called
+    assert len(p2.agent_errors) == 1
+
+    # Run error callback: same behavior.
+    await pm.run_on_run_error_callback(
+        invocation_context=Mock(spec=InvocationContext),
+        error=RuntimeError("app crash"),
+    )
+    assert p1.run_error_called
+    assert len(p2.run_errors) == 1
+
+  @pytest.mark.asyncio
+  async def test_original_exception_propagates_despite_agent_plugin_failure(
+      self,
+  ):
+    """End-to-end: a crashing plugin error callback does not mask
+    the original agent exception seen by the caller."""
+
+    class _FailingPlugin(BasePlugin):
+      __test__ = False
+
+      def __init__(self, name):
+        super().__init__(name)
+
+      async def on_agent_error_callback(self, **kwargs):
+        raise ValueError("plugin internal error")
+
+    plugin = _FailingPlugin(name="bad_plugin")
+    agent = _CrashingAgent(name="crash_agent")
+    ctx = await _create_ctx(agent, plugins=[plugin])
+
+    # The caller must see the original RuntimeError("agent crashed"),
+    # NOT the plugin's ValueError.
+    with pytest.raises(RuntimeError, match="agent crashed"):
+      _ = [e async for e in agent.run_async(ctx)]
+
+  @pytest.mark.asyncio
+  async def test_original_exception_propagates_despite_run_plugin_failure(
+      self,
+  ):
+    """End-to-end: a crashing plugin on_run_error_callback does not mask
+    the original agent exception seen by the runner caller."""
+    from google.adk.runners import Runner
+
+    class _FailingRunPlugin(BasePlugin):
+      __test__ = False
+
+      def __init__(self, name):
+        super().__init__(name)
+
+      async def on_run_error_callback(self, **kwargs):
+        raise ValueError("plugin internal error")
+
+    plugin = _FailingRunPlugin(name="bad_plugin")
+    agent = _CrashingAgent(name="crash_agent")
+    runner = Runner(
+        agent=agent,
+        app_name="test_app",
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    # The caller must see the original RuntimeError("agent crashed"),
+    # NOT the plugin's ValueError.
+    with pytest.raises(RuntimeError, match="agent crashed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+
+# ---------------------------------------------------------------------------
+# Node-runtime run-error coverage
+# ---------------------------------------------------------------------------
+
+
+class TestNodeRuntimeRunErrorCallback:
+  """on_run_error_callback fires for the node runtime path (_run_node_async).
+
+  Runner(node=...) with a non-agent BaseNode root routes through
+  _run_node_async rather than the legacy _exec_with_plugin path.
+  """
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_fires_via_node_runtime(self):
+    from google.adk.runners import Runner
+
+    plugin = _ErrorTrackingPlugin()
+    node = _CrashingNode(name="crash_node")
+    runner = Runner(
+        app_name="test_app",
+        node=node,
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError, match="node crashed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    # Exactly one run-error notification from the node runtime path.
+    assert len(plugin.run_errors) == 1
+    assert str(plugin.run_errors[0]) == "node crashed"
+    # after_run stays success-only.
+    assert not plugin.after_run_called
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_not_fired_on_node_success(self):
+    """No run-error notification for a successful node-runtime run."""
+    from google.adk.runners import Runner
+
+    plugin = _ErrorTrackingPlugin()
+
+    class _OkNode(BaseNode):
+      __test__ = False
+
+      @override
+      async def _run_impl(self, *, ctx, node_input):
+        yield Event(
+            author=self.name,
+            invocation_id=ctx.get_invocation_context().invocation_id,
+            content=types.Content(parts=[types.Part(text="ok")]),
+        )
+
+    runner = Runner(
+        app_name="test_app",
+        node=_OkNode(name="ok_node"),
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    _ = [
+        e
+        async for e in runner.run_async(
+            user_id="test_user",
+            session_id=session.id,
+            new_message=types.Content(parts=[types.Part(text="hello")]),
+        )
+    ]
+
+    assert len(plugin.run_errors) == 0
+    assert plugin.after_run_called
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_fires_on_node_before_run_failure(self):
+    """A before_run_callback failure on the node path notifies on_run_error.
+
+    The setup hooks in _run_node_async run before the main event loop; their
+    failures must still be surfaced to on_run_error_callback.
+    """
+    from google.adk.runners import Runner
+
+    class _FailingBeforeRunPlugin(BasePlugin):
+
+      async def before_run_callback(
+          self, *, invocation_context: InvocationContext
+      ) -> Optional[types.Content]:
+        raise RuntimeError("before_run failed")
+
+    tracker = _ErrorTrackingPlugin()
+    runner = Runner(
+        app_name="test_app",
+        node=_CrashingNode(name="never_runs"),
+        session_service=InMemorySessionService(),
+        plugins=[_FailingBeforeRunPlugin(name="failing_before_run"), tracker],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError, match="before_run failed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    assert len(tracker.run_errors) == 1
+    # PluginManager wraps plugin exceptions with plugin/callback context.
+    assert "before_run failed" in str(tracker.run_errors[0])
+    # after_run stays success-only.
+    assert not tracker.after_run_called
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_fires_on_node_user_message_failure(self):
+    """An on_user_message_callback failure on the node path notifies on_run_error."""
+    from google.adk.runners import Runner
+
+    class _FailingUserMessagePlugin(BasePlugin):
+
+      async def on_user_message_callback(
+          self,
+          *,
+          invocation_context: InvocationContext,
+          user_message: types.Content,
+      ) -> Optional[types.Content]:
+        raise RuntimeError("user_message failed")
+
+    tracker = _ErrorTrackingPlugin()
+    runner = Runner(
+        app_name="test_app",
+        node=_CrashingNode(name="never_runs"),
+        session_service=InMemorySessionService(),
+        plugins=[_FailingUserMessagePlugin(name="failing_user_msg"), tracker],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError, match="user_message failed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    assert len(tracker.run_errors) == 1
+    # PluginManager wraps plugin exceptions with plugin/callback context.
+    assert "user_message failed" in str(tracker.run_errors[0])
+    assert not tracker.after_run_called
+
+  @pytest.mark.asyncio
+  async def test_run_error_callback_fires_on_node_after_run_failure(self):
+    """An after_run_callback failure on the node path notifies on_run_error.
+
+    after_run runs on the node-runtime success path, outside the main-loop
+    catch. A failing after_run plugin (which PluginManager surfaces as a
+    RuntimeError) is an unhandled runner error, so it must still notify
+    on_run_error_callback exactly once while the original error propagates.
+    """
+    from google.adk.runners import Runner
+
+    class _OkNode(BaseNode):
+      __test__ = False
+
+      @override
+      async def _run_impl(self, *, ctx, node_input):
+        yield Event(
+            author=self.name,
+            invocation_id=ctx.get_invocation_context().invocation_id,
+            content=types.Content(parts=[types.Part(text="ok")]),
+        )
+
+    class _FailingAfterRunPlugin(BasePlugin):
+
+      async def after_run_callback(
+          self, *, invocation_context: InvocationContext
+      ) -> None:
+        raise RuntimeError("after_run failed")
+
+    tracker = _ErrorTrackingPlugin()
+    runner = Runner(
+        app_name="test_app",
+        node=_OkNode(name="ok_node"),
+        session_service=InMemorySessionService(),
+        plugins=[_FailingAfterRunPlugin(name="failing_after_run"), tracker],
+    )
+    session = await runner.session_service.create_session(
+        app_name="test_app", user_id="test_user"
+    )
+
+    with pytest.raises(RuntimeError, match="after_run failed"):
+      _ = [
+          e
+          async for e in runner.run_async(
+              user_id="test_user",
+              session_id=session.id,
+              new_message=types.Content(parts=[types.Part(text="hello")]),
+          )
+      ]
+
+    # Exactly one run-error notification for the after_run failure.
+    assert len(tracker.run_errors) == 1
+    # PluginManager wraps plugin exceptions with plugin/callback context.
+    assert "after_run failed" in str(tracker.run_errors[0])

tests/unittests/plugins/test_plugin_manager.py

@@ -0,0 +1,382 @@
+# Copyright 2026 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.
+
+"""Unit tests for the PluginManager."""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+# Assume the following path to your modules
+# You might need to adjust this based on your project structure.
+from google.adk.plugins.plugin_manager import PluginCallbackName
+from google.adk.plugins.plugin_manager import PluginManager
+import pytest
+
+
+# A helper class to use in tests instead of mocks.
+# This makes tests more explicit and easier to debug.
+class TestPlugin(BasePlugin):
+  __test__ = False
+  """
+  A test plugin that can be configured to return specific values or raise
+  exceptions for any callback, and it logs which callbacks were invoked.
+  """
+
+  def __init__(self, name: str):
+    super().__init__(name)
+    # A log to track the names of callbacks that have been called.
+    self.call_log: list[PluginCallbackName] = []
+    # A map to configure return values for specific callbacks.
+    self.return_values: dict[PluginCallbackName, any] = {}
+    # A map to configure exceptions to be raised by specific callbacks.
+    self.exceptions_to_raise: dict[PluginCallbackName, Exception] = {}
+
+  async def _handle_callback(self, name: PluginCallbackName):
+    """Generic handler for all callback methods."""
+    self.call_log.append(name)
+    if name in self.exceptions_to_raise:
+      raise self.exceptions_to_raise[name]
+    return self.return_values.get(name)
+
+  # Implement all callback methods from the BasePlugin interface.
+  async def on_user_message_callback(self, **kwargs):
+    return await self._handle_callback("on_user_message_callback")
+
+  async def before_run_callback(self, **kwargs):
+    return await self._handle_callback("before_run_callback")
+
+  async def after_run_callback(self, **kwargs):
+    return await self._handle_callback("after_run_callback")
+
+  async def on_event_callback(self, **kwargs):
+    return await self._handle_callback("on_event_callback")
+
+  async def before_agent_callback(self, **kwargs):
+    return await self._handle_callback("before_agent_callback")
+
+  async def after_agent_callback(self, **kwargs):
+    return await self._handle_callback("after_agent_callback")
+
+  async def before_tool_callback(self, **kwargs):
+    return await self._handle_callback("before_tool_callback")
+
+  async def after_tool_callback(self, **kwargs):
+    return await self._handle_callback("after_tool_callback")
+
+  async def on_tool_error_callback(self, **kwargs):
+    return await self._handle_callback("on_tool_error_callback")
+
+  async def before_model_callback(self, **kwargs):
+    return await self._handle_callback("before_model_callback")
+
+  async def after_model_callback(self, **kwargs):
+    return await self._handle_callback("after_model_callback")
+
+  async def on_model_error_callback(self, **kwargs):
+    return await self._handle_callback("on_model_error_callback")
+
+  async def on_agent_error_callback(self, **kwargs):
+    return await self._handle_callback("on_agent_error_callback")
+
+  async def on_run_error_callback(self, **kwargs):
+    return await self._handle_callback("on_run_error_callback")
+
+
+@pytest.fixture
+def service() -> PluginManager:
+  """Provides a clean PluginManager instance for each test."""
+  return PluginManager()
+
+
+@pytest.fixture
+def plugin1() -> TestPlugin:
+  """Provides a clean instance of our test plugin named 'plugin1'."""
+  return TestPlugin(name="plugin1")
+
+
+@pytest.fixture
+def plugin2() -> TestPlugin:
+  """Provides a clean instance of our test plugin named 'plugin2'."""
+  return TestPlugin(name="plugin2")
+
+
+def test_register_and_get_plugin(service: PluginManager, plugin1: TestPlugin):
+  """Tests successful registration and retrieval of a plugin."""
+  service.register_plugin(plugin1)
+
+  assert len(service.plugins) == 1
+  assert service.plugins[0] is plugin1
+  assert service.get_plugin("plugin1") is plugin1
+
+
+def test_register_duplicate_plugin_name_raises_value_error(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that registering a plugin with a duplicate name raises an error."""
+  plugin1_duplicate = TestPlugin(name="plugin1")
+  service.register_plugin(plugin1)
+
+  with pytest.raises(
+      ValueError, match="Plugin with name 'plugin1' already registered."
+  ):
+    service.register_plugin(plugin1_duplicate)
+
+
+@pytest.mark.asyncio
+async def test_early_exit_stops_subsequent_plugins(
+    service: PluginManager, plugin1: TestPlugin, plugin2: TestPlugin
+):
+  """Tests the core "early exit" logic: if a plugin returns a value,
+
+  subsequent plugins for that callback should not be executed.
+  """
+  # Configure plugin1 to return a value, simulating a cache hit.
+  mock_response = Mock(spec=LlmResponse)
+  plugin1.return_values["before_run_callback"] = mock_response
+
+  service.register_plugin(plugin1)
+  service.register_plugin(plugin2)
+
+  # Execute the callback chain.
+  result = await service.run_before_run_callback(invocation_context=Mock())
+
+  # Assert that the final result is the one returned by the first plugin.
+  assert result is mock_response
+  # Assert that the first plugin was called.
+  assert "before_run_callback" in plugin1.call_log
+  # CRITICAL: Assert that the second plugin was never called.
+  assert "before_run_callback" not in plugin2.call_log
+
+
+@pytest.mark.asyncio
+async def test_normal_flow_all_plugins_are_called(
+    service: PluginManager, plugin1: TestPlugin, plugin2: TestPlugin
+):
+  """Tests that if no plugin returns a value, all plugins in the chain
+
+  are executed in order.
+  """
+  # By default, plugins are configured to return None.
+  service.register_plugin(plugin1)
+  service.register_plugin(plugin2)
+
+  result = await service.run_before_run_callback(invocation_context=Mock())
+
+  # The final result should be None as no plugin interrupted the flow.
+  assert result is None
+  # Both plugins must have been called.
+  assert "before_run_callback" in plugin1.call_log
+  assert "before_run_callback" in plugin2.call_log
+
+
+@pytest.mark.asyncio
+async def test_plugin_exception_is_wrapped_in_runtime_error(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that if a plugin callback raises an exception, the PluginManager
+
+  catches it and raises a descriptive RuntimeError.
+  """
+  # Configure the plugin to raise an error during a specific callback.
+  original_exception = ValueError("Something went wrong inside the plugin!")
+  plugin1.exceptions_to_raise["before_run_callback"] = original_exception
+  service.register_plugin(plugin1)
+
+  with pytest.raises(RuntimeError) as excinfo:
+    await service.run_before_run_callback(invocation_context=Mock())
+
+  # Check that the error message is informative.
+  assert "Error in plugin 'plugin1'" in str(excinfo.value)
+  assert "before_run_callback" in str(excinfo.value)
+  # Check that the original exception is chained for better tracebacks.
+  assert excinfo.value.__cause__ is original_exception
+
+
+@pytest.mark.asyncio
+async def test_all_callbacks_are_supported(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that all callbacks defined in the BasePlugin interface are supported
+
+  by the PluginManager.
+  """
+  service.register_plugin(plugin1)
+  mock_context = Mock()
+  mock_user_message = Mock()
+
+  # Test all callbacks
+  await service.run_on_user_message_callback(
+      user_message=mock_user_message, invocation_context=mock_context
+  )
+  await service.run_before_run_callback(invocation_context=mock_context)
+  await service.run_after_run_callback(invocation_context=mock_context)
+  await service.run_on_event_callback(
+      invocation_context=mock_context, event=mock_context
+  )
+  await service.run_before_agent_callback(
+      agent=mock_context, callback_context=mock_context
+  )
+  await service.run_after_agent_callback(
+      agent=mock_context, callback_context=mock_context
+  )
+  await service.run_before_tool_callback(
+      tool=mock_context, tool_args={}, tool_context=mock_context
+  )
+  await service.run_after_tool_callback(
+      tool=mock_context, tool_args={}, tool_context=mock_context, result={}
+  )
+  await service.run_on_tool_error_callback(
+      tool=mock_context,
+      tool_args={},
+      tool_context=mock_context,
+      error=mock_context,
+  )
+  await service.run_before_model_callback(
+      callback_context=mock_context, llm_request=mock_context
+  )
+  await service.run_after_model_callback(
+      callback_context=mock_context, llm_response=mock_context
+  )
+  await service.run_on_model_error_callback(
+      callback_context=mock_context,
+      llm_request=mock_context,
+      error=mock_context,
+  )
+  await service.run_on_agent_error_callback(
+      agent=mock_context,
+      callback_context=mock_context,
+      error=mock_context,
+  )
+  await service.run_on_run_error_callback(
+      invocation_context=mock_context,
+      error=mock_context,
+  )
+
+  # Verify all callbacks were logged
+  expected_callbacks = [
+      "on_user_message_callback",
+      "before_run_callback",
+      "after_run_callback",
+      "on_event_callback",
+      "before_agent_callback",
+      "after_agent_callback",
+      "before_tool_callback",
+      "after_tool_callback",
+      "on_tool_error_callback",
+      "before_model_callback",
+      "after_model_callback",
+      "on_model_error_callback",
+      "on_agent_error_callback",
+      "on_run_error_callback",
+  ]
+  assert set(plugin1.call_log) == set(expected_callbacks)
+
+
+@pytest.mark.asyncio
+async def test_close_calls_plugin_close(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that close calls the close method on registered plugins."""
+  plugin1.close = AsyncMock()
+  service.register_plugin(plugin1)
+
+  await service.close()
+
+  plugin1.close.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_close_raises_runtime_error_on_plugin_exception(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that close raises a RuntimeError if a plugin's close fails."""
+  plugin1.close = AsyncMock(side_effect=ValueError("Shutdown error"))
+  service.register_plugin(plugin1)
+
+  with pytest.raises(
+      RuntimeError, match="Failed to close plugins: 'plugin1': ValueError"
+  ):
+    await service.close()
+
+  plugin1.close.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_close_with_timeout(plugin1: TestPlugin):
+  """Tests that close respects the timeout and raises on failure."""
+  service = PluginManager(close_timeout=0.1)
+
+  async def slow_close():
+    await asyncio.sleep(0.2)
+
+  plugin1.close = slow_close
+  service.register_plugin(plugin1)
+
+  with pytest.raises(RuntimeError) as excinfo:
+    await service.close()
+
+  assert "Failed to close plugins: 'plugin1': TimeoutError" in str(
+      excinfo.value
+  )
+
+
+@pytest.mark.asyncio
+async def test_close_is_noop_after_set_skip_closing_plugins(
+    plugin1: TestPlugin,
+):
+  """Tests that close is a no-op after set_skip_closing_plugins(True)."""
+  plugin1.close = AsyncMock()
+  service = PluginManager(plugins=[plugin1])
+  service.set_skip_closing_plugins(True)
+
+  await service.close()
+
+  plugin1.close.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_set_skip_closing_plugins_bypasses_close_timeout(
+    plugin1: TestPlugin,
+):
+  """Tests that set_skip_closing_plugins(True) bypasses close timeout."""
+
+  async def slow_close():
+    await asyncio.sleep(10)  # Would otherwise time out.
+
+  plugin1.close = slow_close
+  service = PluginManager(plugins=[plugin1], close_timeout=0.01)
+  service.set_skip_closing_plugins(True)
+
+  # Should return immediately without raising.
+  await service.close()
+
+
+@pytest.mark.asyncio
+async def test_set_skip_closing_plugins_false_reverts_to_closing(
+    plugin1: TestPlugin,
+):
+  """Tests that set_skip_closing_plugins(False) re-enables closing."""
+  plugin1.close = AsyncMock()
+  service = PluginManager(plugins=[plugin1])
+  service.set_skip_closing_plugins(True)
+  service.set_skip_closing_plugins(False)
+
+  await service.close()
+
+  plugin1.close.assert_awaited_once()

tests/unittests/plugins/test_reflect_retry_tool_plugin.py

@@ -0,0 +1,668 @@
+# Copyright 2026 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 unittest import IsolatedAsyncioTestCase
+from unittest.mock import Mock
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.plugins.reflect_retry_tool_plugin import REFLECT_AND_RETRY_RESPONSE_TYPE
+from google.adk.plugins.reflect_retry_tool_plugin import ReflectAndRetryToolPlugin
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+
+from .. import testing_utils
+
+
+class MockTool(BaseTool):
+  """Mock tool for testing purposes."""
+
+  def __init__(self, name: str = "mock_tool"):
+    self.name = name
+    self.description = f"Mock tool named {name}"
+
+  async def run(self, **kwargs) -> Any:
+    return "mock result"
+
+
+class CustomErrorExtractionPlugin(ReflectAndRetryToolPlugin):
+  """Custom plugin for testing error extraction from tool responses."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.error_conditions = {}
+
+  def set_error_condition(self, condition_func):
+    """Set a custom error condition function for testing."""
+    self.error_condition = condition_func
+
+  async def extract_error_from_result(
+      self, *, tool, tool_args, tool_context, result
+  ):
+    """Extract error based on custom conditions set for testing."""
+    if hasattr(self, "error_condition"):
+      return self.error_condition(result)
+    return None
+
+
+# Inheriting from IsolatedAsyncioTestCase ensures consistent behavior.
+# See https://github.com/pytest-dev/pytest-asyncio/issues/1039
+class TestReflectAndRetryToolPlugin(IsolatedAsyncioTestCase):
+  """Comprehensive tests for ReflectAndRetryToolPlugin focusing on behavior."""
+
+  def get_plugin(self):
+    """Create a default plugin instance for testing."""
+    return ReflectAndRetryToolPlugin()
+
+  def get_custom_plugin(self):
+    """Create a plugin with custom parameters."""
+    return ReflectAndRetryToolPlugin(
+        name="custom_plugin",
+        max_retries=5,
+        throw_exception_if_retry_exceeded=False,
+    )
+
+  def get_mock_tool(self):
+    """Create a mock tool for testing."""
+    return MockTool("test_tool_id")
+
+  def get_mock_tool_context(self):
+    """Create a mock tool context."""
+    return Mock(spec=ToolContext)
+
+  def get_custom_error_plugin(self):
+    """Create a custom error extraction plugin for testing."""
+    return CustomErrorExtractionPlugin(max_retries=3)
+
+  def get_sample_tool_args(self):
+    """Sample tool arguments for testing."""
+    return {"param1": "value1", "param2": 42, "param3": True}
+
+  async def test_plugin_initialization_default(self):
+    """Test plugin initialization with default parameters."""
+    plugin = self.get_plugin()
+
+    self.assertEqual(plugin.name, "reflect_retry_tool_plugin")
+    self.assertEqual(plugin.max_retries, 3)
+    self.assertIs(plugin.throw_exception_if_retry_exceeded, True)
+
+  async def test_plugin_initialization_custom(self):
+    """Test plugin initialization with custom parameters."""
+    plugin = ReflectAndRetryToolPlugin(
+        name="custom_name",
+        max_retries=10,
+        throw_exception_if_retry_exceeded=False,
+    )
+
+    self.assertEqual(plugin.name, "custom_name")
+    self.assertEqual(plugin.max_retries, 10)
+    self.assertIsNot(plugin.throw_exception_if_retry_exceeded, True)
+
+  async def test_after_tool_callback_successful_call(self):
+    """Test after_tool_callback with successful tool call."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    result = {"success": True, "data": "test_data"}
+
+    callback_result = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=result,
+    )
+
+    # Should return None for successful calls
+    self.assertIsNone(callback_result)
+
+  async def test_after_tool_callback_ignore_retry_response(self):
+    """Test that retry responses are ignored in after_tool_callback."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    retry_result = {"response_type": REFLECT_AND_RETRY_RESPONSE_TYPE}
+
+    callback_result = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=retry_result,
+    )
+
+    # Retry responses should be ignored
+    self.assertIsNone(callback_result)
+
+  async def test_on_tool_error_callback_max_retries_zero(self):
+    """Test error callback when max_retries is 0.
+
+    This should return None so that the exception is rethrown
+    """
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(max_retries=0)
+    error = ValueError("Test error")
+
+    with self.assertRaises(ValueError) as cm:
+      await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+
+    # Should re-raise the original exception when max_retries is 0
+    self.assertIs(cm.exception, error)
+
+  async def test_on_tool_error_callback_max_retries_zero_without_exception(
+      self,
+  ):
+    """Test error callback when max_retries is 0 and exception is disabled."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=False
+    )
+    error = ValueError("Test error")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Should return a retry exceeded message instead of raising
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "ValueError")
+    self.assertEqual(result["retry_count"], 0)
+    self.assertIn(
+        "the retry limit has been exceeded", result["reflection_guidance"]
+    )
+
+  async def test_on_tool_error_callback_max_retries_zero_with_dict_error(self):
+    """Test error callback when max_retries is 0 and error is a dict."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Should raise an Exception wrapping the dict
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
+  async def test_on_tool_error_callback_first_failure(self):
+    """Test first tool failure creates reflection response."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    error = ValueError("Test error message")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "ValueError")
+    self.assertEqual(result["error_details"], "Test error message")
+    self.assertEqual(result["retry_count"], 1)
+    self.assertIn("test_tool_id", result["reflection_guidance"])
+    self.assertIn("Test error message", result["reflection_guidance"])
+
+  async def test_retry_behavior_with_consecutive_failures(self):
+    """Test the retry behavior with consecutive failures."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    error = RuntimeError("Runtime error")
+
+    # First failure
+    result1 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Second failure - should have different retry count based on plugin logic
+    result2 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    # The plugin's internal logic determines the exact retry count
+    self.assertIsNotNone(result2)
+    self.assertEqual(result2["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result2["retry_count"], 2)
+
+  async def test_different_tools_behavior(self):
+    """Test behavior when using different tools."""
+    plugin = self.get_plugin()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    tool1 = MockTool("tool1")
+    tool2 = MockTool("tool2")
+    error = ValueError("Test error")
+
+    # First failure on tool1
+    result1 = await plugin.on_tool_error_callback(
+        tool=tool1,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Failure on tool2
+    result2 = await plugin.on_tool_error_callback(
+        tool=tool2,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    # Since tool is different, retry count should start over.
+    self.assertIsNotNone(result2)
+    self.assertEqual(result2["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result2["retry_count"], 1)
+
+  async def test_max_retries_exceeded_with_exception(self):
+    """Test that original exception is raised when max retries exceeded."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=True
+    )
+    error = ConnectionError("Connection failed")
+
+    # First call should succeed and return a retry response
+    await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Second call should exceed max_retries and raise
+    with self.assertRaises(ConnectionError) as cm:
+      await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+
+    # Verify exception properties
+    self.assertIs(cm.exception, error)
+
+  async def test_max_retries_exceeded_with_dict_error(self):
+    """Test that Exception is raised when max retries exceeded with dict error."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    # First call should fail and return a retry response
+    result1 = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result={"some": "result"},
+    )
+    self.assertIsNotNone(result1)
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Second call should exceed max_retries and raise
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Verify exception properties
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
+  async def test_max_retries_exceeded_without_exception(self):
+    """Test max retries exceeded returns failure message when exception is disabled."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=2, throw_exception_if_retry_exceeded=False
+    )
+    error = TimeoutError("Timeout occurred")
+
+    # Call until we exceed the retry limit
+    result = None
+    for _ in range(3):
+      result = await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+
+    # Should get a retry exceeded message on the last call
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "TimeoutError")
+    self.assertIn(
+        "the retry limit has been exceeded", result["reflection_guidance"]
+    )
+    self.assertIn("Do not attempt to use the", result["reflection_guidance"])
+
+  async def test_successful_call_resets_retry_behavior(self):
+    """Test that successful calls reset the retry behavior."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    error = ValueError("Test error")
+
+    # First failure
+    result1 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Successful call
+    await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result={"success": True},
+    )
+
+    # Next failure should start fresh
+    result2 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result2["retry_count"], 1)  # Should restart from 1
+
+  async def test_none_result_handling(self):
+    """Test handling of None results in after_tool_callback."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+
+    # None result should be handled gracefully
+    callback_result = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=None,
+    )
+
+    self.assertIsNone(callback_result)
+
+  async def test_empty_tool_args_handling(self):
+    """Test handling of empty tool arguments."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    empty_args = {}
+    error = ValueError("Test error")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=empty_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    self.assertIsNotNone(result)
+    # Empty args should be represented in the response
+    self.assertIn("{}", result["reflection_guidance"])
+
+  async def test_retry_count_progression(self):
+    """Test that retry counts progress correctly for the same tool."""
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(max_retries=5)
+    error = ValueError("Test error")
+    tool = MockTool("single_tool")
+
+    for i in range(1, 4):
+      result = await plugin.on_tool_error_callback(
+          tool=tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+      self.assertEqual(result["retry_count"], i)
+
+  async def test_max_retries_parameter_behavior(self):
+    """Test that max_retries parameter affects behavior correctly."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    # Test with very low max_retries
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=False
+    )
+    error = ValueError("Test error")
+
+    # First call is fine
+    await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Second call exceeds limit
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Should hit max retries quickly with max_retries=1
+    self.assertIn(
+        "the retry limit has been exceeded.", result["reflection_guidance"]
+    )
+
+  async def test_default_extract_error_returns_none(self):
+    """Test that default extract_error_from_result returns None."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    result = {"status": "success", "data": "some data"}
+
+    error = await plugin.extract_error_from_result(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=result,
+    )
+    self.assertIsNone(error)
+
+  async def test_custom_error_detection_and_success_handling(self):
+    """Test custom error detection, success handling, and retry progression."""
+    custom_error_plugin = self.get_custom_error_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    custom_error_plugin.set_error_condition(
+        lambda result: result if result.get("status") == "error" else None
+    )
+
+    # Test error detection
+    error_result = {"status": "error", "message": "Something went wrong"}
+    callback_result = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=error_result,
+    )
+    self.assertIsNotNone(callback_result)
+    self.assertEqual(
+        callback_result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE
+    )
+    self.assertEqual(callback_result["retry_count"], 1)
+
+    # Test success handling
+    success_result = {"status": "success", "data": "operation completed"}
+    callback_result = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=success_result,
+    )
+    self.assertIsNone(callback_result)
+
+  async def test_retry_state_management(self):
+    """Test retry state management with custom errors and mixed error types."""
+    custom_error_plugin = self.get_custom_error_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    custom_error_plugin.set_error_condition(
+        lambda result: result if result.get("failed") else None
+    )
+
+    # Custom error followed by exception
+    custom_error = {"failed": True, "reason": "Network timeout"}
+    result1 = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=custom_error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Exception should increment retry count
+    exception = ValueError("Invalid parameter")
+    result2 = await custom_error_plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=exception,
+    )
+    self.assertEqual(result2["retry_count"], 2)
+
+    # Success should reset
+    success = {"result": "success"}
+    result3 = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=success,
+    )
+    self.assertIsNone(result3)
+
+    # Next error should start fresh
+    result4 = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=custom_error,
+    )
+    self.assertEqual(result4["retry_count"], 1)
+
+  async def test_hallucinating_tool_name(self):
+    """Test that hallucinating tool name is handled correctly."""
+    wrong_function_call = types.Part.from_function_call(
+        name="increase_by_one", args={"x": 1}
+    )
+    correct_function_call = types.Part.from_function_call(
+        name="increase", args={"x": 1}
+    )
+    responses: list[types.Content] = [
+        wrong_function_call,
+        correct_function_call,
+        "response1",
+    ]
+    mock_model = testing_utils.MockModel.create(responses=responses)
+
+    function_called = 0
+
+    def increase(x: int) -> int:
+      nonlocal function_called
+      function_called += 1
+      return x + 1
+
+    agent = LlmAgent(name="root_agent", model=mock_model, tools=[increase])
+    runner = testing_utils.TestInMemoryRunner(
+        agent=agent, plugins=[self.get_plugin()]
+    )
+
+    events = await runner.run_async_with_new_session("test")
+    # Filter out agent_state events (no content).
+    events = [e for e in events if e.content is not None]
+
+    # Assert that the first event is a function call with the wrong name
+    assert events[0].content.parts[0].function_call.name == "increase_by_one"
+
+    # Assert that the second event is a function response with the
+    # reflection_guidance
+    assert (
+        events[1].content.parts[0].function_response.response["error_type"]
+        == "ValueError"
+    )
+    assert (
+        events[1].content.parts[0].function_response.response["retry_count"]
+        == 1
+    )
+    assert (
+        "Wrong Function Name"
+        in events[1]
+        .content.parts[0]
+        .function_response.response["reflection_guidance"]
+    )
+
+    # Assert that the third event is a function call with the correct name
+    assert events[2].content.parts[0].function_call.name == "increase"
+    self.assertEqual(function_called, 1)

tests/unittests/plugins/test_save_files_as_artifacts.py

@@ -0,0 +1,396 @@
+# Copyright 2026 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 __future__ import annotations
+
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.artifacts.base_artifact_service import ArtifactVersion
+from google.adk.plugins.save_files_as_artifacts_plugin import SaveFilesAsArtifactsPlugin
+from google.genai import types
+import pytest
+
+
+class TestSaveFilesAsArtifactsPlugin:
+  """Test suite for SaveFilesAsArtifactsPlugin."""
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    self.plugin = SaveFilesAsArtifactsPlugin()
+
+    # Mock invocation context
+    self.mock_context = Mock(spec=InvocationContext)
+    self.mock_context.app_name = "test_app"
+    self.mock_context.user_id = "test_user"
+    self.mock_context.invocation_id = "test_invocation_123"
+    self.mock_context.session = Mock()
+    self.mock_context.session.id = "test_session"
+    self.mock_context.session.state = {}
+
+    artifact_service = Mock()
+    artifact_service.save_artifact = AsyncMock(return_value=0)
+
+    async def _mock_get_artifact_version(**kwargs):
+      filename = kwargs.get("filename", "unknown_file")
+      version = kwargs.get("version", 0)
+      return ArtifactVersion(
+          version=version,
+          canonical_uri=f"gs://mock-bucket/{filename}/versions/{version}",
+          mime_type="application/pdf",
+      )
+
+    artifact_service.get_artifact_version = AsyncMock(
+        side_effect=_mock_get_artifact_version
+    )
+    self.mock_context.artifact_service = artifact_service
+
+  @pytest.mark.asyncio
+  async def test_save_files_with_display_name(self):
+    """Test saving files when inline_data has display_name."""
+    inline_data = types.Blob(
+        display_name="test_document.pdf",
+        data=b"test data",
+        mime_type="application/pdf",
+    )
+
+    original_part = types.Part(inline_data=inline_data)
+    user_message = types.Content(parts=[original_part])
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    self.mock_context.artifact_service.save_artifact.assert_called_once_with(
+        app_name="test_app",
+        user_id="test_user",
+        session_id="test_session",
+        filename="test_document.pdf",
+        artifact=original_part,
+    )
+
+    assert result
+    assert len(result.parts) == 2
+    assert result.parts[0].text == '[Uploaded Artifact: "test_document.pdf"]'
+    assert result.parts[1].file_data
+    assert (
+        result.parts[1].file_data.file_uri
+        == "gs://mock-bucket/test_document.pdf/versions/0"
+    )
+    assert result.parts[1].file_data.display_name == "test_document.pdf"
+    assert result.parts[1].file_data.mime_type == "application/pdf"
+
+  @pytest.mark.asyncio
+  async def test_attach_file_reference_false(self):
+    """Test that file reference is not attached when attach_file_reference is False."""
+    plugin = SaveFilesAsArtifactsPlugin(attach_file_reference=False)
+
+    inline_data = types.Blob(
+        display_name="test_document.pdf",
+        data=b"test data",
+        mime_type="application/pdf",
+    )
+
+    original_part = types.Part(inline_data=inline_data)
+    user_message = types.Content(parts=[original_part])
+
+    result = await plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    self.mock_context.artifact_service.save_artifact.assert_called_once_with(
+        app_name="test_app",
+        user_id="test_user",
+        session_id="test_session",
+        filename="test_document.pdf",
+        artifact=original_part,
+    )
+
+    assert result
+    assert len(result.parts) == 1
+    assert result.parts[0].text == '[Uploaded Artifact: "test_document.pdf"]'
+
+  @pytest.mark.asyncio
+  async def test_save_files_without_display_name(self):
+    """Test saving files when inline_data has no display_name."""
+    inline_data = types.Blob(
+        display_name=None, data=b"test data", mime_type="application/pdf"
+    )
+
+    original_part = types.Part(inline_data=inline_data)
+    user_message = types.Content(parts=[original_part])
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    expected_filename = "artifact_test_invocation_123_0"
+    self.mock_context.artifact_service.save_artifact.assert_called_once_with(
+        app_name="test_app",
+        user_id="test_user",
+        session_id="test_session",
+        filename=expected_filename,
+        artifact=original_part,
+    )
+
+    assert result
+    assert len(result.parts) == 2
+    assert result.parts[0].text == f'[Uploaded Artifact: "{expected_filename}"]'
+    assert result.parts[1].file_data
+    assert (
+        result.parts[1].file_data.file_uri
+        == "gs://mock-bucket/artifact_test_invocation_123_0/versions/0"
+    )
+    assert result.parts[1].file_data.display_name == expected_filename
+
+  @pytest.mark.asyncio
+  async def test_multiple_files_in_message(self):
+    """Test handling multiple files in a single message."""
+    inline_data1 = types.Blob(
+        display_name="file1.txt", data=b"file1 content", mime_type="text/plain"
+    )
+    inline_data2 = types.Blob(
+        display_name="file2.jpg", data=b"file2 content", mime_type="image/jpeg"
+    )
+
+    user_message = types.Content(
+        parts=[
+            types.Part(inline_data=inline_data1),
+            types.Part(text="Some text between files"),
+            types.Part(inline_data=inline_data2),
+        ]
+    )
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    assert self.mock_context.artifact_service.save_artifact.call_count == 2
+    first_call = (
+        self.mock_context.artifact_service.save_artifact.call_args_list[0]
+    )
+    second_call = (
+        self.mock_context.artifact_service.save_artifact.call_args_list[1]
+    )
+    assert first_call[1]["filename"] == "file1.txt"
+    assert second_call[1]["filename"] == "file2.jpg"
+
+    assert result
+    assert len(result.parts) == 5
+    assert result.parts[0].text == '[Uploaded Artifact: "file1.txt"]'
+    assert result.parts[1].file_data
+    assert (
+        result.parts[1].file_data.file_uri
+        == "gs://mock-bucket/file1.txt/versions/0"
+    )
+    assert result.parts[1].file_data.display_name == "file1.txt"
+    assert result.parts[2].text == "Some text between files"
+    assert result.parts[3].text == '[Uploaded Artifact: "file2.jpg"]'
+    assert result.parts[4].file_data
+    assert (
+        result.parts[4].file_data.file_uri
+        == "gs://mock-bucket/file2.jpg/versions/0"
+    )
+    assert result.parts[4].file_data.display_name == "file2.jpg"
+
+  @pytest.mark.asyncio
+  async def test_no_artifact_service(self):
+    """Test behavior when artifact service is not available."""
+    self.mock_context.artifact_service = None
+
+    inline_data = types.Blob(
+        display_name="test.pdf", data=b"test data", mime_type="application/pdf"
+    )
+    user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    assert result == user_message
+    assert result.parts[0].inline_data == inline_data
+
+  @pytest.mark.asyncio
+  async def test_no_parts_in_message(self):
+    """Test behavior when message has no parts."""
+    user_message = types.Content(parts=[])
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    assert result is None
+    self.mock_context.artifact_service.save_artifact.assert_not_called()
+
+  @pytest.mark.asyncio
+  async def test_parts_without_inline_data(self):
+    """Test behavior with parts that don't have inline_data."""
+    user_message = types.Content(
+        parts=[types.Part(text="Hello world"), types.Part(text="No files here")]
+    )
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    assert result is None
+    self.mock_context.artifact_service.save_artifact.assert_not_called()
+
+  @pytest.mark.asyncio
+  async def test_save_artifact_failure(self):
+    """Test behavior when saving artifact fails."""
+    self.mock_context.artifact_service.save_artifact.side_effect = Exception(
+        "Storage error"
+    )
+
+    inline_data = types.Blob(
+        display_name="test.pdf", data=b"test data", mime_type="application/pdf"
+    )
+    user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_mixed_success_and_failure(self):
+    """Test behavior when some files save successfully and others fail."""
+    save_calls = 0
+
+    async def _save_side_effect(*_args, **_kwargs):
+      nonlocal save_calls
+      save_calls += 1
+      if save_calls == 2:
+        raise Exception("Storage error on second file")
+      return 0
+
+    self.mock_context.artifact_service.save_artifact.side_effect = (
+        _save_side_effect
+    )
+
+    inline_data1 = types.Blob(
+        display_name="success.pdf",
+        data=b"success data",
+        mime_type="application/pdf",
+    )
+    inline_data2 = types.Blob(
+        display_name="failure.pdf",
+        data=b"failure data",
+        mime_type="application/pdf",
+    )
+
+    original_part2 = types.Part(inline_data=inline_data2)
+    user_message = types.Content(
+        parts=[types.Part(inline_data=inline_data1), original_part2]
+    )
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    assert result
+    assert len(result.parts) == 3
+    assert result.parts[0].text == '[Uploaded Artifact: "success.pdf"]'
+    assert result.parts[1].file_data
+    assert result.parts[2] == original_part2
+    assert result.parts[2].inline_data == inline_data2
+
+  @pytest.mark.asyncio
+  async def test_placeholder_text_format(self):
+    """Test that placeholder text is formatted correctly."""
+    inline_data = types.Blob(
+        display_name="test file with spaces.docx",
+        data=b"document data",
+        mime_type=(
+            "application/vnd.openxmlformats-officedocument."
+            "wordprocessingml.document"
+        ),
+    )
+
+    user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    expected_text = '[Uploaded Artifact: "test file with spaces.docx"]'
+    assert result.parts[0].text == expected_text
+    assert result.parts[1].file_data
+
+  def test_plugin_name_default(self):
+    """Test that plugin has correct default name."""
+    plugin = SaveFilesAsArtifactsPlugin()
+    assert plugin.name == "save_files_as_artifacts_plugin"
+
+  @pytest.mark.asyncio
+  async def test_artifact_delta_reporting(self):
+    """Test that the artifact delta is written to state then event actions."""
+
+    # 1. First Turn - Trigger user message callback
+    blob = types.Blob(
+        display_name="blob.pdf",
+        data=b"test data",
+        mime_type="application/pdf",
+    )
+    user_message = types.Content(parts=[types.Part(inline_data=blob)])
+    await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    # Verify state is updated
+    key = "save_files_as_artifacts_plugin:pending_delta"
+    assert key in self.mock_context.session.state
+    assert self.mock_context.session.state[key] == {"blob.pdf": 0}
+
+    # 2. First Turn - Trigger before agent callback
+    callback_context = Mock()
+    callback_context.state = self.mock_context.session.state
+    callback_context.actions = Mock()
+    callback_context.actions.artifact_delta = {}
+    await self.plugin.before_agent_callback(
+        agent=Mock(), callback_context=callback_context
+    )
+
+    # Verify artifact_delta is updated and state is cleared
+    assert callback_context.actions.artifact_delta == {"blob.pdf": 0}
+    assert self.mock_context.session.state[key] == {}
+
+    # 3. Second Turn - Trigger user message callback
+    blob_2 = types.Blob(
+        display_name="blob_2.pdf",
+        data=b"test data 2",
+        mime_type="application/pdf",
+    )
+    user_message_2 = types.Content(parts=[types.Part(inline_data=blob_2)])
+    await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message_2
+    )
+
+    # Verify state is updated
+    assert self.mock_context.session.state[key] == {"blob_2.pdf": 0}
+
+    # 4. Second Turn - Trigger before agent callback
+    callback_context_2 = Mock()
+    callback_context_2.state = self.mock_context.session.state
+    callback_context_2.actions = Mock()
+    callback_context_2.actions.artifact_delta = {}
+    await self.plugin.before_agent_callback(
+        agent=Mock(), callback_context=callback_context_2
+    )
+
+    # Verify artifact_delta is updated and state is cleared
+    assert callback_context_2.actions.artifact_delta == {"blob_2.pdf": 0}
+    assert self.mock_context.session.state[key] == {}