diff --git a/pyproject.toml b/pyproject.toml index b0b413a..503bda8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-runtime" -version = "0.12.2" +version = "0.12.3" description = "Runtime abstractions and interfaces for building agents and automation scripts in the UiPath ecosystem" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" @@ -9,6 +9,7 @@ dependencies = [ "vaderSentiment>=3.3.2, <4.0", "chardet>=5.2.0, <8.0", ] + classifiers = [ "Intended Audience :: Developers", "Topic :: Software Development :: Build Tools", @@ -20,6 +21,9 @@ maintainers = [ { name = "Cristian Pufu", email = "cristian.pufu@uipath.com" }, ] +[project.optional-dependencies] +governance-rego = ["opa-wasmtime>=0.1.1"] + [project.urls] Homepage = "https://uipath.com" Repository = "https://github.com/UiPath/uipath-runtime-python" diff --git a/src/uipath/runtime/governance/_audit/console.py b/src/uipath/runtime/governance/_audit/console.py new file mode 100644 index 0000000..2b04aa1 --- /dev/null +++ b/src/uipath/runtime/governance/_audit/console.py @@ -0,0 +1,118 @@ +"""Console audit sink for human-readable governance output. + +Writes audit events via Python logging (not stderr print), useful for +debugging and development. +""" +from __future__ import annotations + +import json +import logging + +from .base import AuditEvent, AuditSink, EventType + +_logger = logging.getLogger("uipath.governance.audit.console") + + +class ConsoleAuditSink(AuditSink): + """Audit sink that writes governance events via logging. + + Args: + verbose: If True, show all events. If False, only show matches. + """ + + def __init__(self, verbose: bool = False) -> None: + """Configure the sink's verbosity (verbose shows every event).""" + self._verbose = verbose + + @property + def name(self) -> str: + """Constant sink identifier.""" + return "console" + + def accepts(self, event: AuditEvent) -> bool: + """Filter to matched rules and lifecycle events unless verbose.""" + if self._verbose: + return True + if event.event_type == EventType.RULE_EVALUATION: + return event.data.get("matched", False) + return event.event_type in ( + EventType.SESSION_START, + EventType.SESSION_END, + EventType.HOOK_END, + EventType.POLICY_VIOLATION, + ) + + def emit(self, event: AuditEvent) -> None: + """Write the event to the log using the appropriate formatter.""" + if event.event_type == EventType.RULE_EVALUATION: + self._emit_rule_evaluation(event) + elif event.event_type == EventType.HOOK_END: + self._emit_hook_summary(event) + elif event.event_type == EventType.SESSION_START: + self._emit_session_start(event) + elif event.event_type == EventType.SESSION_END: + self._emit_session_end(event) + else: + self._emit_generic(event) + + def _emit_rule_evaluation(self, event: AuditEvent) -> None: + data = event.data + matched = data.get("matched", False) + status = "MATCHED" if matched else "PASS" + policy_id = data.get("policy_id", "?") + rule_name = data.get("rule_name", "?") + action = data.get("action", "?").upper() + detail = data.get("detail", "") + + if matched: + _logger.warning( + "[GOVERNANCE] [%s] %s | %s | action=%s | %s", + status, policy_id, rule_name, action, detail, + ) + elif self._verbose: + _logger.info("[GOVERNANCE] [%s] %s | %s", status, policy_id, rule_name) + + def _emit_hook_summary(self, event: AuditEvent) -> None: + data = event.data + hook = event.hook + total = data.get("total_rules", 0) + matched = data.get("matched_rules", 0) + action = data.get("final_action", "allow").upper() + mode = data.get("enforcement_mode") + mode_str = mode.value if hasattr(mode, "value") else str(mode or "audit") + + if mode_str == "audit" and action == "DENY": + action = "AUDIT (would deny)" + + level = logging.WARNING if matched else logging.INFO + _logger.log( + level, + "[GOVERNANCE] HOOK: %s | rules=%d | matched=%d | action=%s", + hook, total, matched, action, + ) + + def _emit_session_start(self, event: AuditEvent) -> None: + data = event.data + packs = data.get("packs", []) + mode = data.get("enforcement_mode") + mode_str = mode.value if hasattr(mode, "value") else str(mode or "audit") + _logger.info( + "[GOVERNANCE] Session started | agent=%s | packs=%s | mode=%s", + event.agent_name, ",".join(packs), mode_str, + ) + + def _emit_session_end(self, event: AuditEvent) -> None: + data = event.data + total = data.get("total_evaluations", 0) + matched = data.get("rules_matched", 0) + denied = data.get("rules_denied", 0) + _logger.info( + "[GOVERNANCE] Session ended | evaluations=%d | matched=%d | denied=%d", + total, matched, denied, + ) + + def _emit_generic(self, event: AuditEvent) -> None: + _logger.info( + "[GOVERNANCE] %s | %s | %s", + event.event_type, event.agent_name, json.dumps(event.data, default=str), + ) diff --git a/src/uipath/runtime/governance/_audit/factory.py b/src/uipath/runtime/governance/_audit/factory.py index 334f867..db7fab0 100644 --- a/src/uipath/runtime/governance/_audit/factory.py +++ b/src/uipath/runtime/governance/_audit/factory.py @@ -29,5 +29,10 @@ def create_sink(name: str) -> AuditSink | None: return TracesAuditSink() + if name == "console": + from .console import ConsoleAuditSink + + return ConsoleAuditSink() + logger.warning("Unknown audit sink: %s", name) return None diff --git a/src/uipath/runtime/governance/config.py b/src/uipath/runtime/governance/config.py new file mode 100644 index 0000000..b2fbee0 --- /dev/null +++ b/src/uipath/runtime/governance/config.py @@ -0,0 +1,69 @@ +"""Runtime-level governance enforcement-mode state. + +``EnforcementMode`` is defined in :mod:`uipath.core.governance` and +re-exported here for backward compatibility. The get/set/reset helpers +below are *per-process* state — the native policy loader sets the mode +on each successful backend fetch. +""" + +from __future__ import annotations + +import logging +import os + +from uipath.core.governance import EnforcementMode + +logger = logging.getLogger(__name__) + +ENV_ENFORCEMENT_MODE = "UIPATH_GOVERNANCE_MODE" + +__all__ = [ + "EnforcementMode", + "ENV_ENFORCEMENT_MODE", + "get_enforcement_mode", + "set_enforcement_mode", + "reset_enforcement_mode", +] + +_enforcement_mode: EnforcementMode | None = None + + +def get_enforcement_mode() -> EnforcementMode: + """Return the current enforcement mode. + + Resolution order: + + 1. A value previously set via :func:`set_enforcement_mode` (the + policy loader calls this with the backend-supplied mode on every + successful fetch — that's the canonical source). + 2. ``UIPATH_GOVERNANCE_MODE`` env var (developer override). + 3. Default :attr:`EnforcementMode.AUDIT` — evaluate and log without + blocking. + """ + global _enforcement_mode + if _enforcement_mode is not None: + return _enforcement_mode + + mode_str = os.getenv(ENV_ENFORCEMENT_MODE, "audit").lower() + try: + _enforcement_mode = EnforcementMode(mode_str) + except ValueError: + _enforcement_mode = EnforcementMode.AUDIT + + return _enforcement_mode + + +def set_enforcement_mode(mode: EnforcementMode) -> None: + """Set the enforcement mode programmatically. + + The policy loader calls this with the backend-supplied mode on each + fetch so the evaluator picks up the platform-controlled value. + """ + global _enforcement_mode + _enforcement_mode = mode + + +def reset_enforcement_mode() -> None: + """Clear cached enforcement mode (intended for tests).""" + global _enforcement_mode + _enforcement_mode = None diff --git a/src/uipath/runtime/governance/delegation_guard.py b/src/uipath/runtime/governance/delegation_guard.py new file mode 100644 index 0000000..5a79dea --- /dev/null +++ b/src/uipath/runtime/governance/delegation_guard.py @@ -0,0 +1,205 @@ +"""Delegation depth guard. + +Patches an agent's ``invoke`` method to track recursion depth and raise +a ``GovernanceBlockException`` when the configured maximum is exceeded. +This prevents runaway sub-agent chains. +""" + +from __future__ import annotations + +import asyncio +import functools +import logging +import os +from contextvars import ContextVar, Token +from typing import Any + +from uipath.core.governance.exceptions import ( + GovernanceBlockException, + GovernanceViolation, +) + +logger = logging.getLogger(__name__) + +_DEFAULT_MAX_DELEGATION_DEPTH = 25 +_ENV_MAX_DELEGATION_DEPTH = "UIPATH_GOVERNANCE_MAX_DELEGATION_DEPTH" + +# Single module-level ContextVar holding per-agent delegation depths +# keyed by ``id(agent)``. Each install / uninstall pair shares this one +# ContextVar instead of allocating a new one per agent — the interpreter +# interns ContextVars and never GCs them, so per-agent allocation was an +# unbounded leak in long-running hosts (every `install_delegation_guard` +# call permanently grew the interpreter's ContextVar registry). +# +# Per-context isolation (asyncio task / thread) still works the standard +# ContextVar way: each context sees its own copy of the depths dict, and +# nested invokes use ``set`` / ``reset`` for LIFO depth tracking. The +# dict itself is copied on every increment (copy-on-write) so concurrent +# contexts don't share state through a mutable mapping. +_DELEGATION_DEPTHS: ContextVar[dict[int, int]] = ContextVar( + "_uipath_delegation_depths" +) + + +def _current_depth(agent_key: int) -> int: + """Return the current depth for ``agent_key`` in this context.""" + try: + return _DELEGATION_DEPTHS.get().get(agent_key, 0) + except LookupError: + return 0 + + +def _enter_depth_if_under( + agent_key: int, max_depth: int +) -> tuple[int, Token[dict[int, int]] | None]: + """Attempt to increment depth for ``agent_key``. + + Returns ``(new_depth, token)`` where ``token`` is ``None`` if the + new depth would exceed ``max_depth`` — caller raises and does not + need to clean up. On success, caller must reset via ``token``. + """ + try: + depths = _DELEGATION_DEPTHS.get() + except LookupError: + depths = {} + new_depth = depths.get(agent_key, 0) + 1 + if new_depth > max_depth: + return new_depth, None + new_depths = dict(depths) + new_depths[agent_key] = new_depth + token = _DELEGATION_DEPTHS.set(new_depths) + return new_depth, token + + +def _exit_depth(token: Token[dict[int, int]]) -> None: + """Undo a successful :func:`_enter_depth_if_under` call.""" + try: + _DELEGATION_DEPTHS.reset(token) + except (ValueError, LookupError): + logger.debug("Delegation depth reset from foreign context") + + +def _resolve_max_depth() -> int: + """Read max-depth from env at call time, falling back to default on parse error.""" + raw = os.getenv(_ENV_MAX_DELEGATION_DEPTH) + if raw is None: + return _DEFAULT_MAX_DELEGATION_DEPTH + try: + return int(raw) + except ValueError: + logger.warning( + "Invalid %s=%r; using default %d", + _ENV_MAX_DELEGATION_DEPTH, + raw, + _DEFAULT_MAX_DELEGATION_DEPTH, + ) + return _DEFAULT_MAX_DELEGATION_DEPTH + + +def _build_violation(current: int, resolved_max: int) -> GovernanceBlockException: + """Build the depth-exceeded exception (shared by sync and async guards).""" + return GovernanceBlockException.from_violation( + GovernanceViolation( + rule_id="ASI-02", + rule_name="Excessive Agency", + detail=f"Delegation depth {current} exceeds max {resolved_max}", + ) + ) + + +def _wrap_invoke(original: Any, agent_key: int, resolved_max: int) -> Any: + """Return a depth-guarded wrapper matching the sync/async shape of ``original``.""" + if asyncio.iscoroutinefunction(original): + + @functools.wraps(original) + async def _guarded_async(input_data: Any, **kwargs: Any) -> Any: + current, token = _enter_depth_if_under(agent_key, resolved_max) + if token is None: + raise _build_violation(current, resolved_max) + try: + return await original(input_data, **kwargs) + finally: + _exit_depth(token) + + return _guarded_async + + @functools.wraps(original) + def _guarded_sync(input_data: Any, **kwargs: Any) -> Any: + current, token = _enter_depth_if_under(agent_key, resolved_max) + if token is None: + raise _build_violation(current, resolved_max) + try: + return original(input_data, **kwargs) + finally: + _exit_depth(token) + + return _guarded_sync + + +_GUARDED_METHODS = ("invoke", "ainvoke") + + +def install_delegation_guard(agent: Any, max_depth: int | None = None) -> None: + """Patch the agent's invoke methods to enforce a maximum delegation depth. + + Patches both ``invoke`` and ``ainvoke`` when present. No-op when + neither attribute exists or the agent has already been guarded. + """ + if max_depth is None: + max_depth = _resolve_max_depth() + if getattr(agent, "_delegation_wrapped", False): + return + + originals = { + name: getattr(agent, name, None) + for name in _GUARDED_METHODS + if callable(getattr(agent, name, None)) + } + if not originals: + return + + agent_key = id(agent) + resolved_max = max_depth + + for name, original in originals.items(): + try: + setattr(agent, name, _wrap_invoke(original, agent_key, resolved_max)) + setattr(agent, f"_uipath_original_{name}", original) + except (AttributeError, TypeError) as exc: + logger.debug("Could not patch %s on agent: %s", name, exc) + agent._delegation_wrapped = True + logger.debug( + "Delegation guard installed (max=%d, methods=%s)", + resolved_max, + list(originals), + ) + + +def uninstall_delegation_guard(agent: Any) -> None: + """Restore the agent's invoke methods if a delegation guard was installed. + + Safe to call on agents that were never guarded. + """ + if not getattr(agent, "_delegation_wrapped", False): + return + for name in _GUARDED_METHODS: + attr = f"_uipath_original_{name}" + original = getattr(agent, attr, None) + if original is not None: + try: + setattr(agent, name, original) + except Exception as exc: # noqa: BLE001 + logger.debug("Could not restore original %s: %s", name, exc) + try: + delattr(agent, attr) + except AttributeError: + pass + agent._delegation_wrapped = False + agent_key = id(agent) + try: + depths = _DELEGATION_DEPTHS.get() + except LookupError: + return + if agent_key in depths: + new_depths = {k: v for k, v in depths.items() if k != agent_key} + _DELEGATION_DEPTHS.set(new_depths) diff --git a/src/uipath/runtime/governance/features/__init__.py b/src/uipath/runtime/governance/features/__init__.py new file mode 100644 index 0000000..807b74f --- /dev/null +++ b/src/uipath/runtime/governance/features/__init__.py @@ -0,0 +1,18 @@ +"""Pre-computed feature signals for the Rego WASM evaluator. + +Each feature module registers its functions via ``_registry.register()``. +Importing this package imports all feature modules, triggering registration. + +Usage:: + + from uipath.runtime.governance.features import compute_features, FEATURE_NAMES +""" +from __future__ import annotations + +# Importing each module triggers its @register() calls. +from . import commitment, encoding, incident, sentiment, text_stats # noqa: F401 +from ._registry import _REGISTRY, compute_features + +FEATURE_NAMES: frozenset[str] = frozenset(_REGISTRY) + +__all__ = ["compute_features", "FEATURE_NAMES"] diff --git a/src/uipath/runtime/governance/features/_registry.py b/src/uipath/runtime/governance/features/_registry.py new file mode 100644 index 0000000..f90ce79 --- /dev/null +++ b/src/uipath/runtime/governance/features/_registry.py @@ -0,0 +1,40 @@ +"""Feature registry for Rego WASM pre-computed signals. + +Feature modules call ``register("name")`` at module level. +``compute_features`` is the only function callers outside this package need. +""" +from __future__ import annotations + +from typing import Any, Callable + +from uipath.runtime.governance.native.models import CheckContext + +_REGISTRY: dict[str, Callable[[CheckContext], Any]] = {} + + +def register(name: str) -> Callable[[Callable[[CheckContext], Any]], Callable[[CheckContext], Any]]: + """Decorator that registers a feature function under *name*.""" + def decorator(fn: Callable[[CheckContext], Any]) -> Callable[[CheckContext], Any]: + _REGISTRY[name] = fn + return fn + return decorator + + +def compute_features(context: CheckContext, plan: list[str] | None) -> dict[str, Any]: + """Compute only the features listed in *plan*. + + Unknown names are silently skipped. + Any exception during a feature function excludes that feature from the result. + """ + if not plan: + return {} + result: dict[str, Any] = {} + for name in plan: + fn = _REGISTRY.get(name) + if fn is None: + continue + try: + result[name] = fn(context) + except Exception: # noqa: BLE001 + pass + return result diff --git a/src/uipath/runtime/governance/features/_text.py b/src/uipath/runtime/governance/features/_text.py new file mode 100644 index 0000000..f2e0277 --- /dev/null +++ b/src/uipath/runtime/governance/features/_text.py @@ -0,0 +1,34 @@ +"""Text extraction helpers for feature functions.""" +from __future__ import annotations + +from uipath.runtime.governance.native.models import CheckContext + + +def primary_text(context: CheckContext) -> str: + """Return the most content-rich populated text field for the active hook.""" + for field in ( + context.model_output, + context.model_input, + context.agent_output, + context.agent_input, + context.tool_result, + ): + if field and isinstance(field, str) and field.strip(): + return field + return "" + + +def messages_text(context: CheckContext) -> str: + """Concatenate all message content blocks into one string.""" + if not context.messages: + return "" + parts: list[str] = [] + for msg in context.messages: + content = msg.get("content", "") + if isinstance(content, str): + parts.append(content) + elif isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(block.get("text", "")) + return " ".join(filter(None, parts)) diff --git a/src/uipath/runtime/governance/features/commitment.py b/src/uipath/runtime/governance/features/commitment.py new file mode 100644 index 0000000..ad56ab6 --- /dev/null +++ b/src/uipath/runtime/governance/features/commitment.py @@ -0,0 +1,63 @@ +"""Commitment language features: commitment_verb, commitment_amount, commitment_deadline.""" +from __future__ import annotations + +import re + +from uipath.runtime.governance.native.models import CheckContext + +from ._registry import register +from ._text import primary_text + +_COMMITMENT_VERB_PATTERN = re.compile( + r"(?i)(" + r"\brefund\b|\breimburse\b|" + r"\bwarranty\b|\bwarrant(?:y|ed|ies)\b|\bguarante[ed]+\b|" + r"\bsla\b|" + r"\bwaive[d]?\b|" + r"\b(?:we|i)\s+(?:will|shall|promise|commit|guarantee)\b|" + r"\b(?:we|i|i'?ll)\s+(?:deliver|provide|complete|finish|" + r"handover|hand\s+over|ship)\b|" + r"\bfixed\s+(?:price|cost|fee|scope|bid|rate)\b|" + r"\bcost\s*:\s*\$?\d|" + r"\bquote\s*:\s*\$?\d|" + r"\bdeliverables?\b|" + r"\btimeline\s*:\s*\d+\s*(?:second|minute|hour|day|week|month|year)s?\b|" + r"\bI\s+propose\b" + r")" +) +_COMMITMENT_AMOUNT_FALLBACK = re.compile( + r"(?:\$|€|£|¥|₹|USD|EUR|GBP|JPY|INR)\s*\d[\d,]*(?:\.\d+)?" + r"|\b\d[\d,]*(?:\.\d+)?\s*(?:USD|EUR|GBP|JPY|INR|" + r"dollars?|euros?|pounds?|yen|rupees?)\b" +) +_COMMITMENT_DEADLINE_PATTERN = re.compile( + r"(?i)\bwithin\s+\d+\s*(?:second|minute|hour|day|week|month|year)s?\b" + r"|\bby\s+(?:tomorrow|next\s+\w+|\d+/\d+(?:/\d+)?)\b" +) + + +@register("commitment_verb") +def commitment_verb(context: CheckContext) -> bool: + """True if the primary text contains a commitment verb or proposal marker.""" + text = primary_text(context) + if not text: + return False + return bool(_COMMITMENT_VERB_PATTERN.search(text)) + + +@register("commitment_amount") +def commitment_amount(context: CheckContext) -> bool: + """True if the primary text contains a currency-anchored monetary amount.""" + text = primary_text(context) + if not text: + return False + return bool(_COMMITMENT_AMOUNT_FALLBACK.search(text)) + + +@register("commitment_deadline") +def commitment_deadline(context: CheckContext) -> bool: + """True if the primary text contains a deadline phrase.""" + text = primary_text(context) + if not text: + return False + return bool(_COMMITMENT_DEADLINE_PATTERN.search(text)) diff --git a/src/uipath/runtime/governance/features/encoding.py b/src/uipath/runtime/governance/features/encoding.py new file mode 100644 index 0000000..11ce89c --- /dev/null +++ b/src/uipath/runtime/governance/features/encoding.py @@ -0,0 +1,64 @@ +"""Encoding integrity features: encoding_concern_ratio, encoding_concern_events.""" +from __future__ import annotations + +import re + +from uipath.runtime.governance.native.models import CheckContext + +from ._registry import register +from ._text import primary_text + +_MOJIBAKE_BIGRAMS: tuple[str, ...] = ( + "é", "è", "â", "à ", "ù", "î", "ô", "ç", + "Ä", "Ö", "Ãœ", "ß", + "’", "“", "â€\x9d", "–", "—", "•", + "£", "°", "§", "¶", "©", "®", + "ï¿", "¿½", "ï»", "»¿", +) +_HEX_ESCAPE_PATTERN = re.compile(r"\\x[0-9a-fA-F]{2}") + + +def _corruption_counts(text: str) -> tuple[int, int]: + """Return (events, weighted_chars) for encoding corruption in text.""" + replacement_chars = text.count("�") + literal_ufffd = text.count("\\ufffd") + hex_escapes = len(_HEX_ESCAPE_PATTERN.findall(text)) + mojibake = sum(text.count(b) for b in _MOJIBAKE_BIGRAMS) + + events = replacement_chars + literal_ufffd + hex_escapes + mojibake + weighted = ( + replacement_chars + + 6 * literal_ufffd + + 4 * hex_escapes + + 2 * mojibake + ) + return events, weighted + + +@register("encoding_concern_events") +def encoding_concern_events(context: CheckContext) -> int: + r"""Absolute count of encoding corruption events in the primary text field. + + Each U+FFFD, literal � escape, \xHH hex escape, or mojibake bigram + counts as one event. A value >= 2 is a strong signal in production output. + """ + text = primary_text(context) + if not text: + return 0 + events, _ = _corruption_counts(text) + return events + + +@register("encoding_concern_ratio") +def encoding_concern_ratio(context: CheckContext) -> float: + r"""Weighted corruption ratio (0.0–1.0) for the primary text field. + + Weights: U+FFFD=1, literal �=6, \xHH=4, mojibake bigram=2. + Typical threshold for a policy deny rule: > 0.05. + Returns 0.0 for empty input. + """ + text = primary_text(context) + if not text: + return 0.0 + _, weighted = _corruption_counts(text) + return weighted / max(len(text), 1) diff --git a/src/uipath/runtime/governance/features/incident.py b/src/uipath/runtime/governance/features/incident.py new file mode 100644 index 0000000..63a2007 --- /dev/null +++ b/src/uipath/runtime/governance/features/incident.py @@ -0,0 +1,60 @@ +"""Incident detection feature: incident_categories.""" +from __future__ import annotations + +import re + +from uipath.runtime.governance.native.models import CheckContext + +from ._registry import register +from ._text import primary_text + +_INCIDENT_PATTERNS: dict[str, list[re.Pattern[str]]] = { + "safety_refusal": [ + re.compile( + r"(?i)\b(i\s+(?:cannot|can'?t|am\s+unable\s+to|won'?t\s+be\s+able\s+to)" + r"\s+(?:help|assist|provide|answer|do\s+that))\b" + ), + re.compile(r"(?i)\b(i'?m\s+sorry,?\s+but\s+i\s+(?:cannot|can'?t))\b"), + re.compile(r"(?i)\b(against\s+my\s+(?:guidelines|policies|programming))\b"), + ], + "tool_failure": [ + re.compile( + r"\b(5\d{2})\b\s*(?:internal\s+server\s+error|service\s+unavailable)" + ), + re.compile(r"(?i)\b(ERR_[A-Z_]+|connection\s+refused|ECONNREFUSED)\b"), + re.compile(r"(?i)\b(timed?\s*out|timeout)\b"), + ], + "auth_failure": [ + re.compile(r"\b(401|403)\b\s*(?:unauthori[sz]ed|forbidden)"), + re.compile( + r"(?i)\b(authentication\s+failed|invalid\s+(?:token|credentials))\b" + ), + ], + "quota_exceeded": [ + re.compile(r"\b(429)\b"), + re.compile( + r"(?i)\b(rate\s+limit\s+exceeded|quota\s+exceeded|too\s+many\s+requests)\b" + ), + ], + "hallucination": [ + re.compile(r"(?i)\b(i\s+(?:made\s+(?:that|this)\s+up|am\s+just\s+guessing))\b"), + re.compile(r"(?i)\b(i\s+don'?t\s+actually\s+know|i\s+fabricat(?:ed|ing))\b"), + ], +} + + +@register("incident_categories") +def incident_categories(context: CheckContext) -> dict[str, bool]: + """Categorical incident detection over the primary text field. + + Returns a dict mapping each category name to True/False. + Categories: safety_refusal, tool_failure, auth_failure, quota_exceeded, hallucination. + """ + text = primary_text(context) + result: dict[str, bool] = {} + for category, patterns in _INCIDENT_PATTERNS.items(): + if not text: + result[category] = False + else: + result[category] = any(p.search(text) for p in patterns) + return result diff --git a/src/uipath/runtime/governance/features/sentiment.py b/src/uipath/runtime/governance/features/sentiment.py new file mode 100644 index 0000000..3403f2b --- /dev/null +++ b/src/uipath/runtime/governance/features/sentiment.py @@ -0,0 +1,36 @@ +"""Sentiment feature: vader_compound.""" +from __future__ import annotations + +from typing import Any + +from uipath.runtime.governance.native.models import CheckContext + +from ._registry import register +from ._text import primary_text + +_analyzer: Any = None + + +def _get_analyzer() -> Any: + global _analyzer + if _analyzer is None: + from vaderSentiment.vaderSentiment import ( + SentimentIntensityAnalyzer, # type: ignore[import] + ) + _analyzer = SentimentIntensityAnalyzer() + return _analyzer + + +@register("vader_compound") +def vader_compound(context: CheckContext) -> float: + """VADER compound sentiment score (-1.0 to 1.0) of the primary text field. + + Returns 0.0 for empty input or when vaderSentiment is not installed. + """ + text = primary_text(context) + if not text: + return 0.0 + try: + return float(_get_analyzer().polarity_scores(text)["compound"]) + except Exception: # noqa: BLE001 + return 0.0 diff --git a/src/uipath/runtime/governance/features/text_stats.py b/src/uipath/runtime/governance/features/text_stats.py new file mode 100644 index 0000000..c71856e --- /dev/null +++ b/src/uipath/runtime/governance/features/text_stats.py @@ -0,0 +1,37 @@ +"""Text statistics features: word_count, char_count, shannon_entropy.""" +from __future__ import annotations + +import math +from collections import Counter + +from uipath.runtime.governance.native.models import CheckContext + +from ._registry import register +from ._text import primary_text + + +@register("word_count") +def word_count(context: CheckContext) -> int: + """Number of whitespace-separated tokens in the primary text field.""" + return len(primary_text(context).split()) + + +@register("char_count") +def char_count(context: CheckContext) -> int: + """Character count of the primary text field.""" + return len(primary_text(context)) + + +@register("shannon_entropy") +def shannon_entropy(context: CheckContext) -> float: + """Shannon entropy (bits/symbol) over the primary text field. + + English prose is typically 3.5–4.5 bits/char. Binary noise approaches 8. + Returns 0.0 for empty input. + """ + text = primary_text(context) + if not text: + return 0.0 + counts = Counter(text) + total = len(text) + return -sum((c / total) * math.log2(c / total) for c in counts.values()) diff --git a/src/uipath/runtime/governance/native/backend_client.py b/src/uipath/runtime/governance/native/backend_client.py new file mode 100644 index 0000000..6c40a2c --- /dev/null +++ b/src/uipath/runtime/governance/native/backend_client.py @@ -0,0 +1,384 @@ +"""Governance backend client. + +Hosts the shared infrastructure used by every governance-backend call: + +- :func:`get_backend_base_url` — resolves the cloud host (with the + org/tenant path segments stripped) so each endpoint builder can + append its own scoped path. +- :func:`governance_request_headers` — composes the headers shared by + the policy fetch and the ``/runtime/govern`` compensating POST + (Accept, User-Agent, optional Content-Type, optional Bearer auth). +- :func:`build_governance_url` — composes an org-scoped URL against + the ``agenticgovernance_`` ingress. +- :func:`resolve_organization_id` / :func:`resolve_tenant_id` — read + the active org/tenant from ``UiPathConfig`` with an env-var fallback + for installations that don't have ``uipath-platform``. +- :func:`safe_call` — fail-open helper that catches every non-block + exception so governance hooks never crash an agent run. +- Module-level constants — request timeout, service path prefix, + compensation pool size — all the tunables an operator might care + about. Defined once here so the policy fetch, the compensating + ``/runtime/govern`` call, and the loader share one definition. + +The endpoint clients live next door: + +- :mod:`uipath.runtime.governance.native.policy_api_client` — policy fetch +- :mod:`uipath.runtime.governance.native.guardrail_compensation` — /runtime/govern +""" + +from __future__ import annotations + +import logging +import os +from functools import lru_cache +from typing import Callable +from urllib.parse import urlparse + +logger = logging.getLogger(__name__) + +# ---------------------------------------------------------------------------- +# Env-var names (consumed by the helpers below + diagnostic messages) +# ---------------------------------------------------------------------------- + +# Explicit dev/test override — used verbatim, no path-stripping. +ENV_BACKEND_BASE_URL = "UIPATH_GOVERNANCE_BACKEND_URL" +# The canonical platform URL env var (also backs ``UiPathConfig.base_url``). +ENV_PLATFORM_BASE_URL = "UIPATH_URL" +# Bearer token; missing means the policy fetch and compensating call are +# skipped (and that fact is logged) rather than producing 401s on every call. +ENV_ACCESS_TOKEN = "UIPATH_ACCESS_TOKEN" +# Org / tenant scoping for the agenticgovernance_ ingress. +ENV_ORGANIZATION_ID = "UIPATH_ORGANIZATION_ID" +ENV_TENANT_ID = "UIPATH_TENANT_ID" +# Job-execution context forwarded in the /runtime/govern payload so the +# server can populate the LLMOps trace record (Doc-2 audit structure). +# Each falls back to the named env var when uipath-platform isn't present. +ENV_FOLDER_KEY = "UIPATH_FOLDER_KEY" +ENV_JOB_KEY = "UIPATH_JOB_KEY" +ENV_PROCESS_KEY = "UIPATH_PROCESS_UUID" +ENV_REFERENCE_ID = "UIPATH_AGENT_ID" +ENV_AGENT_VERSION = "UIPATH_PROCESS_VERSION" + +# ---------------------------------------------------------------------------- +# Endpoint shape — all governance calls hit the org-scoped agenticgovernance_ +# service. Centralised so adding a third endpoint is "one new path constant" +# instead of "a new path template that someone forgets to keep in sync." +# ---------------------------------------------------------------------------- + +GOVERNANCE_SERVICE_PREFIX = "agenticgovernance_" +POLICY_API_PATH = "api/v1/runtime/policy" +GOVERN_API_PATH = "api/v1/runtime/govern" +ALL_POLICIES_API_PATH = "api/v1/all-policies" +TENANT_HEADER = "x-uipath-internal-tenantid" +# Query param on the policy fetch that selects the agent-type view of the +# policy: the server's clause-resolver reads the matching container key +# (``*-in-flight-conversational-agents`` vs ``*-in-flight-agents``). It's a +# representation selector (it changes the returned policy), so it travels as a +# query param — cache-correct and part of resource identification — not a +# header. Values: "conversational" | "autonomous". +AGENT_TYPE_PARAM = "agentType" +AGENT_TYPE_CONVERSATIONAL = "conversational" +AGENT_TYPE_AUTONOMOUS = "autonomous" + +# Default base URL when no override and no UiPathConfig / UIPATH_URL value is +# available. Used only on developer machines doing fully-offline work; real +# deployments always have UIPATH_URL injected by the host. +_DEFAULT_BACKEND_BASE_URL = "https://alpha.uipath.com" + +# ---------------------------------------------------------------------------- +# Tunables — one place so an ops change is one edit. The values that bound +# how long a single agent run can spend on governance traffic. +# ---------------------------------------------------------------------------- + +# Per-request timeout for any governance backend HTTP call (policy fetch, +# /runtime/govern compensating POST). Same value used everywhere so an agent +# can't accidentally end up with a "long" timeout on one call and "short" on +# another. +BACKEND_REQUEST_TIMEOUT_SECONDS = 10.0 + +# Bound on concurrent /runtime/govern requests in flight. A misbehaving +# agent that fires `before_model` 100 times in a session with three matched +# fallback rules each would otherwise spawn 100 daemon threads; this pool +# caps the concurrency. Saturated submissions are logged and dropped — the +# server still receives traces from the requests that did land. +COMPENSATION_MAX_WORKERS = 4 + +# Browser-shaped User-Agent. Required because the alpha/production +# governance ingress runs a WAF whose default scanner rule set blocks +# ``Python-urllib/``. Identifying as a real browser keeps the +# request from being rejected before any auth/tenant logic runs. +USER_AGENT = ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/148.0.0.0 Safari/537.36" +) + + +# ---------------------------------------------------------------------------- +# Headers +# ---------------------------------------------------------------------------- + + +def governance_request_headers(*, json_body: bool = False) -> dict[str, str]: + """Return the common HTTP headers for governance backend requests. + + Centralises the headers shared between the policy fetch and the + compensating ``/runtime/govern`` POST so the UA and auth shape are + declared once. + + Args: + json_body: When ``True`` (POST/PATCH/etc. with a JSON payload), + adds ``Content-Type: application/json``. GETs leave it off + so origin servers that 415 on unexpected Content-Type stay + happy. + + Returns: + A new dict with: + + - ``Accept: application/json`` + - ``User-Agent`` (the browser-shaped string above) + - ``Content-Type: application/json`` when ``json_body=True`` + - ``Authorization: Bearer `` when the env + var is set; omitted otherwise (caller decides whether the + missing token is fatal). + + Endpoint-specific headers (e.g. ``x-uipath-internal-tenantid``) are + added by the caller after this helper returns. + """ + headers: dict[str, str] = { + "Accept": "application/json", + "User-Agent": USER_AGENT, + } + if json_body: + headers["Content-Type"] = "application/json" + token = os.environ.get(ENV_ACCESS_TOKEN) + if token: + headers["Authorization"] = f"Bearer {token}" + return headers + + +# ---------------------------------------------------------------------------- +# URL composition +# ---------------------------------------------------------------------------- + + +def _strip_to_origin(raw_url: str) -> str: + """Return ``scheme://host[:port]`` for ``raw_url``, dropping any path. + + Platform URLs are commonly ``https://cloud.uipath.com//``; + the governance endpoints construct their own + ``/{org}/agenticgovernance_/...`` suffix, so the org/tenant segments + in the base must be stripped to avoid a duplicated org path. + """ + parsed = urlparse(raw_url) + if not parsed.scheme or not parsed.netloc: + # Not a parseable absolute URL — leave it to the caller. + return raw_url.rstrip("/") + return f"{parsed.scheme}://{parsed.netloc}" + + +def get_backend_base_url() -> str: + """Resolve the governance backend base URL on each call. + + Resolution order (first hit wins): + + 1. ``UIPATH_GOVERNANCE_BACKEND_URL`` — explicit dev/test override, + used verbatim. + 2. ``UiPathConfig.base_url`` from ``uipath-platform`` — the + canonical platform URL. Org/tenant path segments are stripped + so the caller can append its own org-scoped path. + 3. ``UIPATH_URL`` env var — same as (2) but works when + ``uipath-platform`` is not installed. + 4. ``https://alpha.uipath.com`` — last-resort default for offline + development; real deployments always have ``UIPATH_URL`` set. + + Reading on each call (not at import) lets the runtime entrypoint + configure the env vars after this module is already loaded. + """ + explicit_override = os.environ.get(ENV_BACKEND_BASE_URL) + if explicit_override: + return explicit_override.rstrip("/") + + # Lazy import — uipath-platform is optional; falls through to the + # env-var path when only uipath-core / uipath-runtime are installed. + platform_url: str | None = None + try: + from uipath.platform.common import UiPathConfig + + platform_url = UiPathConfig.base_url + except (ImportError, AttributeError): + pass + + raw = platform_url or os.environ.get(ENV_PLATFORM_BASE_URL) + if raw: + return _strip_to_origin(raw) + + return _DEFAULT_BACKEND_BASE_URL + + +def build_governance_url(org_id: str, path: str) -> str: + """Compose an org-scoped governance backend URL. + + Final shape: ``{backend_base}/{org_id}/{GOVERNANCE_SERVICE_PREFIX}/{path}``. + + Args: + org_id: Active organization id; the URL is meaningless without it. + path: API suffix WITHOUT the org/service prefix + (e.g. :data:`POLICY_API_PATH` or :data:`GOVERN_API_PATH`). + """ + base = get_backend_base_url() + return f"{base}/{org_id}/{GOVERNANCE_SERVICE_PREFIX}/{path}" + + +# ---------------------------------------------------------------------------- +# Org / tenant resolution +# ---------------------------------------------------------------------------- + + +def _resolve_uipath_config_field(attr: str, env_var: str) -> str | None: + """Read a single ``UiPathConfig`` attribute with an env-var fallback. + + Lazy-imports ``UiPathConfig`` so ``uipath-runtime`` doesn't require + ``uipath-platform`` at install time. When the platform package is + missing (``ImportError``) or the attribute isn't yet exposed + (``AttributeError``), falls back to reading the named env var. + """ + try: + from uipath.platform.common import UiPathConfig + + return getattr(UiPathConfig, attr, None) or os.environ.get(env_var) + except ImportError: + return os.environ.get(env_var) + + +# ---------------------------------------------------------------------------- +# Agent-type selector (conversational vs autonomous) +# +# Set once by the governance wrapper at runtime init (before the background +# policy prefetch is kicked off) and read by the policy fetch when composing +# the request URL. A process-level holder — not a ContextVar — because the +# prefetch runs on a separate thread that wouldn't inherit a ContextVar, and a +# coded-agent process hosts a single agent so the value is stable per process. +# ---------------------------------------------------------------------------- + +_agent_is_conversational: bool | None = None + + +def set_agent_conversational(value: bool | None) -> None: + """Record whether the hosted agent is conversational. + + ``None`` clears the selector (used by tests / direct callers); the policy + fetch then omits the param and the server applies its default. + """ + global _agent_is_conversational + _agent_is_conversational = value + + +def agent_type_param() -> str | None: + """Return the ``agentType`` query value, or ``None`` when unknown. + + ``"conversational"`` / ``"autonomous"`` map to the server's + conversational-vs-autonomous container keys; ``None`` (selector never set) + omits the param so the server's default applies. + """ + if _agent_is_conversational is None: + return None + return AGENT_TYPE_CONVERSATIONAL if _agent_is_conversational else AGENT_TYPE_AUTONOMOUS + + +def resolve_organization_id() -> str | None: + """Return the current organization id from ``UiPathConfig`` / env. + + Returns ``None`` when neither source yields a value — callers skip + the backend interaction (no URL can be built without an org id) + and the agent runs with no policies / no compensation. + """ + return _resolve_uipath_config_field("organization_id", ENV_ORGANIZATION_ID) + + +def resolve_tenant_id() -> str | None: + """Return the current tenant id from ``UiPathConfig`` / env. + + Returns ``None`` when neither source yields a value — callers skip + the backend interaction since the ``x-uipath-internal-tenantid`` + header would be missing. + """ + return _resolve_uipath_config_field("tenant_id", ENV_TENANT_ID) + + +@lru_cache(maxsize=1) +def _resolved_job_context() -> tuple[tuple[str, str], ...]: + """Resolve and freeze the job context once per process. + + Returned as a tuple of ``(key, value)`` pairs so the cached value is + immutable — callers materialize a fresh dict each call. Tests that + mutate env vars can invalidate via ``resolve_job_context.cache_clear()``. + """ + candidates = { + "folderKey": _resolve_uipath_config_field("folder_key", ENV_FOLDER_KEY), + "jobKey": _resolve_uipath_config_field("job_key", ENV_JOB_KEY), + "processKey": _resolve_uipath_config_field("process_uuid", ENV_PROCESS_KEY), + "referenceId": _resolve_uipath_config_field("agent_id", ENV_REFERENCE_ID), + "agentVersion": _resolve_uipath_config_field( + "process_version", ENV_AGENT_VERSION + ), + } + return tuple((k, v) for k, v in candidates.items() if v) + + +def resolve_job_context() -> dict[str, str]: + """Return the agent's job-execution context for the govern payload. + + Each field is read from ``UiPathConfig`` (env-var fallback) and only + included when it resolves to a truthy value, so the server receives + exactly the keys the agent actually knows. Cached per-process — the + underlying values are immutable for the agent's lifetime. The server + maps these onto the LLMOps trace record: + + - ``folderKey`` → ``FolderKey`` / ``uipath.folder_key`` + - ``jobKey`` → ``JobKey`` / ``uipath.job_key`` + - ``processKey`` → ``ProcessKey`` + - ``referenceId`` → ``ReferenceId`` (typically the agent id) + - ``agentVersion`` → ``AgentVersion`` + """ + return dict(_resolved_job_context()) + + +resolve_job_context.cache_clear = _resolved_job_context.cache_clear # type: ignore[attr-defined] + + +# ---------------------------------------------------------------------------- +# Generic safe-call helper. Used by callers that want "log and continue" on +# any unexpected failure path without spelling out the same try/except every +# time. The intentional GovernanceBlockException ALWAYS propagates — only +# this exception type carries policy intent; anything else is a bug. +# ---------------------------------------------------------------------------- + + +def safe_call( + fn: Callable[..., None], + *args: object, + what: str, + **kwargs: object, +) -> None: + """Call ``fn(*args, **kwargs)`` and swallow any non-block exception. + + ``GovernanceBlockException`` propagates (intentional policy block); + everything else is logged at WARNING with the ``what`` label and + swallowed so the agent can continue. Designed for fire-and-forget + governance paths that should never fail an agent run. + + Args: + fn: Callable to invoke. + what: Short label used in the log line on failure + (e.g. ``"BEFORE_AGENT governance check"``). + """ + # Lazy import to avoid pulling uipath-core into module load. + from uipath.core.governance.exceptions import GovernanceBlockException + + try: + fn(*args, **kwargs) + except GovernanceBlockException: + raise + except Exception as exc: # noqa: BLE001 - fail-open by contract + logger.warning("%s failed (continuing): %s", what, exc) diff --git a/src/uipath/runtime/governance/native/evaluator.py b/src/uipath/runtime/governance/native/evaluator.py index a9c60a6..7e49505 100644 --- a/src/uipath/runtime/governance/native/evaluator.py +++ b/src/uipath/runtime/governance/native/evaluator.py @@ -16,7 +16,10 @@ from collections import Counter from datetime import datetime, timezone from functools import cache, lru_cache -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from uipath.runtime.governance.rego.evaluator import RegoEvaluator from uipath.core.governance import EnforcementMode from uipath.core.governance.exceptions import GovernanceBlockException @@ -303,6 +306,22 @@ def __init__( self._enforcement_mode = enforcement_mode self._audit_manager = audit_manager self._compensator = compensator + self._rego_evaluator: RegoEvaluator | None = None + + def set_rego_evaluator(self, rego_evaluator: RegoEvaluator | None) -> None: + """Wire a Rego evaluator to run alongside per-step hooks. + + Called by :class:`~uipath.runtime.governance.runtime.UiPathGovernedRuntime` + after the callback handler has already been constructed with a + reference to this evaluator. Mutating in place means the handler + sees Rego evaluation on the next hook invocation without + requiring any changes to the framework adapter. + + Only fires on per-step hooks (BEFORE_MODEL, AFTER_MODEL, + TOOL_CALL, AFTER_TOOL). BEFORE_AGENT / AFTER_AGENT are owned by + the runtime wrapper and are skipped here to avoid double-firing. + """ + self._rego_evaluator = rego_evaluator @property def policy_index(self) -> PolicyIndex: @@ -712,7 +731,21 @@ def evaluate_before_model( messages=messages or [], metadata=kwargs.get("metadata", {}), ) - return self.evaluate(context) + result = self.evaluate(context) + if self._rego_evaluator is not None: + try: + self._rego_evaluator.evaluate_before_model( + model_input=model_input, + agent_name=agent_name, + runtime_id=runtime_id, + messages=messages, + model_name=model_name, + ) + except GovernanceBlockException: + raise + except Exception as exc: + logger.warning("Rego BEFORE_MODEL evaluation failed: %s", exc) + return result def evaluate_after_model( self, @@ -729,7 +762,19 @@ def evaluate_after_model( model_output=model_output, metadata=kwargs.get("metadata", {}), ) - return self.evaluate(context) + result = self.evaluate(context) + if self._rego_evaluator is not None: + try: + self._rego_evaluator.evaluate_after_model( + model_output=model_output, + agent_name=agent_name, + runtime_id=runtime_id, + ) + except GovernanceBlockException: + raise + except Exception as exc: + logger.warning("Rego AFTER_MODEL evaluation failed: %s", exc) + return result def evaluate_tool_call( self, @@ -750,7 +795,21 @@ def evaluate_tool_call( session_state=session_state or {}, metadata=kwargs.get("metadata", {}), ) - return self.evaluate(context) + result = self.evaluate(context) + if self._rego_evaluator is not None: + try: + self._rego_evaluator.evaluate_tool_call( + tool_name=tool_name, + tool_args=tool_args, + agent_name=agent_name, + runtime_id=runtime_id, + session_state=session_state, + ) + except GovernanceBlockException: + raise + except Exception as exc: + logger.warning("Rego TOOL_CALL evaluation failed: %s", exc) + return result def evaluate_after_tool( self, @@ -769,7 +828,20 @@ def evaluate_after_tool( tool_result=tool_result, metadata=kwargs.get("metadata", {}), ) - return self.evaluate(context) + result = self.evaluate(context) + if self._rego_evaluator is not None: + try: + self._rego_evaluator.evaluate_after_tool( + tool_name=tool_name, + tool_result=tool_result, + agent_name=agent_name, + runtime_id=runtime_id, + ) + except GovernanceBlockException: + raise + except Exception as exc: + logger.warning("Rego AFTER_TOOL evaluation failed: %s", exc) + return result def _evaluate_rule(self, rule: Rule, context: CheckContext) -> RuleEvaluation: """Evaluate a single rule against the context.""" diff --git a/src/uipath/runtime/governance/rego/__init__.py b/src/uipath/runtime/governance/rego/__init__.py new file mode 100644 index 0000000..fddbb41 --- /dev/null +++ b/src/uipath/runtime/governance/rego/__init__.py @@ -0,0 +1,10 @@ +"""Rego/WASM-based governance evaluator.""" +from __future__ import annotations + +from .evaluator import RegoEvaluator +from .loader import build_rego_evaluator_async + +__all__ = [ + "RegoEvaluator", + "build_rego_evaluator_async", +] diff --git a/src/uipath/runtime/governance/rego/bundle_cache.py b/src/uipath/runtime/governance/rego/bundle_cache.py new file mode 100644 index 0000000..f888c94 --- /dev/null +++ b/src/uipath/runtime/governance/rego/bundle_cache.py @@ -0,0 +1,65 @@ +"""Disk cache for per-hook Rego WASM bundles. + +Layout under ``$TMPDIR/uipath-governance/{key}/hooks/{hook_type}/``: + - ``bundle.tar.gz`` — raw OPA bundle bytes + - ``etag.txt`` — the ETag string from the last successful download +""" +from __future__ import annotations + +import hashlib +import logging +import tempfile +from pathlib import Path + +logger = logging.getLogger(__name__) + +_BUNDLE_FILE = "bundle.tar.gz" +_ETAG_FILE = "etag.txt" + + +def get_cache_dir(org_id: str, tenant_id: str) -> Path: + """Return (and create) the cache root for the given org/tenant pair.""" + key = hashlib.sha256(f"{org_id}{tenant_id}".encode()).hexdigest()[:16] + base = Path(tempfile.gettempdir()) / "uipath-governance" / key + base.mkdir(parents=True, exist_ok=True) + return base + + +def _hook_dir(cache_dir: Path, hook_type: str) -> Path: + return cache_dir / "hooks" / hook_type + + +def get_cached_etag(cache_dir: Path, hook_type: str) -> str | None: + """Return the cached ETag for ``hook_type``, or ``None`` if absent.""" + etag_path = _hook_dir(cache_dir, hook_type) / _ETAG_FILE + try: + return etag_path.read_text(encoding="utf-8").strip() or None + except FileNotFoundError: + return None + except OSError as exc: + logger.warning("Failed to read cached ETag for %s: %s", hook_type, exc) + return None + + +def get_cached_bundle_path(cache_dir: Path, hook_type: str) -> Path | None: + """Return the cached bundle path for ``hook_type``, or ``None`` if absent.""" + bundle_path = _hook_dir(cache_dir, hook_type) / _BUNDLE_FILE + return bundle_path if bundle_path.exists() else None + + +def save_bundle(cache_dir: Path, hook_type: str, data: bytes, etag: str) -> Path: + """Write ``data`` and ``etag`` to the cache for ``hook_type``. + + Returns the path to the written bundle file. + """ + hook_dir = _hook_dir(cache_dir, hook_type) + hook_dir.mkdir(parents=True, exist_ok=True) + + bundle_path = hook_dir / _BUNDLE_FILE + bundle_path.write_bytes(data) + + etag_path = hook_dir / _ETAG_FILE + etag_path.write_text(etag, encoding="utf-8") + + logger.debug("Cached bundle for hook=%s (%d bytes, etag=%s)", hook_type, len(data), etag) + return bundle_path diff --git a/src/uipath/runtime/governance/rego/evaluator.py b/src/uipath/runtime/governance/rego/evaluator.py new file mode 100644 index 0000000..0c58a76 --- /dev/null +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -0,0 +1,545 @@ +"""WASM-based Rego governance evaluator using opa-wasmtime.""" +from __future__ import annotations + +import io +import logging +import os +import tarfile +import tempfile +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from uipath.core.governance import EnforcementMode +from uipath.core.governance.exceptions import GovernanceBlockException +from uipath.core.governance.models import ( + Action, + AuditRecord, + LifecycleHook, + RuleEvaluation, +) + +from uipath.runtime.governance._audit.base import AuditManager +from uipath.runtime.governance.config import get_enforcement_mode +from uipath.runtime.governance.features import compute_features +from uipath.runtime.governance.native.models import CheckContext + +logger = logging.getLogger(__name__) + +_WARMUP_INPUT: dict[str, Any] = { + "hook": "before_model", + "agent_input": "", "agent_output": "", + "model_input": "", "model_output": "", + "model_name": "", "agent_name": "", + "tool_name": "", "tool_args": {}, "tool_result": "", + "session_state": {"tool_calls": 0, "llm_calls": 0}, +} + + +def context_to_input( + context: CheckContext, + feature_plan: list[str] | None = None, +) -> dict[str, Any]: + """Serialize a CheckContext to the flat input dict Rego rules expect.""" + session = context.session_state if isinstance(context.session_state, dict) else {} + features = compute_features(context, feature_plan) if feature_plan else {} + # For BEFORE_MODEL, the user's message is passed as model_input. WASM + # rules that check agent_input (e.g. "block_ssn_in_prompts") would miss + # it, so normalize: copy model_input → agent_input when agent_input is + # empty. This is safe — for BEFORE_MODEL, both fields represent the same + # user content. + agent_input = context.agent_input + if context.hook == LifecycleHook.BEFORE_MODEL and not agent_input: + agent_input = context.model_input + # WASM rules check input.messages[].content (plain string) rather than + # the flat model_input field. When messages is empty for a BEFORE_MODEL + # hook, synthesize a single user message from model_input so that rules + # like block_ssn_in_prompts can match. The callback-handler path never + # passes messages, so this is the only way to populate the field. + messages = context.messages + if context.hook == LifecycleHook.BEFORE_MODEL and not messages and context.model_input: + messages = [{"role": "user", "content": context.model_input}] + return { + "hook": context.hook.value, + "agent_input": agent_input, + "agent_output": context.agent_output, + "model_input": context.model_input, + "model_output": context.model_output, + "model_name": context.model_name, + "agent_name": context.agent_name, + "tool_name": context.tool_name, + "tool_args": context.tool_args, + "tool_result": context.tool_result, + "session_state": { + "tool_calls": session.get("tool_calls", 0), + "llm_calls": session.get("llm_calls", 0), + }, + "ring": context.ring, + "messages": messages, + "features": features, + } + + +_WASM_MAGIC = b"\x00asm" +_GZIP_MAGIC = b"\x1f\x8b" + + +def _find_wasm_in_tar(tf: Any) -> bytes | None: + """Search a TarFile for a valid WASM member; return bytes or None.""" + members = tf.getmembers() + for candidate in ("/policy.wasm", "policy.wasm", "./policy.wasm"): + try: + fobj = tf.extractfile(tf.getmember(candidate)) + except KeyError: + continue + if fobj is None: + continue + data = fobj.read() + if data[:4] == _WASM_MAGIC: + return data + # Fall back: any .wasm member (skip macOS resource forks) + for m in members: + if m.name.endswith(".wasm") and not m.name.startswith("._"): + fobj = tf.extractfile(m) + if fobj is None: + continue + data = fobj.read() + if data[:4] == _WASM_MAGIC: + logger.info("Found valid WASM at non-standard path %r", m.name) + return data + return None + + +def _extract_wasm_from_bundle(bundle_path: Path) -> bytes: + r"""Extract ``policy.wasm`` bytes from an OPA ``.tar.gz`` bundle. + + OPA's backend produces a nested structure: + outer bundle.tar.gz + └── policy.wasm ← a gzip-compressed inner tar bundle + └── /policy.wasm ← the actual WASM binary (``\x00asm``) + """ + import gzip + + with open(bundle_path, "rb") as f: + bundle_bytes = f.read() + + with tarfile.open(fileobj=io.BytesIO(bundle_bytes), mode="r:gz") as outer_tf: + members = outer_tf.getmembers() + logger.debug("Outer bundle members: %s", [m.name for m in members]) + + for candidate in ("/policy.wasm", "policy.wasm", "./policy.wasm"): + try: + fobj = outer_tf.extractfile(outer_tf.getmember(candidate)) + except KeyError: + continue + if fobj is None: + continue + + raw = fobj.read() + + # Direct WASM (rare — bundle served without inner wrapping) + if raw[:4] == _WASM_MAGIC: + return raw + + # Gzip-compressed: could be the WASM itself or an inner bundle + if raw[:2] == _GZIP_MAGIC: + inner = gzip.decompress(raw) + if inner[:4] == _WASM_MAGIC: + return inner + # Inner content is itself a tar (nested OPA bundle format) + try: + with tarfile.open(fileobj=io.BytesIO(inner), mode="r") as inner_tf: + logger.debug( + "Inner bundle members: %s", + [m.name for m in inner_tf.getmembers()], + ) + wasm = _find_wasm_in_tar(inner_tf) + if wasm is not None: + return wasm + except tarfile.TarError: + pass + + logger.warning( + "Bundle member %r yielded no valid WASM (first 8 bytes: %s)", + candidate, raw[:8].hex(), + ) + + names = [m.name for m in members] + raise ValueError( + f"No valid policy.wasm found in bundle {bundle_path}. Members: {names}" + ) + + +def _extract_data_json_from_bundle(bundle_path: Path) -> dict[str, Any] | None: + """Extract and parse ``data.json`` from an OPA ``.tar.gz`` bundle. + + Returns None when the bundle predates the data.json format. + """ + import json as _json + + with open(bundle_path, "rb") as f: + bundle_bytes = f.read() + with tarfile.open(fileobj=io.BytesIO(bundle_bytes), mode="r:gz") as tf: + for candidate in ("/data.json", "data.json"): + try: + member = tf.getmember(candidate) + except KeyError: + continue + fobj = tf.extractfile(member) + if fobj is None: + continue + return _json.loads(fobj.read().decode("utf-8")) + return None + + +def _load_engine(bundle_path: Path) -> Any: + """Load a WASM bundle via opa-wasmtime. Raises on any failure.""" + from opa_wasmtime import OPAPolicy # type: ignore[import] + + wasm_bytes = _extract_wasm_from_bundle(bundle_path) + with tempfile.NamedTemporaryFile(suffix=".wasm", delete=False) as tmp: + tmp.write(wasm_bytes) + tmp_path = tmp.name + try: + return OPAPolicy(tmp_path) + finally: + os.unlink(tmp_path) + + +def _pack_name_from_rule_id(rule_id: str) -> str: + """Extract the pack/policy name prefix from ``{policyId}/{ruleId}``.""" + return rule_id.split("/")[0] if "/" in rule_id else "custom" + + +class RegoEvaluator: + """WASM-based Rego evaluator with one engine per lifecycle hook. + + Args: + hook_wasm_paths: Map of LifecycleHook → bundle path on disk. + hook_data: Optional map of LifecycleHook → data.json dict. + audit_manager: Optional instance-scoped :class:`AuditManager`. + When provided, rule evaluation and hook summary events are + emitted to it. When ``None``, audit emission is skipped. + enforcement_mode: Optional enforcement mode override. When + ``None``, :func:`get_enforcement_mode` is called on each + evaluate() to read the process-level value set by the native + policy loader. + """ + + def __init__( + self, + hook_wasm_paths: dict[LifecycleHook, Path], + hook_data: dict[LifecycleHook, dict[str, Any]] | None = None, + audit_manager: AuditManager | None = None, + enforcement_mode: EnforcementMode | None = None, + ) -> None: + """Load WASM engines from paths; wire audit manager and enforcement mode.""" + self._engines: dict[LifecycleHook, Any] = {} + self._feature_plans: dict[LifecycleHook, list[str]] = {} + if audit_manager is None: + audit_manager = AuditManager(register_default_sinks=True) + self._audit_manager = audit_manager + self._enforcement_mode = enforcement_mode + hook_data = hook_data or {} + + for hook, path in hook_wasm_paths.items(): + try: + engine = _load_engine(path) + data = hook_data.get(hook) + if data: + try: + engine.set_data(data) + except Exception as exc: + logger.warning( + "Rego WASM set_data failed for hook=%s: %s (continuing without data)", + hook.value, exc, + ) + self._feature_plans[hook] = list(data.get("required_features") or []) + try: + engine.evaluate(_WARMUP_INPUT) + except Exception as warmup_exc: + logger.warning( + "Rego WASM warmup failed for hook=%s: %s (engine kept)", + hook.value, warmup_exc, + ) + self._engines[hook] = engine + logger.info("Loaded Rego WASM for hook=%s from %s", hook.value, path) + except ImportError: + logger.warning( + "opa-wasmtime not installed; Rego evaluation for hook=%s disabled. " + "Install with: pip install uipath-runtime[governance-rego]", + hook.value, + ) + except Exception as exc: + logger.warning( + "Failed to load Rego WASM for hook=%s from %s: %s", + hook.value, path, exc, + ) + + @property + def loaded_hooks(self) -> list[LifecycleHook]: + """Lifecycle hooks for which WASM engines were loaded successfully.""" + return list(self._engines) + + def _resolve_mode(self) -> EnforcementMode: + """Return the enforcement mode for this evaluation.""" + return self._enforcement_mode if self._enforcement_mode is not None else get_enforcement_mode() + + def evaluate(self, context: CheckContext) -> AuditRecord: + """Evaluate a lifecycle hook against loaded Rego policies.""" + mode = self._resolve_mode() + + if mode == EnforcementMode.DISABLED: + return AuditRecord( + timestamp=datetime.now(timezone.utc), + agent_name=context.agent_name, + runtime_id=context.runtime_id, + hook=context.hook, + evaluations=[], + final_action=Action.ALLOW, + metadata={"enforcement_mode": mode.value, "rego_engine": "skipped"}, + ) + + engine = self._engines.get(context.hook) + if engine is None: + return AuditRecord( + timestamp=datetime.now(timezone.utc), + agent_name=context.agent_name, + runtime_id=context.runtime_id, + hook=context.hook, + evaluations=[], + final_action=Action.ALLOW, + metadata={"enforcement_mode": mode.value, "rego_engine": "missing"}, + ) + + try: + plan = self._feature_plans.get(context.hook, []) + wasm_input = context_to_input(context, feature_plan=plan) + raw = engine.evaluate(wasm_input) + logger.info( + "Rego WASM hook=%s agent_input_len=%d model_input_len=%d raw=%r", + context.hook.value, + len(wasm_input.get("agent_input", "")), + len(wasm_input.get("model_input", "")), + raw, + ) + except Exception as exc: + logger.warning("Rego WASM evaluation failed for hook=%s: %s", context.hook.value, exc) + return AuditRecord( + timestamp=datetime.now(timezone.utc), + agent_name=context.agent_name, + runtime_id=context.runtime_id, + hook=context.hook, + evaluations=[], + final_action=Action.ALLOW, + metadata={"enforcement_mode": mode.value, "rego_engine": "error"}, + ) + + result: dict[str, Any] = {} + if isinstance(raw, list) and raw: + result = raw[0].get("result", {}) + elif isinstance(raw, dict): + result = raw + + fired_deny: set[str] = set(result.get("fired_deny") or []) + fired_allow: set[str] = set(result.get("fired_allow") or []) + rule_messages: dict[str, str] = dict(result.get("messages") or {}) + + evaluations: list[RuleEvaluation] = [] + for rule_id in fired_deny: + evaluations.append(RuleEvaluation( + rule_id=rule_id, + rule_name=rule_id, + matched=True, + pack_name=_pack_name_from_rule_id(rule_id), + action=Action.DENY, + detail=rule_messages.get(rule_id, f"Rego rule '{rule_id}' denied the request"), + description="", + )) + for rule_id in fired_allow: + evaluations.append(RuleEvaluation( + rule_id=rule_id, + rule_name=rule_id, + matched=True, + pack_name=_pack_name_from_rule_id(rule_id), + action=Action.ALLOW, + detail=rule_messages.get(rule_id, f"Rego rule '{rule_id}' explicitly allowed"), + description="", + )) + + raw_action = Action.DENY if bool(result.get("deny", False)) else Action.ALLOW + final_action = self._apply_mode(raw_action, mode) + + metadata: dict[str, Any] = {"enforcement_mode": mode.value} + if raw_action == Action.DENY and mode == EnforcementMode.AUDIT: + metadata["audit_mode_would_deny"] = True + + audit = AuditRecord( + timestamp=datetime.now(timezone.utc), + agent_name=context.agent_name, + runtime_id=context.runtime_id, + hook=context.hook, + evaluations=evaluations, + final_action=final_action, + metadata=metadata, + ) + + self._emit_audit(audit, mode) + + if final_action == Action.DENY: + raise GovernanceBlockException.from_audit_record(audit) + + return audit + + def _apply_mode(self, raw_action: Action, mode: EnforcementMode) -> Action: + """Downgrade DENY to AUDIT when mode is not ENFORCE.""" + if mode == EnforcementMode.AUDIT and raw_action in (Action.DENY, Action.ESCALATE): + return Action.AUDIT + return raw_action + + def _emit_audit(self, audit: AuditRecord, mode: EnforcementMode) -> None: + """Emit per-rule and hook-summary events to the injected AuditManager.""" + if self._audit_manager is None: + return + hook_name = audit.hook.name + for ev in audit.evaluations: + try: + self._audit_manager.emit_rule_evaluation( + policy_id=ev.rule_id, + rule_name=ev.rule_name, + pack_name=ev.pack_name, + hook=hook_name, + matched=ev.matched, + action=ev.action.value if ev.matched else "allow", + enforcement_mode=mode, + detail=ev.detail, + agent_name=audit.agent_name, + description=ev.description, + ) + except Exception: # noqa: BLE001 + pass + try: + self._audit_manager.emit_hook_summary( + hook=hook_name, + agent_name=audit.agent_name, + total_rules=len(audit.evaluations), + matched_rules=sum(1 for e in audit.evaluations if e.matched), + final_action=audit.final_action.value, + enforcement_mode=mode, + ) + except Exception: # noqa: BLE001 + pass + + def evaluate_before_agent( + self, + agent_input: str, + agent_name: str, + runtime_id: str, + model_name: str = "", + **kwargs: Any, + ) -> AuditRecord: + """Evaluate the BEFORE_AGENT lifecycle hook.""" + ctx = CheckContext( + hook=LifecycleHook.BEFORE_AGENT, + agent_name=agent_name, + runtime_id=runtime_id, + agent_input=agent_input, + model_name=model_name, + metadata=kwargs.get("metadata", {}), + ) + return self.evaluate(ctx) + + def evaluate_after_agent( + self, + agent_output: str, + agent_name: str, + runtime_id: str, + **kwargs: Any, + ) -> AuditRecord: + """Evaluate the AFTER_AGENT lifecycle hook.""" + ctx = CheckContext( + hook=LifecycleHook.AFTER_AGENT, + agent_name=agent_name, + runtime_id=runtime_id, + agent_output=agent_output, + metadata=kwargs.get("metadata", {}), + ) + return self.evaluate(ctx) + + def evaluate_before_model( + self, + model_input: str, + agent_name: str, + runtime_id: str, + messages: list[dict[str, Any]] | None = None, + model_name: str = "", + **kwargs: Any, + ) -> AuditRecord: + """Evaluate the BEFORE_MODEL lifecycle hook.""" + ctx = CheckContext( + hook=LifecycleHook.BEFORE_MODEL, + agent_name=agent_name, + runtime_id=runtime_id, + model_input=model_input, + model_name=model_name, + messages=messages or [], + metadata=kwargs.get("metadata", {}), + ) + return self.evaluate(ctx) + + def evaluate_after_model( + self, + model_output: str, + agent_name: str, + runtime_id: str, + **kwargs: Any, + ) -> AuditRecord: + """Evaluate the AFTER_MODEL lifecycle hook.""" + ctx = CheckContext( + hook=LifecycleHook.AFTER_MODEL, + agent_name=agent_name, + runtime_id=runtime_id, + model_output=model_output, + metadata=kwargs.get("metadata", {}), + ) + return self.evaluate(ctx) + + def evaluate_tool_call( + self, + tool_name: str, + tool_args: dict[str, Any], + agent_name: str, + runtime_id: str, + session_state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> AuditRecord: + """Evaluate the TOOL_CALL lifecycle hook.""" + ctx = CheckContext( + hook=LifecycleHook.TOOL_CALL, + agent_name=agent_name, + runtime_id=runtime_id, + tool_name=tool_name, + tool_args=tool_args, + session_state=session_state or {}, + metadata=kwargs.get("metadata", {}), + ) + return self.evaluate(ctx) + + def evaluate_after_tool( + self, + tool_name: str, + tool_result: str, + agent_name: str, + runtime_id: str, + **kwargs: Any, + ) -> AuditRecord: + """Evaluate the AFTER_TOOL lifecycle hook.""" + ctx = CheckContext( + hook=LifecycleHook.AFTER_TOOL, + agent_name=agent_name, + runtime_id=runtime_id, + tool_name=tool_name, + tool_result=tool_result, + metadata=kwargs.get("metadata", {}), + ) + return self.evaluate(ctx) diff --git a/src/uipath/runtime/governance/rego/loader.py b/src/uipath/runtime/governance/rego/loader.py new file mode 100644 index 0000000..6d8a19f --- /dev/null +++ b/src/uipath/runtime/governance/rego/loader.py @@ -0,0 +1,131 @@ +"""Rego bundle loader: async bootstrap builder.""" +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from uipath.core.governance.models import LifecycleHook + +from uipath.runtime.governance.rego.bundle_cache import ( + get_cache_dir, + get_cached_bundle_path, + get_cached_etag, + save_bundle, +) + +if TYPE_CHECKING: + from uipath.runtime.governance._audit.base import AuditManager + +logger = logging.getLogger(__name__) + +# Hook type string → LifecycleHook mapping +_HOOK_MAP: dict[str, LifecycleHook] = { + "before_agent": LifecycleHook.BEFORE_AGENT, + "after_agent": LifecycleHook.AFTER_AGENT, + "before_model": LifecycleHook.BEFORE_MODEL, + "after_model": LifecycleHook.AFTER_MODEL, + "tool_call": LifecycleHook.TOOL_CALL, + "after_tool": LifecycleHook.AFTER_TOOL, +} + + +async def build_rego_evaluator_async( + service: object, + audit_manager: "AuditManager | None" = None, +) -> object: + """Fetch bundles via platform service and build a RegoEvaluator. + + Drop-in counterpart to the native evaluator's inline bootstrap: + ``get_policy_async`` → ``build_policy_index_from_yaml`` → ``GovernanceEvaluator``. + Takes any object that exposes ``retrieve_all_policies_async()`` and + ``download_bundle_async(url)`` — typically ``UiPath().governance``. + + Args: + service: Platform service object exposing ``retrieve_all_policies_async()`` + and ``download_bundle_async(url)``. + audit_manager: Optional :class:`~uipath.runtime.governance._audit.base.AuditManager` + to wire into the evaluator. When provided, Rego rule violations and + hook summaries are emitted as trace spans alongside the native checker's + events. When ``None``, audit emission is skipped. + + Returns a :class:`~uipath.runtime.governance.rego.evaluator.RegoEvaluator` + on success, ``None`` when no bundles are available. Never raises. + """ + try: + return await _build_rego_evaluator_async_inner(service, audit_manager) + except Exception as exc: # noqa: BLE001 + logger.warning("Rego evaluator build failed: %s", exc) + return None + + +async def _build_rego_evaluator_async_inner( + service: object, + audit_manager: "AuditManager | None" = None, +) -> object: + from pathlib import Path + + from uipath.runtime.governance.native.backend_client import ( + resolve_organization_id, + resolve_tenant_id, + ) + from uipath.runtime.governance.rego.evaluator import ( + RegoEvaluator, + _extract_data_json_from_bundle, + ) + + org_id = resolve_organization_id() + tenant_id = resolve_tenant_id() + if not org_id or not tenant_id: + logger.warning( + "Rego build skipped: org_id or tenant_id unavailable; " + "agent will run without custom Rego evaluation." + ) + return None + + response = await service.retrieve_all_policies_async() # type: ignore[union-attr] + if not response.hook_bundles: + logger.warning( + "Rego build skipped: all-policies returned no WASM bundles for " + "org=%s tenant=%s — no Rego policies are published to this tenant yet.", + org_id, + tenant_id, + ) + return None + + cache_dir = get_cache_dir(org_id, tenant_id) + hook_wasm_paths: dict[LifecycleHook, Path] = {} + hook_data: dict[LifecycleHook, dict] = {} + + for bundle in response.hook_bundles: + hook_type = bundle.hook_type + cached_etag = get_cached_etag(cache_dir, hook_type) + cached_path = get_cached_bundle_path(cache_dir, hook_type) + + if cached_etag != bundle.etag or cached_path is None: + raw = await service.download_bundle_async(bundle.bundle_url) # type: ignore[union-attr] + save_bundle(cache_dir, hook_type, raw, bundle.etag) + + path = get_cached_bundle_path(cache_dir, hook_type) + if path is None: + logger.warning("Rego bundle unavailable for hook=%s after download", hook_type) + continue + + lifecycle_hook = _HOOK_MAP.get(hook_type) + if lifecycle_hook is None: + logger.warning("Rego build: unknown hook type %r — skipping", hook_type) + continue + + hook_wasm_paths[lifecycle_hook] = path + data_json = _extract_data_json_from_bundle(path) + if data_json is not None: + hook_data[lifecycle_hook] = data_json + + if not hook_wasm_paths: + logger.warning( + "Rego build: no WASM bundles available; " + "agent will run without custom Rego evaluation." + ) + return None + + logger.info("Rego evaluator built (hooks=%s)", [h.value for h in hook_wasm_paths]) + return RegoEvaluator(hook_wasm_paths, hook_data=hook_data or None, audit_manager=audit_manager) diff --git a/src/uipath/runtime/governance/runtime.py b/src/uipath/runtime/governance/runtime.py index c9cf699..92e9fd0 100644 --- a/src/uipath/runtime/governance/runtime.py +++ b/src/uipath/runtime/governance/runtime.py @@ -47,7 +47,7 @@ import json import logging from contextlib import contextmanager -from typing import Any, AsyncGenerator, Iterator +from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterator from uipath.core.governance import EnforcementMode from uipath.core.governance.exceptions import GovernanceBlockException @@ -58,12 +58,16 @@ UiPathRuntimeProtocol, UiPathStreamOptions, ) +from uipath.runtime.errors import UiPathErrorCategory, UiPathErrorContract from uipath.runtime.events import UiPathRuntimeEvent from uipath.runtime.governance.native.evaluator import GovernanceEvaluator from uipath.runtime.governance.native.models import PolicyIndex -from uipath.runtime.result import UiPathRuntimeResult +from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus from uipath.runtime.schema import UiPathRuntimeSchema +if TYPE_CHECKING: + from uipath.runtime.governance.rego.evaluator import RegoEvaluator + logger = logging.getLogger(__name__) @@ -119,6 +123,38 @@ def _governance_root_span(agent_name: str, runtime_id: str) -> Iterator[None]: yield +def _find_governance_block(exc: BaseException) -> GovernanceBlockException | None: + """Walk the exception chain to find a wrapped GovernanceBlockException. + + Framework runtimes (e.g. LangGraph) wrap exceptions in their own error + types. This walks ``__cause__`` and ``__context__`` to find a + GovernanceBlockException even when it's not the top-level exception. + """ + seen: set[int] = set() + current: BaseException | None = exc + while current is not None: + if id(current) in seen: + break + seen.add(id(current)) + if isinstance(current, GovernanceBlockException): + return current + current = current.__cause__ or current.__context__ + return None + + +def _governance_faulted_result(exc: GovernanceBlockException) -> UiPathRuntimeResult: + """Build a FAULTED result from a GovernanceBlockException.""" + return UiPathRuntimeResult( + status=UiPathRuntimeStatus.FAULTED, + error=UiPathErrorContract( + code="Governance.PolicyViolation", + title="Policy violation blocked execution", + detail=str(exc), + category=UiPathErrorCategory.USER, + ), + ) + + def _serialize_payload(payload: Any) -> str: """Serialize an agent input / output to a string for evaluator checks. @@ -165,6 +201,7 @@ def __init__( enforcement_mode: EnforcementMode, *, evaluator: GovernanceEvaluator | None = None, + rego_evaluator: RegoEvaluator | None = None, agent_name: str = "", runtime_id: str = "", ): @@ -187,6 +224,10 @@ def __init__( :meth:`execute` / :meth:`stream`. When ``None`` the wrapper is a pure passthrough — the caller is expected to fire those evaluations itself. + rego_evaluator: Optional :class:`RegoEvaluator` for + WASM-based Rego policy evaluation. Runs sequentially + after the native evaluator on each lifecycle hook. + When ``None``, Rego evaluation is skipped. agent_name: Name of the agent (the runtime's entrypoint). Passed through to the evaluator's hook methods. runtime_id: Runtime-instance id (conversation id, job id, @@ -197,46 +238,94 @@ def __init__( self._policy_index = policy_index self._enforcement_mode = enforcement_mode self._evaluator = evaluator + self._rego_evaluator = rego_evaluator self._agent_name = agent_name self._runtime_id = runtime_id + # Wire the Rego evaluator into the native evaluator for per-step + # hooks (BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, AFTER_TOOL). The + # framework adapter (e.g. uipath-langchain) holds a reference to + # the same GovernanceEvaluator object, so this post-construction + # mutation is visible to its callback handler without requiring + # any changes to the framework adapter's code. + # BEFORE_AGENT / AFTER_AGENT are owned by this runtime and are + # already dispatched separately in _fire_before_agent / _fire_after_agent. + if evaluator is not None and rego_evaluator is not None: + evaluator.set_rego_evaluator(rego_evaluator) def _fire_before_agent(self, input: Any) -> None: - """Fire BEFORE_AGENT when an evaluator is wired; otherwise no-op. + """Fire BEFORE_AGENT through native then Rego evaluators. - ``GovernanceBlockException`` propagates — that's how - ENFORCE-mode DENY rules halt a run. Anything else is logged - and swallowed so a governance bug never breaks the agent. + Both evaluators run regardless of whether the first one flags a + violation (sequential-merge), so the caller sees all violations + rather than only the first. The first ``GovernanceBlockException`` + collected is raised; non-block exceptions are logged and swallowed. """ - if self._evaluator is None: - return - try: - self._evaluator.evaluate_before_agent( - agent_input=_serialize_payload(input), - agent_name=self._agent_name, - runtime_id=self._runtime_id, - ) - except GovernanceBlockException: - raise - except Exception as exc: # noqa: BLE001 — never break a run on audit failure - logger.warning("BEFORE_AGENT governance evaluation failed: %s", exc) + violations: list[GovernanceBlockException] = [] + serialized = _serialize_payload(input) + + if self._evaluator is not None: + try: + self._evaluator.evaluate_before_agent( + agent_input=serialized, + agent_name=self._agent_name, + runtime_id=self._runtime_id, + ) + except GovernanceBlockException as exc: + violations.append(exc) + except Exception as exc: # noqa: BLE001 + logger.warning("BEFORE_AGENT native governance evaluation failed: %s", exc) + + if self._rego_evaluator is not None: + try: + self._rego_evaluator.evaluate_before_agent( + agent_input=serialized, + agent_name=self._agent_name, + runtime_id=self._runtime_id, + ) + except GovernanceBlockException as exc: + violations.append(exc) + except Exception as exc: # noqa: BLE001 + logger.warning("BEFORE_AGENT rego governance evaluation failed: %s", exc) + + if violations: + raise violations[0] def _fire_after_agent(self, result: UiPathRuntimeResult) -> None: - """Fire AFTER_AGENT against ``result.output``. + """Fire AFTER_AGENT through native then Rego evaluators. - Same exception policy as :meth:`_fire_before_agent`. + Same sequential-merge exception policy as :meth:`_fire_before_agent`. """ - if self._evaluator is None: + if self._evaluator is None and self._rego_evaluator is None: return - try: - self._evaluator.evaluate_after_agent( - agent_output=_serialize_payload(result.output), - agent_name=self._agent_name, - runtime_id=self._runtime_id, - ) - except GovernanceBlockException: - raise - except Exception as exc: # noqa: BLE001 - logger.warning("AFTER_AGENT governance evaluation failed: %s", exc) + violations: list[GovernanceBlockException] = [] + serialized = _serialize_payload(getattr(result, "output", result)) + + if self._evaluator is not None: + try: + self._evaluator.evaluate_after_agent( + agent_output=serialized, + agent_name=self._agent_name, + runtime_id=self._runtime_id, + ) + except GovernanceBlockException as exc: + violations.append(exc) + except Exception as exc: # noqa: BLE001 + logger.warning("AFTER_AGENT native governance evaluation failed: %s", exc) + + if self._rego_evaluator is not None: + try: + self._rego_evaluator.evaluate_after_agent( + agent_output=serialized, + agent_name=self._agent_name, + runtime_id=self._runtime_id, + ) + except GovernanceBlockException as exc: + violations.append(exc) + except Exception as exc: # noqa: BLE001 + logger.warning("AFTER_AGENT rego governance evaluation failed: %s", exc) + + if violations: + raise violations[0] async def execute( self, @@ -255,10 +344,18 @@ async def execute( raises, there's no output to evaluate. """ with _governance_root_span(self._agent_name, self._runtime_id): - self._fire_before_agent(input) - result = await self._delegate.execute(input, options=options) - self._fire_after_agent(result) - return result + try: + self._fire_before_agent(input) + result = await self._delegate.execute(input, options=options) + self._fire_after_agent(result) + return result + except GovernanceBlockException as exc: + return _governance_faulted_result(exc) + except Exception as exc: + block = _find_governance_block(exc) + if block is not None: + return _governance_faulted_result(block) + raise async def stream( self, @@ -278,10 +375,18 @@ async def stream( pass through untouched. """ with _governance_root_span(self._agent_name, self._runtime_id): - self._fire_before_agent(input) + try: + self._fire_before_agent(input) + except GovernanceBlockException as exc: + yield _governance_faulted_result(exc) + return async for event in self._delegate.stream(input, options=options): if isinstance(event, UiPathRuntimeResult): - self._fire_after_agent(event) + try: + self._fire_after_agent(event) + except GovernanceBlockException as exc: + yield _governance_faulted_result(exc) + return yield event async def get_schema(self) -> UiPathRuntimeSchema: diff --git a/tests/test_governance_runtime.py b/tests/test_governance_runtime.py index b2ca3e3..6b38472 100644 --- a/tests/test_governance_runtime.py +++ b/tests/test_governance_runtime.py @@ -480,13 +480,15 @@ async def test_execute_fires_before_and_after_agent_when_evaluator_wired() -> No assert evaluator.after_calls[0]["agent_output"] == "agent-output" -async def test_execute_propagates_governance_block_exception() -> None: - """A DENY in ``ENFORCE`` mode must halt the run — the evaluator - raises :class:`GovernanceBlockException` and the wrapper propagates - it rather than swallowing. +async def test_execute_governance_block_returns_faulted() -> None: + """A DENY in ``ENFORCE`` mode must halt the run — the wrapper catches + the :class:`GovernanceBlockException` and returns a FAULTED result + with error code ``Governance.PolicyViolation``. """ from uipath.core.governance.exceptions import GovernanceBlockException + from uipath.runtime.result import UiPathRuntimeStatus + evaluator = _CapturingEvaluator() evaluator.before_raises = GovernanceBlockException("policy denied") @@ -497,12 +499,14 @@ async def test_execute_propagates_governance_block_exception() -> None: evaluator=evaluator, # type: ignore[arg-type] ) - with pytest.raises(GovernanceBlockException): - await runtime.execute({"x": 1}) + result = await runtime.execute({"x": 1}) + assert result.status == UiPathRuntimeStatus.FAULTED + assert result.error is not None + assert result.error.code == "Governance.PolicyViolation" -async def test_after_agent_block_exception_also_propagates() -> None: - """The re-raise contract holds for AFTER_AGENT too, not just +async def test_after_agent_block_also_returns_faulted() -> None: + """The FAULTED-result contract holds for AFTER_AGENT too, not just BEFORE_AGENT. Even though DENY on output is a rarer configuration (most policies @@ -512,6 +516,8 @@ async def test_after_agent_block_exception_also_propagates() -> None: """ from uipath.core.governance.exceptions import GovernanceBlockException + from uipath.runtime.result import UiPathRuntimeStatus + evaluator = _CapturingEvaluator() evaluator.after_raises = GovernanceBlockException("output policy denied") @@ -522,8 +528,10 @@ async def test_after_agent_block_exception_also_propagates() -> None: evaluator=evaluator, # type: ignore[arg-type] ) - with pytest.raises(GovernanceBlockException): - await runtime.execute({"x": 1}) + result = await runtime.execute({"x": 1}) + assert result.status == UiPathRuntimeStatus.FAULTED + assert result.error is not None + assert result.error.code == "Governance.PolicyViolation" async def test_execute_swallows_unexpected_evaluator_errors() -> None: diff --git a/tests/test_rego_bundle_cache.py b/tests/test_rego_bundle_cache.py new file mode 100644 index 0000000..10b5f04 --- /dev/null +++ b/tests/test_rego_bundle_cache.py @@ -0,0 +1,109 @@ +"""Tests for the Rego WASM bundle disk cache.""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +from uipath.runtime.governance.rego.bundle_cache import ( + get_cache_dir, + get_cached_bundle_path, + get_cached_etag, + save_bundle, +) + + +@pytest.fixture +def cache_dir(tmp_path: Path) -> Path: + return get_cache_dir.__wrapped__(tmp_path) if hasattr(get_cache_dir, "__wrapped__") else tmp_path / "cache" + + +# --------------------------------------------------------------------------- +# get_cache_dir +# --------------------------------------------------------------------------- + +def test_get_cache_dir_creates_directory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + import tempfile + monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path)) + result = get_cache_dir("org1", "tenant1") + assert result.exists() + assert result.is_dir() + + +def test_get_cache_dir_same_org_tenant_returns_same_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + import tempfile + monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path)) + a = get_cache_dir("org1", "tenant1") + b = get_cache_dir("org1", "tenant1") + assert a == b + + +def test_get_cache_dir_different_pairs_return_different_paths(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + import tempfile + monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path)) + a = get_cache_dir("org1", "tenant1") + b = get_cache_dir("org1", "tenant2") + assert a != b + + +# --------------------------------------------------------------------------- +# save_bundle / get_cached_bundle_path / get_cached_etag +# --------------------------------------------------------------------------- + +def test_save_and_retrieve_bundle(tmp_path: Path) -> None: + bundle_path = save_bundle(tmp_path, "before_model", b"fake-wasm", "etag-abc") + assert bundle_path.exists() + assert bundle_path.read_bytes() == b"fake-wasm" + + +def test_get_cached_bundle_path_returns_path_after_save(tmp_path: Path) -> None: + save_bundle(tmp_path, "before_model", b"wasm-data", "etag-1") + result = get_cached_bundle_path(tmp_path, "before_model") + assert result is not None + assert result.exists() + + +def test_get_cached_bundle_path_returns_none_when_missing(tmp_path: Path) -> None: + result = get_cached_bundle_path(tmp_path, "before_model") + assert result is None + + +def test_get_cached_etag_returns_etag_after_save(tmp_path: Path) -> None: + save_bundle(tmp_path, "after_agent", b"wasm", "etag-xyz") + result = get_cached_etag(tmp_path, "after_agent") + assert result == "etag-xyz" + + +def test_get_cached_etag_returns_none_when_missing(tmp_path: Path) -> None: + result = get_cached_etag(tmp_path, "after_agent") + assert result is None + + +def test_save_bundle_overwrites_existing(tmp_path: Path) -> None: + save_bundle(tmp_path, "before_model", b"old", "etag-old") + save_bundle(tmp_path, "before_model", b"new", "etag-new") + path = get_cached_bundle_path(tmp_path, "before_model") + assert path is not None + assert path.read_bytes() == b"new" + assert get_cached_etag(tmp_path, "before_model") == "etag-new" + + +def test_different_hook_types_are_isolated(tmp_path: Path) -> None: + save_bundle(tmp_path, "before_model", b"wasm-bm", "etag-bm") + save_bundle(tmp_path, "after_agent", b"wasm-aa", "etag-aa") + assert get_cached_etag(tmp_path, "before_model") == "etag-bm" + assert get_cached_etag(tmp_path, "after_agent") == "etag-aa" + assert get_cached_bundle_path(tmp_path, "tool_call") is None + + +def test_get_cached_etag_handles_oserror(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + save_bundle(tmp_path, "before_model", b"wasm", "etag-1") + # Make the etag file unreadable + hook_dir = tmp_path / "hooks" / "before_model" + etag_file = hook_dir / "etag.txt" + etag_file.chmod(0o000) + try: + result = get_cached_etag(tmp_path, "before_model") + assert result is None + finally: + etag_file.chmod(0o644) diff --git a/tests/test_rego_evaluator.py b/tests/test_rego_evaluator.py new file mode 100644 index 0000000..5dd3117 --- /dev/null +++ b/tests/test_rego_evaluator.py @@ -0,0 +1,423 @@ +"""Tests for RegoEvaluator: enforcement modes, audit emission, WASM mock.""" +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from uipath.core.governance import EnforcementMode +from uipath.core.governance.exceptions import GovernanceBlockException +from uipath.core.governance.models import Action, LifecycleHook + +from uipath.runtime.governance.rego.evaluator import ( + RegoEvaluator, + _extract_data_json_from_bundle, + _extract_wasm_from_bundle, + _pack_name_from_rule_id, + context_to_input, +) +from uipath.runtime.governance.native.models import CheckContext + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +import io +import json +import tarfile + + +def _make_bundle(wasm: bytes = b"\x00asm", data: dict | None = None) -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for name, content in [("policy.wasm", wasm)] + ( + [("data.json", json.dumps(data).encode())] if data else [] + ): + info = tarfile.TarInfo(name=name) + info.size = len(content) + tf.addfile(info, io.BytesIO(content)) + return buf.getvalue() + + +def _write_bundle(path: Path, **kwargs: Any) -> Path: + path.write_bytes(_make_bundle(**kwargs)) + return path + + +def _make_context( + hook: LifecycleHook = LifecycleHook.BEFORE_MODEL, + agent_name: str = "test-agent", + runtime_id: str = "run-1", + model_input: str = "hello", +) -> CheckContext: + return CheckContext( + hook=hook, + agent_name=agent_name, + runtime_id=runtime_id, + model_input=model_input, + ) + + +def _make_engine(result: dict | None = None) -> MagicMock: + """Return a fake OPA engine whose evaluate() returns a single-item list.""" + engine = MagicMock() + engine.evaluate.return_value = [{"result": result or {"deny": False}}] + return engine + + +def _make_evaluator( + hook: LifecycleHook = LifecycleHook.BEFORE_MODEL, + engine_result: dict | None = None, + enforcement_mode: EnforcementMode = EnforcementMode.ENFORCE, + bundle_path: Path | None = None, +) -> tuple[RegoEvaluator, MagicMock]: + engine = _make_engine(engine_result) + with patch("uipath.runtime.governance.rego.evaluator._load_engine", return_value=engine): + ev = RegoEvaluator( + hook_wasm_paths={hook: bundle_path or Path("/fake/bundle.tar.gz")}, + enforcement_mode=enforcement_mode, + ) + return ev, engine + + +# --------------------------------------------------------------------------- +# _pack_name_from_rule_id +# --------------------------------------------------------------------------- + +def test_pack_name_with_slash() -> None: + assert _pack_name_from_rule_id("my_pack/block_ssn") == "my_pack" + + +def test_pack_name_without_slash() -> None: + assert _pack_name_from_rule_id("block_ssn") == "custom" + + +# --------------------------------------------------------------------------- +# context_to_input +# --------------------------------------------------------------------------- + +def test_context_to_input_basic_fields() -> None: + ctx = _make_context(hook=LifecycleHook.BEFORE_MODEL, model_input="test input") + result = context_to_input(ctx) + assert result["hook"] == "before_model" + assert result["model_input"] == "test input" + assert result["agent_name"] == "test-agent" + + +def test_context_to_input_session_state_defaults() -> None: + ctx = _make_context() + result = context_to_input(ctx) + assert result["session_state"] == {"tool_calls": 0, "llm_calls": 0} + + +def test_context_to_input_before_model_normalizes_agent_input() -> None: + # When agent_input is empty for BEFORE_MODEL, model_input is copied so + # WASM rules checking either field receive the user's message. + ctx = CheckContext( + hook=LifecycleHook.BEFORE_MODEL, + agent_name="a", + runtime_id="r", + model_input="identify 999-99-9999", + agent_input="", + ) + result = context_to_input(ctx) + assert result["agent_input"] == "identify 999-99-9999" + assert result["model_input"] == "identify 999-99-9999" + + +def test_context_to_input_before_model_keeps_explicit_agent_input() -> None: + # When agent_input is already set, it is NOT overwritten. + ctx = CheckContext( + hook=LifecycleHook.BEFORE_MODEL, + agent_name="a", + runtime_id="r", + model_input="model text", + agent_input="explicit agent text", + ) + result = context_to_input(ctx) + assert result["agent_input"] == "explicit agent text" + + +def test_context_to_input_other_hooks_no_normalization() -> None: + # Normalization must not apply to other hooks. + ctx = CheckContext( + hook=LifecycleHook.BEFORE_AGENT, + agent_name="a", + runtime_id="r", + agent_input="", + model_input="should not copy", + ) + result = context_to_input(ctx) + assert result["agent_input"] == "" + + +def test_context_to_input_before_model_synthesizes_messages() -> None: + # When messages is empty for BEFORE_MODEL, a synthetic user message is + # built from model_input so WASM rules checking input.messages fire. + ctx = CheckContext( + hook=LifecycleHook.BEFORE_MODEL, + agent_name="a", + runtime_id="r", + model_input="identify 999-99-9999", + messages=[], + ) + result = context_to_input(ctx) + assert result["messages"] == [{"role": "user", "content": "identify 999-99-9999"}] + + +def test_context_to_input_before_model_keeps_explicit_messages() -> None: + # When messages is already populated, it is NOT overwritten. + existing = [{"role": "user", "content": "original"}] + ctx = CheckContext( + hook=LifecycleHook.BEFORE_MODEL, + agent_name="a", + runtime_id="r", + model_input="other text", + messages=existing, + ) + result = context_to_input(ctx) + assert result["messages"] == existing + + +def test_context_to_input_other_hooks_no_messages_synthesis() -> None: + # Messages synthesis must not apply to other hooks. + ctx = CheckContext( + hook=LifecycleHook.BEFORE_AGENT, + agent_name="a", + runtime_id="r", + model_input="should not synthesize", + messages=[], + ) + result = context_to_input(ctx) + assert result["messages"] == [] + + +# --------------------------------------------------------------------------- +# _extract_wasm_from_bundle / _extract_data_json_from_bundle +# --------------------------------------------------------------------------- + +def test_extract_wasm_from_bundle(tmp_path: Path) -> None: + p = _write_bundle(tmp_path / "b.tar.gz", wasm=b"fake-wasm") + assert _extract_wasm_from_bundle(p) == b"fake-wasm" + + +def test_extract_data_json_from_bundle(tmp_path: Path) -> None: + p = _write_bundle(tmp_path / "b.tar.gz", data={"required_features": ["sentiment"]}) + result = _extract_data_json_from_bundle(p) + assert result == {"required_features": ["sentiment"]} + + +def test_extract_data_json_returns_none_when_absent(tmp_path: Path) -> None: + p = _write_bundle(tmp_path / "b.tar.gz") + assert _extract_data_json_from_bundle(p) is None + + +# --------------------------------------------------------------------------- +# RegoEvaluator — DISABLED mode +# --------------------------------------------------------------------------- + +def test_evaluate_disabled_mode_returns_allow_without_engine() -> None: + ev = RegoEvaluator( + hook_wasm_paths={}, + enforcement_mode=EnforcementMode.DISABLED, + ) + ctx = _make_context() + record = ev.evaluate(ctx) + assert record.final_action == Action.ALLOW + assert record.metadata.get("rego_engine") == "skipped" + + +# --------------------------------------------------------------------------- +# RegoEvaluator — missing engine +# --------------------------------------------------------------------------- + +def test_evaluate_missing_engine_returns_allow() -> None: + ev, _ = _make_evaluator(hook=LifecycleHook.AFTER_AGENT) + ctx = _make_context(hook=LifecycleHook.BEFORE_AGENT) + record = ev.evaluate(ctx) + assert record.final_action == Action.ALLOW + assert record.metadata.get("rego_engine") == "missing" + + +# --------------------------------------------------------------------------- +# RegoEvaluator — ENFORCE mode, allow +# --------------------------------------------------------------------------- + +def test_evaluate_enforce_allow_returns_audit_record() -> None: + ev, _ = _make_evaluator(engine_result={"deny": False}) + ctx = _make_context() + record = ev.evaluate(ctx) + assert record.final_action == Action.ALLOW + assert record.hook == LifecycleHook.BEFORE_MODEL + + +# --------------------------------------------------------------------------- +# RegoEvaluator — ENFORCE mode, deny raises GovernanceBlockException +# --------------------------------------------------------------------------- + +def test_evaluate_enforce_deny_raises() -> None: + ev, _ = _make_evaluator(engine_result={"deny": True, "fired_deny": ["pack/rule1"]}) + ctx = _make_context() + with pytest.raises(GovernanceBlockException): + ev.evaluate(ctx) + + +# --------------------------------------------------------------------------- +# RegoEvaluator — AUDIT mode, deny does NOT raise +# --------------------------------------------------------------------------- + +def test_evaluate_audit_mode_deny_does_not_raise() -> None: + ev, _ = _make_evaluator( + engine_result={"deny": True, "fired_deny": ["pack/rule1"]}, + enforcement_mode=EnforcementMode.AUDIT, + ) + ctx = _make_context() + record = ev.evaluate(ctx) + assert record.final_action == Action.AUDIT + assert record.metadata.get("audit_mode_would_deny") is True + + +# --------------------------------------------------------------------------- +# RegoEvaluator — engine raises, fail-open +# --------------------------------------------------------------------------- + +def test_evaluate_engine_exception_returns_allow() -> None: + engine = MagicMock() + engine.evaluate.side_effect = RuntimeError("wasm error") + with patch("uipath.runtime.governance.rego.evaluator._load_engine", return_value=engine): + ev = RegoEvaluator( + hook_wasm_paths={LifecycleHook.BEFORE_MODEL: Path("/fake/bundle.tar.gz")}, + enforcement_mode=EnforcementMode.ENFORCE, + ) + ctx = _make_context() + record = ev.evaluate(ctx) + assert record.final_action == Action.ALLOW + assert record.metadata.get("rego_engine") == "error" + + +# --------------------------------------------------------------------------- +# RegoEvaluator — fired_deny / fired_allow rule evaluations +# --------------------------------------------------------------------------- + +def test_evaluate_populates_rule_evaluations() -> None: + ev, _ = _make_evaluator(engine_result={ + "deny": True, + "fired_deny": ["pack/block_ssn"], + "fired_allow": ["pack/safe_rule"], + "messages": {"pack/block_ssn": "SSN detected"}, + }) + ctx = _make_context() + with pytest.raises(GovernanceBlockException) as exc_info: + ev.evaluate(ctx) + # The audit record is embedded in the exception + record = exc_info.value.audit_record + deny_evals = [e for e in record.evaluations if e.action == Action.DENY] + allow_evals = [e for e in record.evaluations if e.action == Action.ALLOW] + assert len(deny_evals) == 1 + assert deny_evals[0].rule_id == "pack/block_ssn" + assert deny_evals[0].detail == "SSN detected" + assert len(allow_evals) == 1 + + +# --------------------------------------------------------------------------- +# RegoEvaluator — opa_wasmtime not installed, load warning +# --------------------------------------------------------------------------- + +def test_load_engine_import_error_skips_hook() -> None: + with patch("uipath.runtime.governance.rego.evaluator._load_engine", + side_effect=ImportError("opa_wasmtime not installed")): + ev = RegoEvaluator( + hook_wasm_paths={LifecycleHook.BEFORE_MODEL: Path("/fake/bundle.tar.gz")}, + enforcement_mode=EnforcementMode.ENFORCE, + ) + assert LifecycleHook.BEFORE_MODEL not in ev._engines + assert ev.loaded_hooks == [] + + +# --------------------------------------------------------------------------- +# RegoEvaluator — audit emission +# --------------------------------------------------------------------------- + +def test_evaluate_emits_audit_events() -> None: + from uipath.runtime.governance._audit.base import AuditManager, AuditSink, AuditEvent + + class _Sink(AuditSink): + def __init__(self) -> None: + self.events: list[AuditEvent] = [] + @property + def name(self) -> str: + return "test-sink" + def emit(self, event: AuditEvent) -> None: + self.events.append(event) + + sink = _Sink() + manager = AuditManager() + manager.register_sink(sink) + + ev, _ = _make_evaluator( + engine_result={"deny": False, "fired_allow": ["pack/safe"]}, + enforcement_mode=EnforcementMode.ENFORCE, + ) + ev._audit_manager = manager + + ctx = _make_context() + ev.evaluate(ctx) + assert len(sink.events) > 0 + + +# --------------------------------------------------------------------------- +# RegoEvaluator — convenience evaluate_* methods +# --------------------------------------------------------------------------- + +def test_evaluate_before_model() -> None: + ev, _ = _make_evaluator(engine_result={"deny": False}) + record = ev.evaluate_before_model("prompt", "agent", "run-1", model_name="gpt-4") + assert record.hook == LifecycleHook.BEFORE_MODEL + + +def test_evaluate_before_agent() -> None: + ev, _ = _make_evaluator( + hook=LifecycleHook.BEFORE_AGENT, + engine_result={"deny": False}, + ) + record = ev.evaluate_before_agent("input", "agent", "run-1") + assert record.hook == LifecycleHook.BEFORE_AGENT + + +def test_evaluate_after_agent() -> None: + ev, _ = _make_evaluator( + hook=LifecycleHook.AFTER_AGENT, + engine_result={"deny": False}, + ) + record = ev.evaluate_after_agent("output", "agent", "run-1") + assert record.hook == LifecycleHook.AFTER_AGENT + + +def test_evaluate_tool_call() -> None: + ev, _ = _make_evaluator( + hook=LifecycleHook.TOOL_CALL, + engine_result={"deny": False}, + ) + record = ev.evaluate_tool_call("my_tool", {"arg": "val"}, "agent", "run-1") + assert record.hook == LifecycleHook.TOOL_CALL + + +def test_evaluate_after_tool() -> None: + ev, _ = _make_evaluator( + hook=LifecycleHook.AFTER_TOOL, + engine_result={"deny": False}, + ) + record = ev.evaluate_after_tool("my_tool", "result", "agent", "run-1") + assert record.hook == LifecycleHook.AFTER_TOOL + + +def test_evaluate_after_model() -> None: + ev, _ = _make_evaluator( + hook=LifecycleHook.AFTER_MODEL, + engine_result={"deny": False}, + ) + record = ev.evaluate_after_model("output", "agent", "run-1") + assert record.hook == LifecycleHook.AFTER_MODEL diff --git a/tests/test_rego_loader.py b/tests/test_rego_loader.py new file mode 100644 index 0000000..e645a6d --- /dev/null +++ b/tests/test_rego_loader.py @@ -0,0 +1,209 @@ +"""Tests for the Rego async evaluator bootstrap (loader.py).""" +from __future__ import annotations + +import io +import tarfile +import tempfile +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from uipath.core.governance.models import LifecycleHook + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_tar_gz(files: dict[str, bytes]) -> bytes: + """Build an in-memory .tar.gz with the given filename → bytes mapping.""" + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for name, data in files.items(): + info = tarfile.TarInfo(name=name) + info.size = len(data) + tf.addfile(info, io.BytesIO(data)) + return buf.getvalue() + + +def _fake_bundle(wasm: bytes = b"\x00asm\x01\x00\x00\x00", data: dict | None = None) -> bytes: + """Build a minimal OPA .tar.gz bundle.""" + import json + files: dict[str, bytes] = {"policy.wasm": wasm} + if data is not None: + files["data.json"] = json.dumps(data).encode() + return _make_tar_gz(files) + + +def _hook_bundle(hook_type: str, etag: str = "etag-1") -> Any: + b = MagicMock() + b.hook_type = hook_type + b.bundle_url = f"https://cdn.example.com/{hook_type}.tar.gz" + b.etag = etag + return b + + +def _all_policies_response(*hook_bundles: Any) -> Any: + r = MagicMock() + r.hook_bundles = list(hook_bundles) + return r + + +# --------------------------------------------------------------------------- +# build_rego_evaluator_async — no bundles +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_returns_none_when_no_hook_bundles() -> None: + from uipath.runtime.governance.rego.loader import build_rego_evaluator_async + + service = MagicMock() + service.retrieve_all_policies_async = AsyncMock(return_value=_all_policies_response()) + + with patch("uipath.runtime.governance.native.backend_client.resolve_organization_id", return_value="org1"), \ + patch("uipath.runtime.governance.native.backend_client.resolve_tenant_id", return_value="tenant1"): + result = await build_rego_evaluator_async(service) + + assert result is None + service.download_bundle_async.assert_not_called() + + +# --------------------------------------------------------------------------- +# build_rego_evaluator_async — missing org/tenant +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_returns_none_when_org_id_missing() -> None: + from uipath.runtime.governance.rego.loader import build_rego_evaluator_async + + service = MagicMock() + with patch("uipath.runtime.governance.native.backend_client.resolve_organization_id", return_value=None), \ + patch("uipath.runtime.governance.native.backend_client.resolve_tenant_id", return_value="tenant1"): + result = await build_rego_evaluator_async(service) + + assert result is None + + +@pytest.mark.asyncio +async def test_returns_none_when_tenant_id_missing() -> None: + from uipath.runtime.governance.rego.loader import build_rego_evaluator_async + + service = MagicMock() + with patch("uipath.runtime.governance.native.backend_client.resolve_organization_id", return_value="org1"), \ + patch("uipath.runtime.governance.native.backend_client.resolve_tenant_id", return_value=None): + result = await build_rego_evaluator_async(service) + + assert result is None + + +# --------------------------------------------------------------------------- +# build_rego_evaluator_async — service raises +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_returns_none_when_service_raises() -> None: + from uipath.runtime.governance.rego.loader import build_rego_evaluator_async + + service = MagicMock() + service.retrieve_all_policies_async = AsyncMock(side_effect=RuntimeError("network error")) + + with patch("uipath.runtime.governance.native.backend_client.resolve_organization_id", return_value="org1"), \ + patch("uipath.runtime.governance.native.backend_client.resolve_tenant_id", return_value="tenant1"): + result = await build_rego_evaluator_async(service) + + assert result is None + + +# --------------------------------------------------------------------------- +# build_rego_evaluator_async — unknown hook type skipped +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_unknown_hook_type_is_skipped(tmp_path: Path) -> None: + from uipath.runtime.governance.rego.loader import build_rego_evaluator_async + + bundle_data = _fake_bundle() + service = MagicMock() + service.retrieve_all_policies_async = AsyncMock( + return_value=_all_policies_response(_hook_bundle("unknown_hook_xyz")) + ) + service.download_bundle_async = AsyncMock(return_value=bundle_data) + + with patch("uipath.runtime.governance.native.backend_client.resolve_organization_id", return_value="org1"), \ + patch("uipath.runtime.governance.native.backend_client.resolve_tenant_id", return_value="tenant1"), \ + patch("uipath.runtime.governance.rego.loader.get_cache_dir", return_value=tmp_path), \ + patch("uipath.runtime.governance.rego.loader.get_cached_etag", return_value=None), \ + patch("uipath.runtime.governance.rego.loader.get_cached_bundle_path", return_value=tmp_path / "b.tar.gz"), \ + patch("uipath.runtime.governance.rego.loader.save_bundle"): + result = await build_rego_evaluator_async(service) + + assert result is None + + +# --------------------------------------------------------------------------- +# build_rego_evaluator_async — etag cache hit skips download +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_etag_cache_hit_skips_download(tmp_path: Path) -> None: + from uipath.runtime.governance.rego.loader import build_rego_evaluator_async + + bundle_path = tmp_path / "bundle.tar.gz" + bundle_path.write_bytes(_fake_bundle()) + + service = MagicMock() + service.retrieve_all_policies_async = AsyncMock( + return_value=_all_policies_response(_hook_bundle("before_model", etag="etag-cached")) + ) + service.download_bundle_async = AsyncMock() + + mock_evaluator = MagicMock() + mock_evaluator.loaded_hooks = [LifecycleHook.BEFORE_MODEL] + + with patch("uipath.runtime.governance.native.backend_client.resolve_organization_id", return_value="org1"), \ + patch("uipath.runtime.governance.native.backend_client.resolve_tenant_id", return_value="tenant1"), \ + patch("uipath.runtime.governance.rego.loader.get_cache_dir", return_value=tmp_path), \ + patch("uipath.runtime.governance.rego.loader.get_cached_etag", return_value="etag-cached"), \ + patch("uipath.runtime.governance.rego.loader.get_cached_bundle_path", return_value=bundle_path), \ + patch("uipath.runtime.governance.rego.evaluator._extract_data_json_from_bundle", return_value=None), \ + patch("uipath.runtime.governance.rego.evaluator.RegoEvaluator", return_value=mock_evaluator): + result = await build_rego_evaluator_async(service) + + service.download_bundle_async.assert_not_called() + assert result is mock_evaluator + + +# --------------------------------------------------------------------------- +# build_rego_evaluator_async — etag mismatch triggers download +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_etag_mismatch_triggers_download(tmp_path: Path) -> None: + from uipath.runtime.governance.rego.loader import build_rego_evaluator_async + + bundle_data = _fake_bundle() + bundle_path = tmp_path / "bundle.tar.gz" + bundle_path.write_bytes(bundle_data) + + service = MagicMock() + service.retrieve_all_policies_async = AsyncMock( + return_value=_all_policies_response(_hook_bundle("before_model", etag="etag-new")) + ) + service.download_bundle_async = AsyncMock(return_value=bundle_data) + + mock_evaluator = MagicMock() + + with patch("uipath.runtime.governance.native.backend_client.resolve_organization_id", return_value="org1"), \ + patch("uipath.runtime.governance.native.backend_client.resolve_tenant_id", return_value="tenant1"), \ + patch("uipath.runtime.governance.rego.loader.get_cache_dir", return_value=tmp_path), \ + patch("uipath.runtime.governance.rego.loader.get_cached_etag", return_value="etag-old"), \ + patch("uipath.runtime.governance.rego.loader.get_cached_bundle_path", return_value=bundle_path), \ + patch("uipath.runtime.governance.rego.loader.save_bundle"), \ + patch("uipath.runtime.governance.rego.evaluator._extract_data_json_from_bundle", return_value=None), \ + patch("uipath.runtime.governance.rego.evaluator.RegoEvaluator", return_value=mock_evaluator): + result = await build_rego_evaluator_async(service) + + service.download_bundle_async.assert_called_once() + assert result is mock_evaluator diff --git a/uv.lock b/uv.lock index afe9408..1aa86de 100644 --- a/uv.lock +++ b/uv.lock @@ -1153,7 +1153,7 @@ wheels = [ [[package]] name = "uipath-runtime" -version = "0.12.2" +version = "0.12.3" source = { editable = "." } dependencies = [ { name = "chardet" },