diff --git a/packages/uipath-agent-framework/pyproject.toml b/packages/uipath-agent-framework/pyproject.toml index cecd9918..b0cbc040 100644 --- a/packages/uipath-agent-framework/pyproject.toml +++ b/packages/uipath-agent-framework/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "aiosqlite>=0.20.0", "openinference-instrumentation-agent-framework>=0.1.0", "uipath>=2.10.0, <2.11.0", + "uipath-core>=0.5.18, <0.7.0", "uipath-runtime>=0.11.0, <0.12.0", ] classifiers = [ diff --git a/packages/uipath-agent-framework/src/uipath_agent_framework/governance/__init__.py b/packages/uipath-agent-framework/src/uipath_agent_framework/governance/__init__.py new file mode 100644 index 00000000..f9d6dbe0 --- /dev/null +++ b/packages/uipath-agent-framework/src/uipath_agent_framework/governance/__init__.py @@ -0,0 +1,28 @@ +"""Governance integration for ``uipath-agent-framework``. + +Exposes :func:`install_governance` — appends governance middleware +(:class:`GovernanceChatMiddleware` + :class:`GovernanceFunctionMiddleware`) to +each ``agent_framework`` agent's ``middleware`` list, governing model/tool +events (BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, AFTER_TOOL). Wired into a run by +passing an ``evaluator`` to :class:`UiPathAgentFrameworkRuntimeFactory`; 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 + +from .middleware import ( + GovernanceCallbacks, + GovernanceChatMiddleware, + GovernanceFunctionMiddleware, + install_governance, +) + +__all__ = [ + "GovernanceCallbacks", + "GovernanceChatMiddleware", + "GovernanceFunctionMiddleware", + "install_governance", +] \ No newline at end of file diff --git a/packages/uipath-agent-framework/src/uipath_agent_framework/governance/middleware.py b/packages/uipath-agent-framework/src/uipath_agent_framework/governance/middleware.py new file mode 100644 index 00000000..afe99145 --- /dev/null +++ b/packages/uipath-agent-framework/src/uipath_agent_framework/governance/middleware.py @@ -0,0 +1,404 @@ +"""Microsoft Agent Framework governance middleware for UiPath. + +Provides governance for ``agent_framework`` agents (``Agent`` and +``WorkflowAgent`` graphs). The framework runs agents through a middleware +pipeline that it rebuilds from ``agent.middleware`` on **every** ``run`` call +("Re-categorize self.middleware at runtime to support dynamic changes"). So, +like the Google ADK and OpenAI Agents integrations — and unlike the LangChain +one, which wraps the ``Runnable`` — governance is installed by appending +middleware to each agent's ``middleware`` list in place: + +- :class:`GovernanceChatMiddleware` (a ``ChatMiddleware``) brackets the LLM + call → BEFORE_MODEL before ``call_next`` / AFTER_MODEL after it. +- :class:`GovernanceFunctionMiddleware` (a ``FunctionMiddleware``) brackets a + tool call → TOOL_CALL before ``call_next`` / AFTER_TOOL after it. + +Both subclass the framework's middleware base classes because the framework's +``categorize_middleware`` sorts middleware into chat/function/agent pipelines +by ``isinstance`` — a duck-typed object would be silently dropped. + +Because the mutation is in place, :func:`install_governance` returns the +**original agent**. For a ``WorkflowAgent`` the inner agents reachable via +``workflow.executors[*]._agent`` are governed too, so a multi-agent app is +covered end to end. + +Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are owned by the +governance host, so they are not fired here. The framework's +``AgentMiddleware`` slot is therefore left untouched. + +The evaluator protocol comes from ``uipath-core``; this package contributes +only the Agent-Framework-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. Each middleware only extracts the relevant +payload and calls the matching ``evaluate_*`` method; +:class:`GovernanceBlockException` is allowed to propagate (it aborts the run), +anything else is logged and swallowed so a governance bug never breaks a run. +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Mapping +from typing import Any, Awaitable, Callable, Dict, List + +from agent_framework._middleware import ( + ChatContext, + ChatMiddleware, + FunctionInvocationContext, + FunctionMiddleware, +) +from uipath.core.adapters import 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 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 message, not the full +# prompt — see :meth:`GovernanceCallbacks._latest_message_text`. +_BEFORE_MODEL_TEXT_CAP = 64000 + +# Hard cap on how many nodes the workflow walk visits, guarding against cyclic +# or pathologically deep (nested) workflows. Hitting it is logged, not silent. +_MAX_GRAPH_NODES = 1000 + + +def install_governance( + agent: Any, + evaluator: EvaluatorProtocol, + *, + agent_name: str, + session_id: str, +) -> Any: + """Append governance middleware to the agent graph (mutated in place). + + Returns the original ``agent`` — the framework rebuilds the middleware + pipeline from ``agent.middleware`` on each ``run``, so the in-place append + is what wires governance into execution. Idempotent: an already-governed + agent is skipped. For a ``WorkflowAgent`` the inner agents reachable via + ``workflow.executors[*]._agent`` are governed too. + + Called by :class:`UiPathAgentFrameworkRuntimeFactory` when an ``evaluator`` + is supplied to ``new_runtime``. + """ + callbacks = GovernanceCallbacks( + evaluator=evaluator, agent_name=agent_name, session_id=session_id + ) + targets = _iter_agents(agent) + if not targets: + logger.warning( + "install_governance found no agent in %s — hooks will not fire", + type(agent).__name__, + ) + return agent + + installed = 0 + for node in targets: + existing = list(getattr(node, "middleware", None) or []) + if any(isinstance(m, _GOVERNANCE_MIDDLEWARE) for m in existing): + continue # idempotent — already governed + # Governance runs first so it can BLOCK before user middleware. + node.middleware = [ + GovernanceChatMiddleware(callbacks), + GovernanceFunctionMiddleware(callbacks), + *existing, + ] + installed += 1 + logger.debug("Installed governance middleware on %d agent(s)", installed) + return agent + + +def _iter_agents(root: Any) -> List[Any]: + """Return every agent node carrying a ``middleware`` slot. + + A plain ``Agent`` is itself the target. A ``WorkflowAgent`` exposes its + inner agents through ``workflow.executors[*]._agent``. Those inner agents + can themselves be ``WorkflowAgent``s (workflow-of-workflows), so the walk + **recurses** through nested workflows rather than stopping one level down — + otherwise a nested workflow's agents would run ungoverned. 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] + 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 + seen.add(id(node)) + if hasattr(node, "middleware"): + found.append(node) + workflow = getattr(node, "workflow", None) + executors = getattr(workflow, "executors", None) + if isinstance(executors, Mapping): + for executor in executors.values(): + inner = getattr(executor, "_agent", None) + if inner is not None: + stack.append(inner) + if capped: + logger.warning( + "install_governance stopped walking the agent graph at the %d-node " + "cap; agents beyond it will not be governed", + _MAX_GRAPH_NODES, + ) + return found + + +class GovernanceCallbacks: + """Holds the governance evaluator + per-attach state shared by the two + middleware classes. + + Each method extracts the relevant payload and calls the matching + ``evaluate_*`` method. :class:`GovernanceBlockException` is re-raised (it + aborts the run); anything else is logged and swallowed so a governance bug + never breaks an agent run. + """ + + 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 + # ``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} + + # ----- Model -------------------------------------------------------- + + def before_model(self, messages: Any) -> None: + """Evaluate BEFORE_MODEL on the latest message only (see ADK rationale).""" + try: + self._evaluator.evaluate_before_model( + model_input=self._latest_message_text(messages), + agent_name=self._agent_name, + runtime_id=self._session_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 + except Exception as e: # noqa: BLE001 - governance must not break the run + logger.warning("before_model governance check failed (continuing): %s", e) + + def after_model(self, result: Any) -> None: + """Evaluate AFTER_MODEL on the chat response text.""" + try: + self._evaluator.evaluate_after_model( + model_output=self._response_text(result), + agent_name=self._agent_name, + runtime_id=self._session_id, + ) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 + logger.warning("after_model governance check failed (continuing): %s", e) + + # ----- Tools -------------------------------------------------------- + + def before_tool(self, function: Any, arguments: Any) -> None: + """Evaluate TOOL_CALL with the tool name + arguments.""" + try: + self._evaluator.evaluate_tool_call( + tool_name=getattr(function, "name", None) or "unknown", + tool_args=_coerce_args(arguments), + agent_name=self._agent_name, + runtime_id=self._session_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: # noqa: BLE001 + logger.warning("before_tool governance check failed (continuing): %s", e) + + def after_tool(self, function: Any, result: Any) -> None: + """Evaluate AFTER_TOOL with the tool result.""" + try: + self._evaluator.evaluate_after_tool( + tool_name=getattr(function, "name", None) or "unknown", + tool_result="" if result is None else _stringify(result), + agent_name=self._agent_name, + runtime_id=self._session_id, + ) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 + logger.warning("after_tool governance check failed (continuing): %s", e) + + # ----- Text extraction --------------------------------------------- + # Payload text is read defensively via ``.text`` rather than + # isinstance-checking agent-framework message/response models: those + # shapes are still pre-release (rc) and not stable public types, so we + # avoid coupling extraction to types that may move. + + @classmethod + def _latest_message_text(cls, messages: Any) -> str: + """Text of the most-recent message in a chat request.""" + if not messages: + return "" + if isinstance(messages, (list, tuple)): + return cls._message_text(messages[-1]) + return cls._message_text(messages) + + @classmethod + def _message_text(cls, message: Any) -> str: + """Pull text from a ``Message`` (``.text``) or a bare string.""" + if message is None: + return "" + if isinstance(message, str): + return message[:_BEFORE_MODEL_TEXT_CAP] + text = getattr(message, "text", None) + if isinstance(text, str): + return text[:_BEFORE_MODEL_TEXT_CAP] + return _stringify(message)[:_BEFORE_MODEL_TEXT_CAP] + + @classmethod + def _response_text(cls, result: Any) -> str: + """Pull text from a ``ChatResponse`` (``.text``) or fallbacks.""" + if result is None: + return "" + text = getattr(result, "text", None) + if isinstance(text, str) and text: + return text[:_BEFORE_MODEL_TEXT_CAP] + messages = getattr(result, "messages", None) + if isinstance(messages, (list, tuple)) and messages: + return cls._message_text(messages[-1]) + return _stringify(result)[:_BEFORE_MODEL_TEXT_CAP] + + +class GovernanceChatMiddleware(ChatMiddleware): + """Brackets each LLM call: BEFORE_MODEL, then ``call_next``, then AFTER_MODEL.""" + + def __init__(self, callbacks: GovernanceCallbacks) -> None: + self._cb = callbacks + + async def process( + self, context: ChatContext, call_next: Callable[[], Awaitable[None]] + ) -> None: + self._cb.before_model(getattr(context, "messages", None)) + + if getattr(context, "stream", False): + # Streaming: after ``call_next`` ``context.result`` is a + # ResponseStream, not finalized text — reading it for AFTER_MODEL + # yields nothing. Register a result hook so AFTER_MODEL runs on the + # finalized ``ChatResponse`` the framework assembles once the stream + # is consumed. + hooks = getattr(context, "stream_result_hooks", None) + if isinstance(hooks, list): + hooks.append(self._govern_streamed_result) + else: # pragma: no cover - defensive: framework always provides it + logger.debug( + "ChatContext has no stream_result_hooks; AFTER_MODEL will " + "not run for this streamed response" + ) + await call_next() + return + + try: + await call_next() + finally: + # AFTER_MODEL must run even if the model call raised, so audit and + # rules still observe whatever result is present. + self._cb.after_model(getattr(context, "result", None)) + + def _govern_streamed_result(self, response: Any) -> Any: + """``stream_result_hook``: govern the finalized streamed ``ChatResponse``. + + Returns the response unchanged (governance observes, it does not + rewrite). A DENY raised here still propagates to abort the run. + """ + self._cb.after_model(response) + return response + + +class GovernanceFunctionMiddleware(FunctionMiddleware): + """Brackets each tool call: TOOL_CALL, then ``call_next``, then AFTER_TOOL.""" + + def __init__(self, callbacks: GovernanceCallbacks) -> None: + self._cb = callbacks + + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + function = getattr(context, "function", None) + self._cb.before_tool(function, getattr(context, "arguments", None)) + try: + await call_next() + finally: + # AFTER_TOOL must run even if the tool call raised. + self._cb.after_tool(function, getattr(context, "result", None)) + + +# Tuple used for isinstance idempotency / detach checks. +_GOVERNANCE_MIDDLEWARE = (GovernanceChatMiddleware, GovernanceFunctionMiddleware) + + +# -------------------------------------------------------------------------- +# Helpers +# -------------------------------------------------------------------------- + + +def _coerce_args(arguments: Any) -> Dict[str, Any]: + """Normalise tool arguments (Mapping / pydantic model / None) to a dict.""" + if arguments is None: + return {} + if isinstance(arguments, Mapping): + return dict(arguments) + model_dump = getattr(arguments, "model_dump", None) + if callable(model_dump): + try: + dumped = model_dump() + if isinstance(dumped, dict): + return dumped + except Exception as e: # noqa: BLE001 + # Don't silently drop the args from governance visibility — surface + # that they couldn't be coerced. + logger.warning( + "governance: could not coerce %s tool args to a dict (%s); " + "TOOL_CALL will see empty args", + type(arguments).__name__, + e, + ) + return {} + + +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 or message payload can't + hand a multi-megabyte string to the evaluator. Callers that slice the + result again (the ``_message_text`` / ``_response_text`` fallbacks) are + unaffected. + """ + if isinstance(value, str): + return value[:cap] + try: + return json.dumps(value, default=str, ensure_ascii=False)[:cap] + except (TypeError, ValueError): + return str(value)[:cap] diff --git a/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py index 1b3bfb42..4576e1d4 100644 --- a/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py +++ b/packages/uipath-agent-framework/src/uipath_agent_framework/runtime/factory.py @@ -11,6 +11,7 @@ ) from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider +from uipath.core.adapters import EvaluatorProtocol from uipath.platform.resume_triggers import UiPathResumeTriggerHandler from uipath.runtime import ( UiPathResumableRuntime, @@ -21,6 +22,7 @@ ) from uipath.runtime.errors import UiPathErrorCategory +from uipath_agent_framework.governance import install_governance from uipath_agent_framework.runtime.config import AgentFrameworkConfig from uipath_agent_framework.runtime.errors import ( UiPathAgentFrameworkErrorCode, @@ -202,13 +204,25 @@ async def _create_runtime_instance( agent: WorkflowAgent, runtime_id: str, entrypoint: str, + evaluator: EvaluatorProtocol | None = None, ) -> UiPathRuntimeProtocol: """Create a runtime instance from an agent. Creates the runtime with a shared SqliteResumableStorage for persistent conversation history and HITL trigger management. Wraps with UiPathResumableRuntime for resume trigger lifecycle handling. + + When ``evaluator`` is supplied, governance middleware is installed on + the agent graph in place via :func:`install_governance`. """ + if evaluator is not None: + install_governance( + agent, + evaluator, + agent_name=entrypoint, + session_id=runtime_id, + ) + storage = await self._get_storage() assert storage.checkpoint_storage is not None checkpoint_storage = ScopedCheckpointStorage( @@ -239,7 +253,9 @@ async def new_runtime( Args: entrypoint: Agent name from agent_framework.json runtime_id: Unique identifier for the runtime instance - **kwargs: Additional keyword arguments (unused) + **kwargs: Forwarded factory kwargs. Recognized: ``evaluator`` + (``EvaluatorProtocol``) — when present, governance middleware + is installed on the agent via :func:`install_governance`. Returns: Configured runtime instance with agent @@ -250,6 +266,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-agent-framework/tests/governance/__init__.py b/packages/uipath-agent-framework/tests/governance/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/uipath-agent-framework/tests/governance/test_middleware.py b/packages/uipath-agent-framework/tests/governance/test_middleware.py new file mode 100644 index 00000000..d031c78b --- /dev/null +++ b/packages/uipath-agent-framework/tests/governance/test_middleware.py @@ -0,0 +1,485 @@ +"""Unit tests for the Microsoft Agent Framework governance middleware. + +The middleware classes subclass ``agent_framework`` base classes (the +framework routes middleware by ``isinstance``), so importing the adapter +requires ``agent-framework-core`` — but the messages / responses / tools / +contexts under test are lightweight duck-typed fakes. + +The package is configured with ``asyncio_mode = "auto"``, so ``async def`` +tests run without an explicit marker. +""" + +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_agent_framework.governance.middleware import ( + _BEFORE_MODEL_TEXT_CAP, + GovernanceCallbacks, + GovernanceChatMiddleware, + GovernanceFunctionMiddleware, + install_governance, +) + +# -------------------------------------------------------------------------- +# 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[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, *args: Any, **kwargs: Any) -> Any: + self._record("before_agent", **kwargs) + + def evaluate_after_agent(self, *args: Any, **kwargs: Any) -> Any: + self._record("after_agent", **kwargs) + + def evaluate_before_model(self, *args: Any, **kwargs: Any) -> Any: + self._record("before_model", **kwargs) + + def evaluate_after_model(self, *args: Any, **kwargs: Any) -> Any: + self._record("after_model", **kwargs) + + def evaluate_tool_call(self, *args: Any, **kwargs: Any) -> Any: + self._record("tool_call", **kwargs) + + def evaluate_after_tool(self, *args: Any, **kwargs: Any) -> Any: + self._record("after_tool", **kwargs) + + +class FakeAgent: + """Minimal stand-in for an ``agent_framework`` Agent (duck-typed).""" + + def __init__(self, name: str = "agent"): + self.name = name + self.middleware: Any = None + + async def run(self, *_a: Any, **_k: Any) -> None: # marks it as an agent + return None + + +class FakeWorkflowAgent: + """Stand-in for ``WorkflowAgent`` exposing inner agents via executors.""" + + def __init__(self, inner_agents: List[Any]): + self.middleware: Any = None + executors = { + f"e{i}": SimpleNamespace(_agent=a) for i, a in enumerate(inner_agents) + } + self.workflow = SimpleNamespace(executors=executors) + + +class FakeTool: + def __init__(self, name: str): + self.name = name + + +def _msg(text: str) -> SimpleNamespace: + return SimpleNamespace(text=text) + + +def _make_callbacks(ev: FakeEvaluator) -> GovernanceCallbacks: + return GovernanceCallbacks(evaluator=ev, agent_name="agent-1", session_id="sess-1") + + +async def _noop_next() -> None: + return None + + +# -------------------------------------------------------------------------- +# install_governance +# -------------------------------------------------------------------------- + + +def test_install_governance_appends_both_middleware(): + agent = FakeAgent() + returned = install_governance( + agent, FakeEvaluator(), agent_name="x", session_id="s" + ) + assert returned is agent + kinds = [type(m) for m in agent.middleware] + assert GovernanceChatMiddleware in kinds + assert GovernanceFunctionMiddleware in kinds + + +def test_install_governance_installs_on_workflow_inner_agents(): + a, b = FakeAgent("a"), FakeAgent("b") + root = FakeWorkflowAgent([a, b]) + install_governance(root, FakeEvaluator(), agent_name="x", session_id="s") + for node in (a, b): + assert any(isinstance(m, GovernanceChatMiddleware) for m in node.middleware) + + +def test_install_governance_recurses_into_nested_workflow(): + """A WorkflowAgent inside a WorkflowAgent (workflow-of-workflows): the deep + leaf agent must still be governed, not left one level below the walk.""" + leaf = FakeAgent("leaf") + inner = FakeWorkflowAgent([leaf]) + root = FakeWorkflowAgent([inner]) + install_governance(root, FakeEvaluator(), agent_name="x", session_id="s") + assert any(isinstance(m, GovernanceChatMiddleware) for m in leaf.middleware) + + +def test_install_governance_is_cycle_safe(): + """A workflow whose executor points back at itself must not loop forever.""" + w = FakeWorkflowAgent([]) + w.workflow.executors = {"self": SimpleNamespace(_agent=w)} + # completes (id-visited set breaks the cycle) and governs w exactly once + install_governance(w, FakeEvaluator(), agent_name="x", session_id="s") + assert sum(isinstance(m, GovernanceChatMiddleware) for m in w.middleware) == 1 + + +def test_install_governance_is_idempotent(): + agent = FakeAgent() + ev = FakeEvaluator() + install_governance(agent, ev, agent_name="x", session_id="s") + install_governance(agent, ev, agent_name="x", session_id="s") + assert sum(isinstance(m, GovernanceChatMiddleware) for m in agent.middleware) == 1 + + +def test_install_governance_preserves_existing_middleware_and_runs_first(): + user_mw = object() + agent = FakeAgent() + agent.middleware = [user_mw] + install_governance(agent, FakeEvaluator(), agent_name="x", session_id="s") + # governance prepended → runs first; user middleware preserved at the end + assert isinstance(agent.middleware[0], GovernanceChatMiddleware) + assert agent.middleware[-1] is user_mw + + +def test_install_governance_warns_when_no_agent(caplog): + with caplog.at_level(logging.WARNING): + install_governance(object(), FakeEvaluator(), agent_name="x", session_id="s") + assert any("no agent" in r.message for r in caplog.records) + + +# -------------------------------------------------------------------------- +# Factory wiring — the evaluator kwarg drives install_governance +# -------------------------------------------------------------------------- + + +def _factory_without_init(): + """A factory instance that skips __init__ (avoids config/IO).""" + from uipath_agent_framework.runtime.factory import ( + UiPathAgentFrameworkRuntimeFactory, + ) + + return UiPathAgentFrameworkRuntimeFactory.__new__( + UiPathAgentFrameworkRuntimeFactory + ) + + +def _stub_factory_runtime(monkeypatch, factory_mod): + """Stub storage + runtime constructions so only the governance branch runs.""" + monkeypatch.setattr(factory_mod, "ScopedCheckpointStorage", lambda *a, **k: None) + monkeypatch.setattr( + factory_mod, "UiPathAgentFrameworkRuntime", lambda **kw: SimpleNamespace(**kw) + ) + monkeypatch.setattr( + factory_mod, "UiPathResumableRuntime", lambda **kw: SimpleNamespace(**kw) + ) + monkeypatch.setattr(factory_mod, "UiPathResumeTriggerHandler", lambda *a, **k: None) + + async def _storage(self): + return SimpleNamespace(checkpoint_storage=object()) + + monkeypatch.setattr( + factory_mod.UiPathAgentFrameworkRuntimeFactory, "_get_storage", _storage + ) + + +async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch): + from uipath_agent_framework.runtime import factory as factory_mod + + _stub_factory_runtime(monkeypatch, factory_mod) + agent = FakeAgent() + await _factory_without_init()._create_runtime_instance( + agent=agent, runtime_id="r", entrypoint="e", evaluator=FakeEvaluator() + ) + assert any(isinstance(m, GovernanceChatMiddleware) for m in agent.middleware) + + +async def test_factory_skips_governance_without_evaluator(monkeypatch): + from uipath_agent_framework.runtime import factory as factory_mod + + _stub_factory_runtime(monkeypatch, factory_mod) + agent = FakeAgent() + await _factory_without_init()._create_runtime_instance( + agent=agent, runtime_id="r", entrypoint="e" + ) + assert agent.middleware is None + + +# -------------------------------------------------------------------------- +# ChatMiddleware → BEFORE_MODEL / AFTER_MODEL +# -------------------------------------------------------------------------- + + +async def test_chat_middleware_brackets_call_with_before_and_after(): + ev = FakeEvaluator() + mw = GovernanceChatMiddleware(_make_callbacks(ev)) + order: List[str] = [] + + async def call_next() -> None: + order.append("model_call") + + context: Any = SimpleNamespace( + messages=[_msg("old"), _msg("the question")], + result=SimpleNamespace(text="the answer"), + ) + await mw.process(context, call_next) + + hooks = [h for h, _ in ev.calls] + assert hooks == ["before_model", "after_model"] + assert order == ["model_call"] + assert ev.calls[0][1]["model_input"] == "the question" # latest only + assert ev.calls[1][1]["model_output"] == "the answer" + + +async def test_chat_middleware_caps_text(): + ev = FakeEvaluator() + mw = GovernanceChatMiddleware(_make_callbacks(ev)) + huge = "x" * (_BEFORE_MODEL_TEXT_CAP + 5000) + context: Any = SimpleNamespace( + messages=[_msg(huge)], result=SimpleNamespace(text="") + ) + await mw.process(context, _noop_next) + assert len(ev.calls[0][1]["model_input"]) <= _BEFORE_MODEL_TEXT_CAP + + +async def test_after_model_runs_even_when_model_call_raises(): + """AFTER_MODEL must fire from the finally so audit/rules still observe the + turn, and the underlying error must still propagate.""" + ev = FakeEvaluator() + mw = GovernanceChatMiddleware(_make_callbacks(ev)) + + async def boom_next() -> None: + raise RuntimeError("model exploded") + + context: Any = SimpleNamespace( + messages=[_msg("hi")], result=SimpleNamespace(text="") + ) + with pytest.raises(RuntimeError, match="model exploded"): + await mw.process(context, boom_next) + assert [h for h, _ in ev.calls] == ["before_model", "after_model"] + + +async def test_streaming_governs_finalized_response_via_result_hook(): + """Streaming: context.result is a ResponseStream after call_next, so + AFTER_MODEL runs from a stream_result_hook on the finalized ChatResponse.""" + ev = FakeEvaluator() + mw = GovernanceChatMiddleware(_make_callbacks(ev)) + context: Any = SimpleNamespace( + messages=[_msg("the question")], + stream=True, + stream_result_hooks=[], + result=None, + ) + await mw.process(context, _noop_next) + + # BEFORE_MODEL fired; AFTER_MODEL deferred to the registered hook. + assert [h for h, _ in ev.calls] == ["before_model"] + assert len(context.stream_result_hooks) == 1 + + finalized = SimpleNamespace(text="the streamed answer") + returned = context.stream_result_hooks[0](finalized) + assert returned is finalized # hook returns the response unchanged + assert [h for h, _ in ev.calls] == ["before_model", "after_model"] + assert ev.calls[-1][1]["model_output"] == "the streamed answer" + + +# -------------------------------------------------------------------------- +# FunctionMiddleware → TOOL_CALL / AFTER_TOOL +# -------------------------------------------------------------------------- + + +async def test_function_middleware_passes_name_args_and_result(): + ev = FakeEvaluator() + mw = GovernanceFunctionMiddleware(_make_callbacks(ev)) + order: List[str] = [] + + async def call_next() -> None: + order.append("tool_call") + + context: Any = SimpleNamespace( + function=FakeTool("transfer"), + arguments={"amount": 50}, + result={"status": "ok"}, + ) + await mw.process(context, call_next) + + hooks = [h for h, _ in ev.calls] + assert hooks == ["tool_call", "after_tool"] + assert order == ["tool_call"] + assert ev.calls[0][1]["tool_name"] == "transfer" + assert ev.calls[0][1]["tool_args"] == {"amount": 50} + assert ev.calls[0][1]["session_state"]["tool_calls"] == 1 + assert "ok" in ev.calls[1][1]["tool_result"] + + +async def test_function_middleware_coerces_pydantic_args(): + ev = FakeEvaluator() + mw = GovernanceFunctionMiddleware(_make_callbacks(ev)) + args = SimpleNamespace(model_dump=lambda: {"x": 1}) + context: Any = SimpleNamespace(function=FakeTool("t"), arguments=args, result=None) + await mw.process(context, _noop_next) + assert ev.calls[0][1]["tool_args"] == {"x": 1} + assert ev.calls[1][1]["tool_result"] == "" # None result → "" + + +async def test_after_tool_runs_even_when_tool_call_raises(): + ev = FakeEvaluator() + mw = GovernanceFunctionMiddleware(_make_callbacks(ev)) + + async def boom_next() -> None: + raise RuntimeError("tool exploded") + + context: Any = SimpleNamespace(function=FakeTool("t"), arguments={}, result=None) + with pytest.raises(RuntimeError, match="tool exploded"): + await mw.process(context, boom_next) + assert [h for h, _ in ev.calls] == ["tool_call", "after_tool"] + + +def test_blocked_before_tool_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"), {}) + assert ev.calls[-1][1]["session_state"]["tool_calls"] == 0 + assert cb._session_state["tool_calls"] == 0 + + +# -------------------------------------------------------------------------- +# enforcement semantics +# -------------------------------------------------------------------------- + + +async def test_block_in_before_model_aborts_before_call_next(): + ev = FakeEvaluator(block_on="before_model") + mw = GovernanceChatMiddleware(_make_callbacks(ev)) + called = {"next": False} + + async def call_next() -> None: + called["next"] = True + + context: Any = SimpleNamespace(messages=[_msg("hi")], result=None) + with pytest.raises(GovernanceBlockException): + await mw.process(context, call_next) + assert called["next"] is False # tool/model never ran + + +async def test_block_in_before_tool_aborts_before_call_next(): + ev = FakeEvaluator(block_on="tool_call") + mw = GovernanceFunctionMiddleware(_make_callbacks(ev)) + called = {"next": False} + + async def call_next() -> None: + called["next"] = True + + context: Any = SimpleNamespace(function=FakeTool("t"), arguments={}, result=None) + with pytest.raises(GovernanceBlockException): + await mw.process(context, call_next) + assert called["next"] is False + + +async def test_non_block_exception_is_swallowed(caplog): + class Boom: + def evaluate_before_model(self, **_: Any) -> None: + raise RuntimeError("evaluator bug") + + cb = GovernanceCallbacks(evaluator=Boom(), agent_name="a", session_id="s") # type: ignore[arg-type] + mw = GovernanceChatMiddleware(cb) + with caplog.at_level(logging.WARNING): + await mw.process( + SimpleNamespace(messages=[_msg("x")], result=None), # type: ignore[arg-type] + _noop_next, + ) + assert any("governance check failed" in r.message for r in caplog.records) + + +# -------------------------------------------------------------------------- +# coverage: swallow on every callback + extraction / _coerce_args 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.before_model([_msg("x")]), + lambda cb: cb.after_model(SimpleNamespace(text="y")), + lambda cb: cb.before_tool(FakeTool("t"), {}), + lambda cb: cb.after_tool(FakeTool("t"), {"r": 1}), + ], +) +def test_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_text_extraction_and_coerce_edges(): + from uipath_agent_framework.governance.middleware import _coerce_args, _stringify + + M = GovernanceCallbacks + # _latest_message_text: empty + single (non-list) + assert M._latest_message_text([]) == "" + assert M._latest_message_text(SimpleNamespace(text="solo")) == "solo" + # _message_text: None / str / object-without-.text -> _stringify fallback + assert M._message_text(None) == "" + assert M._message_text("plain") == "plain" + assert isinstance(M._message_text(SimpleNamespace()), str) + # _response_text: None / .text / .messages[-1] / _stringify fallback + assert M._response_text(None) == "" + assert M._response_text(SimpleNamespace(text="via")) == "via" + assert "m" in M._response_text( + SimpleNamespace(text=None, messages=[SimpleNamespace(text="m")]) + ) + assert isinstance(M._response_text(SimpleNamespace(text=None, messages=None)), str) + # _coerce_args: None / Mapping / model_dump / non-coercible + assert _coerce_args(None) == {} + assert _coerce_args({"a": 1}) == {"a": 1} + assert _coerce_args(SimpleNamespace(model_dump=lambda: {"b": 2})) == {"b": 2} + assert _coerce_args(object()) == {} + # _stringify: str passthrough + circular-ref fallback (no crash) + assert _stringify("hi") == "hi" + circular: dict[str, Any] = {} + circular["self"] = circular + assert isinstance(_stringify(circular), str) + + +def test_coerce_args_warns_on_model_dump_failure(caplog): + from uipath_agent_framework.governance.middleware import _coerce_args + + def _bad() -> dict[str, Any]: + raise ValueError("boom") + + with caplog.at_level(logging.WARNING): + assert _coerce_args(SimpleNamespace(model_dump=_bad)) == {} + assert any("could not coerce" in r.message for r in caplog.records) diff --git a/packages/uipath-agent-framework/uv.lock b/packages/uipath-agent-framework/uv.lock index e551beac..721f5961 100644 --- a/packages/uipath-agent-framework/uv.lock +++ b/packages/uipath-agent-framework/uv.lock @@ -2486,6 +2486,7 @@ dependencies = [ { name = "aiosqlite" }, { name = "openinference-instrumentation-agent-framework" }, { name = "uipath" }, + { name = "uipath-core" }, { name = "uipath-runtime" }, ] @@ -2515,6 +2516,7 @@ requires-dist = [ { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.43.0" }, { name = "openinference-instrumentation-agent-framework", specifier = ">=0.1.0" }, { 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"] @@ -2532,16 +2534,16 @@ dev = [ [[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]]