From 5ff0568bc503a1af77cda2131c15514ff448acbf Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Mon, 22 Jun 2026 23:09:21 +0530 Subject: [PATCH 1/9] feat(governance): add Google ADK governance adapter Installs governance on each LlmAgent's native callbacks (before/after_model_callback -> BEFORE/AFTER_MODEL, before/after_tool_callback -> TOOL_CALL/AFTER_TOOL) in place, walking the sub_agents tree. Self-registers via the uipath.governance.adapters entry point; unit-tested and cloud-verified end to end (CodedAgent03). BEFORE/AFTER_AGENT remain owned by the uipath-runtime wrapper. Co-Authored-By: Claude Opus 4.8 --- packages/uipath-google-adk/pyproject.toml | 4 + .../uipath_google_adk/governance/__init__.py | 58 +++ .../uipath_google_adk/governance/adapter.py | 434 ++++++++++++++++++ .../tests/governance/__init__.py | 0 .../tests/governance/test_adapter.py | 356 ++++++++++++++ 5 files changed, 852 insertions(+) create mode 100644 packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py create mode 100644 packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py create mode 100644 packages/uipath-google-adk/tests/governance/__init__.py create mode 100644 packages/uipath-google-adk/tests/governance/test_adapter.py diff --git a/packages/uipath-google-adk/pyproject.toml b/packages/uipath-google-adk/pyproject.toml index 73a696d6..48f2f6ea 100644 --- a/packages/uipath-google-adk/pyproject.toml +++ b/packages/uipath-google-adk/pyproject.toml @@ -8,6 +8,7 @@ dependencies = [ "google-adk>=1.25.1", "openinference-instrumentation-google-adk>=0.1.9", "uipath>=2.10.0, <2.11.0", + "uipath-core>=0.5.18, <0.7.0", "uipath-runtime>=0.11.0, <0.12.0", ] classifiers = [ @@ -30,6 +31,9 @@ register = "uipath_google_adk.middlewares:register_middleware" [project.entry-points."uipath.runtime.factories"] google-adk = "uipath_google_adk.runtime:register_runtime_factory" +[project.entry-points."uipath.governance.adapters"] +google-adk = "uipath_google_adk.governance:register_governance_adapter" + [project.urls] Homepage = "https://uipath.com" Repository = "https://github.com/UiPath/uipath-integrations-python" diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py new file mode 100644 index 00000000..f7421421 --- /dev/null +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py @@ -0,0 +1,58 @@ +"""Governance integration for ``uipath-google-adk``. + +Registers :class:`GoogleADKAdapter` with the global adapter registry in +``uipath.core.adapters`` so ``uipath.runtime.governance.GovernanceRuntime`` +can attach the ADK-specific inner hooks (BEFORE_MODEL, AFTER_MODEL, +TOOL_CALL, AFTER_TOOL) when it sees a Google ADK agent. + +Registration is **idempotent**: calling :func:`register_governance_adapter` +twice is a no-op on the second call. + +Wiring: + 1. Importing this module triggers registration as a side-effect, so + any caller that does ``import uipath_google_adk.governance`` is + opted in. + 2. The package also exposes :func:`register_governance_adapter` as an + entry point under ``uipath.governance.adapters`` so the registry's + entry-point discovery can plug us in without an explicit import. +""" + +from __future__ import annotations + +import logging + +from uipath.core.adapters import get_adapter_registry + +from .adapter import GoogleADKAdapter, GovernanceCallbacks + +logger = logging.getLogger(__name__) + +_registered: bool = False + + +def register_governance_adapter() -> None: + """Register :class:`GoogleADKAdapter` with the global registry. + + Idempotent — safe to call multiple times. + """ + global _registered + if _registered: + return + registry = get_adapter_registry() + if any(a.name == "GoogleADK" for a in registry.get_all()): + _registered = True + return + registry.register(GoogleADKAdapter()) + _registered = True + logger.debug("Registered uipath-google-adk governance adapter") + + +# Side-effect registration on module import. +register_governance_adapter() + + +__all__ = [ + "GoogleADKAdapter", + "GovernanceCallbacks", + "register_governance_adapter", +] diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py new file mode 100644 index 00000000..6364e807 --- /dev/null +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py @@ -0,0 +1,434 @@ +"""Google ADK adapter for UiPath governance. + +Provides governance for Google ADK agents (``google.adk.agents.LlmAgent`` +and any ``BaseAgent`` tree containing them). Unlike the LangChain adapter +— which wraps a ``Runnable`` and intercepts ``invoke`` / ``ainvoke`` — ADK +agents are executed by a ``Runner`` that holds its **own** reference to +the agent object. Replacing ``runtime.agent`` with a proxy would never +reach the ``Runner``. So this adapter installs governance directly onto +each ``LlmAgent``'s native callback attributes, mutating them in place: + +- ``before_model_callback`` → BEFORE_MODEL +- ``after_model_callback`` → AFTER_MODEL +- ``before_tool_callback`` → TOOL_CALL +- ``after_tool_callback`` → AFTER_TOOL + +Because the mutation is in place, :meth:`GoogleADKAdapter.attach` returns +the **original agent** (hooks installed) rather than a wrapping proxy. +Returning a proxy here would also break ADK's own ``isinstance(agent, +LlmAgent)`` checks in output-schema / graph resolution, since ``LlmAgent`` +is a Pydantic model. + +Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are intentionally +*not* fired from here — they are owned by the runtime wrapper layer in +``uipath-runtime`` (``GovernanceRuntime.execute`` / ``.stream``). Firing +them here too would duplicate every boundary evaluation. + +Contracts and the evaluator protocol come from ``uipath-core``; this +package contributes only the ADK-specific implementation and +self-registers it with the global adapter registry when +``uipath_google_adk.governance`` is imported. + +Audit emission and enforcement (raising :class:`GovernanceBlockException` +on DENY) are owned by the evaluator itself. Each callback only extracts +the relevant payload and calls the matching ``evaluate_*`` method; +:class:`GovernanceBlockException` is allowed to propagate (it aborts the +``Runner`` run), anything else is logged and swallowed so a governance +bug never breaks an agent run. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, List +from uuid import uuid4 + +from uipath.core.adapters import BaseAdapter, EvaluatorProtocol +from uipath.core.governance.exceptions import GovernanceBlockException + +logger = logging.getLogger(__name__) + +# Cap on the text blob passed to BEFORE_MODEL / AFTER_MODEL governance +# evaluation. Sized to match the runtime side (``_GOVERNANCE_TEXT_CAP`` +# in ``uipath.runtime.governance.wrapper``) and the LangChain adapter so +# scan-time budgets are consistent across hooks. A long conversation +# history is governed at the LLM layer by scanning only the latest +# request content, not the full prompt — see +# :meth:`GovernanceCallbacks._latest_request_text`. +_BEFORE_MODEL_TEXT_CAP = 64000 + +# Native LlmAgent callback attribute names this adapter manages. +_MODEL_BEFORE = "before_model_callback" +_MODEL_AFTER = "after_model_callback" +_TOOL_BEFORE = "before_tool_callback" +_TOOL_AFTER = "after_tool_callback" + + +def _is_governance_callable(fn: Any) -> bool: + """True if ``fn`` is a bound method of a :class:`GovernanceCallbacks`.""" + return isinstance(getattr(fn, "__self__", None), GovernanceCallbacks) + + +def _install_callback(agent: Any, attr: str, fn: Any) -> None: + """Prepend ``fn`` to an ADK callback slot, preserving existing handlers. + + ADK accepts a single callable or a ``list`` of callables for each + ``*_callback`` field and runs them in order, stopping early if one + returns a value (a short-circuit). Governance is prepended (runs + first) so it always evaluates — and can BLOCK — before any + user-supplied callback gets a chance to short-circuit the model / + tool call. + + Idempotent: if a governance callback is already present in the slot, + this is a no-op (so a double ``attach`` does not stack duplicates). + """ + existing = getattr(agent, attr, None) + if existing is None: + handlers: List[Any] = [] + elif isinstance(existing, list): + handlers = list(existing) + else: + handlers = [existing] + if any(_is_governance_callable(h) for h in handlers): + return + setattr(agent, attr, [fn, *handlers]) + + +def _remove_callbacks(agent: Any) -> None: + """Strip this adapter's governance callbacks from every managed slot.""" + for attr in (_MODEL_BEFORE, _MODEL_AFTER, _TOOL_BEFORE, _TOOL_AFTER): + existing = getattr(agent, attr, None) + if existing is None: + continue + if isinstance(existing, list): + kept = [h for h in existing if not _is_governance_callable(h)] + setattr(agent, attr, kept or None) + elif _is_governance_callable(existing): + setattr(agent, attr, None) + + +def _iter_llm_agents(root: Any) -> List[Any]: + """Return every ``LlmAgent``-shaped node in the ``sub_agents`` tree. + + A node qualifies if it exposes the model-callback surface (duck-typed + via :data:`_MODEL_BEFORE` so we don't hard-require ``LlmAgent`` to be + importable). Container agents (``Sequential`` / ``Parallel`` / ``Loop``) + have no model callbacks themselves but their ``sub_agents`` are walked + so a multi-agent app is governed end to end. Cycles and pathological + depth are bounded by an id-visited set and a hard cap. + """ + found: List[Any] = [] + seen: set[int] = set() + stack: List[Any] = [root] + while stack and len(seen) < 1000: + node = stack.pop() + if node is None or id(node) in seen: + continue + seen.add(id(node)) + if hasattr(node, _MODEL_BEFORE): + found.append(node) + sub_agents = getattr(node, "sub_agents", None) + if isinstance(sub_agents, (list, tuple)): + stack.extend(sub_agents) + return found + + +class GoogleADKAdapter(BaseAdapter): + """Adapter for the Google ADK framework. + + Detects ``google.adk`` agents and installs governance callbacks on + every ``LlmAgent`` reachable through the ``sub_agents`` tree. + """ + + @property + def name(self) -> str: + return "GoogleADK" + + def can_handle(self, agent: Any) -> bool: + """Return True if this adapter knows how to hook into the agent.""" + try: + from google.adk.agents import BaseAgent + + if isinstance(agent, BaseAgent): + return True + except ImportError: + pass + + # Duck-typed fallback: an ADK agent exposes a name plus either the + # model-callback surface (LlmAgent) or a sub_agents container. + if hasattr(agent, "name") and ( + hasattr(agent, _MODEL_BEFORE) or hasattr(agent, "sub_agents") + ): + return True + return False + + def attach( + self, + agent: Any, + agent_id: str, + session_id: str, + evaluator: EvaluatorProtocol, + ) -> Any: + """Install governance callbacks on the agent (mutated in place). + + Returns the original ``agent`` — the ``Runner`` already holds this + reference, so in-place mutation is what actually wires governance + into execution. A wrapping proxy would not reach the ``Runner`` + and would break ADK's ``isinstance(agent, LlmAgent)`` checks. + """ + callbacks = GovernanceCallbacks( + evaluator=evaluator, + agent_name=agent_id, + session_id=session_id, + ) + llm_agents = _iter_llm_agents(agent) + for node in llm_agents: + _install_callback(node, _MODEL_BEFORE, callbacks.before_model) + _install_callback(node, _MODEL_AFTER, callbacks.after_model) + _install_callback(node, _TOOL_BEFORE, callbacks.before_tool) + _install_callback(node, _TOOL_AFTER, callbacks.after_tool) + if not llm_agents: + logger.warning( + "GoogleADKAdapter found no LlmAgent in %s — deep hooks will not fire", + type(agent).__name__, + ) + else: + logger.debug( + "Installed governance callbacks on %d ADK LlmAgent(s)", + len(llm_agents), + ) + return agent + + def detach(self, governed: Any) -> Any: + """Remove governance callbacks from the agent tree and return it.""" + for node in _iter_llm_agents(governed): + _remove_callbacks(node) + return governed + + +class GovernanceCallbacks: + """Holds the four ADK callbacks bound to one governance evaluator. + + The evaluator owns audit emission and DENY-raising. Each callback + extracts the relevant payload, calls the matching ``evaluate_*`` + method, and returns ``None`` (never short-circuiting the model / tool + on its own). :class:`GovernanceBlockException` is allowed to + propagate — it aborts the ``Runner`` run — anything else is logged + and swallowed. + """ + + def __init__( + self, + evaluator: EvaluatorProtocol, + agent_name: str, + session_id: str, + ) -> None: + self._evaluator = evaluator + self._agent_name = agent_name + self._session_id = session_id + self._trace_id = str(uuid4()) + self._session_state: Dict[str, Any] = {"tool_calls": 0, "llm_calls": 0} + + # ----- Model callbacks ------------------------------------------------- + + def before_model(self, callback_context: Any, llm_request: Any) -> None: + """Evaluate BEFORE_MODEL rules at model start. + + Scans only the **latest request content** — not the full history. + The model still receives the entire history (this callback does + not mutate ``llm_request``); the evaluator focuses on the new + content the agent is about to respond to. Without this scoping, a + violation in an earlier turn would re-fire on every subsequent + model call because that text stays in the prompt for context. + + Returns ``None`` so ADK proceeds with the model call. + """ + try: + self._session_state["llm_calls"] = ( + self._session_state.get("llm_calls", 0) + 1 + ) + model_input = self._latest_request_text(llm_request) + self._evaluator.evaluate_before_model( + model_input=model_input, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning("before_model governance check failed (continuing): %s", e) + return None + + def after_model(self, callback_context: Any, llm_response: Any) -> None: + """Evaluate AFTER_MODEL rules at model end. + + Partial (streamed) responses are skipped — ADK fires + ``after_model_callback`` for each chunk with ``partial=True`` and + once more for the aggregated final response. Governing only the + final response avoids re-scanning the same text token-by-token. + + Returns ``None`` so ADK keeps the model's response unchanged. + """ + try: + if getattr(llm_response, "partial", False): + return None + content = getattr(llm_response, "content", None) + model_output = self._content_text(content) + self._evaluator.evaluate_after_model( + model_output=model_output, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning("after_model governance check failed (continuing): %s", e) + return None + + # ----- Tool callbacks -------------------------------------------------- + + def before_tool(self, tool: Any, args: Dict[str, Any], tool_context: Any) -> None: + """Evaluate TOOL_CALL rules at tool start. + + Returns ``None`` so ADK proceeds with the tool call (a non-None + return would short-circuit it with a substitute result). + """ + try: + self._session_state["tool_calls"] = ( + self._session_state.get("tool_calls", 0) + 1 + ) + tool_name = getattr(tool, "name", None) or "unknown" + self._evaluator.evaluate_tool_call( + tool_name=tool_name, + tool_args=args or {}, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + session_state=self._session_state, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning("before_tool governance check failed (continuing): %s", e) + return None + + def after_tool( + self, + tool: Any, + args: Dict[str, Any], + tool_context: Any, + tool_response: Any, + ) -> None: + """Evaluate AFTER_TOOL rules at tool end. + + Returns ``None`` so ADK keeps the tool's result unchanged. + """ + try: + tool_name = getattr(tool, "name", None) or "unknown" + tool_result = ( + "" if tool_response is None else self._stringify(tool_response) + ) + self._evaluator.evaluate_after_tool( + tool_name=tool_name, + tool_result=tool_result, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning("after_tool governance check failed (continuing): %s", e) + return None + + # ----- Text extraction ------------------------------------------------- + + def _latest_request_text(self, llm_request: Any) -> str: + """Extract text from the most-recent content in an ``LlmRequest``. + + ``llm_request.contents`` is the full ``list[Content]`` sent to the + model. We take the last entry — the new user message, or the tool + ``function_response`` being fed back — and pull its text cleanly + via :meth:`_content_text`. Returns ``""`` when there is nothing + extractable. + """ + contents = getattr(llm_request, "contents", None) + if not contents: + return "" + return self._content_text(contents[-1]) + + @classmethod + def _content_text(cls, content: Any) -> str: + """Return governance-relevant text from a ``Content`` (or part list). + + Walks ``content.parts`` and pulls, per part: + + - ``part.text`` — plain text. + - ``part.function_call`` — the tool name plus JSON-encoded + ``args``; ADK / Gemini routinely carry the user-visible reply in + a function call (e.g. a "submit final answer" tool). + - ``part.function_response`` — the tool result fed back to the + model; relevant when it is the latest content for BEFORE_MODEL. + + Capped at :data:`_BEFORE_MODEL_TEXT_CAP` so a runaway response or + large tool payload can't blow scan budgets. + """ + if content is None: + return "" + parts = getattr(content, "parts", None) + if parts is None: + # Some shapes hand us a bare string or a list of parts. + if isinstance(content, str): + return content[:_BEFORE_MODEL_TEXT_CAP] + if isinstance(content, (list, tuple)): + parts = content + else: + return "" + collected: List[str] = [] + remaining = _BEFORE_MODEL_TEXT_CAP + for part in parts: + if remaining <= 0: + break + piece = cls._part_text(part) + if piece: + collected.append(piece) + remaining -= len(piece) + 1 + return "\n".join(collected)[:_BEFORE_MODEL_TEXT_CAP] + + @classmethod + def _part_text(cls, part: Any) -> str: + """Return text / function-call args / function-response from one part.""" + pieces: List[str] = [] + text = getattr(part, "text", None) + if isinstance(text, str) and text: + pieces.append(text) + + function_call = getattr(part, "function_call", None) + if function_call is not None: + name = getattr(function_call, "name", "") or "" + fc_args = getattr(function_call, "args", None) + if name: + pieces.append(name) + if fc_args: + pieces.append(cls._stringify(fc_args)) + + function_response = getattr(part, "function_response", None) + if function_response is not None: + response = getattr(function_response, "response", None) + if response: + pieces.append(cls._stringify(response)) + + return "\n".join(p for p in pieces if p) + + @staticmethod + def _stringify(value: Any) -> str: + """Render a dict / object payload as compact, scannable text.""" + if isinstance(value, str): + return value + try: + return json.dumps(value, default=str, ensure_ascii=False) + except (TypeError, ValueError): + return str(value) diff --git a/packages/uipath-google-adk/tests/governance/__init__.py b/packages/uipath-google-adk/tests/governance/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/uipath-google-adk/tests/governance/test_adapter.py b/packages/uipath-google-adk/tests/governance/test_adapter.py new file mode 100644 index 00000000..f9bf18e4 --- /dev/null +++ b/packages/uipath-google-adk/tests/governance/test_adapter.py @@ -0,0 +1,356 @@ +"""Unit tests for the Google ADK governance adapter. + +These tests deliberately avoid importing ``google.adk`` — the adapter +duck-types every Google type (it only hard-imports ``uipath.core``), so +lightweight fakes for ``Part`` / ``Content`` / ``LlmRequest`` / +``LlmResponse`` / tool / agent exercise the real code paths without the +heavy ADK dependency. +""" + +from __future__ import annotations + +import logging +from types import SimpleNamespace +from typing import Any, List + +import pytest +from uipath.core.governance.exceptions import GovernanceBlockException + +from uipath_google_adk.governance.adapter import ( + _BEFORE_MODEL_TEXT_CAP, + GoogleADKAdapter, + GovernanceCallbacks, +) + +# -------------------------------------------------------------------------- +# Fakes +# -------------------------------------------------------------------------- + + +class FakeEvaluator: + """Records evaluate_* calls; optionally BLOCKs on a named hook.""" + + def __init__(self, block_on: str | None = None) -> None: + self.block_on = block_on + self.calls: List[tuple[str, dict]] = [] + + def _record(self, hook: str, **kwargs: Any) -> None: + self.calls.append((hook, kwargs)) + if self.block_on == hook: + raise GovernanceBlockException("blocked") # type: ignore[call-arg] + + def evaluate_before_agent(self, **kwargs: Any) -> None: + self._record("before_agent", **kwargs) + + def evaluate_after_agent(self, **kwargs: Any) -> None: + self._record("after_agent", **kwargs) + + def evaluate_before_model(self, **kwargs: Any) -> None: + self._record("before_model", **kwargs) + + def evaluate_after_model(self, **kwargs: Any) -> None: + self._record("after_model", **kwargs) + + def evaluate_tool_call(self, **kwargs: Any) -> None: + self._record("tool_call", **kwargs) + + def evaluate_after_tool(self, **kwargs: Any) -> None: + self._record("after_tool", **kwargs) + + +class FakeLlmAgent: + """Minimal stand-in for ``google.adk.agents.LlmAgent``.""" + + def __init__(self, name: str = "agent", sub_agents: List[Any] | None = None): + self.name = name + self.before_model_callback: Any = None + self.after_model_callback: Any = None + self.before_tool_callback: Any = None + self.after_tool_callback: Any = None + self.sub_agents = sub_agents or [] + + +class FakeContainerAgent: + """Container agent (Sequential/Parallel) with no model callbacks.""" + + def __init__(self, name: str, sub_agents: List[Any]): + self.name = name + self.sub_agents = sub_agents + + +class FakeTool: + def __init__(self, name: str): + self.name = name + + +def _part( + text: str | None = None, + function_call: Any = None, + function_response: Any = None, +) -> SimpleNamespace: + return SimpleNamespace( + text=text, + function_call=function_call, + function_response=function_response, + ) + + +def _content(parts: List[Any], role: str = "user") -> SimpleNamespace: + return SimpleNamespace(role=role, parts=parts) + + +def _make_callbacks(evaluator: FakeEvaluator) -> GovernanceCallbacks: + return GovernanceCallbacks( + evaluator=evaluator, agent_name="agent-1", session_id="sess-1" + ) + + +# -------------------------------------------------------------------------- +# can_handle +# -------------------------------------------------------------------------- + + +def test_can_handle_llm_agent(): + assert GoogleADKAdapter().can_handle(FakeLlmAgent()) is True + + +def test_can_handle_container_agent(): + container = FakeContainerAgent("root", [FakeLlmAgent()]) + assert GoogleADKAdapter().can_handle(container) is True + + +def test_can_handle_rejects_plain_object(): + assert GoogleADKAdapter().can_handle(object()) is False + + +# -------------------------------------------------------------------------- +# attach / detach +# -------------------------------------------------------------------------- + + +def test_attach_installs_on_all_llm_agents_in_tree(): + leaf_a = FakeLlmAgent("a") + leaf_b = FakeLlmAgent("b") + root = FakeContainerAgent("root", [leaf_a, leaf_b]) + + returned = GoogleADKAdapter().attach( + root, agent_id="x", session_id="s", evaluator=FakeEvaluator() + ) + + assert returned is root # original returned, not a proxy + for leaf in (leaf_a, leaf_b): + assert isinstance(leaf.before_model_callback, list) + assert len(leaf.before_model_callback) == 1 + assert leaf.after_model_callback and leaf.before_tool_callback + assert leaf.after_tool_callback + + +def test_attach_is_idempotent(): + agent = FakeLlmAgent() + adapter = GoogleADKAdapter() + ev = FakeEvaluator() + adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) + adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) + assert len(agent.before_model_callback) == 1 + + +def test_attach_preserves_existing_callback_and_runs_governance_first(): + def user_cb(*_a, **_k): + return None + + agent = FakeLlmAgent() + agent.before_model_callback = user_cb + GoogleADKAdapter().attach( + agent, agent_id="x", session_id="s", evaluator=FakeEvaluator() + ) + cbs = agent.before_model_callback + assert isinstance(cbs, list) and len(cbs) == 2 + # governance prepended → runs first + assert getattr(cbs[0], "__self__", None).__class__ is GovernanceCallbacks + assert cbs[1] is user_cb + + +def test_detach_removes_governance_callbacks(): + def user_cb(*_a, **_k): + return None + + agent = FakeLlmAgent() + agent.after_tool_callback = user_cb + adapter = GoogleADKAdapter() + adapter.attach(agent, agent_id="x", session_id="s", evaluator=FakeEvaluator()) + adapter.detach(agent) + assert agent.before_model_callback is None + # unrelated user callback survives + assert agent.after_tool_callback == [user_cb] + + +def test_attach_warns_when_no_llm_agent(caplog): + container = FakeContainerAgent("root", []) + with caplog.at_level(logging.WARNING): + GoogleADKAdapter().attach( + container, agent_id="x", session_id="s", evaluator=FakeEvaluator() + ) + assert any("no LlmAgent" in r.message for r in caplog.records) + + +# -------------------------------------------------------------------------- +# before_model +# -------------------------------------------------------------------------- + + +def test_before_model_scopes_to_latest_content(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + req = SimpleNamespace( + contents=[ + _content([_part(text="OLD turn — secret leak here")]), + _content([_part(text="the new question")]), + ] + ) + cb.before_model(callback_context=None, llm_request=req) + hook, kwargs = ev.calls[-1] + assert hook == "before_model" + assert kwargs["model_input"] == "the new question" + assert "OLD turn" not in kwargs["model_input"] + + +def test_before_model_extracts_function_response_when_latest(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + fr = SimpleNamespace(name="lookup", response={"balance": "1000"}) + req = SimpleNamespace(contents=[_content([_part(function_response=fr)])]) + cb.before_model(callback_context=None, llm_request=req) + assert "1000" in ev.calls[-1][1]["model_input"] + + +def test_before_model_caps_text(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + huge = "x" * (_BEFORE_MODEL_TEXT_CAP + 5000) + req = SimpleNamespace(contents=[_content([_part(text=huge)])]) + cb.before_model(callback_context=None, llm_request=req) + assert len(ev.calls[-1][1]["model_input"]) <= _BEFORE_MODEL_TEXT_CAP + + +def test_before_model_empty_contents(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + cb.before_model(callback_context=None, llm_request=SimpleNamespace(contents=[])) + assert ev.calls[-1][1]["model_input"] == "" + + +# -------------------------------------------------------------------------- +# after_model +# -------------------------------------------------------------------------- + + +def test_after_model_skips_partial(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + resp = SimpleNamespace(partial=True, content=_content([_part(text="chunk")])) + cb.after_model(callback_context=None, llm_response=resp) + assert ev.calls == [] + + +def test_after_model_extracts_text_and_function_call(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + fc = SimpleNamespace(name="submit_answer", args={"content": "final reply"}) + resp = SimpleNamespace( + partial=False, + content=_content( + [_part(text="thinking"), _part(function_call=fc)], role="model" + ), + ) + cb.after_model(callback_context=None, llm_response=resp) + out = ev.calls[-1][1]["model_output"] + assert "thinking" in out and "submit_answer" in out and "final reply" in out + + +# -------------------------------------------------------------------------- +# tools +# -------------------------------------------------------------------------- + + +def test_before_tool_passes_args_and_session_state(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + cb.before_tool(FakeTool("transfer"), {"amount": 50}, tool_context=None) + hook, kwargs = ev.calls[-1] + assert hook == "tool_call" + assert kwargs["tool_name"] == "transfer" + assert kwargs["tool_args"] == {"amount": 50} + assert kwargs["session_state"]["tool_calls"] == 1 + + +def test_after_tool_stringifies_dict_response(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + cb.after_tool(FakeTool("lookup"), {}, tool_context=None, tool_response={"x": 1}) + out = ev.calls[-1][1]["tool_result"] + assert "x" in out and "1" in out + + +def test_after_tool_none_response(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + cb.after_tool(FakeTool("noop"), {}, tool_context=None, tool_response=None) + assert ev.calls[-1][1]["tool_result"] == "" + + +# -------------------------------------------------------------------------- +# enforcement semantics +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "hook,invoke", + [ + ( + "before_model", + lambda cb: cb.before_model( + None, SimpleNamespace(contents=[_content([_part(text="hi")])]) + ), + ), + ( + "after_model", + lambda cb: cb.after_model( + None, + SimpleNamespace(partial=False, content=_content([_part(text="o")])), + ), + ), + ("tool_call", lambda cb: cb.before_tool(FakeTool("t"), {}, None)), + ( + "after_tool", + lambda cb: cb.after_tool(FakeTool("t"), {}, None, {"r": 1}), + ), + ], +) +def test_block_exception_propagates(hook, invoke): + cb = _make_callbacks(FakeEvaluator(block_on=hook)) + with pytest.raises(GovernanceBlockException): + invoke(cb) + + +def test_non_block_exception_is_swallowed(caplog): + class Boom: + def evaluate_before_model(self, **_): + raise RuntimeError("evaluator bug") + + cb = GovernanceCallbacks( + evaluator=Boom(), + agent_name="a", + session_id="s", # type: ignore[arg-type] + ) + with caplog.at_level(logging.WARNING): + # must NOT raise — a governance bug can't break the agent run + cb.before_model(None, SimpleNamespace(contents=[_content([_part(text="x")])])) + assert any("governance check failed" in r.message for r in caplog.records) + + +def test_callbacks_return_none(): + cb = _make_callbacks(FakeEvaluator()) + assert cb.before_model(None, SimpleNamespace(contents=[])) is None + assert cb.after_model(None, SimpleNamespace(partial=False, content=None)) is None + assert cb.before_tool(FakeTool("t"), {}, None) is None + assert cb.after_tool(FakeTool("t"), {}, None, {}) is None From e0500c9b603b116ed7e45dcfe7bb850dd14e11c1 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 17:52:33 +0530 Subject: [PATCH 2/9] chore(governance): apply review feedback (no import-time registration, framework-only can_handle) Mirror radu's LangChain-adapter review across the Google ADK adapter: - __init__: drop the import-time registration side-effect; registration only via the uipath.governance.adapters entry point. - can_handle: claim only a real google.adk BaseAgent; remove the duck-typed (name + model-callback / sub_agents) fallback. - docstring: 'governance host' instead of uipath-runtime internals. - tests: can_handle uses a real LlmAgent; duck-typed look-alikes are now rejected. Co-Authored-By: Claude Opus 4.8 --- .../uipath_google_adk/governance/__init__.py | 23 ++++++---------- .../uipath_google_adk/governance/adapter.py | 27 ++++++------------- .../tests/governance/test_adapter.py | 15 ++++++----- 3 files changed, 24 insertions(+), 41 deletions(-) diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py index f7421421..03e22dee 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py @@ -1,20 +1,17 @@ """Governance integration for ``uipath-google-adk``. -Registers :class:`GoogleADKAdapter` with the global adapter registry in -``uipath.core.adapters`` so ``uipath.runtime.governance.GovernanceRuntime`` -can attach the ADK-specific inner hooks (BEFORE_MODEL, AFTER_MODEL, -TOOL_CALL, AFTER_TOOL) when it sees a Google ADK agent. +Registers :class:`GoogleADKAdapter` with the adapter registry in +``uipath.core.adapters`` so the governance host can attach the ADK-specific +inner hooks (BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, AFTER_TOOL) when it sees a +Google ADK agent. Registration is **idempotent**: calling :func:`register_governance_adapter` twice is a no-op on the second call. -Wiring: - 1. Importing this module triggers registration as a side-effect, so - any caller that does ``import uipath_google_adk.governance`` is - opted in. - 2. The package also exposes :func:`register_governance_adapter` as an - entry point under ``uipath.governance.adapters`` so the registry's - entry-point discovery can plug us in without an explicit import. +Wiring: the package exposes :func:`register_governance_adapter` as an entry +point under ``uipath.governance.adapters``. The governance adapter discovery +path calls it to register the adapter. Importing this module does not, by +itself, mutate the global registry. """ from __future__ import annotations @@ -47,10 +44,6 @@ def register_governance_adapter() -> None: logger.debug("Registered uipath-google-adk governance adapter") -# Side-effect registration on module import. -register_governance_adapter() - - __all__ = [ "GoogleADKAdapter", "GovernanceCallbacks", diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py index 6364e807..476872e6 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py @@ -20,14 +20,13 @@ is a Pydantic model. Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are intentionally -*not* fired from here — they are owned by the runtime wrapper layer in -``uipath-runtime`` (``GovernanceRuntime.execute`` / ``.stream``). Firing -them here too would duplicate every boundary evaluation. +*not* fired from here — they are owned by the governance host. Firing them +here too would duplicate every boundary evaluation. Contracts and the evaluator protocol come from ``uipath-core``; this -package contributes only the ADK-specific implementation and -self-registers it with the global adapter registry when -``uipath_google_adk.governance`` is imported. +package contributes only the ADK-specific implementation and registers it +with the adapter registry via the ``uipath.governance.adapters`` entry +point. Audit emission and enforcement (raising :class:`GovernanceBlockException` on DENY) are owned by the evaluator itself. Each callback only extracts @@ -146,22 +145,12 @@ def name(self) -> str: return "GoogleADK" def can_handle(self, agent: Any) -> bool: - """Return True if this adapter knows how to hook into the agent.""" + """Return True only for a Google ADK ``BaseAgent`` (incl. LlmAgent trees).""" try: from google.adk.agents import BaseAgent - - if isinstance(agent, BaseAgent): - return True except ImportError: - pass - - # Duck-typed fallback: an ADK agent exposes a name plus either the - # model-callback surface (LlmAgent) or a sub_agents container. - if hasattr(agent, "name") and ( - hasattr(agent, _MODEL_BEFORE) or hasattr(agent, "sub_agents") - ): - return True - return False + return False + return isinstance(agent, BaseAgent) def attach( self, diff --git a/packages/uipath-google-adk/tests/governance/test_adapter.py b/packages/uipath-google-adk/tests/governance/test_adapter.py index f9bf18e4..e6aeb828 100644 --- a/packages/uipath-google-adk/tests/governance/test_adapter.py +++ b/packages/uipath-google-adk/tests/governance/test_adapter.py @@ -110,16 +110,17 @@ def _make_callbacks(evaluator: FakeEvaluator) -> GovernanceCallbacks: # -------------------------------------------------------------------------- -def test_can_handle_llm_agent(): - assert GoogleADKAdapter().can_handle(FakeLlmAgent()) is True +def test_can_handle_real_agent(): + from google.adk.agents import LlmAgent + assert GoogleADKAdapter().can_handle(LlmAgent(name="t")) is True -def test_can_handle_container_agent(): - container = FakeContainerAgent("root", [FakeLlmAgent()]) - assert GoogleADKAdapter().can_handle(container) is True - -def test_can_handle_rejects_plain_object(): +def test_can_handle_rejects_non_adk_agent(): + # Duck-typed look-alikes (name + model-callback / sub_agents) must NOT be + # claimed — only a real google.adk BaseAgent is. + assert GoogleADKAdapter().can_handle(FakeLlmAgent()) is False + assert GoogleADKAdapter().can_handle(FakeContainerAgent("root", [FakeLlmAgent()])) is False assert GoogleADKAdapter().can_handle(object()) is False From 6cba9143d4622b5288074b445c1f17921b5369ff Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 18:41:49 +0530 Subject: [PATCH 3/9] docs(governance): address Copilot review on the Google ADK adapter - Text-cap comment: refer to the governance host, not the uipath-runtime wrapper constant. - Test docstring: note can_handle uses a real google.adk LlmAgent (isinstance detection); only payload shapes are faked. Co-Authored-By: Claude Opus 4.8 --- .../src/uipath_google_adk/governance/adapter.py | 3 +-- .../uipath-google-adk/tests/governance/test_adapter.py | 10 +++++----- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py index 476872e6..7828ef1f 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py @@ -49,8 +49,7 @@ logger = logging.getLogger(__name__) # Cap on the text blob passed to BEFORE_MODEL / AFTER_MODEL governance -# evaluation. Sized to match the runtime side (``_GOVERNANCE_TEXT_CAP`` -# in ``uipath.runtime.governance.wrapper``) and the LangChain adapter so +# evaluation. Sized to match the governance host and the other adapters so # scan-time budgets are consistent across hooks. A long conversation # history is governed at the LLM layer by scanning only the latest # request content, not the full prompt — see diff --git a/packages/uipath-google-adk/tests/governance/test_adapter.py b/packages/uipath-google-adk/tests/governance/test_adapter.py index e6aeb828..4feec4e5 100644 --- a/packages/uipath-google-adk/tests/governance/test_adapter.py +++ b/packages/uipath-google-adk/tests/governance/test_adapter.py @@ -1,10 +1,10 @@ """Unit tests for the Google ADK governance adapter. -These tests deliberately avoid importing ``google.adk`` — the adapter -duck-types every Google type (it only hard-imports ``uipath.core``), so -lightweight fakes for ``Part`` / ``Content`` / ``LlmRequest`` / -``LlmResponse`` / tool / agent exercise the real code paths without the -heavy ADK dependency. +``can_handle`` is tested against a real ``google.adk`` ``LlmAgent`` (the +adapter detects agents with ``isinstance(..., BaseAgent)``). The remaining +tests duck-type the ADK payloads — lightweight fakes for ``Part`` / +``Content`` / ``LlmRequest`` / ``LlmResponse`` / tool / agent — so the +callback code paths are exercised without driving the heavy ADK runtime. """ from __future__ import annotations From e0a0cf1075520114ac2b8b1a3f5ed15d8f0846f8 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Tue, 30 Jun 2026 16:58:50 +0530 Subject: [PATCH 4/9] docs(governance): note intentional duck-typed extraction in Google ADK adapter Record that LlmRequest/LlmResponse/content/parts are read via getattr rather than isinstance on ADK's typed models, to avoid hard-coupling to google-adk internals and to let tests duck-type payloads without the package installed. Co-Authored-By: Claude Opus 4.8 --- .../src/uipath_google_adk/governance/adapter.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py index 7828ef1f..340bb35f 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py @@ -333,6 +333,10 @@ def after_tool( return None # ----- Text extraction ------------------------------------------------- + # Read LlmRequest/LlmResponse/content/parts defensively via getattr + # rather than isinstance on ADK's typed models: this keeps the adapter + # from hard-coupling to google-adk internal types that may shift, and + # lets the tests duck-type the payloads without a google-adk install. def _latest_request_text(self, llm_request: Any) -> str: """Extract text from the most-recent content in an ``LlmRequest``. From 39948ab30437e678dd19a719efb37e8a41c50b5a Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Tue, 30 Jun 2026 20:34:53 +0530 Subject: [PATCH 5/9] refactor(governance): migrate Google ADK adapter to factory-evaluator Core PR #1761 removed BaseAdapter from uipath-core. Migrate to the factory-evaluator pattern (matching #899): - governance/adapter.py -> callbacks.py: replace the BaseAdapter subclass (name/can_handle/attach/detach) with module-level install_governance() that installs governance on each LlmAgent's native *_callback slots (walking sub_agents); keep GovernanceCallbacks + the callback helpers; drop the detach-only _remove_callbacks. File named for its seam (callbacks), like LangChain's callbacks.py. - runtime/factory.py: new_runtime reads `evaluator` from kwargs and calls install_governance before building the Runner. - governance/__init__.py: drop register_governance_adapter + registry import; expose install_governance. No import-time side effects. - pyproject.toml: remove the uipath.governance.adapters entry point. - tests (test_adapter.py -> test_callbacks.py): drop can_handle/attach/ detach; cover install_governance + factory wiring. ruff + mypy clean; 21 governance tests pass. Co-Authored-By: Claude Opus 4.8 --- packages/uipath-google-adk/pyproject.toml | 3 - .../uipath_google_adk/governance/__init__.py | 53 ++----- .../governance/{adapter.py => callbacks.py} | 132 +++++++----------- .../src/uipath_google_adk/runtime/factory.py | 20 ++- .../{test_adapter.py => test_callbacks.py} | 125 ++++++++++------- packages/uipath-google-adk/uv.lock | 2 + 6 files changed, 156 insertions(+), 179 deletions(-) rename packages/uipath-google-adk/src/uipath_google_adk/governance/{adapter.py => callbacks.py} (80%) rename packages/uipath-google-adk/tests/governance/{test_adapter.py => test_callbacks.py} (79%) diff --git a/packages/uipath-google-adk/pyproject.toml b/packages/uipath-google-adk/pyproject.toml index 48f2f6ea..f0d14711 100644 --- a/packages/uipath-google-adk/pyproject.toml +++ b/packages/uipath-google-adk/pyproject.toml @@ -31,9 +31,6 @@ register = "uipath_google_adk.middlewares:register_middleware" [project.entry-points."uipath.runtime.factories"] google-adk = "uipath_google_adk.runtime:register_runtime_factory" -[project.entry-points."uipath.governance.adapters"] -google-adk = "uipath_google_adk.governance:register_governance_adapter" - [project.urls] Homepage = "https://uipath.com" Repository = "https://github.com/UiPath/uipath-integrations-python" diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py index 03e22dee..fefee0ce 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/__init__.py @@ -1,51 +1,20 @@ """Governance integration for ``uipath-google-adk``. -Registers :class:`GoogleADKAdapter` with the adapter registry in -``uipath.core.adapters`` so the governance host can attach the ADK-specific -inner hooks (BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, AFTER_TOOL) when it sees a -Google ADK agent. - -Registration is **idempotent**: calling :func:`register_governance_adapter` -twice is a no-op on the second call. - -Wiring: the package exposes :func:`register_governance_adapter` as an entry -point under ``uipath.governance.adapters``. The governance adapter discovery -path calls it to register the adapter. Importing this module does not, by -itself, mutate the global registry. +Exposes :func:`install_governance` — installs governance callbacks +(BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, AFTER_TOOL) on every ``LlmAgent`` in an +ADK agent tree's native ``*_callback`` slots. Wired into a run by passing an +``evaluator`` to :class:`UiPathGoogleADKRuntimeFactory`; the factory calls +:func:`install_governance` on the resolved agent. + +Importing this module has no side effects: no adapter is registered, no global +state is mutated. """ from __future__ import annotations -import logging - -from uipath.core.adapters import get_adapter_registry - -from .adapter import GoogleADKAdapter, GovernanceCallbacks - -logger = logging.getLogger(__name__) - -_registered: bool = False - - -def register_governance_adapter() -> None: - """Register :class:`GoogleADKAdapter` with the global registry. - - Idempotent — safe to call multiple times. - """ - global _registered - if _registered: - return - registry = get_adapter_registry() - if any(a.name == "GoogleADK" for a in registry.get_all()): - _registered = True - return - registry.register(GoogleADKAdapter()) - _registered = True - logger.debug("Registered uipath-google-adk governance adapter") - +from .callbacks import GovernanceCallbacks, install_governance __all__ = [ - "GoogleADKAdapter", "GovernanceCallbacks", - "register_governance_adapter", -] + "install_governance", +] \ No newline at end of file diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py similarity index 80% rename from packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py rename to packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py index 340bb35f..18b7c531 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/adapter.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py @@ -1,20 +1,21 @@ -"""Google ADK adapter for UiPath governance. +"""Google ADK governance callbacks for UiPath. Provides governance for Google ADK agents (``google.adk.agents.LlmAgent`` -and any ``BaseAgent`` tree containing them). Unlike the LangChain adapter +and any ``BaseAgent`` tree containing them). Unlike the LangChain integration — which wraps a ``Runnable`` and intercepts ``invoke`` / ``ainvoke`` — ADK agents are executed by a ``Runner`` that holds its **own** reference to the agent object. Replacing ``runtime.agent`` with a proxy would never -reach the ``Runner``. So this adapter installs governance directly onto -each ``LlmAgent``'s native callback attributes, mutating them in place: +reach the ``Runner``. So :func:`install_governance` installs governance +directly onto each ``LlmAgent``'s native callback attributes, mutating them +in place: - ``before_model_callback`` → BEFORE_MODEL - ``after_model_callback`` → AFTER_MODEL - ``before_tool_callback`` → TOOL_CALL - ``after_tool_callback`` → AFTER_TOOL -Because the mutation is in place, :meth:`GoogleADKAdapter.attach` returns -the **original agent** (hooks installed) rather than a wrapping proxy. +Because the mutation is in place, :func:`install_governance` returns the +**original agent** (hooks installed) rather than a wrapping proxy. Returning a proxy here would also break ADK's own ``isinstance(agent, LlmAgent)`` checks in output-schema / graph resolution, since ``LlmAgent`` is a Pydantic model. @@ -23,10 +24,11 @@ *not* fired from here — they are owned by the governance host. Firing them here too would duplicate every boundary evaluation. -Contracts and the evaluator protocol come from ``uipath-core``; this -package contributes only the ADK-specific implementation and registers it -with the adapter registry via the ``uipath.governance.adapters`` entry -point. +The evaluator protocol comes from ``uipath-core``; this package contributes +only the ADK-specific wiring. Governance is installed by the runtime +factory: passing an ``evaluator`` to ``new_runtime`` calls +:func:`install_governance` on the resolved agent. No adapter registry, no +entry point, no import-time side effects. Audit emission and enforcement (raising :class:`GovernanceBlockException` on DENY) are owned by the evaluator itself. Each callback only extracts @@ -43,7 +45,7 @@ from typing import Any, Dict, List from uuid import uuid4 -from uipath.core.adapters import BaseAdapter, EvaluatorProtocol +from uipath.core.adapters import EvaluatorProtocol from uipath.core.governance.exceptions import GovernanceBlockException logger = logging.getLogger(__name__) @@ -93,19 +95,6 @@ def _install_callback(agent: Any, attr: str, fn: Any) -> None: setattr(agent, attr, [fn, *handlers]) -def _remove_callbacks(agent: Any) -> None: - """Strip this adapter's governance callbacks from every managed slot.""" - for attr in (_MODEL_BEFORE, _MODEL_AFTER, _TOOL_BEFORE, _TOOL_AFTER): - existing = getattr(agent, attr, None) - if existing is None: - continue - if isinstance(existing, list): - kept = [h for h in existing if not _is_governance_callable(h)] - setattr(agent, attr, kept or None) - elif _is_governance_callable(existing): - setattr(agent, attr, None) - - def _iter_llm_agents(root: Any) -> List[Any]: """Return every ``LlmAgent``-shaped node in the ``sub_agents`` tree. @@ -132,67 +121,46 @@ def _iter_llm_agents(root: Any) -> List[Any]: return found -class GoogleADKAdapter(BaseAdapter): - """Adapter for the Google ADK framework. - - Detects ``google.adk`` agents and installs governance callbacks on - every ``LlmAgent`` reachable through the ``sub_agents`` tree. - """ +def install_governance( + agent: Any, + evaluator: EvaluatorProtocol, + *, + agent_name: str, + session_id: str, +) -> Any: + """Install governance callbacks on the agent tree (mutated in place). - @property - def name(self) -> str: - return "GoogleADK" - - def can_handle(self, agent: Any) -> bool: - """Return True only for a Google ADK ``BaseAgent`` (incl. LlmAgent trees).""" - try: - from google.adk.agents import BaseAgent - except ImportError: - return False - return isinstance(agent, BaseAgent) - - def attach( - self, - agent: Any, - agent_id: str, - session_id: str, - evaluator: EvaluatorProtocol, - ) -> Any: - """Install governance callbacks on the agent (mutated in place). + Walks every ``LlmAgent`` reachable through ``sub_agents`` and prepends + governance to each model/tool callback slot, preserving existing handlers. + Returns the original ``agent`` — the ``Runner`` already holds this + reference, so in-place mutation is what wires governance into execution. + Idempotent: a slot that already carries a governance callback is skipped. - Returns the original ``agent`` — the ``Runner`` already holds this - reference, so in-place mutation is what actually wires governance - into execution. A wrapping proxy would not reach the ``Runner`` - and would break ADK's ``isinstance(agent, LlmAgent)`` checks. - """ - callbacks = GovernanceCallbacks( - evaluator=evaluator, - agent_name=agent_id, - session_id=session_id, + Called by :class:`UiPathGoogleADKRuntimeFactory` when an ``evaluator`` + is supplied to ``new_runtime``. + """ + callbacks = GovernanceCallbacks( + evaluator=evaluator, + agent_name=agent_name, + session_id=session_id, + ) + llm_agents = _iter_llm_agents(agent) + for node in llm_agents: + _install_callback(node, _MODEL_BEFORE, callbacks.before_model) + _install_callback(node, _MODEL_AFTER, callbacks.after_model) + _install_callback(node, _TOOL_BEFORE, callbacks.before_tool) + _install_callback(node, _TOOL_AFTER, callbacks.after_tool) + if not llm_agents: + logger.warning( + "install_governance found no LlmAgent in %s — deep hooks will not fire", + type(agent).__name__, ) - llm_agents = _iter_llm_agents(agent) - for node in llm_agents: - _install_callback(node, _MODEL_BEFORE, callbacks.before_model) - _install_callback(node, _MODEL_AFTER, callbacks.after_model) - _install_callback(node, _TOOL_BEFORE, callbacks.before_tool) - _install_callback(node, _TOOL_AFTER, callbacks.after_tool) - if not llm_agents: - logger.warning( - "GoogleADKAdapter found no LlmAgent in %s — deep hooks will not fire", - type(agent).__name__, - ) - else: - logger.debug( - "Installed governance callbacks on %d ADK LlmAgent(s)", - len(llm_agents), - ) - return agent - - def detach(self, governed: Any) -> Any: - """Remove governance callbacks from the agent tree and return it.""" - for node in _iter_llm_agents(governed): - _remove_callbacks(node) - return governed + else: + logger.debug( + "Installed governance callbacks on %d ADK LlmAgent(s)", + len(llm_agents), + ) + return agent class GovernanceCallbacks: diff --git a/packages/uipath-google-adk/src/uipath_google_adk/runtime/factory.py b/packages/uipath-google-adk/src/uipath_google_adk/runtime/factory.py index 338b6883..2658302e 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/runtime/factory.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/runtime/factory.py @@ -8,6 +8,7 @@ from google.adk.runners import Runner from google.adk.sessions.sqlite_session_service import SqliteSessionService from openinference.instrumentation.google_adk import GoogleADKInstrumentor +from uipath.core.adapters import EvaluatorProtocol from uipath.runtime import ( UiPathRuntimeContext, UiPathRuntimeFactorySettings, @@ -16,6 +17,7 @@ ) from uipath.runtime.errors import UiPathErrorCategory +from uipath_google_adk.governance import install_governance from uipath_google_adk.runtime.config import GoogleADKConfig from uipath_google_adk.runtime.errors import ( UiPathGoogleADKErrorCode, @@ -209,6 +211,7 @@ async def _create_runtime_instance( agent: BaseAgent, runtime_id: str, entrypoint: str, + evaluator: EvaluatorProtocol | None = None, ) -> UiPathRuntimeProtocol: """ Create a runtime instance from an agent. @@ -217,7 +220,19 @@ async def _create_runtime_instance( retrieves or creates a session for the given runtime_id. Sessions persist across calls, enabling multi-turn conversations where only the current user message is sent each time. + + When ``evaluator`` is supplied, governance callbacks are installed on + the agent tree in place via :func:`install_governance` before the + ``Runner`` is created. """ + if evaluator is not None: + install_governance( + agent, + evaluator, + agent_name=entrypoint, + session_id=runtime_id, + ) + session_service = await self._get_session_service() runner = Runner( agent=agent, @@ -256,7 +271,9 @@ async def new_runtime( Args: entrypoint: Agent name from google_adk.json runtime_id: Unique identifier for the runtime instance - **kwargs: Additional keyword arguments (unused) + **kwargs: Forwarded factory kwargs. Recognized: ``evaluator`` + (``EvaluatorProtocol``) — when present, governance callbacks + are installed on the agent via :func:`install_governance`. Returns: Configured runtime instance with agent @@ -267,6 +284,7 @@ async def new_runtime( agent=agent, runtime_id=runtime_id, entrypoint=entrypoint, + evaluator=kwargs.get("evaluator"), ) async def dispose(self) -> None: diff --git a/packages/uipath-google-adk/tests/governance/test_adapter.py b/packages/uipath-google-adk/tests/governance/test_callbacks.py similarity index 79% rename from packages/uipath-google-adk/tests/governance/test_adapter.py rename to packages/uipath-google-adk/tests/governance/test_callbacks.py index 4feec4e5..3bdd06b6 100644 --- a/packages/uipath-google-adk/tests/governance/test_adapter.py +++ b/packages/uipath-google-adk/tests/governance/test_callbacks.py @@ -1,4 +1,4 @@ -"""Unit tests for the Google ADK governance adapter. +"""Unit tests for the Google ADK governance callbacks. ``can_handle`` is tested against a real ``google.adk`` ``LlmAgent`` (the adapter detects agents with ``isinstance(..., BaseAgent)``). The remaining @@ -16,10 +16,10 @@ import pytest from uipath.core.governance.exceptions import GovernanceBlockException -from uipath_google_adk.governance.adapter import ( +from uipath_google_adk.governance.callbacks import ( _BEFORE_MODEL_TEXT_CAP, - GoogleADKAdapter, GovernanceCallbacks, + install_governance, ) # -------------------------------------------------------------------------- @@ -106,37 +106,16 @@ def _make_callbacks(evaluator: FakeEvaluator) -> GovernanceCallbacks: # -------------------------------------------------------------------------- -# can_handle +# install_governance # -------------------------------------------------------------------------- -def test_can_handle_real_agent(): - from google.adk.agents import LlmAgent - - assert GoogleADKAdapter().can_handle(LlmAgent(name="t")) is True - - -def test_can_handle_rejects_non_adk_agent(): - # Duck-typed look-alikes (name + model-callback / sub_agents) must NOT be - # claimed — only a real google.adk BaseAgent is. - assert GoogleADKAdapter().can_handle(FakeLlmAgent()) is False - assert GoogleADKAdapter().can_handle(FakeContainerAgent("root", [FakeLlmAgent()])) is False - assert GoogleADKAdapter().can_handle(object()) is False - - -# -------------------------------------------------------------------------- -# attach / detach -# -------------------------------------------------------------------------- - - -def test_attach_installs_on_all_llm_agents_in_tree(): +def test_install_governance_installs_on_all_llm_agents_in_tree(): leaf_a = FakeLlmAgent("a") leaf_b = FakeLlmAgent("b") root = FakeContainerAgent("root", [leaf_a, leaf_b]) - returned = GoogleADKAdapter().attach( - root, agent_id="x", session_id="s", evaluator=FakeEvaluator() - ) + returned = install_governance(root, FakeEvaluator(), agent_name="x", session_id="s") assert returned is root # original returned, not a proxy for leaf in (leaf_a, leaf_b): @@ -146,24 +125,21 @@ def test_attach_installs_on_all_llm_agents_in_tree(): assert leaf.after_tool_callback -def test_attach_is_idempotent(): +def test_install_governance_is_idempotent(): agent = FakeLlmAgent() - adapter = GoogleADKAdapter() ev = FakeEvaluator() - adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) - adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) + install_governance(agent, ev, agent_name="x", session_id="s") + install_governance(agent, ev, agent_name="x", session_id="s") assert len(agent.before_model_callback) == 1 -def test_attach_preserves_existing_callback_and_runs_governance_first(): +def test_install_governance_preserves_existing_callback_and_runs_first(): def user_cb(*_a, **_k): return None agent = FakeLlmAgent() agent.before_model_callback = user_cb - GoogleADKAdapter().attach( - agent, agent_id="x", session_id="s", evaluator=FakeEvaluator() - ) + install_governance(agent, FakeEvaluator(), agent_name="x", session_id="s") cbs = agent.before_model_callback assert isinstance(cbs, list) and len(cbs) == 2 # governance prepended → runs first @@ -171,27 +147,74 @@ def user_cb(*_a, **_k): assert cbs[1] is user_cb -def test_detach_removes_governance_callbacks(): - def user_cb(*_a, **_k): +def test_install_governance_warns_when_no_llm_agent(caplog): + container = FakeContainerAgent("root", []) + with caplog.at_level(logging.WARNING): + install_governance(container, FakeEvaluator(), agent_name="x", session_id="s") + assert any("no LlmAgent" in r.message for r in caplog.records) + + +# -------------------------------------------------------------------------- +# Factory wiring — the evaluator kwarg drives install_governance +# -------------------------------------------------------------------------- + + +class _FakeRuntime: + APP_NAME = "app" + USER_ID = "user" + + def __init__(self, **kw: Any) -> None: + pass + + +class _FakeSessionService: + async def get_session(self, **kw: Any) -> Any: return None + async def create_session(self, **kw: Any) -> Any: + return object() + + +def _factory_without_init(): + """A factory instance that skips __init__ (avoids config/IO).""" + from uipath_google_adk.runtime.factory import UiPathGoogleADKRuntimeFactory + + return UiPathGoogleADKRuntimeFactory.__new__(UiPathGoogleADKRuntimeFactory) + + +def _stub_factory_runtime(monkeypatch, factory_mod): + """Stub Runner + runtime + session service so only the governance branch runs.""" + monkeypatch.setattr(factory_mod, "Runner", lambda **kw: None) + monkeypatch.setattr(factory_mod, "UiPathGoogleADKRuntime", _FakeRuntime) + + async def _session_service(self): + return _FakeSessionService() + + monkeypatch.setattr( + factory_mod.UiPathGoogleADKRuntimeFactory, "_get_session_service", _session_service + ) + + +async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch): + from uipath_google_adk.runtime import factory as factory_mod + + _stub_factory_runtime(monkeypatch, factory_mod) agent = FakeLlmAgent() - agent.after_tool_callback = user_cb - adapter = GoogleADKAdapter() - adapter.attach(agent, agent_id="x", session_id="s", evaluator=FakeEvaluator()) - adapter.detach(agent) - assert agent.before_model_callback is None - # unrelated user callback survives - assert agent.after_tool_callback == [user_cb] + await _factory_without_init()._create_runtime_instance( + agent=agent, runtime_id="r", entrypoint="e", evaluator=FakeEvaluator() + ) + assert isinstance(agent.before_model_callback, list) -def test_attach_warns_when_no_llm_agent(caplog): - container = FakeContainerAgent("root", []) - with caplog.at_level(logging.WARNING): - GoogleADKAdapter().attach( - container, agent_id="x", session_id="s", evaluator=FakeEvaluator() - ) - assert any("no LlmAgent" in r.message for r in caplog.records) +async def test_factory_skips_governance_without_evaluator(monkeypatch): + from uipath_google_adk.runtime import factory as factory_mod + + _stub_factory_runtime(monkeypatch, factory_mod) + agent = FakeLlmAgent() + await _factory_without_init()._create_runtime_instance( + agent=agent, runtime_id="r", entrypoint="e" + ) + assert agent.before_model_callback is None # -------------------------------------------------------------------------- diff --git a/packages/uipath-google-adk/uv.lock b/packages/uipath-google-adk/uv.lock index 7028436e..50a7a21f 100644 --- a/packages/uipath-google-adk/uv.lock +++ b/packages/uipath-google-adk/uv.lock @@ -3611,6 +3611,7 @@ dependencies = [ { name = "google-adk" }, { name = "openinference-instrumentation-google-adk" }, { name = "uipath" }, + { name = "uipath-core" }, { name = "uipath-runtime" }, ] @@ -3636,6 +3637,7 @@ requires-dist = [ { name = "google-adk", specifier = ">=1.25.1" }, { name = "openinference-instrumentation-google-adk", specifier = ">=0.1.9" }, { name = "uipath", specifier = ">=2.10.0,<2.11.0" }, + { name = "uipath-core", specifier = ">=0.5.18,<0.7.0" }, { name = "uipath-runtime", specifier = ">=0.11.0,<0.12.0" }, ] provides-extras = ["anthropic"] From ef4518ba72494ef38fbfffca2a0e3821226a5980 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 15:43:04 +0530 Subject: [PATCH 6/9] =?UTF-8?q?fix(google-adk):=20address=20governance=20r?= =?UTF-8?q?eview=20=E2=80=94=20trace=5Fid,=20rebind,=20AgentTool=20walk?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings (Viswa) for PR #362: - Drop the per-callbacks uuid trace_id (identical for every call); trace correlation is owned by the layer below, matching LangChain. Requires uipath-core >= 0.5.20, which removed trace_id from EvaluatorProtocol — bumped in the lock. - Count llm/tool calls only after governance passes: a DENY raised before the bump, inflating the counter on blocked calls. - Follow AgentTool-wrapped agents: an agent exposed as a tool lives in ``tools`` (on ``tool.agent``), not ``sub_agents``, so the walk missed it. - Refresh governance metadata on cached-agent reuse: the factory caches agents by entrypoint, so a second new_runtime with a new session_id would otherwise keep the first run's session_id (install was a no-op). Added GovernanceCallbacks.rebind() + _find_governance_callbacks(). - Log (not silently drop) when the tree walk hits the node cap. - Cap _stringify output so an oversized tool result / args / response can't hand a multi-megabyte string to the evaluator. - Trailing newline (W292). Tests: added AgentTool-walk, cached-agent rebind, and no-inflation-on-block coverage. Co-Authored-By: Claude Opus 4.8 --- .../uipath_google_adk/governance/callbacks.py | 132 ++++++++++++++---- .../tests/governance/test_callbacks.py | 57 +++++++- packages/uipath-google-adk/uv.lock | 6 +- 3 files changed, 164 insertions(+), 31 deletions(-) diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py index 18b7c531..ce64fddc 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py @@ -43,7 +43,6 @@ import json import logging from typing import Any, Dict, List -from uuid import uuid4 from uipath.core.adapters import EvaluatorProtocol from uipath.core.governance.exceptions import GovernanceBlockException @@ -58,6 +57,10 @@ # :meth:`GovernanceCallbacks._latest_request_text`. _BEFORE_MODEL_TEXT_CAP = 64000 +# Hard cap on how many nodes the agent-tree walk visits, guarding against +# cyclic or pathologically deep trees. Hitting it is logged, not silent. +_MAX_GRAPH_NODES = 1000 + # Native LlmAgent callback attribute names this adapter manages. _MODEL_BEFORE = "before_model_callback" _MODEL_AFTER = "after_model_callback" @@ -70,6 +73,23 @@ def _is_governance_callable(fn: Any) -> bool: return isinstance(getattr(fn, "__self__", None), GovernanceCallbacks) +def _find_governance_callbacks(agent: Any) -> "GovernanceCallbacks | None": + """Return the :class:`GovernanceCallbacks` already installed on ``agent``. + + Scans the four callback slots for a governance-owned callable and returns + the instance backing it, else ``None``. Used to detect a cached agent that + was governed by a previous ``new_runtime`` so its metadata can be refreshed + rather than left stale. + """ + for attr in (_MODEL_BEFORE, _MODEL_AFTER, _TOOL_BEFORE, _TOOL_AFTER): + existing = getattr(agent, attr, None) + handlers = existing if isinstance(existing, list) else [existing] + for h in handlers: + if _is_governance_callable(h): + return h.__self__ # type: ignore[no-any-return] + return None + + def _install_callback(agent: Any, attr: str, fn: Any) -> None: """Prepend ``fn`` to an ADK callback slot, preserving existing handlers. @@ -96,19 +116,28 @@ def _install_callback(agent: Any, attr: str, fn: Any) -> None: def _iter_llm_agents(root: Any) -> List[Any]: - """Return every ``LlmAgent``-shaped node in the ``sub_agents`` tree. + """Return every ``LlmAgent``-shaped node in the agent tree. A node qualifies if it exposes the model-callback surface (duck-typed via :data:`_MODEL_BEFORE` so we don't hard-require ``LlmAgent`` to be importable). Container agents (``Sequential`` / ``Parallel`` / ``Loop``) have no model callbacks themselves but their ``sub_agents`` are walked - so a multi-agent app is governed end to end. Cycles and pathological - depth are bounded by an id-visited set and a hard cap. + so a multi-agent app is governed end to end. + + ``AgentTool``-wrapped agents are also followed: an agent exposed to another + agent as a tool carries its target on ``tool.agent`` and lives in ``tools`` + (not ``sub_agents``), so it would otherwise be missed. Cycles and + pathological depth are bounded by an id-visited set and a hard cap + (``_MAX_GRAPH_NODES``), which logs rather than silently truncating. """ found: List[Any] = [] seen: set[int] = set() stack: List[Any] = [root] - while stack and len(seen) < 1000: + capped = False + while stack: + if len(seen) >= _MAX_GRAPH_NODES: + capped = True + break node = stack.pop() if node is None or id(node) in seen: continue @@ -118,6 +147,20 @@ def _iter_llm_agents(root: Any) -> List[Any]: sub_agents = getattr(node, "sub_agents", None) if isinstance(sub_agents, (list, tuple)): stack.extend(sub_agents) + # AgentTool wraps its target agent on ``.agent``; follow tools so an + # agent-as-tool is governed too. + tools = getattr(node, "tools", None) + if isinstance(tools, (list, tuple)): + for tool in tools: + wrapped = getattr(tool, "agent", None) + if wrapped is not None: + stack.append(wrapped) + if capped: + logger.warning( + "install_governance stopped walking the agent tree at the %d-node " + "cap; agents beyond it will not be governed", + _MAX_GRAPH_NODES, + ) return found @@ -139,13 +182,25 @@ def install_governance( Called by :class:`UiPathGoogleADKRuntimeFactory` when an ``evaluator`` is supplied to ``new_runtime``. """ - callbacks = GovernanceCallbacks( - evaluator=evaluator, - agent_name=agent_name, - session_id=session_id, - ) llm_agents = _iter_llm_agents(agent) + callbacks: GovernanceCallbacks | None = None for node in llm_agents: + already = _find_governance_callbacks(node) + if already is not None: + # Cached agent reused for a new runtime: refresh the evaluator and + # session/agent so governance attributes to *this* run rather than + # the first one that installed it (the factory caches agents by + # entrypoint across runtime_ids). + already.rebind( + evaluator=evaluator, agent_name=agent_name, session_id=session_id + ) + continue + if callbacks is None: + callbacks = GovernanceCallbacks( + evaluator=evaluator, + agent_name=agent_name, + session_id=session_id, + ) _install_callback(node, _MODEL_BEFORE, callbacks.before_model) _install_callback(node, _MODEL_AFTER, callbacks.after_model) _install_callback(node, _TOOL_BEFORE, callbacks.before_tool) @@ -183,9 +238,29 @@ def __init__( self._evaluator = evaluator self._agent_name = agent_name self._session_id = session_id - self._trace_id = str(uuid4()) + # ``trace_id`` is intentionally NOT held here. A single uuid minted at + # install time would be identical for every call. Trace correlation is + # owned by the layer below (OTel span / HTTP resolve at call time), + # matching the LangChain adapter. self._session_state: Dict[str, Any] = {"tool_calls": 0, "llm_calls": 0} + def rebind( + self, + evaluator: EvaluatorProtocol, + agent_name: str, + session_id: str, + ) -> None: + """Re-point this callback set at a new run. + + Called when a cached agent (already carrying these callbacks) is reused + for a fresh ``new_runtime`` — updates the evaluator and identifiers and + resets the per-run counters so state does not bleed across runtimes. + """ + self._evaluator = evaluator + self._agent_name = agent_name + self._session_id = session_id + self._session_state = {"tool_calls": 0, "llm_calls": 0} + # ----- Model callbacks ------------------------------------------------- def before_model(self, callback_context: Any, llm_request: Any) -> None: @@ -201,15 +276,16 @@ def before_model(self, callback_context: Any, llm_request: Any) -> None: Returns ``None`` so ADK proceeds with the model call. """ try: - self._session_state["llm_calls"] = ( - self._session_state.get("llm_calls", 0) + 1 - ) model_input = self._latest_request_text(llm_request) self._evaluator.evaluate_before_model( model_input=model_input, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, + ) + # Count only calls that passed governance — a DENY raises above, so + # a blocked call must not inflate the counter. + self._session_state["llm_calls"] = ( + self._session_state.get("llm_calls", 0) + 1 ) except GovernanceBlockException: raise @@ -236,7 +312,6 @@ def after_model(self, callback_context: Any, llm_response: Any) -> None: model_output=model_output, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, ) except GovernanceBlockException: raise @@ -253,18 +328,19 @@ def before_tool(self, tool: Any, args: Dict[str, Any], tool_context: Any) -> Non return would short-circuit it with a substitute result). """ try: - self._session_state["tool_calls"] = ( - self._session_state.get("tool_calls", 0) + 1 - ) tool_name = getattr(tool, "name", None) or "unknown" self._evaluator.evaluate_tool_call( tool_name=tool_name, tool_args=args or {}, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, session_state=self._session_state, ) + # Count only calls that passed governance; the evaluator saw the + # count of prior tool calls, and a DENY raises before this bump. + self._session_state["tool_calls"] = ( + self._session_state.get("tool_calls", 0) + 1 + ) except GovernanceBlockException: raise except Exception as e: @@ -292,7 +368,6 @@ def after_tool( tool_result=tool_result, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, ) except GovernanceBlockException: raise @@ -384,11 +459,16 @@ def _part_text(cls, part: Any) -> str: return "\n".join(p for p in pieces if p) @staticmethod - def _stringify(value: Any) -> str: - """Render a dict / object payload as compact, scannable text.""" + def _stringify(value: Any, cap: int = _BEFORE_MODEL_TEXT_CAP) -> str: + """Render a dict / object payload as compact, scannable text, capped. + + Bounded by ``cap`` so an oversized tool result, function-call args + blob, or function-response can't hand a multi-megabyte string to the + evaluator. + """ if isinstance(value, str): - return value + return value[:cap] try: - return json.dumps(value, default=str, ensure_ascii=False) + return json.dumps(value, default=str, ensure_ascii=False)[:cap] except (TypeError, ValueError): - return str(value) + return str(value)[:cap] diff --git a/packages/uipath-google-adk/tests/governance/test_callbacks.py b/packages/uipath-google-adk/tests/governance/test_callbacks.py index 3bdd06b6..b5479f45 100644 --- a/packages/uipath-google-adk/tests/governance/test_callbacks.py +++ b/packages/uipath-google-adk/tests/governance/test_callbacks.py @@ -61,13 +61,27 @@ def evaluate_after_tool(self, **kwargs: Any) -> None: class FakeLlmAgent: """Minimal stand-in for ``google.adk.agents.LlmAgent``.""" - def __init__(self, name: str = "agent", sub_agents: List[Any] | None = None): + def __init__( + self, + name: str = "agent", + sub_agents: List[Any] | None = None, + tools: List[Any] | None = None, + ): self.name = name self.before_model_callback: Any = None self.after_model_callback: Any = None self.before_tool_callback: Any = None self.after_tool_callback: Any = None self.sub_agents = sub_agents or [] + self.tools = tools or [] + + +class FakeAgentTool: + """Stand-in for ``google.adk.tools.agent_tool.AgentTool`` — wraps an agent.""" + + def __init__(self, agent: Any): + self.agent = agent + self.name = getattr(agent, "name", "agent_tool") class FakeContainerAgent: @@ -154,6 +168,33 @@ def test_install_governance_warns_when_no_llm_agent(caplog): assert any("no LlmAgent" in r.message for r in caplog.records) +def test_install_governance_follows_agent_tool_wrapped_agents(): + """An agent exposed to another agent via AgentTool lives in ``tools``, not + ``sub_agents`` — it must still be governed.""" + wrapped = FakeLlmAgent("researcher") + root = FakeLlmAgent("root", tools=[FakeAgentTool(wrapped)]) + install_governance(root, FakeEvaluator(), agent_name="x", session_id="s") + assert isinstance(wrapped.before_model_callback, list) + assert len(wrapped.before_model_callback) == 1 + + +def test_install_governance_rebinds_session_on_cached_agent_reuse(): + """The factory caches agents by entrypoint; a second new_runtime reuses the + same agent, so governance metadata must refresh to the new session.""" + agent = FakeLlmAgent() + install_governance(agent, FakeEvaluator(), agent_name="a", session_id="session-1") + gov = agent.before_model_callback[0].__self__ + assert gov._session_id == "session-1" + + ev2 = FakeEvaluator() + install_governance(agent, ev2, agent_name="a", session_id="session-2") + # same callback object, not re-stacked, but re-pointed at the new run + assert len(agent.before_model_callback) == 1 + assert agent.before_model_callback[0].__self__ is gov + assert gov._session_id == "session-2" + assert gov._evaluator is ev2 + + # -------------------------------------------------------------------------- # Factory wiring — the evaluator kwarg drives install_governance # -------------------------------------------------------------------------- @@ -191,7 +232,9 @@ async def _session_service(self): return _FakeSessionService() monkeypatch.setattr( - factory_mod.UiPathGoogleADKRuntimeFactory, "_get_session_service", _session_service + factory_mod.UiPathGoogleADKRuntimeFactory, + "_get_session_service", + _session_service, ) @@ -322,6 +365,16 @@ def test_after_tool_none_response(): assert ev.calls[-1][1]["tool_result"] == "" +def test_blocked_tool_call_does_not_increment_counter(): + """A DENY raises before the counter bump, so the count is not inflated.""" + ev = FakeEvaluator(block_on="tool_call") + cb = _make_callbacks(ev) + with pytest.raises(GovernanceBlockException): + cb.before_tool(FakeTool("t"), {}, tool_context=None) + assert ev.calls[-1][1]["session_state"]["tool_calls"] == 0 + assert cb._session_state["tool_calls"] == 0 + + # -------------------------------------------------------------------------- # enforcement semantics # -------------------------------------------------------------------------- diff --git a/packages/uipath-google-adk/uv.lock b/packages/uipath-google-adk/uv.lock index 50a7a21f..11a9f288 100644 --- a/packages/uipath-google-adk/uv.lock +++ b/packages/uipath-google-adk/uv.lock @@ -3591,16 +3591,16 @@ wheels = [ [[package]] name = "uipath-core" -version = "0.5.18" +version = "0.5.28" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/14/b1/d4e555a1a2ccf298195a5f2968e538b0cea8592b3e03f43fc12b178d6c69/uipath_core-0.5.18.tar.gz", hash = "sha256:63ebe8bdb818ca30a4bc9ab0ea8171315680691429931282939359ce039401ab", size = 131988, upload-time = "2026-06-08T14:04:49.688Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4c/f9/8d2f1d98cbebbcf059cf4561f38f34ad4cd58423e4f15cad22bd297a2563/uipath_core-0.5.28.tar.gz", hash = "sha256:942987f6b612c64f93d612ad7b242276ed75f129fdd8f25bc71c24ec8887e388", size = 130578, upload-time = "2026-06-30T14:04:48.841Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/57/de/1a820b33f7bff4565d7649772bc54c88480ac7e70f707097f7da37d05157/uipath_core-0.5.18-py3-none-any.whl", hash = "sha256:351d6faeecfc6a0acea93182e01526f39c04a77e09fa0444be5f4fb580463f5a", size = 54572, upload-time = "2026-06-08T14:04:48.22Z" }, + { url = "https://files.pythonhosted.org/packages/e8/1e/385bb166232a57ebe938cc57ad2717f350bc922bb5d2ce31af84306b7569/uipath_core-0.5.28-py3-none-any.whl", hash = "sha256:b952a46a21710073cbc16d6d5684e9aa645c107f57a636b778cfb94aa81a1e48", size = 54980, upload-time = "2026-06-30T14:04:47.374Z" }, ] [[package]] From e24a78381affd2f5e778bc64cfa3abb8cd875f79 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 18:25:30 +0530 Subject: [PATCH 7/9] fix(google-adk): cap before_tool args + review test coverage (page 6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the #362 review pass — the earlier commit covered AgentTool walk, _stringify cap, cached-agent rebind, trace_id drop, and the cap warning; these are the rest of page 6: - callbacks.py:249 (Major) — before_tool no longer passes tool args uncapped. New _cap_args bounds the payload: within budget the dict passes through unchanged (per-key rules still work); once its serialized size exceeds the cap it is replaced with a single capped {"_truncated": ...}. Contrast with after_tool, which already capped its result. - test_callbacks.py (Minor) — the tree test now also asserts the container agent is NOT decorated; added a huge-args cap test. (Double-attach with a pre-existing user callback is already covered by test_install_governance_preserves_existing_callback_and_runs_first.) Acknowledged, unchanged (flagged as follow-ups): - callbacks sync-only (Major): making governance async is a protocol-wide change (the evaluator is sync) — deferred, not one-off here. - _session_state accumulation across cached-agent reuse (Minor): fixed by the rebind added in the earlier commit (rebind resets the counters). Per-task safety for ParallelAgent's concurrent callbacks is a known limitation. - LIFO cap ordering (Minor): the cap now logs; which nodes are dropped past the cap is non-deterministic but all found nodes are governed. Tests: 66 pass. Co-Authored-By: Claude Opus 4.8 --- .../uipath_google_adk/governance/callbacks.py | 20 ++++++++++++++++++- .../tests/governance/test_callbacks.py | 14 +++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py b/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py index ce64fddc..5c1aaef7 100644 --- a/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py +++ b/packages/uipath-google-adk/src/uipath_google_adk/governance/callbacks.py @@ -331,7 +331,7 @@ def before_tool(self, tool: Any, args: Dict[str, Any], tool_context: Any) -> Non tool_name = getattr(tool, "name", None) or "unknown" self._evaluator.evaluate_tool_call( tool_name=tool_name, - tool_args=args or {}, + tool_args=self._cap_args(args or {}), agent_name=self._agent_name, runtime_id=self._session_id, session_state=self._session_state, @@ -458,6 +458,24 @@ def _part_text(cls, part: Any) -> str: return "\n".join(p for p in pieces if p) + @classmethod + def _cap_args(cls, args: Dict[str, Any], cap: int = _BEFORE_MODEL_TEXT_CAP) -> Any: + """Bound the tool-args payload before it reaches the evaluator. + + ``before_tool`` receives args straight from ADK; a huge blob (e.g. a + tool called with a multi-megabyte string) would otherwise be scanned + uncapped — contrast with ``after_tool``, which caps its result. Within + budget the dict is passed through unchanged (so per-key rules still + work); once its serialized size exceeds ``cap`` it is replaced with a + single capped, stringified form. + """ + if not isinstance(args, dict) or not args: + return args + blob = cls._stringify(args, cap + 1) + if len(blob) <= cap: + return args + return {"_truncated": blob[:cap]} + @staticmethod def _stringify(value: Any, cap: int = _BEFORE_MODEL_TEXT_CAP) -> str: """Render a dict / object payload as compact, scannable text, capped. diff --git a/packages/uipath-google-adk/tests/governance/test_callbacks.py b/packages/uipath-google-adk/tests/governance/test_callbacks.py index b5479f45..1e82eea9 100644 --- a/packages/uipath-google-adk/tests/governance/test_callbacks.py +++ b/packages/uipath-google-adk/tests/governance/test_callbacks.py @@ -137,6 +137,8 @@ def test_install_governance_installs_on_all_llm_agents_in_tree(): assert len(leaf.before_model_callback) == 1 assert leaf.after_model_callback and leaf.before_tool_callback assert leaf.after_tool_callback + # the container agent has no model-callback surface → must NOT be decorated + assert not hasattr(root, "before_model_callback") def test_install_governance_is_idempotent(): @@ -350,6 +352,18 @@ def test_before_tool_passes_args_and_session_state(): assert kwargs["session_state"]["tool_calls"] == 1 +def test_before_tool_caps_huge_args(): + """A huge arg blob must not reach the evaluator uncapped (contrast with the + small-args case, which passes through unchanged).""" + ev = FakeEvaluator() + cb = _make_callbacks(ev) + huge = "x" * (_BEFORE_MODEL_TEXT_CAP + 5000) + cb.before_tool(FakeTool("t"), {"blob": huge}, tool_context=None) + tool_args = ev.calls[-1][1]["tool_args"] + assert set(tool_args) == {"_truncated"} + assert len(tool_args["_truncated"]) <= _BEFORE_MODEL_TEXT_CAP + + def test_after_tool_stringifies_dict_response(): ev = FakeEvaluator() cb = _make_callbacks(ev) From 6a42ec7c40a079648e8f6b2a136ea229d224bb64 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 19:24:21 +0530 Subject: [PATCH 8/9] test(google-adk): make test_callbacks mypy-clean (fix CI lint) Same pre-existing CI-lint failure (mypy runs over tests): - FakeEvaluator evaluate_* -> (self, *args, **kwargs) -> Any so it satisfies EvaluatorProtocol; bare dict -> dict[str, Any]. - moved the Boom test-double arg-type ignore onto the reported line; func-returns-value ignores on the return-None pass-through asserts. mypy . clean (26 files); ruff + 66 tests green. Co-Authored-By: Claude Opus 4.8 --- .../tests/governance/test_callbacks.py | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/packages/uipath-google-adk/tests/governance/test_callbacks.py b/packages/uipath-google-adk/tests/governance/test_callbacks.py index 1e82eea9..39b92b04 100644 --- a/packages/uipath-google-adk/tests/governance/test_callbacks.py +++ b/packages/uipath-google-adk/tests/governance/test_callbacks.py @@ -32,29 +32,29 @@ class FakeEvaluator: def __init__(self, block_on: str | None = None) -> None: self.block_on = block_on - self.calls: List[tuple[str, dict]] = [] + self.calls: List[tuple[str, dict[str, Any]]] = [] def _record(self, hook: str, **kwargs: Any) -> None: self.calls.append((hook, kwargs)) if self.block_on == hook: raise GovernanceBlockException("blocked") # type: ignore[call-arg] - def evaluate_before_agent(self, **kwargs: Any) -> None: + def evaluate_before_agent(self, *args: Any, **kwargs: Any) -> Any: self._record("before_agent", **kwargs) - def evaluate_after_agent(self, **kwargs: Any) -> None: + def evaluate_after_agent(self, *args: Any, **kwargs: Any) -> Any: self._record("after_agent", **kwargs) - def evaluate_before_model(self, **kwargs: Any) -> None: + def evaluate_before_model(self, *args: Any, **kwargs: Any) -> Any: self._record("before_model", **kwargs) - def evaluate_after_model(self, **kwargs: Any) -> None: + def evaluate_after_model(self, *args: Any, **kwargs: Any) -> Any: self._record("after_model", **kwargs) - def evaluate_tool_call(self, **kwargs: Any) -> None: + def evaluate_tool_call(self, *args: Any, **kwargs: Any) -> Any: self._record("tool_call", **kwargs) - def evaluate_after_tool(self, **kwargs: Any) -> None: + def evaluate_after_tool(self, *args: Any, **kwargs: Any) -> Any: self._record("after_tool", **kwargs) @@ -429,9 +429,9 @@ def evaluate_before_model(self, **_): raise RuntimeError("evaluator bug") cb = GovernanceCallbacks( - evaluator=Boom(), + evaluator=Boom(), # type: ignore[arg-type] # minimal test double agent_name="a", - session_id="s", # type: ignore[arg-type] + session_id="s", ) with caplog.at_level(logging.WARNING): # must NOT raise — a governance bug can't break the agent run @@ -440,8 +440,11 @@ def evaluate_before_model(self, **_): def test_callbacks_return_none(): + # callbacks return None (ADK: a None return means "don't override the + # model/tool"); the type: ignores silence mypy's func-returns-value on the + # None-returning callbacks while the asserts document that contract. cb = _make_callbacks(FakeEvaluator()) - assert cb.before_model(None, SimpleNamespace(contents=[])) is None - assert cb.after_model(None, SimpleNamespace(partial=False, content=None)) is None - assert cb.before_tool(FakeTool("t"), {}, None) is None - assert cb.after_tool(FakeTool("t"), {}, None, {}) is None + assert cb.before_model(None, SimpleNamespace(contents=[])) is None # type: ignore[func-returns-value] + assert cb.after_model(None, SimpleNamespace(partial=False, content=None)) is None # type: ignore[func-returns-value] + assert cb.before_tool(FakeTool("t"), {}, None) is None # type: ignore[func-returns-value] + assert cb.after_tool(FakeTool("t"), {}, None, {}) is None # type: ignore[func-returns-value] From 14cd09effb123f42fb614ede61eb2db7504b5a3f Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 19:48:07 +0530 Subject: [PATCH 9/9] test(google-adk): cover swallow/extraction branches (Sonar coverage) New-code coverage ~89% -> over the 90% gate. Added non-block swallow on the model/tool callbacks and _content_text / _cap_args / _stringify helper edges. governance/callbacks.py: 89% -> 96%. Co-Authored-By: Claude Opus 4.8 --- .../tests/governance/test_callbacks.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/packages/uipath-google-adk/tests/governance/test_callbacks.py b/packages/uipath-google-adk/tests/governance/test_callbacks.py index 39b92b04..be0b4aa0 100644 --- a/packages/uipath-google-adk/tests/governance/test_callbacks.py +++ b/packages/uipath-google-adk/tests/governance/test_callbacks.py @@ -448,3 +448,56 @@ def test_callbacks_return_none(): assert cb.after_model(None, SimpleNamespace(partial=False, content=None)) is None # type: ignore[func-returns-value] assert cb.before_tool(FakeTool("t"), {}, None) is None # type: ignore[func-returns-value] assert cb.after_tool(FakeTool("t"), {}, None, {}) is None # type: ignore[func-returns-value] + + +# -------------------------------------------------------------------------- +# coverage: swallow on model/tool callbacks + extraction / helper edges +# -------------------------------------------------------------------------- + + +class _Boom: + """Evaluator whose every evaluate_* raises a non-block error.""" + + def __getattr__(self, _name: str) -> Any: + def _raise(*_a: Any, **_k: Any) -> None: + raise RuntimeError("evaluator bug") + + return _raise + + +@pytest.mark.parametrize( + "invoke", + [ + lambda cb: cb.after_model(None, SimpleNamespace(partial=False, content=None)), + lambda cb: cb.before_tool(FakeTool("t"), {}, None), + lambda cb: cb.after_tool(FakeTool("t"), {}, None, {"r": 1}), + ], +) +def test_model_tool_callbacks_swallow_non_block_errors(invoke, caplog): + cb = GovernanceCallbacks(evaluator=_Boom(), agent_name="a", session_id="s") + with caplog.at_level(logging.WARNING): + invoke(cb) # must NOT raise — a governance bug can't break the run + assert any("governance check failed" in r.message for r in caplog.records) + + +def test_content_text_and_helper_edges(): + G = GovernanceCallbacks + # _content_text: None / bare str / list-of-parts / unsupported object + assert G._content_text(None) == "" + assert G._content_text("bare") == "bare" + assert G._content_text(123) == "" + fc = SimpleNamespace(name="lookup", args={"q": "x"}) + fr = SimpleNamespace(response={"ok": 1}) + out = G._content_text( + _content( + [_part(text="hi"), _part(function_call=fc), _part(function_response=fr)] + ) + ) + assert "hi" in out and "lookup" in out and "ok" in out + # _cap_args: non-dict passes through untouched + assert G._cap_args("notdict") == "notdict" # type: ignore[arg-type] + # _stringify: str passthrough + circular-ref fallback (no crash) + assert G._stringify("hi") == "hi" + circular: dict[str, Any] = {} + circular["self"] = circular + assert isinstance(G._stringify(circular), str)