diff --git a/contributing/samples/agent_hooks/README.md b/contributing/samples/agent_hooks/README.md new file mode 100644 index 00000000000..b4d8f35cd63 --- /dev/null +++ b/contributing/samples/agent_hooks/README.md @@ -0,0 +1,86 @@ +# Governing an ADK agent with agent-hooks + +This sample shows how to govern an ADK agent with +[agent-hooks](https://github.com/responsibleai/agent-hooks), a framework-neutral +_control_ contract for AI agent systems. You register one or more +**interceptors** (policy engines, content filters, egress guards, ...) once, and +[`AgentHooksPlugin`](../../../src/google/adk/plugins/_agent_hooks_plugin.py) +enforces their verdicts at every governed point in the ADK lifecycle. + +## What it demonstrates + +The agent is a customer-support assistant with two tools: `lookup_account` and +`delete_account`. A single [`ToolGovernanceInterceptor`](governance.py) applies +two policies: + +- **deny** — `delete_account` is destructive, so the interceptor blocks the tool + call before it runs. The model receives a policy error and tells the user it + cannot perform the action. +- **transform** — `lookup_account` returns an `email` and an `api_key`. The + interceptor redacts those fields _before the model or the transcript sees + them_. + +Every decision is recorded as an auditable `InterceptionRecord`; `main.py` +prints the trail at the end. + +## Interception-point mapping + +| ADK plugin callback | agent-hooks point | +| ----------------------------- | ----------------- | +| `before_run_callback` | `agent_startup` | +| `on_user_message_callback` | `input` | +| `before_model_callback` | `pre_model_call` | +| `after_model_callback` | `post_model_call` | +| `before_tool_callback` | `pre_tool_call` | +| `after_tool_callback` | `post_tool_call` | +| `on_event_callback` (final) | `output` | +| `after_run_callback` | `agent_shutdown` | + +## Prerequisites + +1. Install ADK with the optional `agent-hooks` extra, plus LiteLLM for the local + model: + + ```bash + pip install "google-adk[agent-hooks]" litellm + ``` + +2. Install [Ollama](https://ollama.com/) and pull a tool-capable model: + + ```bash + ollama pull qwen2.5 + ``` + + The example runs against a real local model, so tool-calling behavior is not + scripted. Any tool-capable Ollama model works; edit the `LiteLlm(model=...)` + line in [`agent.py`](agent.py) to change it. + +## Run + +```bash +python -m contributing.samples.agent_hooks.main +``` + +Expected shape of the output: + +- For "look up account 42", the agent calls `lookup_account` and summarizes the + result — with `email` and `api_key` already redacted. +- For "delete account 42", the `delete_account` call is denied and the agent + explains it cannot delete the account. +- The audit trail lists every interception point and its verdict, including the + `pre_tool_call -> deny` and `post_tool_call -> transform` decisions. + +## Enforcement semantics + +`AgentHooksPlugin` **fails closed**: a `deny` blocks the guarded action, a +`transform` rewrites the guarded value, and any engine error, malformed verdict, +or interceptor timeout becomes a fail-closed block — it never fails open. Set +`mode="evaluate_only"` on the plugin to record decisions without enforcing them. + +## Trust model + +agent-hooks is a _cooperative_ control contract, **not** a security boundary: +interceptors run in-process with full data access and the interception points do +not guarantee complete mediation. See the agent-hooks +[`SECURITY.md`](https://github.com/responsibleai/agent-hooks/blob/main/SECURITY.md) +before relying on it for isolation. diff --git a/contributing/samples/agent_hooks/__init__.py b/contributing/samples/agent_hooks/__init__.py new file mode 100644 index 00000000000..4015e47d6e4 --- /dev/null +++ b/contributing/samples/agent_hooks/__init__.py @@ -0,0 +1,15 @@ +# 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 . import agent diff --git a/contributing/samples/agent_hooks/agent.py b/contributing/samples/agent_hooks/agent.py new file mode 100644 index 00000000000..2f37feda1cc --- /dev/null +++ b/contributing/samples/agent_hooks/agent.py @@ -0,0 +1,74 @@ +# 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. + +"""A customer-support agent with two tools, one of them destructive. + +The agent runs on a local Ollama model so the example exercises real model +behavior (see README.md). agent-hooks governance is wired in ``main.py``: the +``delete_account`` tool call is denied, and ``lookup_account`` results are +redacted, before the model ever sees them. +""" + +from __future__ import annotations + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.models.lite_llm import LiteLlm + + +def lookup_account(user_id: str) -> dict: + """Looks up a customer account. + + Args: + user_id: The id of the account to look up. + + Returns: + The account record, including fields the governance policy will redact. + """ + return { + "user_id": user_id, + "name": "Alice Example", + "email": "alice@example.com", + "api_key": "EXAMPLE_NOT_A_REAL_KEY", + "plan": "pro", + } + + +def delete_account(user_id: str) -> dict: + """Permanently deletes a customer account. + + This is a destructive tool; the governance policy denies it before it runs. + + Args: + user_id: The id of the account to delete. + + Returns: + A confirmation record (never reached under the governance policy). + """ + return {"user_id": user_id, "status": "deleted"} + + +root_agent = LlmAgent( + name="support_agent", + model=LiteLlm(model="ollama_chat/qwen2.5:latest"), + description="A customer-support agent guarded by agent-hooks.", + instruction=( + "You are a customer-support assistant. Always use the available tools" + " to fulfill the user's request: call lookup_account to read an account" + " and call delete_account when the user asks to delete one. After a" + " tool returns, summarize its result for the user. If a tool result" + " reports that it was blocked by policy, tell the user you were not" + " allowed to perform that action." + ), + tools=[lookup_account, delete_account], +) diff --git a/contributing/samples/agent_hooks/governance.py b/contributing/samples/agent_hooks/governance.py new file mode 100644 index 00000000000..2214cedf775 --- /dev/null +++ b/contributing/samples/agent_hooks/governance.py @@ -0,0 +1,106 @@ +# 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. + +"""An example agent-hooks interceptor: a small tool-governance policy. + +The interceptor implements the ``intercept(AgentContext) -> Verdict`` contract. +It demonstrates the two enforcement primitives that matter most for tool use: + +- ``deny``: a destructive tool (``delete_account``) is blocked before it runs. +- ``transform``: sensitive fields returned by a tool are redacted before the + model (and the transcript) ever see them. + +An interceptor is framework-neutral: this same class works against any +agent-hooks host (ADK, crewAI, ...), not just ADK. +""" + +from __future__ import annotations + +import re +from typing import Any + +from agent_hooks import AgentContext +from agent_hooks import Decision +from agent_hooks import Transform +from agent_hooks import Verdict + +#: Tools that must never execute under this policy. +_DENIED_TOOLS = frozenset({"delete_account"}) + +#: Result fields whose values are masked before the model sees them. +_SENSITIVE_KEYS = frozenset( + {"email", "api_key", "password", "secret", "token", "ssn"} +) + +_EMAIL_RE = re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+") + + +def _redact(value: Any) -> tuple[Any, bool]: + """Return ``(redacted_value, changed)`` for a tool result. + + Masks the values of sensitive keys and any email address found in a string. + """ + changed = False + + def walk(node: Any) -> Any: + nonlocal changed + if isinstance(node, dict): + out: dict[str, Any] = {} + for key, item in node.items(): + if key in _SENSITIVE_KEYS and isinstance(item, str): + out[key] = "[REDACTED]" + changed = True + else: + out[key] = walk(item) + return out + if isinstance(node, list): + return [walk(item) for item in node] + if isinstance(node, str): + masked = _EMAIL_RE.sub("[REDACTED_EMAIL]", node) + if masked != node: + changed = True + return masked + return node + + return walk(value), changed + + +class ToolGovernanceInterceptor: + """Deny destructive tools and redact sensitive tool results.""" + + name = "tool_governance" + + def intercept(self, ctx: AgentContext) -> Verdict: + point = ctx["interception_point"] + + if point == "pre_tool_call": + tool_name = ctx["tool_call"]["name"] + if tool_name in _DENIED_TOOLS: + return Verdict.deny( + reason="tool_denied", + message=f"Tool '{tool_name}' is disabled by policy.", + ) + return Verdict(decision=Decision.ALLOW) + + if point == "post_tool_call": + # ``ctx["target"]`` is the tool result value at post_tool_call. + redacted, changed = _redact(ctx["target"]) + if changed: + return Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target", value=redacted), + ) + return Verdict(decision=Decision.ALLOW) + + return Verdict(decision=Decision.ALLOW) diff --git a/contributing/samples/agent_hooks/main.py b/contributing/samples/agent_hooks/main.py new file mode 100644 index 00000000000..5f0fdd38282 --- /dev/null +++ b/contributing/samples/agent_hooks/main.py @@ -0,0 +1,92 @@ +# 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. + +"""Runs the support agent with agent-hooks governance enabled. + +Prerequisites (see README.md): + * Ollama running locally with the ``qwen2.5:7b`` model pulled. + * ``pip install "google-adk[agent-hooks]" litellm`` + +Run: + python -m contributing.samples.agent_hooks.main +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from google.adk.apps.app import App +from google.adk.plugins import AgentHooksPlugin +from google.adk.runners import InMemoryRunner +from google.genai import types + +from .agent import root_agent +from .governance import ToolGovernanceInterceptor + +_APP_NAME = "agent_hooks_demo" + + +async def main() -> None: + """Runs two prompts: one benign (redacted), one destructive (denied).""" + # ``record_sink`` receives an auditable InterceptionRecord per decision. + records: list[Any] = [] + plugin = AgentHooksPlugin( + interceptors=[ToolGovernanceInterceptor()], + record_sink=records.append, + ) + + app = App(name=_APP_NAME, root_agent=root_agent, plugins=[plugin]) + runner = InMemoryRunner(app=app) + session = await runner.session_service.create_session( + user_id="user", app_name=_APP_NAME + ) + + prompts = [ + "Look up the account details for user 42.", + "Now delete account 42.", + ] + for prompt in prompts: + print(f"\n=== USER: {prompt} ===") + async for event in runner.run_async( + user_id="user", + session_id=session.id, + new_message=types.Content( + role="user", parts=[types.Part.from_text(text=prompt)] + ), + ): + if event.content and event.content.parts: + for part in event.content.parts: + if part.text: + print(f"[{event.author}] {part.text}") + if part.function_call: + print(f"[{event.author}] -> tool call: {part.function_call.name}") + if part.function_response: + print( + f"[{event.author}] <- tool result:" + f" {part.function_response.response}" + ) + + print("\n=== agent-hooks audit trail ===") + for record in records: + verdict = record.verdict + print( + f"seq={record.sequence:<2} {record.interception_point.value:<16}" + f" -> {verdict.decision.value}" + + (f" ({verdict.reason})" if verdict.reason else "") + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/pyproject.toml b/pyproject.toml index 21cd6a1f11d..29503a49d68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,11 @@ dependencies = [ optional-dependencies.a2a = [ "a2a-sdk>=0.3.4,<2", ] +optional-dependencies.agent-hooks = [ + # Framework-neutral agent lifecycle governance contract with a compiled + # native core; used by google.adk.plugins.AgentHooksPlugin. + "agent-hooks-sdk>=0.1.0a4", +] optional-dependencies.agent-identity = [ "google-cloud-agentidentitycredentials>=0.1,<0.2", "google-cloud-iamconnectorcredentials>=0.1,<0.2", diff --git a/src/google/adk/plugins/__init__.py b/src/google/adk/plugins/__init__.py index 893d3dd7c8c..f33a4877790 100644 --- a/src/google/adk/plugins/__init__.py +++ b/src/google/adk/plugins/__init__.py @@ -21,12 +21,14 @@ from .plugin_manager import PluginManager if TYPE_CHECKING: + from ._agent_hooks_plugin import AgentHooksPlugin from ._reflect_retry_model_plugin import ReflectAndRetryModelPlugin from .debug_logging_plugin import DebugLoggingPlugin from .logging_plugin import LoggingPlugin from .reflect_retry_tool_plugin import ReflectAndRetryToolPlugin __all__ = [ + "AgentHooksPlugin", "BasePlugin", "DebugLoggingPlugin", "LoggingPlugin", @@ -36,6 +38,7 @@ ] _LAZY_MEMBERS: dict[str, str] = { + "AgentHooksPlugin": "_agent_hooks_plugin", "DebugLoggingPlugin": "debug_logging_plugin", "LoggingPlugin": "logging_plugin", "ReflectAndRetryModelPlugin": "_reflect_retry_model_plugin", diff --git a/src/google/adk/plugins/_agent_hooks_plugin.py b/src/google/adk/plugins/_agent_hooks_plugin.py new file mode 100644 index 00000000000..df8ea1ee8b2 --- /dev/null +++ b/src/google/adk/plugins/_agent_hooks_plugin.py @@ -0,0 +1,758 @@ +# 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. + +"""Governance plugin that makes ADK a host for the agent-hooks contract. + +`agent-hooks `_ is a +framework-neutral *control* contract: a fixed set of agent lifecycle +interception points, the ``AgentContext`` a host supplies at each, and the +``Verdict`` an interceptor returns (``allow`` / ``deny`` / ``transform``). This +plugin turns ADK's :class:`~google.adk.plugins.base_plugin.BasePlugin` seam +into a conformant host: you register one or more interceptors (policy engines, +content filters, rate limiters, egress guards) once, and every governed ADK +lifecycle point delegates its decision to the agent-hooks emitter, which runs +the interceptors and records every decision as an auditable +``InterceptionRecord``. + +Interception-point mapping (ADK callback -> agent-hooks point): + ``before_run_callback`` -> ``agent_startup`` + ``on_user_message_callback`` -> ``input`` + ``before_model_callback`` -> ``pre_model_call`` + ``after_model_callback`` -> ``post_model_call`` + ``before_tool_callback`` -> ``pre_tool_call`` + ``after_tool_callback`` -> ``post_tool_call`` + ``on_event_callback`` (final) -> ``output`` + ``after_run_callback`` -> ``agent_shutdown`` + +Because ADK invokes the user-message seam before the run seam, ``input`` is +emitted before ``agent_startup`` in a turn; the agent-hooks ``sequence`` field +reflects that real ADK order. + +Enforcement semantics (fail closed): + - ``deny`` blocks the guarded action. A blocked model call returns a + refusal ``LlmResponse``; a blocked tool returns an error result; a + blocked run/input/output surfaces a refusal ``Content``. + - ``transform`` rewrites the guarded value from the interceptor's + ``$target`` transform: tool args (pre-tool), tool result (post-tool), + user text (input), model text (post-model), or final output text. + - Any engine-internal error, malformed verdict, or interceptor + timeout is turned into a fail-closed *deny* by the emitter, and this + plugin surfaces it as a blocked action — it never fails open. + - ``pre_model_call`` supports ``allow`` / ``deny``; a ``transform`` there + is treated as a fail-closed deny, because rebuilding a provider-native + request from wire messages is not round-trip safe. + +Trust model: agent-hooks is a *cooperative* control contract, **not** a +security boundary. Interceptors run in-process with full data access and the +interception points do not guarantee complete mediation. See the agent-hooks +``SECURITY.md`` before relying on it for isolation. + +``agent-hooks`` is an **optional** dependency with a compiled native core; +install it with ``pip install "google-adk[agent-hooks]"``. Importing this +module never requires it — the import is deferred to +:class:`AgentHooksPlugin` construction, which raises an actionable +``ImportError`` when the package is missing. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import importlib +import json +import logging +import math +from typing import Any +from typing import Optional +from typing import TYPE_CHECKING + +from google.genai import types +from typing_extensions import override + +from ..agents.callback_context import CallbackContext +from ..events.event import Event +from ..models.llm_request import LlmRequest +from ..models.llm_response import LlmResponse +from ..tools.base_tool import BaseTool +from .base_plugin import BasePlugin + +if TYPE_CHECKING: + from collections.abc import Callable + from collections.abc import Sequence + + from agent_hooks import AgentContext + from agent_hooks import AgentContextBuilder + from agent_hooks import CompositionConfig + from agent_hooks import IdentityProvider + from agent_hooks import InterceptionEmitter + from agent_hooks import InterceptionRecord + from agent_hooks import Interceptor + + from ..agents.invocation_context import InvocationContext + from ..tools.tool_context import ToolContext + +logger = logging.getLogger("google_adk." + __name__) + +#: ``framework`` identifier stamped on every emitted ``AgentContext``. +_FRAMEWORK = "google-adk" + +#: Default identity provider name understood by the emitter (agent-hooks §10.2). +_DEFAULT_IDENTITY = "jcs-sha256" + +#: Default bound on the emitter's in-memory record buffer per invocation. +_DEFAULT_MAX_RECORDS = 1000 + +_INSTALL_HINT = ( + "agent-hooks is not installed (or its native core failed to load). It is " + "an optional dependency with a compiled core; install it with:\n" + ' pip install "google-adk[agent-hooks]"\n' + "See https://github.com/responsibleai/agent-hooks for details." +) + + +def _require_agent_hooks() -> Any: + """Import and return the ``agent_hooks`` module, or fail actionably. + + Raises: + ImportError: If the ``agent_hooks`` package (or its compiled core) cannot + be imported. + """ + try: + return importlib.import_module("agent_hooks") + except Exception as exc: # ImportError or a native-core load failure. + raise ImportError(_INSTALL_HINT) from exc + + +#: Maximum nesting depth normalized by :func:`_json_safe`; deeper structures +#: are truncated so an untrusted, deeply nested payload cannot raise +#: ``RecursionError`` before a decision is made. +_MAX_JSON_DEPTH = 100 + + +def _json_safe( + value: Any, _seen: frozenset[int] = frozenset(), _depth: int = 0 +) -> Any: + """Coerce ``value`` into a JSON-serializable form for an ``AgentContext``. + + Tool arguments, tool results, and model output are untrusted and may hold + objects the agent-hooks wire format cannot carry (which would otherwise fail + the emission closed). This normalizes them at the boundary so an interceptor + always sees a stable, inspectable value. Reference cycles are broken with a + ``""`` placeholder and nesting past ``_MAX_JSON_DEPTH`` is truncated to + ``""``, so a self-referential or deeply nested container cannot + raise ``RecursionError`` before a decision is made. + """ + if _depth > _MAX_JSON_DEPTH: + return "" + if value is None or isinstance(value, (bool, int, str)): + return value + if isinstance(value, float): + return value if math.isfinite(value) else str(value) + if isinstance(value, dict): + if id(value) in _seen: + return "" + seen = _seen | {id(value)} + return {str(k): _json_safe(v, seen, _depth + 1) for k, v in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + if id(value) in _seen: + return "" + seen = _seen | {id(value)} + return [_json_safe(v, seen, _depth + 1) for v in value] + if isinstance(value, (bytes, bytearray)): + return bytes(value).decode("utf-8", errors="replace") + dump = getattr(value, "model_dump", None) + if callable(dump): + try: + return _json_safe(dump(mode="json"), _seen, _depth + 1) + except Exception: + logger.debug("model_dump() failed while normalizing %s", type(value)) + return str(value) + + +def _content_text(content: Optional[types.Content]) -> str: + """Join the text parts of a ``types.Content`` into a single string.""" + if content is None or not content.parts: + return "" + return "".join(part.text for part in content.parts if part.text) + + +def _text_content(text: str, *, role: str) -> types.Content: + """Build a single-text-part ``types.Content`` with the given role.""" + return types.Content(role=role, parts=[types.Part.from_text(text=text)]) + + +def _target_text(target: Any) -> str: + """Extract replacement text from a transformed ``input``/``output`` target. + + The ``input``/``output`` L1 envelopes wrap the value as ``{"content": ...}``; + an interceptor may transform the whole envelope or replace ``$target`` with a + bare value. + """ + if isinstance(target, dict) and "content" in target: + target = target["content"] + return target if isinstance(target, str) else json.dumps(target, default=str) + + +def _request_messages(llm_request: LlmRequest) -> list[dict[str, Any]]: + """Project an ``LlmRequest`` into inspectable ``{role, content}`` messages.""" + messages: list[dict[str, Any]] = [] + system = getattr(llm_request.config, "system_instruction", None) + if isinstance(system, str) and system: + messages.append({"role": "system", "content": system}) + for content in llm_request.contents: + messages.append({ + "role": content.role or "user", + "content": _content_text(content), + }) + return messages + + +def _response_tool_calls(response: LlmResponse) -> list[dict[str, Any]]: + """Surface the function calls a model response requested.""" + if response.content is None or not response.content.parts: + return [] + calls: list[dict[str, Any]] = [] + for part in response.content.parts: + call = part.function_call + if call is not None: + calls.append({ + "id": call.id or "", + "name": call.name or "", + "args": _json_safe(dict(call.args) if call.args else {}), + }) + return calls + + +def _response_finish_reason( + response: LlmResponse, *, has_tool_calls: bool +) -> str: + """Best-effort finish reason for the ``post_model_call`` context.""" + reason = response.finish_reason + if reason is not None: + return getattr(reason, "name", str(reason)).lower() + return "tool_calls" if has_tool_calls else "stop" + + +def _synth_call_id(name: str, args: dict[str, Any]) -> str: + """Deterministic tool-call id when ADK supplies none. + + Derived from ``name`` and ``args``. The plugin stashes the pre-tool id so the + matching post-tool record reuses it, because a pre-tool transform may rewrite + the args in place before the post-tool point is reached. + """ + digest = hashlib.sha256() + digest.update(name.encode("utf-8")) + digest.update(json.dumps(args, sort_keys=True, default=str).encode("utf-8")) + return f"tc-{digest.hexdigest()[:16]}" + + +def _verdict_reason(record: Optional[InterceptionRecord]) -> str: + """Human-readable reason for a blocked action (payload-free).""" + if record is None: + return "agent-hooks engine error (failing closed)" + verdict = record.verdict + reason = (verdict.reason or "").strip() + message = (verdict.message or "").strip() + if reason and message and reason != message: + return f"{reason}: {message}" + return reason or message or "blocked by agent-hooks policy" + + +class _InvocationState: + """Per-invocation agent-hooks builder and emitter. + + One instance exists per ADK invocation id; it owns the monotonic + ``sequence`` for that turn's records and is evicted when the invocation + ends (or errors) so long-running processes do not leak state. + """ + + __slots__ = ( + "builder", + "emitter", + "startup_denied", + "startup_record", + "synth_call_ids", + ) + + def __init__( + self, builder: AgentContextBuilder, emitter: InterceptionEmitter + ) -> None: + self.builder = builder + self.emitter = emitter + # Set when agent_startup denies, so the deny is re-enforced at the first + # model call on runner paths that ignore before_run_callback's return. + self.startup_denied: bool = False + self.startup_record: Optional[InterceptionRecord] = None + # Synthesized pre-tool call ids, queued per tool name so the matching + # post-tool record reuses the same id after an in-place args transform. + self.synth_call_ids: dict[str, list[str]] = {} + + +class AgentHooksPlugin(BasePlugin): + """ADK plugin that enforces agent-hooks interceptors across the lifecycle. + + Register the plugin on the ``Runner`` with one or more interceptors; each + governed ADK callback emits an ``AgentContext`` to the agent-hooks emitter + and enforces the returned verdict, failing closed on any error. + + Example: + >>> from google.adk.plugins import AgentHooksPlugin + >>> from agent_hooks import AgentContext, Decision, Verdict + >>> + >>> class BlockDangerousTools: + ... def intercept(self, ctx: AgentContext) -> Verdict: + ... if ( + ... ctx["interception_point"] == "pre_tool_call" + ... and ctx["tool_call"]["name"] == "delete_account" + ... ): + ... return Verdict.deny(reason="tool_denied") + ... return Verdict(decision=Decision.ALLOW) + >>> + >>> plugin = AgentHooksPlugin(interceptors=[BlockDangerousTools()]) + >>> # runner = InMemoryRunner(agent=root_agent, plugins=[plugin]) + """ + + def __init__( + self, + interceptors: Sequence[Interceptor], + *, + name: str = "agent_hooks", + mode: str = "enforce", + timeout: Optional[float] = 5.0, + composition: Optional[CompositionConfig] = None, + identity_provider: Optional[str | IdentityProvider] = _DEFAULT_IDENTITY, + record_sink: Optional[Callable[[InterceptionRecord], None]] = None, + max_records: Optional[int] = _DEFAULT_MAX_RECORDS, + ) -> None: + """Initializes the plugin. + + Args: + interceptors: The interceptors to run at every governed point, in + registration order. Each is an object with an + ``intercept(AgentContext) -> Verdict`` method (sync or async). + name: Unique identifier for this plugin instance. + mode: ``"enforce"`` (act on verdicts) or ``"evaluate_only"`` (record + decisions without blocking or transforming). + timeout: Per-interceptor timeout in seconds; ``None`` disables it. + composition: agent-hooks composition profile; defaults to + ``sequential/first_deny``. + identity_provider: agent-hooks identity provider for audit records; + ``"jcs-sha256"`` by default, or ``None`` for identity-unbound records. + record_sink: Optional callback invoked with every ``InterceptionRecord`` + for audit persistence. A sink exception is swallowed by the emitter. + max_records: Bound on the per-invocation in-memory record buffer. + + Raises: + ImportError: If the optional ``agent-hooks`` package is not installed. + """ + super().__init__(name) + ah = _require_agent_hooks() + self._ah = ah + self._interceptors: list[Any] = list(interceptors) + self._mode = ah.EnforcementMode(mode) + self._enforcing = self._mode == ah.EnforcementMode("enforce") + self._timeout = timeout + self._composition = composition + self._identity_provider = identity_provider + self._record_sink = record_sink + self._max_records = max_records + self._states: dict[str, _InvocationState] = {} + + # --- state lifecycle ------------------------------------------------------- + + def _new_state( + self, *, invocation_id: str, session_id: str, agent_name: str + ) -> _InvocationState: + ah = self._ah + builder = ah.AgentContextBuilder( + agent_id=agent_name, + framework=_FRAMEWORK, + session_id=session_id, + agent_name=agent_name, + ) + emitter = ah.InterceptionEmitter( + mode=self._mode, + timeout=self._timeout, + composition=self._composition, + identity_provider=self._identity_provider, + ) + for interceptor in self._interceptors: + emitter.register(interceptor, type(interceptor).__name__) + if self._record_sink is not None: + emitter.set_record_sink(self._record_sink) + if self._max_records is not None: + emitter.set_max_records(self._max_records) + state = _InvocationState(builder, emitter) + self._states[invocation_id] = state + return state + + def _state_for_invocation( + self, invocation_context: InvocationContext + ) -> _InvocationState: + agent = invocation_context.agent + agent_name = agent.name if agent is not None else "unknown" + state = self._states.get(invocation_context.invocation_id) + if state is None: + state = self._new_state( + invocation_id=invocation_context.invocation_id, + session_id=invocation_context.session.id, + agent_name=agent_name, + ) + return state + + def _state_for_context( + self, context: CallbackContext | ToolContext + ) -> _InvocationState: + state = self._states.get(context.invocation_id) + if state is None: + state = self._new_state( + invocation_id=context.invocation_id, + session_id=context.session.id, + agent_name=context.agent_name, + ) + return state + + def _agent_envelope(self, agent_name: str) -> dict[str, Any]: + return {"id": agent_name, "framework": _FRAMEWORK, "name": agent_name} + + async def _emit( + self, state: _InvocationState, ctx: AgentContext, *, agent_name: str + ) -> Optional[InterceptionRecord]: + """Emit ``ctx`` and return the record, or ``None`` on engine failure. + + A ``None`` return signals the caller to fail closed. ``CancelledError`` + is propagated unchanged so task cancellation is honoured. + """ + ctx["agent"] = self._agent_envelope(agent_name) + try: + return await state.emitter.emit_unchecked(ctx) + except asyncio.CancelledError: + raise + except Exception: + logger.exception( + "agent-hooks emission failed at %s; failing closed", + ctx.get("interception_point"), + ) + return None + + @staticmethod + def _blocked(record: Optional[InterceptionRecord]) -> bool: + """Whether the record denies the guarded action (or is an engine error).""" + return record is None or not record.proceeds + + def _is_transform(self, record: InterceptionRecord) -> bool: + # Only enforce mode actually rewrites ``ctx["target"]``; evaluate_only + # records the verdict without transforming, so applying it here is lossy. + return self._enforcing and record.verdict.transform is not None + + def _pre_tool_call_id( + self, + state: _InvocationState, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + ) -> str: + """Call id for the pre-tool record. + + When ADK supplies no ``function_call_id`` the id is synthesized and stashed + so the matching post-tool record can reuse it even after a pre-tool + transform rewrites the args in place. + """ + function_call_id = tool_context.function_call_id + if function_call_id: + return function_call_id + call_id = _synth_call_id(tool.name, tool_args) + state.synth_call_ids.setdefault(tool.name, []).append(call_id) + return call_id + + def _post_tool_call_id( + self, + state: _InvocationState, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + ) -> str: + """Call id for the post-tool record, correlated with its pre-tool id.""" + function_call_id = tool_context.function_call_id + if function_call_id: + return function_call_id + pending = state.synth_call_ids.get(tool.name) + if pending: + return pending.pop(0) + return _synth_call_id(tool.name, tool_args) + + # --- lifecycle callbacks --------------------------------------------------- + + @override + async def before_run_callback( + self, *, invocation_context: InvocationContext + ) -> Optional[types.Content]: + """agent_startup: deny halts the run with a refusal message.""" + state = self._state_for_invocation(invocation_context) + agent = invocation_context.agent + agent_name = agent.name if agent is not None else "unknown" + ctx = state.builder.agent_startup(tools_registered=self._tool_names(agent)) + record = await self._emit(state, ctx, agent_name=agent_name) + if self._blocked(record): + # Some runner paths ignore this return value; the flag makes the deny + # stick by also blocking the first model call (see before_model_callback). + state.startup_denied = True + state.startup_record = record + return _text_content( + f"[blocked by agent-hooks: {_verdict_reason(record)}]", role="model" + ) + return None + + @override + async def on_user_message_callback( + self, + *, + invocation_context: InvocationContext, + user_message: types.Content, + ) -> Optional[types.Content]: + """input: deny replaces the user message; transform rewrites it.""" + state = self._state_for_invocation(invocation_context) + agent = invocation_context.agent + agent_name = agent.name if agent is not None else "unknown" + ctx = state.builder.input(content=_content_text(user_message)) + record = await self._emit(state, ctx, agent_name=agent_name) + if self._blocked(record): + return _text_content( + f"[input blocked by agent-hooks: {_verdict_reason(record)}]", + role="user", + ) + if record is not None and self._is_transform(record): + return _text_content(_target_text(ctx.get("target")), role="user") + return None + + @override + async def before_model_callback( + self, *, callback_context: CallbackContext, llm_request: LlmRequest + ) -> Optional[LlmResponse]: + """pre_model_call: deny (or transform, treated as deny) blocks the call.""" + state = self._state_for_context(callback_context) + if state.startup_denied: + # A denied agent_startup halts the run even on runner paths that ignore + # before_run_callback's return value. + return self._blocked_response(state.startup_record) + ctx = state.builder.pre_model_call( + model_id=llm_request.model or "unknown", + messages=_request_messages(llm_request), + ) + record = await self._emit( + state, ctx, agent_name=callback_context.agent_name + ) + if self._blocked(record): + return self._blocked_response(record) + if record is not None and self._is_transform(record): + logger.warning( + "agent-hooks pre_model_call transform is not applied to the " + "provider request; failing closed" + ) + return self._blocked_response(record) + return None + + @override + async def after_model_callback( + self, *, callback_context: CallbackContext, llm_response: LlmResponse + ) -> Optional[LlmResponse]: + """post_model_call: deny blocks; transform rewrites the response text.""" + state = self._state_for_context(callback_context) + tool_calls = _response_tool_calls(llm_response) + ctx = state.builder.post_model_call( + model_id=llm_response.model_version or "unknown", + content=_content_text(llm_response.content), + tool_calls=tool_calls, + finish_reason=_response_finish_reason( + llm_response, has_tool_calls=bool(tool_calls) + ), + ) + record = await self._emit( + state, ctx, agent_name=callback_context.agent_name + ) + if self._blocked(record): + return self._blocked_response(record) + if record is not None and self._is_transform(record): + return llm_response.model_copy( + update={ + "content": _text_content( + _target_text(ctx.get("target")), role="model" + ) + } + ) + return None + + @override + async def before_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + ) -> Optional[dict[str, Any]]: + """pre_tool_call: deny blocks the tool; transform rewrites the args.""" + state = self._state_for_context(tool_context) + call_id = self._pre_tool_call_id(state, tool, tool_args, tool_context) + ctx = state.builder.pre_tool_call( + call_id=call_id, name=tool.name, args=_json_safe(tool_args) + ) + record = await self._emit(state, ctx, agent_name=tool_context.agent_name) + if self._blocked(record): + return self._blocked_tool_result(record) + if record is not None and self._is_transform(record): + new_args = ctx.get("target") + if not isinstance(new_args, dict): + logger.warning( + "agent-hooks pre_tool_call transform did not yield an args " + "object; failing closed" + ) + return self._blocked_tool_result(record) + # Mutating tool_args in place propagates to the actual tool call. + tool_args.clear() + tool_args.update(new_args) + return None + + @override + async def after_tool_callback( + self, + *, + tool: BaseTool, + tool_args: dict[str, Any], + tool_context: ToolContext, + result: dict[str, Any], + ) -> Optional[dict[str, Any]]: + """post_tool_call: deny blocks; transform replaces the tool result.""" + state = self._state_for_context(tool_context) + call_id = self._post_tool_call_id(state, tool, tool_args, tool_context) + ctx = state.builder.post_tool_call( + call_id=call_id, + name=tool.name, + args=_json_safe(tool_args), + value=_json_safe(result), + ) + record = await self._emit(state, ctx, agent_name=tool_context.agent_name) + if self._blocked(record): + return self._blocked_tool_result(record) + if record is not None and self._is_transform(record): + new_value = ctx.get("target") + return new_value if isinstance(new_value, dict) else {"result": new_value} + return None + + @override + async def on_event_callback( + self, *, invocation_context: InvocationContext, event: Event + ) -> Optional[Event]: + """output: govern the final response event (deny/transform its content).""" + if not event.is_final_response(): + return None + state = self._state_for_invocation(invocation_context) + agent = invocation_context.agent + agent_name = event.author or ( + agent.name if agent is not None else "unknown" + ) + ctx = state.builder.output(content=_content_text(event.content)) + record = await self._emit(state, ctx, agent_name=agent_name) + if self._blocked(record): + return event.model_copy( + update={ + "content": _text_content( + f"[output blocked by agent-hooks: {_verdict_reason(record)}]", + role="model", + ) + } + ) + if record is not None and self._is_transform(record): + return event.model_copy( + update={ + "content": _text_content( + _target_text(ctx.get("target")), role="model" + ) + } + ) + return None + + @override + async def after_run_callback( + self, *, invocation_context: InvocationContext + ) -> None: + """agent_shutdown: emit for audit, then evict per-invocation state.""" + state = self._states.pop(invocation_context.invocation_id, None) + if state is None: + return None + agent = invocation_context.agent + agent_name = agent.name if agent is not None else "unknown" + try: + ctx = state.builder.agent_shutdown(reason="completed") + await self._emit(state, ctx, agent_name=agent_name) + except Exception: + logger.debug("agent_shutdown emission failed", exc_info=True) + return None + + @override + async def on_run_error_callback( + self, *, invocation_context: InvocationContext, error: Exception + ) -> None: + """Evict per-invocation state on an error path (notification-only).""" + state = self._states.pop(invocation_context.invocation_id, None) + if state is None: + return None + agent = invocation_context.agent + agent_name = agent.name if agent is not None else "unknown" + try: + ctx = state.builder.agent_shutdown( + reason="error", error=type(error).__name__ + ) + await self._emit(state, ctx, agent_name=agent_name) + except Exception: + logger.debug("agent_shutdown emission failed on run error", exc_info=True) + return None + + @override + async def close(self) -> None: + """Drop any residual per-invocation state.""" + self._states.clear() + + # --- block-result builders ------------------------------------------------- + + def _blocked_response( + self, record: Optional[InterceptionRecord] + ) -> LlmResponse: + reason = _verdict_reason(record) + return LlmResponse( + content=_text_content( + f"[blocked by agent-hooks: {reason}]", role="model" + ), + custom_metadata={"agent_hooks_blocked": True, "reason": reason}, + ) + + def _blocked_tool_result( + self, record: Optional[InterceptionRecord] + ) -> dict[str, Any]: + reason = _verdict_reason(record) + return { + "error": f"blocked by agent-hooks: {reason}", + "agent_hooks_blocked": True, + "reason": reason, + } + + @staticmethod + def _tool_names(agent: Any) -> list[str]: + """Best-effort declared tool names for the ``agent_startup`` context.""" + names: list[str] = [] + tools = getattr(agent, "tools", None) or [] + for tool in tools: + name = getattr(tool, "name", None) or getattr(tool, "__name__", None) + if isinstance(name, str) and name not in names: + names.append(name) + return names diff --git a/tests/unittests/plugins/test_agent_hooks_plugin.py b/tests/unittests/plugins/test_agent_hooks_plugin.py new file mode 100644 index 00000000000..677b1c58c98 --- /dev/null +++ b/tests/unittests/plugins/test_agent_hooks_plugin.py @@ -0,0 +1,614 @@ +# 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 AgentHooksPlugin, the agent-hooks governance host. + +The optional ``agent_hooks`` package (with its compiled native core) is +required; the whole module is skipped when it is not importable. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import Mock + +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.genai import types +import pytest + +agent_hooks = pytest.importorskip("agent_hooks") + +from google.adk.plugins import AgentHooksPlugin # noqa: E402 + +AgentContext = agent_hooks.AgentContext +Decision = agent_hooks.Decision +Transform = agent_hooks.Transform +Verdict = agent_hooks.Verdict + + +# --------------------------------------------------------------------------- +# Interceptors (real agent-hooks verdicts) +# --------------------------------------------------------------------------- + + +class _AllowAll: + + def intercept(self, ctx: AgentContext) -> Any: + return Verdict(decision=Decision.ALLOW) + + +class _DenyTool: + + def __init__(self, tool_name: str) -> None: + self._tool_name = tool_name + + def intercept(self, ctx: AgentContext) -> Any: + if ( + ctx["interception_point"] == "pre_tool_call" + and ctx["tool_call"]["name"] == self._tool_name + ): + return Verdict.deny(reason="tool_denied", message="not allowed") + return Verdict(decision=Decision.ALLOW) + + +class _TransformToolArgs: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "pre_tool_call": + new_args = dict(ctx["tool_call"]["args"]) + new_args["redacted"] = True + return Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target", value=new_args), + ) + return Verdict(decision=Decision.ALLOW) + + +class _TransformToolResult: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "post_tool_call": + return Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target", value={"clean": "value"}), + ) + return Verdict(decision=Decision.ALLOW) + + +class _DenyModel: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "pre_model_call": + return Verdict.deny(reason="model_denied") + return Verdict(decision=Decision.ALLOW) + + +class _TransformModel: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "pre_model_call": + return Verdict( + decision=Decision.TRANSFORM, + transform=Transform( + path="$target", value=[{"role": "user", "content": "x"}] + ), + ) + if ctx["interception_point"] == "post_model_call": + return Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target.content", value="SAFE"), + ) + return Verdict(decision=Decision.ALLOW) + + +class _DenyInput: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "input": + return Verdict.deny(reason="input_denied") + return Verdict(decision=Decision.ALLOW) + + +class _TransformInput: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "input": + return Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target.content", value="CLEANED"), + ) + return Verdict(decision=Decision.ALLOW) + + +class _DenyOutput: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "output": + return Verdict.deny(reason="output_denied") + return Verdict(decision=Decision.ALLOW) + + +class _Raiser: + + def intercept(self, ctx: AgentContext) -> Any: + raise RuntimeError("boom") + + +# --------------------------------------------------------------------------- +# Context factories (the plugin only reads a few attributes) +# --------------------------------------------------------------------------- + + +def _invocation_context( + *, + invocation_id: str = "inv-1", + session_id: str = "sess-1", + agent_name: str = "agent", + tools: list[Any] | None = None, +) -> Any: + agent = Mock() + agent.name = agent_name + agent.tools = tools or [] + session = Mock() + session.id = session_id + ic = Mock() + ic.invocation_id = invocation_id + ic.agent = agent + ic.session = session + return ic + + +def _callback_context( + *, + invocation_id: str = "inv-1", + session_id: str = "sess-1", + agent_name: str = "agent", +) -> Any: + session = Mock() + session.id = session_id + ctx = Mock() + ctx.invocation_id = invocation_id + ctx.session = session + ctx.agent_name = agent_name + return ctx + + +def _tool_context( + *, + invocation_id: str = "inv-1", + session_id: str = "sess-1", + agent_name: str = "agent", + function_call_id: str | None = "fc-1", +) -> Any: + ctx = _callback_context( + invocation_id=invocation_id, + session_id=session_id, + agent_name=agent_name, + ) + ctx.function_call_id = function_call_id + return ctx + + +def _tool(name: str = "delete_account") -> Any: + tool = Mock() + tool.name = name + return tool + + +def _final_event(text: str = "answer", author: str = "agent") -> Event: + return Event( + invocation_id="inv-1", + author=author, + content=types.Content( + role="model", parts=[types.Part.from_text(text=text)] + ), + ) + + +# --------------------------------------------------------------------------- +# Tool point +# --------------------------------------------------------------------------- + + +async def test_before_tool_allow_returns_none() -> None: + plugin = AgentHooksPlugin(interceptors=[_AllowAll()]) + result = await plugin.before_tool_callback( + tool=_tool("safe"), + tool_args={"x": 1}, + tool_context=_tool_context(), + ) + assert result is None + + +async def test_before_tool_deny_blocks_with_error() -> None: + plugin = AgentHooksPlugin(interceptors=[_DenyTool("delete_account")]) + args = {"user_id": 42} + result = await plugin.before_tool_callback( + tool=_tool("delete_account"), + tool_args=args, + tool_context=_tool_context(), + ) + assert result is not None + assert result["agent_hooks_blocked"] is True + assert "tool_denied" in result["reason"] + assert "error" in result + # The original args are untouched on a deny. + assert args == {"user_id": 42} + + +async def test_before_tool_transform_mutates_args_in_place() -> None: + plugin = AgentHooksPlugin(interceptors=[_TransformToolArgs()]) + args = {"user_id": 42} + result = await plugin.before_tool_callback( + tool=_tool("lookup"), + tool_args=args, + tool_context=_tool_context(), + ) + # Transform proceeds (returns None) but rewrites args in place so the + # real tool call sees the transformed values. + assert result is None + assert args == {"user_id": 42, "redacted": True} + + +async def test_after_tool_transform_replaces_result() -> None: + plugin = AgentHooksPlugin(interceptors=[_TransformToolResult()]) + result = await plugin.after_tool_callback( + tool=_tool("lookup"), + tool_args={"user_id": 42}, + tool_context=_tool_context(), + result={"secret": "xyz"}, + ) + assert result == {"clean": "value"} + + +async def test_before_tool_fails_closed_on_interceptor_error() -> None: + plugin = AgentHooksPlugin(interceptors=[_Raiser()]) + result = await plugin.before_tool_callback( + tool=_tool("lookup"), + tool_args={"x": 1}, + tool_context=_tool_context(), + ) + assert result is not None + assert result["agent_hooks_blocked"] is True + + +async def test_before_tool_fails_closed_with_no_interceptors() -> None: + plugin = AgentHooksPlugin(interceptors=[]) + result = await plugin.before_tool_callback( + tool=_tool("lookup"), + tool_args={"x": 1}, + tool_context=_tool_context(), + ) + assert result is not None + assert result["agent_hooks_blocked"] is True + + +# --------------------------------------------------------------------------- +# Model point +# --------------------------------------------------------------------------- + + +async def test_before_model_deny_returns_blocked_response() -> None: + plugin = AgentHooksPlugin(interceptors=[_DenyModel()]) + response = await plugin.before_model_callback( + callback_context=_callback_context(), + llm_request=LlmRequest(model="gemini-2.5-flash"), + ) + assert isinstance(response, LlmResponse) + assert response.custom_metadata is not None + assert response.custom_metadata["agent_hooks_blocked"] is True + + +async def test_before_model_transform_fails_closed() -> None: + plugin = AgentHooksPlugin(interceptors=[_TransformModel()]) + response = await plugin.before_model_callback( + callback_context=_callback_context(), + llm_request=LlmRequest(model="gemini-2.5-flash"), + ) + # A transform at pre_model_call is not round-trip safe -> fail closed. + assert isinstance(response, LlmResponse) + assert response.custom_metadata["agent_hooks_blocked"] is True + + +async def test_after_model_transform_rewrites_content() -> None: + plugin = AgentHooksPlugin(interceptors=[_TransformModel()]) + original = LlmResponse( + content=types.Content( + role="model", parts=[types.Part.from_text(text="unsafe")] + ) + ) + response = await plugin.after_model_callback( + callback_context=_callback_context(), + llm_response=original, + ) + assert isinstance(response, LlmResponse) + assert response.content is not None + assert response.content.parts[0].text == "SAFE" + + +async def test_after_model_allow_returns_none() -> None: + plugin = AgentHooksPlugin(interceptors=[_AllowAll()]) + response = await plugin.after_model_callback( + callback_context=_callback_context(), + llm_response=LlmResponse( + content=types.Content( + role="model", parts=[types.Part.from_text(text="hi")] + ) + ), + ) + assert response is None + + +# --------------------------------------------------------------------------- +# Input / output points +# --------------------------------------------------------------------------- + + +async def test_input_deny_replaces_message() -> None: + plugin = AgentHooksPlugin(interceptors=[_DenyInput()]) + result = await plugin.on_user_message_callback( + invocation_context=_invocation_context(), + user_message=types.Content( + role="user", parts=[types.Part.from_text(text="malicious")] + ), + ) + assert result is not None + assert result.role == "user" + assert "input blocked" in result.parts[0].text + + +async def test_input_transform_rewrites_message() -> None: + plugin = AgentHooksPlugin(interceptors=[_TransformInput()]) + result = await plugin.on_user_message_callback( + invocation_context=_invocation_context(), + user_message=types.Content( + role="user", parts=[types.Part.from_text(text="raw")] + ), + ) + assert result is not None + assert result.parts[0].text == "CLEANED" + + +async def test_output_deny_replaces_event_content() -> None: + plugin = AgentHooksPlugin(interceptors=[_DenyOutput()]) + event = await plugin.on_event_callback( + invocation_context=_invocation_context(), + event=_final_event("leaked secret"), + ) + assert event is not None + assert event.content is not None + assert "output blocked" in event.content.parts[0].text + + +async def test_output_ignores_non_final_events() -> None: + plugin = AgentHooksPlugin(interceptors=[_DenyOutput()]) + non_final = Event( + invocation_id="inv-1", + author="agent", + content=types.Content( + role="model", + parts=[ + types.Part(function_call=types.FunctionCall(name="t", args={})) + ], + ), + ) + result = await plugin.on_event_callback( + invocation_context=_invocation_context(), event=non_final + ) + assert result is None + + +# --------------------------------------------------------------------------- +# Startup / lifecycle / state hygiene +# --------------------------------------------------------------------------- + + +async def test_before_run_deny_halts() -> None: + class _DenyStartup: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "agent_startup": + return Verdict.deny(reason="startup_denied") + return Verdict(decision=Decision.ALLOW) + + plugin = AgentHooksPlugin(interceptors=[_DenyStartup()]) + result = await plugin.before_run_callback( + invocation_context=_invocation_context() + ) + assert result is not None + assert "blocked by agent-hooks" in result.parts[0].text + + +async def test_state_is_evicted_after_run() -> None: + plugin = AgentHooksPlugin(interceptors=[_AllowAll()]) + ic = _invocation_context() + await plugin.before_run_callback(invocation_context=ic) + assert ic.invocation_id in plugin._states + await plugin.after_run_callback(invocation_context=ic) + assert ic.invocation_id not in plugin._states + + +async def test_state_is_evicted_on_run_error() -> None: + plugin = AgentHooksPlugin(interceptors=[_AllowAll()]) + ic = _invocation_context() + await plugin.before_run_callback(invocation_context=ic) + assert ic.invocation_id in plugin._states + await plugin.on_run_error_callback( + invocation_context=ic, error=RuntimeError("x") + ) + assert ic.invocation_id not in plugin._states + + +async def test_close_clears_state() -> None: + plugin = AgentHooksPlugin(interceptors=[_AllowAll()]) + await plugin.before_run_callback(invocation_context=_invocation_context()) + assert plugin._states + await plugin.close() + assert not plugin._states + + +async def test_evaluate_only_does_not_block() -> None: + plugin = AgentHooksPlugin( + interceptors=[_DenyTool("delete_account")], mode="evaluate_only" + ) + result = await plugin.before_tool_callback( + tool=_tool("delete_account"), + tool_args={"user_id": 42}, + tool_context=_tool_context(), + ) + # evaluate_only records the deny but does not enforce it. + assert result is None + + +async def test_record_sink_receives_records() -> None: + records: list[Any] = [] + plugin = AgentHooksPlugin( + interceptors=[_AllowAll()], record_sink=records.append + ) + await plugin.before_tool_callback( + tool=_tool("safe"), + tool_args={"x": 1}, + tool_context=_tool_context(), + ) + assert len(records) == 1 + assert records[0].interception_point.value == "pre_tool_call" + + +# --------------------------------------------------------------------------- +# Enforcement-mode fidelity, startup enforcement, and audit correlation +# --------------------------------------------------------------------------- + + +async def test_evaluate_only_does_not_transform_tool_args() -> None: + plugin = AgentHooksPlugin( + interceptors=[_TransformToolArgs()], mode="evaluate_only" + ) + args = {"user_id": 42} + result = await plugin.before_tool_callback( + tool=_tool("lookup"), + tool_args=args, + tool_context=_tool_context(), + ) + # evaluate_only records the transform verdict but must not rewrite the args. + assert result is None + assert args == {"user_id": 42} + + +async def test_evaluate_only_preserves_model_response() -> None: + plugin = AgentHooksPlugin( + interceptors=[_TransformModel()], mode="evaluate_only" + ) + original = LlmResponse( + content=types.Content( + role="model", + parts=[ + types.Part.from_text(text="calling tool"), + types.Part(function_call=types.FunctionCall(name="t", args={})), + ], + ) + ) + response = await plugin.after_model_callback( + callback_context=_callback_context(), + llm_response=original, + ) + # evaluate_only must not rewrite the response or drop the tool-call part. + assert response is None + + +async def test_startup_deny_blocks_model_call() -> None: + class _DenyStartup: + + def intercept(self, ctx: AgentContext) -> Any: + if ctx["interception_point"] == "agent_startup": + return Verdict.deny(reason="startup_denied") + return Verdict(decision=Decision.ALLOW) + + plugin = AgentHooksPlugin(interceptors=[_DenyStartup()]) + await plugin.before_run_callback(invocation_context=_invocation_context()) + # Runner paths that ignore before_run's return value must still be blocked + # at the first model call. + response = await plugin.before_model_callback( + callback_context=_callback_context(), + llm_request=LlmRequest(model="gemini-2.5-flash"), + ) + assert isinstance(response, LlmResponse) + assert response.custom_metadata is not None + assert response.custom_metadata["agent_hooks_blocked"] is True + assert "startup_denied" in response.custom_metadata["reason"] + + +async def test_synth_call_id_correlates_pre_and_post_after_transform() -> None: + class _CaptureAndTransform: + + def __init__(self) -> None: + self.ids: dict[str, str] = {} + + def intercept(self, ctx: AgentContext) -> Any: + point = ctx["interception_point"] + if point == "pre_tool_call": + self.ids["pre"] = ctx["tool_call"]["id"] + new_args = dict(ctx["tool_call"]["args"]) + new_args["redacted"] = True + return Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target", value=new_args), + ) + if point == "post_tool_call": + self.ids["post"] = ctx["tool_call"]["id"] + return Verdict(decision=Decision.ALLOW) + + interceptor = _CaptureAndTransform() + plugin = AgentHooksPlugin(interceptors=[interceptor]) + # No function_call_id -> the plugin synthesizes the id and must correlate the + # pre/post pair even though the pre-tool transform rewrites the args. + tool_context = _tool_context(function_call_id=None) + args = {"user_id": 42} + await plugin.before_tool_callback( + tool=_tool("lookup"), tool_args=args, tool_context=tool_context + ) + assert args == {"user_id": 42, "redacted": True} + await plugin.after_tool_callback( + tool=_tool("lookup"), + tool_args=args, + tool_context=tool_context, + result={"ok": True}, + ) + assert interceptor.ids["pre"] == interceptor.ids["post"] + assert interceptor.ids["pre"].startswith("tc-") + + +def test_json_safe_bounds_recursion_depth() -> None: + from google.adk.plugins._agent_hooks_plugin import _json_safe + + node: dict[str, Any] = {} + cursor = node + for _ in range(5000): + child: dict[str, Any] = {} + cursor["next"] = child + cursor = child + # Deeply nested untrusted input must truncate, not raise RecursionError. + assert "" in str(_json_safe(node)) + + +def test_missing_dependency_raises_actionable_error(monkeypatch) -> None: + import google.adk.plugins._agent_hooks_plugin as mod + + def _boom(_name: str) -> Any: + raise ImportError("no module") + + monkeypatch.setattr(mod.importlib, "import_module", _boom) + with pytest.raises(ImportError, match="google-adk\\[agent-hooks\\]"): + AgentHooksPlugin(interceptors=[])