diff --git a/src/google/adk/agents/__init__.py b/src/google/adk/agents/__init__.py index 6a2f464913..57308ba8b9 100644 --- a/src/google/adk/agents/__init__.py +++ b/src/google/adk/agents/__init__.py @@ -19,6 +19,8 @@ from ..utils import _lazy if TYPE_CHECKING: + from ._callback_metadata import CallbackHook + from ._callback_metadata import CallbackInvocationInfo from ._managed_agent import ManagedAgent from .base_agent import BaseAgent from .base_agent_config import BaseAgentConfig @@ -42,6 +44,8 @@ 'Agent': '.llm_agent', 'BaseAgent': '.base_agent', 'BaseAgentConfig': '.base_agent_config', + 'CallbackHook': '._callback_metadata', + 'CallbackInvocationInfo': '._callback_metadata', 'Context': '.context', 'InvocationContext': '.invocation_context', 'LiveRequest': '.live_request_queue', @@ -61,6 +65,8 @@ __all__ = [ 'Agent', 'BaseAgent', + 'CallbackHook', + 'CallbackInvocationInfo', 'Context', 'LlmAgent', 'LoopAgent', diff --git a/src/google/adk/agents/_callback_metadata.py b/src/google/adk/agents/_callback_metadata.py new file mode 100644 index 0000000000..2b8c1daa21 --- /dev/null +++ b/src/google/adk/agents/_callback_metadata.py @@ -0,0 +1,43 @@ +# 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 dataclasses import dataclass +import enum + + +class CallbackHook(str, enum.Enum): + """Lifecycle hooks that receive callback-scoped context metadata.""" + + BEFORE_AGENT = 'before_agent' + AFTER_AGENT = 'after_agent' + BEFORE_MODEL = 'before_model' + AFTER_MODEL = 'after_model' + BEFORE_TOOL = 'before_tool' + AFTER_TOOL = 'after_tool' + ON_MODEL_ERROR = 'on_model_error' + ON_TOOL_ERROR = 'on_tool_error' + ON_AGENT_ERROR = 'on_agent_error' + + +@dataclass(frozen=True) +class CallbackInvocationInfo: + """Describes the lifecycle callback currently being invoked. + + Attributes: + hook: The lifecycle hook for the active callback. + """ + + hook: CallbackHook diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index d8692a933e..8d1f36c31f 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -51,6 +51,7 @@ from ..utils._callback_pipeline import _stop_on_truthy from ..utils.context_utils import Aclosing from ..workflow import BaseNode +from ._callback_metadata import CallbackHook from .base_agent_config import BaseAgentConfig as BaseAgentConfig from .callback_context import CallbackContext from .context import Context @@ -528,11 +529,12 @@ async def _handle_before_agent_callback( # callbacks. callbacks = self.canonical_before_agent_callbacks if not before_agent_callback_content and callbacks: - before_agent_callback_content = await _run_callbacks( - callbacks, - _stop_on_truthy, - callback_context=callback_context, - ) + with callback_context._callback_scope(CallbackHook.BEFORE_AGENT): + before_agent_callback_content = await _run_callbacks( + callbacks, + _stop_on_truthy, + callback_context=callback_context, + ) # Process the override content if exists, and further process the state # change if exists. @@ -583,11 +585,12 @@ async def _handle_after_agent_callback( # callbacks. callbacks = self.canonical_after_agent_callbacks if not after_agent_callback_content and callbacks: - after_agent_callback_content = await _run_callbacks( - callbacks, - _stop_on_truthy, - callback_context=callback_context, - ) + with callback_context._callback_scope(CallbackHook.AFTER_AGENT): + after_agent_callback_content = await _run_callbacks( + callbacks, + _stop_on_truthy, + callback_context=callback_context, + ) # Process the override content if exists, and further process the state # change if exists. diff --git a/src/google/adk/agents/callback_context.py b/src/google/adk/agents/callback_context.py index e7ffd58b9c..5dee26afa0 100644 --- a/src/google/adk/agents/callback_context.py +++ b/src/google/adk/agents/callback_context.py @@ -14,6 +14,8 @@ from __future__ import annotations +from ._callback_metadata import CallbackHook as CallbackHook +from ._callback_metadata import CallbackInvocationInfo as CallbackInvocationInfo from .context import Context # Keep ReadonlyContext for backward compatibility diff --git a/src/google/adk/agents/context.py b/src/google/adk/agents/context.py index 50bf522fce..8d7a103494 100644 --- a/src/google/adk/agents/context.py +++ b/src/google/adk/agents/context.py @@ -16,8 +16,11 @@ from __future__ import annotations +from collections.abc import Iterator from collections.abc import Mapping from collections.abc import Sequence +from contextlib import contextmanager +from contextvars import ContextVar from typing import Any from typing import cast from typing import TYPE_CHECKING @@ -25,6 +28,8 @@ from opentelemetry import context as context_api from typing_extensions import override +from ._callback_metadata import CallbackHook +from ._callback_metadata import CallbackInvocationInfo from .readonly_context import ReadonlyContext if TYPE_CHECKING: @@ -229,6 +234,9 @@ def __init__( self._output_for_ancestors = [] self._error: Exception | None = None self._error_node_path: str = '' + self._callback_info: ContextVar[CallbackInvocationInfo | None] = ContextVar( + 'callback_info', default=None + ) @property @override @@ -247,6 +255,20 @@ def function_call_id(self, value: str | None) -> None: """Sets the function call id of the current tool call.""" self._function_call_id = value + @property + def callback_info(self) -> CallbackInvocationInfo | None: + """Returns metadata for the callback currently being invoked, if any.""" + return self._callback_info.get() + + @contextmanager + def _callback_scope(self, hook: CallbackHook) -> Iterator[None]: + """Sets task-local callback metadata for the duration of a callback.""" + token = self._callback_info.set(CallbackInvocationInfo(hook=hook)) + try: + yield + finally: + self._callback_info.reset(token) + @property def branch(self) -> str | None: """The branch path of the current invocation context.""" diff --git a/src/google/adk/flows/llm_flows/_tool_caller.py b/src/google/adk/flows/llm_flows/_tool_caller.py index f837b9c76a..38688d76a7 100644 --- a/src/google/adk/flows/llm_flows/_tool_caller.py +++ b/src/google/adk/flows/llm_flows/_tool_caller.py @@ -39,6 +39,7 @@ from google.genai import types from . import _tool_error_handler +from ...agents._callback_metadata import CallbackHook from ...agents.active_streaming_tool import ActiveStreamingTool from ...events.event import Event from ...live.live_request_queue import LiveRequestQueue @@ -679,13 +680,14 @@ async def _prepare_single( # Step 2: If no overrides are provided from the plugins, further run the # canonical callback. if override_response is None: - override_response = await _run_callbacks( - agent.canonical_before_tool_callbacks, # type: ignore[arg-type] - _stop_on_non_none, - tool=tool, - args=function_args, - tool_context=tool_context, - ) + with tool_context._callback_scope(CallbackHook.BEFORE_TOOL): + override_response = await _run_callbacks( + agent.canonical_before_tool_callbacks, # type: ignore[arg-type] + _stop_on_non_none, + tool=tool, + args=function_args, + tool_context=tool_context, + ) # Handle tool lookup failure if before-tool callbacks did not override the # response. @@ -797,14 +799,15 @@ async def _run_with_trace() -> Event | None: # Step 5: If no overrides are provided from the plugins, further run the # canonical after_tool_callbacks. if altered_function_response is None: - altered_function_response = await _run_callbacks( - agent.canonical_after_tool_callbacks, # type: ignore[arg-type] - _stop_on_non_none, - tool=tool, - args=function_args, - tool_context=tool_context, - tool_response=callback_tool_response, - ) + with tool_context._callback_scope(CallbackHook.AFTER_TOOL): + altered_function_response = await _run_callbacks( + agent.canonical_after_tool_callbacks, # type: ignore[arg-type] + _stop_on_non_none, + tool=tool, + args=function_args, + tool_context=tool_context, + tool_response=callback_tool_response, + ) # Step 6: If alternative response exists from after_tool_callback, use it # instead of the original function response. diff --git a/src/google/adk/flows/llm_flows/_tool_error_handler.py b/src/google/adk/flows/llm_flows/_tool_error_handler.py index 76f88a74b3..eb09e3d7fe 100644 --- a/src/google/adk/flows/llm_flows/_tool_error_handler.py +++ b/src/google/adk/flows/llm_flows/_tool_error_handler.py @@ -22,6 +22,7 @@ from typing import Optional from typing import TYPE_CHECKING +from ...agents._callback_metadata import CallbackHook from ...tools.base_tool import BaseTool from ...tools.tool_context import ToolContext from ...utils._callback_pipeline import _run_callbacks @@ -117,11 +118,12 @@ async def run_on_tool_error_callbacks( if error_response is not None: return error_response - return await _run_callbacks( - agent.canonical_on_tool_error_callbacks, # type: ignore[arg-type] - _stop_on_non_none, - tool=tool, - args=tool_args, - tool_context=tool_context, - error=error, - ) + with tool_context._callback_scope(CallbackHook.ON_TOOL_ERROR): + return await _run_callbacks( + agent.canonical_on_tool_error_callbacks, # type: ignore[arg-type] + _stop_on_non_none, + tool=tool, + args=tool_args, + tool_context=tool_context, + error=error, + ) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 975f0e528c..bd48faeba2 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -30,6 +30,7 @@ from . import _live_llm_flow from . import _output_schema_processor from . import functions +from ...agents._callback_metadata import CallbackHook from ...agents._streaming_mode import StreamingMode from ...agents.base_agent import BaseAgent from ...agents.callback_context import CallbackContext @@ -257,12 +258,13 @@ async def _handle_before_model_callback( # If no overrides are provided from the plugins, further run the canonical # callbacks. - callback_response = await _run_callbacks( - agent.canonical_before_model_callbacks, - _stop_on_truthy, - callback_context=callback_context, - llm_request=llm_request, - ) + with callback_context._callback_scope(CallbackHook.BEFORE_MODEL): + callback_response = await _run_callbacks( + agent.canonical_before_model_callbacks, + _stop_on_truthy, + callback_context=callback_context, + llm_request=llm_request, + ) if callback_response: return callback_response return None @@ -327,12 +329,13 @@ async def _maybe_add_grounding_metadata( # If no overrides are provided from the plugins, further run the canonical # callbacks. - callback_response = await _run_callbacks( - agent.canonical_after_model_callbacks, - _stop_on_truthy, - callback_context=callback_context, - llm_response=llm_response, - ) + with callback_context._callback_scope(CallbackHook.AFTER_MODEL): + callback_response = await _run_callbacks( + agent.canonical_after_model_callbacks, + _stop_on_truthy, + callback_context=callback_context, + llm_response=llm_response, + ) if callback_response: return await _maybe_add_grounding_metadata(callback_response) return await _maybe_add_grounding_metadata() @@ -390,13 +393,14 @@ async def _run_on_model_error_callbacks( if error_response is not None: return error_response - return await _run_callbacks( - agent.canonical_on_model_error_callbacks, - _stop_on_non_none, - callback_context=callback_context, - llm_request=llm_request, - error=error, - ) + with callback_context._callback_scope(CallbackHook.ON_MODEL_ERROR): + return await _run_callbacks( + agent.canonical_on_model_error_callbacks, + _stop_on_non_none, + callback_context=callback_context, + llm_request=llm_request, + error=error, + ) try: async with _instrumentation.record_inference_telemetry( diff --git a/src/google/adk/plugins/plugin_manager.py b/src/google/adk/plugins/plugin_manager.py index 11c62937cc..79d99cf9dd 100644 --- a/src/google/adk/plugins/plugin_manager.py +++ b/src/google/adk/plugins/plugin_manager.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio +from contextlib import nullcontext import logging import sys from typing import Any @@ -25,6 +26,7 @@ from google.genai import types +from ..agents._callback_metadata import CallbackHook from .base_plugin import BasePlugin if TYPE_CHECKING: @@ -298,28 +300,41 @@ async def _run_callbacks( RuntimeError: If a plugin encounters an unhandled exception during execution. The original exception is chained. """ - for plugin in self.plugins: - # Each plugin might not implement all callbacks. The base class provides - # default `pass` implementations, so `getattr` will always succeed. - callback_method = getattr(plugin, callback_name) - try: - result = await callback_method(**kwargs) - if result is not None: - # Early exit: A plugin has returned a value. We stop - # processing further plugins and return this value immediately. - logger.debug( - "Plugin '%s' returned a value for callback '%s', exiting early.", - plugin.name, - callback_name, - ) - return result - except Exception as e: - error_message = ( - f"Error in plugin '{plugin.name}' during '{callback_name}'" - f" callback: {e}" + callback_context = kwargs.get("callback_context") + if callback_context is None: + callback_context = kwargs.get("tool_context") + callback_scope = ( + callback_context._callback_scope( + CallbackHook(callback_name.removesuffix("_callback")) ) - logger.error(error_message, exc_info=True) - raise RuntimeError(error_message) from e + if callback_context is not None + else nullcontext() + ) + + with callback_scope: + for plugin in self.plugins: + # Each plugin might not implement all callbacks. The base class provides + # default `pass` implementations, so `getattr` will always succeed. + callback_method = getattr(plugin, callback_name) + try: + result = await callback_method(**kwargs) + if result is not None: + # Early exit: A plugin has returned a value. We stop + # processing further plugins and return this value immediately. + logger.debug( + "Plugin '%s' returned a value for callback '%s', exiting" + " early.", + plugin.name, + callback_name, + ) + return result + except Exception as e: + error_message = ( + f"Error in plugin '{plugin.name}' during '{callback_name}'" + f" callback: {e}" + ) + logger.error(error_message, exc_info=True) + raise RuntimeError(error_message) from e return None @@ -365,18 +380,28 @@ async def _run_notification_callbacks( callback_name: The name of the callback method to execute. **kwargs: Keyword arguments to be passed to the callback method. """ - for plugin in self.plugins: - callback_method = getattr(plugin, callback_name) - try: - await callback_method(**kwargs) - except Exception as e: - logger.error( - "Error in plugin '%s' during '%s' callback: %s", - plugin.name, - callback_name, - e, - exc_info=True, + callback_context = kwargs.get("callback_context") + callback_scope = ( + callback_context._callback_scope( + CallbackHook(callback_name.removesuffix("_callback")) ) + if callback_context is not None + else nullcontext() + ) + + with callback_scope: + for plugin in self.plugins: + callback_method = getattr(plugin, callback_name) + try: + await callback_method(**kwargs) + except Exception as e: + logger.error( + "Error in plugin '%s' during '%s' callback: %s", + plugin.name, + callback_name, + e, + exc_info=True, + ) async def close(self) -> None: """Calls the close method on all registered plugins concurrently. diff --git a/tests/unittests/agents/test_base_agent.py b/tests/unittests/agents/test_base_agent.py index 698520cd7a..7d0c4eacd3 100644 --- a/tests/unittests/agents/test_base_agent.py +++ b/tests/unittests/agents/test_base_agent.py @@ -25,6 +25,7 @@ from unittest import mock import warnings +from google.adk.agents import CallbackHook from google.adk.agents.base_agent import BaseAgent from google.adk.agents.base_agent import BaseAgentState from google.adk.agents.callback_context import CallbackContext @@ -244,6 +245,44 @@ async def test_run_async_before_agent_callback_noop( spy_run_async_impl.assert_called_once() +@pytest.mark.asyncio +async def test_agent_callbacks_receive_current_hook_metadata( + request: pytest.FixtureRequest, +): + """Agent callback lists observe their active hook only during execution.""" + observations = [] + callback_contexts = [] + + def first_before_callback(callback_context: CallbackContext) -> None: + observations.append(callback_context.callback_info.hook) + callback_contexts.append(callback_context) + + def second_before_callback(callback_context: CallbackContext) -> None: + observations.append(callback_context.callback_info.hook) + + def after_callback(callback_context: CallbackContext) -> None: + observations.append(callback_context.callback_info.hook) + callback_contexts.append(callback_context) + + agent = _TestingAgent( + name=f'{request.function.__name__}_test_agent', + before_agent_callback=[first_before_callback, second_before_callback], + after_agent_callback=after_callback, + ) + parent_ctx = await _create_parent_invocation_context( + request.function.__name__, agent + ) + + _ = [event async for event in agent.run_async(parent_ctx)] + + assert observations == [ + CallbackHook.BEFORE_AGENT, + CallbackHook.BEFORE_AGENT, + CallbackHook.AFTER_AGENT, + ] + assert all(context.callback_info is None for context in callback_contexts) + + @pytest.mark.asyncio async def test_run_async_before_agent_callback_use_plugin( request: pytest.FixtureRequest, diff --git a/tests/unittests/agents/test_callback_metadata.py b/tests/unittests/agents/test_callback_metadata.py new file mode 100644 index 0000000000..0e0c0d13a8 --- /dev/null +++ b/tests/unittests/agents/test_callback_metadata.py @@ -0,0 +1,50 @@ +# 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 dataclasses import FrozenInstanceError + +from google.adk.agents import CallbackHook +from google.adk.agents import CallbackInvocationInfo +from google.adk.agents.callback_context import CallbackHook as ContextCallbackHook +from google.adk.agents.callback_context import CallbackInvocationInfo as ContextCallbackInvocationInfo +import pytest + + +def test_callback_metadata_is_exported_from_public_agent_modules(): + """Callback metadata has stable public import paths.""" + assert ContextCallbackHook is CallbackHook + assert ContextCallbackInvocationInfo is CallbackInvocationInfo + + +def test_callback_hook_values_are_stable(): + """Callback hook values match the documented lifecycle names.""" + assert [hook.value for hook in CallbackHook] == [ + 'before_agent', + 'after_agent', + 'before_model', + 'after_model', + 'before_tool', + 'after_tool', + 'on_model_error', + 'on_tool_error', + 'on_agent_error', + ] + + +def test_callback_invocation_info_is_immutable(): + """Callback invocation metadata cannot be changed by consumers.""" + callback_info = CallbackInvocationInfo(hook=CallbackHook.BEFORE_MODEL) + + with pytest.raises(FrozenInstanceError): + callback_info.hook = CallbackHook.AFTER_MODEL diff --git a/tests/unittests/flows/llm_flows/test_model_callbacks.py b/tests/unittests/flows/llm_flows/test_model_callbacks.py index 833ffd17ff..9382cd42d0 100644 --- a/tests/unittests/flows/llm_flows/test_model_callbacks.py +++ b/tests/unittests/flows/llm_flows/test_model_callbacks.py @@ -16,6 +16,7 @@ from typing import Optional from unittest import mock +from google.adk.agents import CallbackHook from google.adk.agents.callback_context import CallbackContext from google.adk.agents.llm_agent import Agent from google.adk.models.llm_request import LlmRequest @@ -175,6 +176,40 @@ def test_after_model_callback_lambda_with_arbitrary_param_names(): ] +def test_model_callbacks_receive_current_hook_metadata(): + """Model callbacks observe their active hook only during execution.""" + observations = [] + callback_contexts = [] + + def before_callback( + callback_context: CallbackContext, llm_request: LlmRequest + ) -> None: + observations.append(callback_context.callback_info.hook) + callback_contexts.append(callback_context) + + def after_callback( + callback_context: CallbackContext, llm_response: LlmResponse + ) -> None: + observations.append(callback_context.callback_info.hook) + callback_contexts.append(callback_context) + + agent = Agent( + name='root_agent', + model=testing_utils.MockModel.create(responses=['model_response']), + before_model_callback=before_callback, + after_model_callback=after_callback, + ) + + runner = testing_utils.InMemoryRunner(agent) + _ = runner.run('test') + + assert observations == [ + CallbackHook.BEFORE_MODEL, + CallbackHook.AFTER_MODEL, + ] + assert all(context.callback_info is None for context in callback_contexts) + + @pytest.mark.asyncio async def test_on_model_callback_model_error_noop(): """Test that the on_model_error_callback is a no-op when the model returns an error.""" @@ -212,6 +247,40 @@ async def test_on_model_callback_model_error_modify_model_response(): ) == [('root_agent', 'on_model_error_callback_response')] +@pytest.mark.asyncio +async def test_model_error_callback_receives_current_hook_metadata(): + """A model error callback observes the error hook only while it runs.""" + observations = [] + callback_contexts = [] + + def error_callback( + callback_context: CallbackContext, + llm_request: LlmRequest, + error: Exception, + ) -> LlmResponse: + observations.append(callback_context.callback_info.hook) + callback_contexts.append(callback_context) + return LlmResponse( + content=testing_utils.ModelContent( + [types.Part.from_text(text='recovered')] + ) + ) + + agent = Agent( + name='root_agent', + model=testing_utils.MockModel.create( + responses=[], error=SystemError('error') + ), + on_model_error_callback=error_callback, + ) + runner = testing_utils.TestInMemoryRunner(agent) + + await runner.run_async_with_new_session('test') + + assert observations == [CallbackHook.ON_MODEL_ERROR] + assert all(context.callback_info is None for context in callback_contexts) + + @pytest.mark.asyncio async def test_on_model_error_callback_chain_stops_on_recovery_response(): """Test that model error recovery stops after a non-None response.""" diff --git a/tests/unittests/flows/llm_flows/test_tool_callbacks.py b/tests/unittests/flows/llm_flows/test_tool_callbacks.py index 5c048e30d7..a03f4c284e 100644 --- a/tests/unittests/flows/llm_flows/test_tool_callbacks.py +++ b/tests/unittests/flows/llm_flows/test_tool_callbacks.py @@ -15,6 +15,7 @@ from typing import Any from unittest import mock +from google.adk.agents import CallbackHook from google.adk.agents.llm_agent import Agent from google.adk.tools.base_tool import BaseTool from google.adk.tools.tool_context import ToolContext @@ -265,6 +266,45 @@ def test_after_tool_callback_noop(): ] +def test_tool_callbacks_receive_current_hook_metadata(): + """Tool callbacks observe their active hook only during execution.""" + observations = [] + tool_contexts = [] + + def context_aware_function(tool_context: ToolContext) -> str: + observations.append(tool_context.callback_info) + return 'tool_response' + + def before_callback(tool, args, tool_context): + observations.append(tool_context.callback_info.hook) + tool_contexts.append(tool_context) + + def after_callback(tool, args, tool_context, tool_response): + observations.append(tool_context.callback_info.hook) + tool_contexts.append(tool_context) + + responses = [ + types.Part.from_function_call(name='context_aware_function', args={}), + 'response1', + ] + agent = Agent( + name='root_agent', + model=testing_utils.MockModel.create(responses=responses), + before_tool_callback=before_callback, + after_tool_callback=after_callback, + tools=[context_aware_function], + ) + + _ = testing_utils.InMemoryRunner(agent).run('test') + + assert observations == [ + CallbackHook.BEFORE_TOOL, + None, + CallbackHook.AFTER_TOOL, + ] + assert all(context.callback_info is None for context in tool_contexts) + + def test_after_tool_callback_modify_tool_response(): """Test that the after_tool_callback modifies the tool response.""" responses = [ @@ -478,6 +518,33 @@ async def async_on_tool_error_callback( ] +def test_tool_error_callback_receives_current_hook_metadata(): + """A tool error callback observes the error hook only while it runs.""" + observations = [] + tool_contexts = [] + + def error_callback(tool, args, tool_context, error): + observations.append(tool_context.callback_info.hook) + tool_contexts.append(tool_context) + return {'result': 'recovered'} + + responses = [ + types.Part.from_function_call(name='simple_function_with_error', args={}), + 'response1', + ] + agent = Agent( + name='root_agent', + model=testing_utils.MockModel.create(responses=responses), + on_tool_error_callback=error_callback, + tools=[simple_function_with_error], + ) + + _ = testing_utils.InMemoryRunner(agent).run('test') + + assert observations == [CallbackHook.ON_TOOL_ERROR] + assert all(context.callback_info is None for context in tool_contexts) + + def test_before_tool_callback_lambda_with_arbitrary_param_names(): """Test that before_tool_callback works with lambda having non-matching param names.""" captured = [] diff --git a/tests/unittests/plugins/test_notification_error_callbacks.py b/tests/unittests/plugins/test_notification_error_callbacks.py index 21ed745578..81a596fba6 100644 --- a/tests/unittests/plugins/test_notification_error_callbacks.py +++ b/tests/unittests/plugins/test_notification_error_callbacks.py @@ -164,6 +164,10 @@ async def _create_ctx( ) +async def _create_callback_context(agent: BaseAgent) -> CallbackContext: + return CallbackContext(await _create_ctx(agent)) + + # --------------------------------------------------------------------------- # Agent-level error callback tests # --------------------------------------------------------------------------- @@ -510,14 +514,13 @@ async def test_run_on_agent_error_callback_dispatches(self): plugin2 = _ErrorTrackingPlugin(name="p2") pm = PluginManager(plugins=[plugin1, plugin2]) - mock_agent = Mock(spec=BaseAgent) - mock_agent.name = "test_agent" - mock_ctx = Mock(spec=CallbackContext) + agent = _SuccessAgent(name="test_agent") + callback_context = await _create_callback_context(agent) err = RuntimeError("boom") await pm.run_on_agent_error_callback( - agent=mock_agent, - callback_context=mock_ctx, + agent=agent, + callback_context=callback_context, error=err, ) @@ -561,10 +564,11 @@ async def on_agent_error_callback(self, **kwargs): p1 = _ReturningPlugin(name="p1") p2 = _ReturningPlugin(name="p2") pm = PluginManager(plugins=[p1, p2]) + agent = _SuccessAgent(name="test_agent") await pm.run_on_agent_error_callback( - agent=Mock(spec=BaseAgent), - callback_context=Mock(spec=CallbackContext), + agent=agent, + callback_context=await _create_callback_context(agent), error=RuntimeError("x"), ) @@ -627,11 +631,10 @@ async def on_run_error_callback(self, **kwargs): 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" + agent = _SuccessAgent(name="test_agent") await pm.run_on_agent_error_callback( - agent=mock_agent, - callback_context=Mock(spec=CallbackContext), + agent=agent, + callback_context=await _create_callback_context(agent), error=RuntimeError("app crash"), ) assert p1.agent_error_called diff --git a/tests/unittests/plugins/test_plugin_manager.py b/tests/unittests/plugins/test_plugin_manager.py index 5b92e6c3df..cac731b55a 100644 --- a/tests/unittests/plugins/test_plugin_manager.py +++ b/tests/unittests/plugins/test_plugin_manager.py @@ -20,6 +20,9 @@ from unittest.mock import AsyncMock from unittest.mock import Mock +from google.adk.agents import CallbackHook +from google.adk.agents import CallbackInvocationInfo +from google.adk.agents.callback_context import CallbackContext from google.adk.models.llm_response import LlmResponse from google.adk.plugins.base_plugin import BasePlugin # Assume the following path to your modules @@ -46,56 +49,64 @@ def __init__(self, name: str): self.return_values: dict[PluginCallbackName, any] = {} # A map to configure exceptions to be raised by specific callbacks. self.exceptions_to_raise: dict[PluginCallbackName, Exception] = {} + self.callback_info_log: list[ + tuple[PluginCallbackName, CallbackInvocationInfo] + ] = [] - async def _handle_callback(self, name: PluginCallbackName): + async def _handle_callback(self, name: PluginCallbackName, **kwargs): """Generic handler for all callback methods.""" self.call_log.append(name) + callback_context = kwargs.get("callback_context") + if callback_context is None: + callback_context = kwargs.get("tool_context") + if callback_context is not None: + self.callback_info_log.append((name, callback_context.callback_info)) 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") + return await self._handle_callback("on_user_message_callback", **kwargs) async def before_run_callback(self, **kwargs): - return await self._handle_callback("before_run_callback") + return await self._handle_callback("before_run_callback", **kwargs) async def after_run_callback(self, **kwargs): - return await self._handle_callback("after_run_callback") + return await self._handle_callback("after_run_callback", **kwargs) async def on_event_callback(self, **kwargs): - return await self._handle_callback("on_event_callback") + return await self._handle_callback("on_event_callback", **kwargs) async def before_agent_callback(self, **kwargs): - return await self._handle_callback("before_agent_callback") + return await self._handle_callback("before_agent_callback", **kwargs) async def after_agent_callback(self, **kwargs): - return await self._handle_callback("after_agent_callback") + return await self._handle_callback("after_agent_callback", **kwargs) async def before_tool_callback(self, **kwargs): - return await self._handle_callback("before_tool_callback") + return await self._handle_callback("before_tool_callback", **kwargs) async def after_tool_callback(self, **kwargs): - return await self._handle_callback("after_tool_callback") + return await self._handle_callback("after_tool_callback", **kwargs) async def on_tool_error_callback(self, **kwargs): - return await self._handle_callback("on_tool_error_callback") + return await self._handle_callback("on_tool_error_callback", **kwargs) async def before_model_callback(self, **kwargs): - return await self._handle_callback("before_model_callback") + return await self._handle_callback("before_model_callback", **kwargs) async def after_model_callback(self, **kwargs): - return await self._handle_callback("after_model_callback") + return await self._handle_callback("after_model_callback", **kwargs) async def on_model_error_callback(self, **kwargs): - return await self._handle_callback("on_model_error_callback") + return await self._handle_callback("on_model_error_callback", **kwargs) async def on_agent_error_callback(self, **kwargs): - return await self._handle_callback("on_agent_error_callback") + return await self._handle_callback("on_agent_error_callback", **kwargs) async def on_run_error_callback(self, **kwargs): - return await self._handle_callback("on_run_error_callback") + return await self._handle_callback("on_run_error_callback", **kwargs) @pytest.fixture @@ -219,6 +230,10 @@ async def test_all_callbacks_are_supported( service.register_plugin(plugin1) mock_context = Mock() mock_user_message = Mock() + mock_invocation_context = Mock() + mock_invocation_context.session.state = {} + mock_invocation_context._state_schema = None + callback_context = CallbackContext(mock_invocation_context) # Test all callbacks await service.run_on_user_message_callback( @@ -230,37 +245,37 @@ async def test_all_callbacks_are_supported( invocation_context=mock_context, event=mock_context ) await service.run_before_agent_callback( - agent=mock_context, callback_context=mock_context + agent=mock_context, callback_context=callback_context ) await service.run_after_agent_callback( - agent=mock_context, callback_context=mock_context + agent=mock_context, callback_context=callback_context ) await service.run_before_tool_callback( - tool=mock_context, tool_args={}, tool_context=mock_context + tool=mock_context, tool_args={}, tool_context=callback_context ) await service.run_after_tool_callback( - tool=mock_context, tool_args={}, tool_context=mock_context, result={} + tool=mock_context, tool_args={}, tool_context=callback_context, result={} ) await service.run_on_tool_error_callback( tool=mock_context, tool_args={}, - tool_context=mock_context, + tool_context=callback_context, error=mock_context, ) await service.run_before_model_callback( - callback_context=mock_context, llm_request=mock_context + callback_context=callback_context, llm_request=mock_context ) await service.run_after_model_callback( - callback_context=mock_context, llm_response=mock_context + callback_context=callback_context, llm_response=mock_context ) await service.run_on_model_error_callback( - callback_context=mock_context, + callback_context=callback_context, llm_request=mock_context, error=mock_context, ) await service.run_on_agent_error_callback( agent=mock_context, - callback_context=mock_context, + callback_context=callback_context, error=mock_context, ) await service.run_on_run_error_callback( @@ -286,6 +301,109 @@ async def test_all_callbacks_are_supported( "on_run_error_callback", ] assert set(plugin1.call_log) == set(expected_callbacks) + assert plugin1.callback_info_log == [ + ( + f"{hook.value}_callback", + CallbackInvocationInfo(hook=hook), + ) + for hook in ( + CallbackHook.BEFORE_AGENT, + CallbackHook.AFTER_AGENT, + CallbackHook.BEFORE_TOOL, + CallbackHook.AFTER_TOOL, + CallbackHook.ON_TOOL_ERROR, + CallbackHook.BEFORE_MODEL, + CallbackHook.AFTER_MODEL, + CallbackHook.ON_MODEL_ERROR, + CallbackHook.ON_AGENT_ERROR, + ) + ] + + +@pytest.mark.asyncio +async def test_plugin_callback_receives_current_hook_metadata(): + """Callback metadata identifies the active plugin hook only while it runs.""" + + class _RecordingPlugin(BasePlugin): + + def __init__(self) -> None: + super().__init__(name="recording_plugin") + self.callback_info: CallbackInvocationInfo | None = None + + async def before_model_callback( + self, *, callback_context, llm_request + ) -> None: + self.callback_info = callback_context.callback_info + + invocation_context = Mock() + invocation_context.session.state = {} + invocation_context._state_schema = None + callback_context = CallbackContext(invocation_context) + plugin = _RecordingPlugin() + service = PluginManager(plugins=[plugin]) + + assert callback_context.callback_info is None + + await service.run_before_model_callback( + callback_context=callback_context, + llm_request=Mock(), + ) + + assert plugin.callback_info == CallbackInvocationInfo( + hook=CallbackHook.BEFORE_MODEL + ) + assert callback_context.callback_info is None + + +@pytest.mark.asyncio +async def test_concurrent_plugin_callbacks_keep_hook_metadata_isolated(): + """Concurrent callbacks sharing a context observe only their own hook.""" + model_started = asyncio.Event() + tool_started = asyncio.Event() + model_observed = asyncio.Event() + observations = [] + + class _ConcurrentPlugin(BasePlugin): + + async def before_model_callback( + self, *, callback_context, llm_request + ) -> None: + model_started.set() + await tool_started.wait() + observations.append(callback_context.callback_info.hook) + model_observed.set() + + async def before_tool_callback( + self, *, tool, tool_args, tool_context + ) -> None: + await model_started.wait() + tool_started.set() + await model_observed.wait() + observations.append(tool_context.callback_info.hook) + + invocation_context = Mock() + invocation_context.session.state = {} + invocation_context._state_schema = None + callback_context = CallbackContext(invocation_context) + service = PluginManager(plugins=[_ConcurrentPlugin("concurrent_plugin")]) + + await asyncio.gather( + service.run_before_model_callback( + callback_context=callback_context, + llm_request=Mock(), + ), + service.run_before_tool_callback( + tool=Mock(), + tool_args={}, + tool_context=callback_context, + ), + ) + + assert observations == [ + CallbackHook.BEFORE_MODEL, + CallbackHook.BEFORE_TOOL, + ] + assert callback_context.callback_info is None @pytest.mark.asyncio