From c895b3728734d3a9c6f44627ac14b045f3fce62b Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Thu, 9 Jul 2026 15:20:36 +0530 Subject: [PATCH 01/12] feat: port Rego evaluator to feat/rego-evaluator-v2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds WASM-based Rego evaluation alongside the native (YAML/regex) evaluator, aligned with main's DI-first, FF-free architecture: - governance/features/ — shared NLP/regex signal extractors (moved from rego/features/ so both evaluators can reuse them); commitment, encoding, incident, sentiment, text_stats extractors registered via @register() - governance/rego/ — RegoEvaluator (opa-wasmtime), bundle cache, api_client, loader with ETag-driven disk cache + 30s background refresh - governance/native/ — backend_client, policy_api_client, _yaml_to_index, loader (prefetch + background load, no FF gating) - governance/config.py — re-exports EnforcementMode from uipath.core.governance, keeps get/set/reset_enforcement_mode for process-level state - governance/delegation_guard.py — ASI-02 recursion depth guard - governance/_audit/console.py — human-readable logging sink - governance/_audit/factory.py — add "console" sink registration Architecture alignment with main: - RegoEvaluator accepts audit_manager + enforcement_mode via DI (no globals) - emit_rule_evaluation uses policy_id= (main API), enforcement_mode as object - No FF gating anywhere; enable/disable is server-side via EnforcementMode - sequential-merge: both evaluators run, violations collected before raise runtime.py changes: - rego_evaluator kwarg added to UiPathGovernedRuntime - _fire_before/after_agent: sequential-merge pattern (native then rego) - execute()/stream(): catch GovernanceBlockException → FAULTED result with error code Governance.PolicyViolation - _find_governance_block(): walks __cause__/__context__ chain to unwrap framework-wrapped GovernanceBlockExceptions Tests updated: two tests now assert FAULTED result instead of raised exception. Generated with Claude Code Co-Authored-By: Claude --- .../runtime/governance/_audit/console.py | 118 +++++ .../runtime/governance/_audit/factory.py | 5 + src/uipath/runtime/governance/config.py | 69 +++ .../runtime/governance/delegation_guard.py | 205 ++++++++ .../runtime/governance/features/__init__.py | 18 + .../runtime/governance/features/_registry.py | 40 ++ .../runtime/governance/features/_text.py | 34 ++ .../runtime/governance/features/commitment.py | 63 +++ .../runtime/governance/features/encoding.py | 64 +++ .../runtime/governance/features/incident.py | 60 +++ .../runtime/governance/features/sentiment.py | 36 ++ .../runtime/governance/features/text_stats.py | 37 ++ .../governance/native/_yaml_to_index.py | 459 ++++++++++++++++++ .../governance/native/backend_client.py | 384 +++++++++++++++ .../runtime/governance/native/loader.py | 284 +++++++++++ .../governance/native/policy_api_client.py | 227 +++++++++ .../runtime/governance/rego/__init__.py | 12 + .../runtime/governance/rego/api_client.py | 145 ++++++ .../runtime/governance/rego/bundle_cache.py | 65 +++ .../runtime/governance/rego/evaluator.py | 448 +++++++++++++++++ src/uipath/runtime/governance/rego/loader.py | 235 +++++++++ src/uipath/runtime/governance/rego/models.py | 18 + src/uipath/runtime/governance/runtime.py | 169 +++++-- tests/test_governance_runtime.py | 28 +- 24 files changed, 3176 insertions(+), 47 deletions(-) create mode 100644 src/uipath/runtime/governance/_audit/console.py create mode 100644 src/uipath/runtime/governance/config.py create mode 100644 src/uipath/runtime/governance/delegation_guard.py create mode 100644 src/uipath/runtime/governance/features/__init__.py create mode 100644 src/uipath/runtime/governance/features/_registry.py create mode 100644 src/uipath/runtime/governance/features/_text.py create mode 100644 src/uipath/runtime/governance/features/commitment.py create mode 100644 src/uipath/runtime/governance/features/encoding.py create mode 100644 src/uipath/runtime/governance/features/incident.py create mode 100644 src/uipath/runtime/governance/features/sentiment.py create mode 100644 src/uipath/runtime/governance/features/text_stats.py create mode 100644 src/uipath/runtime/governance/native/_yaml_to_index.py create mode 100644 src/uipath/runtime/governance/native/backend_client.py create mode 100644 src/uipath/runtime/governance/native/loader.py create mode 100644 src/uipath/runtime/governance/native/policy_api_client.py create mode 100644 src/uipath/runtime/governance/rego/__init__.py create mode 100644 src/uipath/runtime/governance/rego/api_client.py create mode 100644 src/uipath/runtime/governance/rego/bundle_cache.py create mode 100644 src/uipath/runtime/governance/rego/evaluator.py create mode 100644 src/uipath/runtime/governance/rego/loader.py create mode 100644 src/uipath/runtime/governance/rego/models.py 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/_yaml_to_index.py b/src/uipath/runtime/governance/native/_yaml_to_index.py new file mode 100644 index 0000000..2deb463 --- /dev/null +++ b/src/uipath/runtime/governance/native/_yaml_to_index.py @@ -0,0 +1,459 @@ +"""Runtime YAML → PolicyIndex parser. + +Mirrors the shape produced by ``packs/compile_packs.py`` but builds the +PolicyIndex directly from parsed YAML data rather than generating Python +source. Used by :mod:`uipath.runtime.governance.native.loader` when policies are fetched +from the governance backend at startup. + +Accepts either a single YAML document (one pack) or a multi-document +stream (``---``-separated packs). Unknown check types and malformed +rules are skipped with a warning — partial packs are preferred over +failing the whole load. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import yaml +from uipath.core.governance.models import Action, LifecycleHook + +from uipath.runtime.governance.native.models import ( + Check, + Condition, + PolicyIndex, + PolicyPack, + Rule, + Severity, +) + +logger = logging.getLogger(__name__) + + +_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, + "wrap_tool_call": LifecycleHook.TOOL_CALL, + "tool_call": LifecycleHook.TOOL_CALL, + "after_tool": LifecycleHook.AFTER_TOOL, +} + +_ACTION_MAP: dict[str, Action] = { + "block": Action.DENY, + "deny": Action.DENY, + "log": Action.AUDIT, + "audit": Action.AUDIT, + "allow": Action.ALLOW, + "require_approval": Action.ESCALATE, + "escalate": Action.ESCALATE, +} + +_SEVERITY_MAP: dict[str, Severity] = { + "low": Severity.LOW, + "medium": Severity.MEDIUM, + "high": Severity.HIGH, + "critical": Severity.CRITICAL, +} + + +def build_policy_index_from_yaml(yaml_text: str) -> PolicyIndex: + """Parse YAML policy packs into a PolicyIndex. + + Args: + yaml_text: YAML body, either a single document or ``---``-separated + multi-document stream. Each document is one pack. + + Returns: + PolicyIndex with all successfully parsed packs added. Empty when + the input has no parseable packs. + + Raises: + yaml.YAMLError: If the YAML itself is malformed. Callers are + expected to fall back to the compiled index on this error. + """ + index = PolicyIndex() + documents = list(yaml.safe_load_all(yaml_text)) + + for doc in documents: + if not isinstance(doc, dict): + continue + pack = _build_pack(doc) + if pack is not None and pack.rules: + index.add_pack(pack) + + logger.debug( + "Built PolicyIndex from YAML: packs=%s, rules=%d", + index.pack_names, + index.total_rules, + ) + return index + + +def _build_pack(data: dict[str, Any]) -> PolicyPack | None: + """Build a PolicyPack from one YAML document.""" + name = data.get("standard") or data.get("name") + if not name: + logger.warning("Skipping pack: missing 'standard'/'name' field") + return None + + default_action_str = data.get("default_action", "block") + default_action = _ACTION_MAP.get(default_action_str, Action.DENY) + + rules: list[Rule] = [] + for i, rule_data in enumerate(data.get("rules", []) or []): + if not isinstance(rule_data, dict): + continue + rule = _build_rule(rule_data, default_action, i) + if rule is not None: + rules.append(rule) + + return PolicyPack( + name=str(name), + version=str(data.get("version", "1.0.0")), + description=str(data.get("description", "")), + rules=rules, + ) + + +def _build_rule( + data: dict[str, Any], default_action: Action, index: int +) -> Rule | None: + """Build a single Rule from a YAML rule entry.""" + hook = _HOOK_MAP.get(data.get("hook", "before_model")) + if hook is None: + logger.warning( + "Skipping rule %s: unknown hook %r", data.get("id"), data.get("hook") + ) + return None + + action_str = data.get("action") + action = ( + _ACTION_MAP.get(action_str, default_action) if action_str else default_action + ) + + default_sev = "high" if action == Action.DENY else "medium" + severity = _SEVERITY_MAP.get(data.get("severity", default_sev), Severity.HIGH) + + checks = _build_checks( + data.get("checks", []) or [], + action, + mapped_to_uipath=bool(data.get("mapped_to_uipath", False)), + policy_enabled=bool(data.get("policy_enabled", True)), + ) + + # If checks were declared but none could be parsed (e.g. all unknown + # types), skip the rule. A rule with zero checks "always matches" in + # the evaluator, so keeping it would make it fire on every request. + declared = data.get("checks", []) or [] + if declared and not checks: + logger.warning( + "Skipping rule %s: none of its %d declared check(s) could be parsed", + data.get("id"), + len(declared), + ) + return None + + return Rule( + rule_id=str(data.get("id", f"RULE-{index}")), + name=str(data.get("name", data.get("id", f"RULE-{index}"))), + clause=str(data.get("clause", data.get("owasp_ref", ""))), + hook=hook, + action=action, + severity=severity, + checks=checks, + enabled=bool(data.get("enabled", True)), + description=str(data.get("description", "")), + ) + + +def _build_checks( + checks_data: list[dict[str, Any]], + default_action: Action, + *, + mapped_to_uipath: bool = False, + policy_enabled: bool = True, +) -> list[Check]: + """Build the checks list for a rule. + + ``mapped_to_uipath`` / ``policy_enabled`` are rule-level flags read + by ``guardrail_fallback`` checks so the per-check condition can + decide whether to fire the compensating governance call. + """ + checks: list[Check] = [] + for check_data in checks_data: + if not isinstance(check_data, dict): + continue + check = _build_check( + check_data, + default_action, + mapped_to_uipath=mapped_to_uipath, + policy_enabled=policy_enabled, + ) + if check is not None: + checks.append(check) + return checks + + +def _build_check( + data: dict[str, Any], + default_action: Action, + *, + mapped_to_uipath: bool = False, + policy_enabled: bool = True, +) -> Check | None: + """Build one Check from a YAML check entry. + + Supports the same check types as ``compile_packs.py``: explicit + conditions, regex, budget, tool_allowlist, parameter_validation, + rate_limit, field_regex, sentiment_concern, data_quality_score, + incident_taxonomy, commitment_extractor, plus ``guardrail_fallback`` + (reads the rule-level ``mapped_to_uipath`` / ``policy_enabled`` flags + threaded in from ``_build_rule``). + """ + conditions: list[Condition] = [] + message = "" + + raw_conditions = data.get("conditions") + has_explicit_conditions = ( + isinstance(raw_conditions, list) + and raw_conditions + and isinstance(raw_conditions[0], dict) + and "operator" in raw_conditions[0] + ) + + check_type = data.get("type", "regex") + + if has_explicit_conditions: + assert isinstance(raw_conditions, list) # narrowed by has_explicit_conditions + conditions.extend(_make_conditions(raw_conditions)) + message = str(data.get("message", "")) + + elif check_type == "regex": + patterns = data.get("patterns", []) or [] + scope = data.get("scope", ["human", "ai"]) + field = _field_for_scope(scope) + for pattern in patterns: + conditions.append(Condition(operator="regex", field=field, value=pattern)) + message = f"Pattern matched in {scope}" + + elif check_type == "budget": + if "max_tool_calls_per_session" in data: + conditions.append( + Condition( + operator="gt", + field="session_state.tool_calls", + value=data["max_tool_calls_per_session"], + ) + ) + if "max_tool_calls_per_minute" in data: + conditions.append( + Condition( + operator="gt", + field="session_state.tool_calls_per_minute", + value=data["max_tool_calls_per_minute"], + ) + ) + if "max_consecutive_tool_calls" in data: + conditions.append( + Condition( + operator="gt", + field="session_state.consecutive_tool_calls", + value=data["max_consecutive_tool_calls"], + ) + ) + message = "Tool budget exceeded" + + elif check_type == "tool_allowlist": + blocked_tools = data.get("blocked_tools", []) or [] + if blocked_tools: + conditions.append( + Condition(operator="in_list", field="tool_name", value=blocked_tools) + ) + message = "Tool not allowed" + + elif check_type == "parameter_validation": + for pattern in data.get("additional_patterns", []) or []: + conditions.append( + Condition(operator="regex", field="tool_args", value=pattern) + ) + message = "Suspicious pattern in tool parameters" + + elif check_type == "rate_limit": + if "max_llm_calls_per_session" in data: + conditions.append( + Condition( + operator="gt", + field="session_state.llm_calls", + value=data["max_llm_calls_per_session"], + ) + ) + if "max_llm_calls_per_minute" in data: + conditions.append( + Condition( + operator="gt", + field="session_state.llm_calls_per_minute", + value=data["max_llm_calls_per_minute"], + ) + ) + message = "Rate limit exceeded" + + elif check_type == "field_regex": + conditions.extend(_make_conditions(data.get("conditions", []) or [])) + message = str(data.get("message", "Field regex check failed")) + + elif check_type == "data_quality_score": + field = data.get("field", "tool_result") + if data.get("check_encoding", True): + conditions.append( + Condition( + operator="encoding_concern", + field=field, + value={ + "min_confidence": float(data.get("min_confidence", 0.5)), + "max_replacement_ratio": float( + data.get("max_replacement_ratio", 0.05) + ), + "min_corruption_events": int( + data.get("min_corruption_events", 2) + ), + }, + ) + ) + if data.get("check_entropy", True): + conditions.append( + Condition( + operator="entropy_concern", + field=field, + value={ + "min": float(data.get("entropy_min", 1.5)), + "max": float(data.get("entropy_max", 7.5)), + }, + ) + ) + message = str( + data.get("message", "A.7.4: Data quality signal (encoding or entropy)") + ) + + elif check_type == "incident_taxonomy": + field = data.get("field", "model_output") + categories = data.get("categories") + value: dict[str, Any] = {} + if categories: + value["categories"] = list(categories) + conditions.append( + Condition(operator="incident_concern", field=field, value=value) + ) + message = str(data.get("message", "A.8.4: Incident signal detected")) + + elif check_type == "commitment_extractor": + field = data.get("field", "model_output") + conditions.append( + Condition( + operator="commitment_concern", + field=field, + value={ + "require_amount": bool(data.get("require_amount", True)), + "require_deadline": bool(data.get("require_deadline", False)), + }, + ) + ) + message = str( + data.get("message", "A.10.4: Customer commitment language detected") + ) + + elif check_type == "sentiment_concern": + field = data.get("field", "model_input") + threshold = float(data.get("threshold", -0.3)) + conditions.append( + Condition( + operator="vader_concern", + field=field, + value={"threshold": threshold}, + ) + ) + message = str( + data.get( + "message", + f"Negative sentiment detected (VADER compound <= {threshold})", + ) + ) + + elif check_type == "guardrail_fallback": + # Centralized guardrail compensating control. The on/off state + # lives at the RULE level (mapped_to_uipath / policy_enabled), + # threaded in from ``_build_rule``; ``validator`` names which + # guardrail check the server should run on behalf of the agent. + # The condition matches only when the guardrail is mapped to + # UiPath but disabled — see the ``guardrail_fallback`` operator + # in :class:`GovernanceEvaluator`. + conditions.append( + Condition( + operator="guardrail_fallback", + field="", + value={ + "validator": str(data.get("validator", "")), + "mapped_to_uipath": mapped_to_uipath, + "policy_enabled": policy_enabled, + }, + ) + ) + message = str( + data.get("message", "Guardrail disabled — compensating check needed.") + ) + + else: + logger.debug("Skipping check: unknown type %r", check_type) + return None + + if not conditions: + return None + + action_str = data.get("action") + action = ( + _ACTION_MAP.get(action_str, default_action) if action_str else default_action + ) + + message = str(data.get("message", message)) + + # Multi-pattern regex/parameter_validation defaults to OR semantics + # (any pattern indicates a hit); explicit `logic` in YAML wins. + if check_type in ("parameter_validation", "regex") and len(conditions) > 1: + default_logic = "any" + else: + default_logic = "all" + logic = str(data.get("logic", default_logic)) + + return Check(conditions=conditions, action=action, message=message, logic=logic) + + +def _make_conditions(raw: list[dict[str, Any]]) -> list[Condition]: + """Translate a list of YAML condition dicts into Condition objects.""" + out: list[Condition] = [] + for cond in raw: + if not isinstance(cond, dict): + continue + out.append( + Condition( + operator=str(cond.get("operator", "regex")), + field=str(cond.get("field", "model_input")), + value=cond.get("value", ""), + negate=bool(cond.get("negate", False)), + ) + ) + return out + + +def _field_for_scope(scope: list[str] | str) -> str: + """Map a YAML `scope` value to the CheckContext field it targets.""" + if isinstance(scope, str): + scope = [scope] + if "system" in scope or "human" in scope: + return "model_input" + if "ai" in scope: + return "model_output" + if "tool_result" in scope: + return "tool_result" + return "model_input" 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/loader.py b/src/uipath/runtime/governance/native/loader.py new file mode 100644 index 0000000..4d8271f --- /dev/null +++ b/src/uipath/runtime/governance/native/loader.py @@ -0,0 +1,284 @@ +"""Policy pack loader. + +Resolves the active PolicyIndex at startup. Policies are fetched +exclusively from the governance backend (``api/v1/policy``); there is +no local compiled fallback. When the backend is unavailable, the +access token is unset, or the fetch times out, the loader returns an +empty PolicyIndex and the agent runs without any rules. +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +from collections import Counter + +import yaml +from uipath.core.governance import EnforcementMode + +from uipath.runtime.governance.config import set_enforcement_mode +from uipath.runtime.governance.native._yaml_to_index import build_policy_index_from_yaml +from uipath.runtime.governance.native.backend_client import ENV_ACCESS_TOKEN +from uipath.runtime.governance.native.models import PolicyIndex +from uipath.runtime.governance.native.policy_api_client import ( + ENV_ORGANIZATION_ID, + ENV_TENANT_ID, + POLICY_API_TIMEOUT_SECONDS, + fetch_policy_response, + resolve_organization_id, + resolve_tenant_id, +) + +logger = logging.getLogger(__name__) + +# Pack name aliases for backward compatibility +PACK_ALIASES: dict[str, str] = { + "owasp": "owasp_agentic", + "hipaa": "hipaa_runtime", + "soc2": "soc2_runtime", + "nist": "nist_ai_rmf_runtime", + "eu_ai": "eu_ai_act_runtime", + "iso": "iso42001_runtime", +} + + +# Module-level cache +_policy_index: PolicyIndex | None = None + +# Background-prefetch coordination. ``_prefetch_event`` is set once the +# background load_policy_index() call finishes (success OR failure); +# callers of ``get_policy_index()`` wait on it. ``_prefetch_lock`` +# protects the start-once semantics so concurrent ``prefetch`` calls +# don't kick off duplicate threads. +_prefetch_event: threading.Event | None = None +_prefetch_lock = threading.Lock() + +# Default wait when ``get_policy_index()`` blocks on an in-flight +# prefetch. Matched to the policy-API HTTP timeout so a stuck backend +# bounds the total time spent waiting at first hook fire to +# ~POLICY_API_TIMEOUT_SECONDS. If the wait expires we return an empty +# PolicyIndex — the agent runs without any policies rather than +# blocking further or retrying. +_PREFETCH_WAIT_SECONDS = POLICY_API_TIMEOUT_SECONDS + + +def prefetch_policy_index() -> None: + """Kick off a background load of the policy index. + + Non-blocking. Designed to be called as early as possible (at + ``GovernanceRuntime.__init__``) so the HTTP call to the governance + backend overlaps with the rest of agent setup. Idempotent. + """ + global _prefetch_event + + with _prefetch_lock: + if _policy_index is not None: + return # already loaded + if _prefetch_event is not None: + return # already in flight + event = threading.Event() + _prefetch_event = event + + def _worker() -> None: + global _policy_index + try: + loaded = load_policy_index() + except Exception as exc: # noqa: BLE001 + logger.warning("Policy prefetch failed: %s", exc) + else: + with _prefetch_lock: + _policy_index = loaded + finally: + event.set() + + threading.Thread( + target=_worker, + name="governance-policy-prefetch", + daemon=True, + ).start() + + +def get_policy_index() -> PolicyIndex: + """Get the cached policy index, loading if necessary. + + Resolution order on first call: + 1. If a prefetch (see :func:`prefetch_policy_index`) is in flight, + wait for it to complete (bounded by ``_PREFETCH_WAIT_SECONDS``). + 2. Governance backend at ``api/v1/policy`` (one HTTP GET, cached). + 3. Empty PolicyIndex when the backend is unavailable or times out. + + Result is cached for the process lifetime; per-hook evaluation never + touches the network. Call :func:`clear_policy_cache` to force a + refetch (mainly for tests). + """ + global _policy_index + + if _policy_index is not None: + return _policy_index + + event = _prefetch_event + if event is not None: + completed = event.wait(timeout=_PREFETCH_WAIT_SECONDS) + if completed and _policy_index is not None: + return _policy_index + if not completed: + logger.warning( + "Policy prefetch did not complete in %.1fs; " + "agent will run without any policies", + _PREFETCH_WAIT_SECONDS, + ) + else: + logger.warning( + "Policy prefetch completed but produced no PolicyIndex " + "(see prior WARN for the root cause); agent will run " + "without any policies" + ) + _policy_index = PolicyIndex() + return _policy_index + + # No prefetch was started (direct callers / tests). Sync load. + _policy_index = load_policy_index() + return _policy_index + + +def load_policy_index(pack_name: str | None = None) -> PolicyIndex: + """Load the active PolicyIndex from the governance backend. + + Args: + pack_name: Ignored. Pack selection is controlled entirely by the + backend. + + Returns: + PolicyIndex parsed from the backend response. Empty PolicyIndex + when the backend is unavailable, the token is unset, the YAML + is malformed, or the response yields zero rules. + """ + start = time.perf_counter() + + api_index = _load_from_api() + if api_index is not None: + _log_index_summary(api_index) + logger.info( + "Policy index ready: source=backend, total_ms=%.1f", + (time.perf_counter() - start) * 1000, + ) + return api_index + + reason = _empty_index_reason() + logger.info( + "Policy index ready: source=empty (%s), total_ms=%.1f", + reason, + (time.perf_counter() - start) * 1000, + ) + return PolicyIndex() + + +def _empty_index_reason() -> str: + """Diagnose why the policy fetch produced nothing.""" + if not resolve_organization_id(): + return ( + f"UiPathConfig.organization_id unavailable — set {ENV_ORGANIZATION_ID} " + "or install uipath-platform; backend API not contacted" + ) + if not resolve_tenant_id(): + return ( + f"UiPathConfig.tenant_id unavailable — set {ENV_TENANT_ID} " + "or install uipath-platform; backend API not contacted" + ) + if not os.environ.get(ENV_ACCESS_TOKEN): + return f"{ENV_ACCESS_TOKEN} unset — backend API not contacted" + return "backend returned no policies (timeout / error / empty body)" + + +def _apply_enforcement_mode(mode_str: str | None) -> None: + """Map a backend-supplied mode string onto :class:`EnforcementMode`.""" + if not mode_str: + return + try: + mode = EnforcementMode(mode_str.lower()) + except ValueError: + logger.warning( + "Backend returned unknown enforcement mode %r; keeping current mode", + mode_str, + ) + return + set_enforcement_mode(mode) + logger.info("Enforcement mode set from backend: %s", mode.value) + + +def _load_from_api() -> PolicyIndex | None: + """Fetch and parse the policy index from the governance backend.""" + start = time.perf_counter() + response = fetch_policy_response() + if response is None: + return None + + _apply_enforcement_mode(response.mode) + + if not response.policy: + logger.warning( + "Policy fetch returned empty policy field; " + "agent will run without any policies" + ) + return None + + try: + index = build_policy_index_from_yaml(response.policy) + except yaml.YAMLError as exc: + logger.warning("Policy YAML from backend was malformed: %s", exc) + return None + except Exception as exc: # noqa: BLE001 + logger.warning("Failed to build PolicyIndex from backend YAML: %s", exc) + return None + + if index.total_rules == 0: + logger.warning( + "Policy YAML from backend yielded zero rules; " + "agent will run without any policies" + ) + return None + + elapsed_ms = (time.perf_counter() - start) * 1000 + logger.info( + "Loaded policy index from backend: packs=%s, rules=%d, elapsed_ms=%.1f", + index.pack_names, + index.total_rules, + elapsed_ms, + ) + return index + + +def _log_index_summary(index: PolicyIndex) -> None: + """Log summary of loaded policy index.""" + hook_counts: Counter[str] = Counter() + for rule in index.all_rules: + hook_counts[rule.hook.value] += 1 + + logger.debug( + "Policy packs: %s, total rules: %d, by hook: %s", + index.pack_names, + index.total_rules, + dict(hook_counts), + ) + + +def get_available_packs() -> list[str]: + """Get list of pack names from the currently loaded policy index.""" + if _policy_index is None: + return [] + return _policy_index.pack_names + + +def clear_policy_cache() -> None: + """Clear the cached policy index and any in-flight prefetch state.""" + global _policy_index, _prefetch_event + with _prefetch_lock: + _policy_index = None + _prefetch_event = None + logger.debug("Policy index cache cleared") + + +# Backward compatibility alias +reset_policy_index = clear_policy_cache diff --git a/src/uipath/runtime/governance/native/policy_api_client.py b/src/uipath/runtime/governance/native/policy_api_client.py new file mode 100644 index 0000000..325b4e0 --- /dev/null +++ b/src/uipath/runtime/governance/native/policy_api_client.py @@ -0,0 +1,227 @@ +"""Governance policy API client. + +Fetches the governance backend response so policies can be controlled +centrally without redeploying agents. Called once at process startup +from :mod:`uipath.runtime.governance.native.loader`; per-hook evaluation +stays in-process. + +Response shape (JSON):: + + { + "mode": "audit" | "enforce" | "disabled", + "policies": "" + } + +``mode`` is the platform-controlled enforcement mode for the tenant; +the loader applies it via +:func:`uipath.runtime.governance.config.set_enforcement_mode`. ``policies`` +is the YAML the evaluator compiles into a :class:`PolicyIndex`. + +Failure mode is fail-open: when the organization id is unknown, the +access token is missing, the backend errors, or the body can't be +parsed, the caller falls back to an empty PolicyIndex. The fetch is +single-shot (no retry by design — see :func:`_get_once`) so a slow +backend can't extend agent startup beyond +:data:`BACKEND_REQUEST_TIMEOUT_SECONDS`. Nothing in this module ever +raises to the caller. +""" + +from __future__ import annotations + +import json +import logging +import os +import urllib.error +import urllib.request +from dataclasses import dataclass +from urllib.parse import urlencode + +from uipath.runtime.governance.native.backend_client import ( + AGENT_TYPE_PARAM, + BACKEND_REQUEST_TIMEOUT_SECONDS, + ENV_ACCESS_TOKEN, + ENV_ORGANIZATION_ID, + ENV_TENANT_ID, + POLICY_API_PATH, + TENANT_HEADER, + agent_type_param, + build_governance_url, + governance_request_headers, + resolve_organization_id, + resolve_tenant_id, +) + +logger = logging.getLogger(__name__) + +# Re-exported alias kept for callers that imported the old name. +POLICY_API_TIMEOUT_SECONDS = BACKEND_REQUEST_TIMEOUT_SECONDS + + +@dataclass(frozen=True) +class PolicyResponse: + """Parsed governance backend response. + + Attributes: + mode: Enforcement mode string the backend returned + (``"audit"`` / ``"enforce"`` / ``"disabled"``), or ``None`` + when the backend omitted it. Loader applies this via + :func:`uipath.runtime.governance.config.set_enforcement_mode`. + policy: Policy pack YAML to compile into a ``PolicyIndex``. May + be an empty string if the backend returned no rules. + """ + + mode: str | None + policy: str + + +def build_policy_url(org_id: str) -> str: + """Build the policy endpoint URL for the given organization id. + + The tenant id is not part of the URL; it travels in the + ``x-uipath-internal-tenantid`` request header (see + :func:`fetch_policy_response`). + + When the hosted agent's type is known (see + :func:`uipath.runtime.governance.native.backend_client.set_agent_conversational`), + an ``agentType`` query param is appended so the server resolves the + conversational-vs-autonomous container key. Omitted when unknown — the + server then applies its default. + """ + url = build_governance_url(org_id, POLICY_API_PATH) + agent_type = agent_type_param() + if agent_type: + url = f"{url}?{urlencode({AGENT_TYPE_PARAM: agent_type})}" + return url + + +def fetch_policy_response() -> PolicyResponse | None: + """Fetch the governance backend's policy response. + + Single shot, no retry: a failed fetch (timeout / network error / + HTTP error / malformed body) returns ``None`` and the caller falls + back to an empty PolicyIndex. The agent must not spend time on a + second attempt — keeping governance off the critical path is more + important than maximising policy availability. + + Returns: + :class:`PolicyResponse` on success. ``None`` on any failure + path — caller falls back to an empty PolicyIndex. + + Never raises. + """ + try: + return _fetch_policy_response_inner() + except Exception as exc: # noqa: BLE001 - loader path must never raise + logger.warning("Policy fetch failed unexpectedly: %s", exc) + return None + + +def _fetch_policy_response_inner() -> PolicyResponse | None: + org_id = resolve_organization_id() + if not org_id: + logger.warning( + "Policy fetch skipped: UiPathConfig.organization_id is not " + "available (set %s in the environment, or ensure uipath-platform " + "is installed); governance will run with no policies. The " + "backend API was NOT contacted.", + ENV_ORGANIZATION_ID, + ) + return None + + tenant_id = resolve_tenant_id() + if not tenant_id: + logger.warning( + "Policy fetch skipped: UiPathConfig.tenant_id is not " + "available (set %s in the environment, or ensure uipath-platform " + "is installed); governance will run with no policies. The " + "backend API was NOT contacted.", + ENV_TENANT_ID, + ) + return None + + policy_url = build_policy_url(org_id) + + token = os.environ.get(ENV_ACCESS_TOKEN) + if not token: + logger.warning( + "Policy fetch skipped: %s is not set in the environment; " + "governance will run with no policies.", + ENV_ACCESS_TOKEN, + ) + return None + + # Policy fetch is a GET; ``json_body=False`` so ``Content-Type`` is + # omitted. Strict origin servers may 415 on unexpected Content-Type + # for GETs (see :func:`governance_request_headers` docstring). + headers = governance_request_headers(json_body=False) + headers[TENANT_HEADER] = tenant_id + logger.info("Policy fetch starting (org=%s, tenant=%s)", org_id, tenant_id) + + body = _get_once(policy_url, headers) + if body is None: + return None + return _parse_policy_body(body) + + +def _get_once(url: str, headers: dict[str, str]) -> bytes | None: + """GET ``url`` once. Returns body bytes, or ``None`` on any failure. + + No retry by design — see :func:`fetch_policy_response` for the + rationale. Every failure path logs a single WARNING and returns + ``None`` so the caller (the loader) falls back to an empty + PolicyIndex without delay. + """ + request = urllib.request.Request(url, headers=headers, method="GET") + try: + with urllib.request.urlopen( # noqa: S310 - URL is built from config + request, timeout=BACKEND_REQUEST_TIMEOUT_SECONDS + ) as response: + return response.read() + except urllib.error.HTTPError as exc: + logger.warning("Policy fetch returned HTTP %d: %s", exc.code, exc) + except (urllib.error.URLError, TimeoutError, OSError) as exc: + logger.warning("Policy fetch failed: %s", exc) + return None + + +def _parse_policy_body(body: bytes) -> PolicyResponse | None: + """Parse the JSON envelope into a :class:`PolicyResponse`.""" + if not body: + logger.warning("Policy fetch returned empty body") + return None + + try: + payload = json.loads(body.decode("utf-8")) + except UnicodeDecodeError as exc: + logger.warning("Policy fetch returned non-UTF8 body: %s", exc) + return None + except json.JSONDecodeError as exc: + logger.warning( + "Policy fetch returned malformed JSON " + "(server may have returned an HTML error page): %s", + exc, + ) + return None + + if not isinstance(payload, dict): + logger.warning( + "Policy fetch returned unexpected JSON shape (expected object, got %s)", + type(payload).__name__, + ) + return None + + raw_mode = payload.get("mode") + mode = raw_mode if isinstance(raw_mode, str) and raw_mode else None + + raw_policy = payload.get("policies", "") + if not isinstance(raw_policy, str): + logger.warning( + "Policy fetch returned non-string 'policies' field (got %s)", + type(raw_policy).__name__, + ) + return None + + logger.info( + "Policy fetch ok: mode=%s, policy_bytes=%d", mode, len(raw_policy) + ) + return PolicyResponse(mode=mode, policy=raw_policy) diff --git a/src/uipath/runtime/governance/rego/__init__.py b/src/uipath/runtime/governance/rego/__init__.py new file mode 100644 index 0000000..3d0b6a1 --- /dev/null +++ b/src/uipath/runtime/governance/rego/__init__.py @@ -0,0 +1,12 @@ +"""Rego/WASM-based governance evaluator.""" +from __future__ import annotations + +from .evaluator import RegoEvaluator +from .loader import clear_rego_cache, get_rego_evaluator, prefetch_rego_bundles + +__all__ = [ + "RegoEvaluator", + "clear_rego_cache", + "get_rego_evaluator", + "prefetch_rego_bundles", +] diff --git a/src/uipath/runtime/governance/rego/api_client.py b/src/uipath/runtime/governance/rego/api_client.py new file mode 100644 index 0000000..0b82176 --- /dev/null +++ b/src/uipath/runtime/governance/rego/api_client.py @@ -0,0 +1,145 @@ +"""Governance Rego bundle API client. + +Fetches per-hook WASM bundle metadata from ``GET /all-policies/:tenantId`` +and downloads individual bundles from their CDN URLs. + +Failure mode is fail-open throughout: every public function returns +``None`` / empty on any error so callers fall back to stale cache or +run without Rego evaluation. +""" +from __future__ import annotations + +import json +import logging +import os +import urllib.error +import urllib.request + +from uipath.runtime.governance.native.backend_client import ( + ALL_POLICIES_API_PATH, + BACKEND_REQUEST_TIMEOUT_SECONDS, + ENV_ACCESS_TOKEN, + TENANT_HEADER, + build_governance_url, + governance_request_headers, + resolve_organization_id, + resolve_tenant_id, +) +from uipath.runtime.governance.rego.models import AllPoliciesResponse, HookBundle + +logger = logging.getLogger(__name__) + + +def build_all_policies_url(org_id: str, tenant_id: str) -> str: + """Build the ``/all-policies/:tenantId`` endpoint URL.""" + return build_governance_url(org_id, f"{ALL_POLICIES_API_PATH}/{tenant_id}") + + +def fetch_all_policies() -> AllPoliciesResponse | None: + """Fetch per-hook bundle metadata from the governance backend. + + Returns AllPoliciesResponse on success, None on any failure. Never raises. + """ + try: + return _fetch_all_policies_inner() + except Exception as exc: # noqa: BLE001 + logger.warning("fetch_all_policies failed unexpectedly: %s", exc) + return None + + +def _fetch_all_policies_inner() -> AllPoliciesResponse | None: + org_id = resolve_organization_id() + if not org_id: + logger.warning( + "Rego bundle fetch skipped: organization_id unavailable; " + "agent will run without custom Rego policies." + ) + return None + + tenant_id = resolve_tenant_id() + if not tenant_id: + logger.warning( + "Rego bundle fetch skipped: tenant_id unavailable; " + "agent will run without custom Rego policies." + ) + return None + + token = os.environ.get(ENV_ACCESS_TOKEN) + if not token: + logger.warning( + "Rego bundle fetch skipped: %s not set; " + "agent will run without custom Rego policies.", + ENV_ACCESS_TOKEN, + ) + return None + + url = build_all_policies_url(org_id, tenant_id) + headers = governance_request_headers(json_body=False) + headers[TENANT_HEADER] = tenant_id + + body = _get_once(url, headers) + if body is None: + return None + return _parse_response(body) + + +def _get_once(url: str, headers: dict[str, str]) -> bytes | None: + request = urllib.request.Request(url, headers=headers, method="GET") + try: + with urllib.request.urlopen( # noqa: S310 + request, timeout=BACKEND_REQUEST_TIMEOUT_SECONDS + ) as response: + return response.read() + except urllib.error.HTTPError as exc: + logger.warning("Rego bundle fetch returned HTTP %d: %s", exc.code, exc) + except (urllib.error.URLError, TimeoutError, OSError) as exc: + logger.warning("Rego bundle fetch failed: %s", exc) + return None + + +def _parse_response(body: bytes) -> AllPoliciesResponse | None: + if not body: + logger.warning("Rego bundle fetch: empty response body") + return None + try: + payload = json.loads(body.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError) as exc: + logger.warning("Rego bundle fetch: malformed response: %s", exc) + return None + + if not isinstance(payload, dict): + logger.warning("Rego bundle fetch: unexpected JSON shape") + return None + + raw_bundles = payload.get("hookBundles", []) + if not isinstance(raw_bundles, list): + logger.warning("Rego bundle fetch: hookBundles is not a list") + return None + + bundles: list[HookBundle] = [] + for item in raw_bundles: + if not isinstance(item, dict): + continue + hook_type = item.get("hookType", "") + url_str = item.get("bundleUrl", "") + etag = item.get("etag", "") + if hook_type and url_str: + bundles.append(HookBundle(hook_type=hook_type, bundle_url=url_str, etag=etag)) + + return AllPoliciesResponse(hook_bundles=bundles) + + +def download_bundle(url: str) -> bytes | None: + """Download a WASM bundle from a pre-signed CDN URL. + + No auth header is sent — the URL is pre-signed. Returns raw bytes or None. Never raises. + """ + try: + request = urllib.request.Request(url, method="GET") + with urllib.request.urlopen( # noqa: S310 + request, timeout=BACKEND_REQUEST_TIMEOUT_SECONDS + ) as response: + return response.read() + except Exception as exc: # noqa: BLE001 + logger.warning("Bundle download from %s failed: %s", url, exc) + return None 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..e9925bc --- /dev/null +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -0,0 +1,448 @@ +"""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 {} + return { + "hook": context.hook.value, + "agent_input": context.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": context.messages, + "features": features, + } + + +def _extract_wasm_from_bundle(bundle_path: Path) -> bytes: + """Extract ``policy.wasm`` bytes from an OPA ``.tar.gz`` bundle.""" + with open(bundle_path, "rb") as f: + bundle_bytes = f.read() + with tarfile.open(fileobj=io.BytesIO(bundle_bytes), mode="r:gz") as tf: + try: + member = tf.getmember("/policy.wasm") + except KeyError: + member = tf.getmember("policy.wasm") + fobj = tf.extractfile(member) + if fobj is None: + raise ValueError(f"policy.wasm is not a regular file in {bundle_path}") + return fobj.read() + + +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]] = {} + 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, + trace_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, + trace_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, []) + raw = engine.evaluate(context_to_input(context, feature_plan=plan)) + 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, + trace_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, + trace_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..e74e5fc --- /dev/null +++ b/src/uipath/runtime/governance/rego/loader.py @@ -0,0 +1,235 @@ +"""Rego bundle loader: startup prefetch, ETag-driven download, background refresh.""" +from __future__ import annotations + +import logging +import os +import threading +import time + +from uipath.core.governance.models import LifecycleHook + +from uipath.runtime.governance.native.backend_client import ( + BACKEND_REQUEST_TIMEOUT_SECONDS, + resolve_organization_id, + resolve_tenant_id, +) +from uipath.runtime.governance.rego.api_client import ( + download_bundle, + fetch_all_policies, +) +from uipath.runtime.governance.rego.bundle_cache import ( + get_cache_dir, + get_cached_bundle_path, + get_cached_etag, + save_bundle, +) + +logger = logging.getLogger(__name__) + +# Module-level state +_rego_evaluator = None # RegoEvaluator | None +_rego_prefetch_event: threading.Event | None = None +_rego_prefetch_lock = threading.Lock() +_rego_refresh_started = False +_rego_refresh_lock = threading.Lock() + +_DEFAULT_REFRESH_SECONDS = 30 + +# 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, +} + + +def prefetch_rego_bundles() -> None: + """Kick off a background download of changed WASM bundles. Non-blocking. Idempotent.""" + global _rego_prefetch_event + + with _rego_prefetch_lock: + if _rego_evaluator is not None: + return + if _rego_prefetch_event is not None: + return + event = threading.Event() + _rego_prefetch_event = event + + def _worker() -> None: + try: + _sync_bundles_to_disk() + except Exception as exc: # noqa: BLE001 + logger.warning("Rego bundle prefetch failed: %s", exc) + finally: + event.set() + + threading.Thread( + target=_worker, + name="governance-rego-prefetch", + daemon=True, + ).start() + + _start_background_refresh() + + +def _sync_bundles_to_disk() -> None: + """Fetch /all-policies and download any bundle whose ETag has changed.""" + org_id = resolve_organization_id() + tenant_id = resolve_tenant_id() + if not org_id or not tenant_id: + logger.warning( + "Rego bundle sync skipped: org_id or tenant_id unavailable; " + "agent will run without custom Rego policies." + ) + return + + cache_dir = get_cache_dir(org_id, tenant_id) + response = fetch_all_policies() + if response is None: + logger.warning( + "Rego bundle sync: /all-policies fetch returned nothing; " + "using cached bundles if available." + ) + return + + for bundle in response.hook_bundles: + hook_type = bundle.hook_type + cached_etag = get_cached_etag(cache_dir, hook_type) + + if cached_etag == bundle.etag: + logger.debug("Rego bundle for hook=%s is up to date (etag=%s)", hook_type, bundle.etag) + continue + + logger.info( + "Rego bundle for hook=%s changed (cached=%s new=%s); downloading", + hook_type, cached_etag, bundle.etag, + ) + data = download_bundle(bundle.bundle_url) + if data is None: + stale = get_cached_bundle_path(cache_dir, hook_type) + if stale: + logger.warning( + "Rego bundle download failed for hook=%s; using stale cache.", hook_type + ) + else: + logger.warning( + "Rego bundle download failed for hook=%s and no cache exists; " + "hook will pass through without Rego evaluation.", + hook_type, + ) + continue + + save_bundle(cache_dir, hook_type, data, bundle.etag) + logger.info("Rego bundle saved for hook=%s (etag=%s)", hook_type, bundle.etag) + # Invalidate the in-memory evaluator so the next run picks up the new bundle. + global _rego_evaluator + _rego_evaluator = None + + +def get_rego_evaluator(): + """Return the cached RegoEvaluator, building it from disk if needed. + + Waits for the startup prefetch (bounded by BACKEND_REQUEST_TIMEOUT_SECONDS). + Returns None when no bundles are available (fail-open). + """ + global _rego_evaluator + + if _rego_evaluator is not None: + return _rego_evaluator + + event = _rego_prefetch_event + if event is not None: + completed = event.wait(timeout=BACKEND_REQUEST_TIMEOUT_SECONDS) + if not completed: + logger.warning( + "Rego bundle prefetch timed out after %.1fs; " + "using whatever is on disk (may be empty).", + BACKEND_REQUEST_TIMEOUT_SECONDS, + ) + + return _build_evaluator_from_disk() + + +def _build_evaluator_from_disk(): + """Read the disk cache and construct a RegoEvaluator. Returns None if no bundles exist.""" + from pathlib import Path + + 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: + return None + + cache_dir = get_cache_dir(org_id, tenant_id) + hook_wasm_paths: dict[LifecycleHook, Path] = {} + hook_data: dict[LifecycleHook, dict] = {} + + for hook_type_str, lifecycle_hook in _HOOK_MAP.items(): + path = get_cached_bundle_path(cache_dir, hook_type_str) + if path is not None: + hook_wasm_paths[lifecycle_hook] = path + data = _extract_data_json_from_bundle(path) + if data is not None: + hook_data[lifecycle_hook] = data + + if not hook_wasm_paths: + logger.warning( + "No Rego WASM bundles found on disk; " + "agent will run without custom Rego evaluation." + ) + return None + + logger.info( + "Building RegoEvaluator from disk: hooks=%s", + [h.value for h in hook_wasm_paths], + ) + return RegoEvaluator(hook_wasm_paths, hook_data=hook_data or None) + + +def _start_background_refresh() -> None: + """Start the background refresh daemon (once per process).""" + global _rego_refresh_started + + refresh_seconds = int( + os.environ.get("UIPATH_GOVERNANCE_BUNDLE_REFRESH_SECONDS", _DEFAULT_REFRESH_SECONDS) + ) + if refresh_seconds <= 0: + return + + with _rego_refresh_lock: + if _rego_refresh_started: + return + _rego_refresh_started = True + + def _refresh_loop() -> None: + while True: + time.sleep(refresh_seconds) + try: + _sync_bundles_to_disk() + except Exception as exc: # noqa: BLE001 + logger.warning("Rego bundle background refresh failed: %s", exc) + + threading.Thread( + target=_refresh_loop, + name="governance-rego-refresh", + daemon=True, + ).start() + logger.debug("Rego bundle background refresh started (interval=%ds)", refresh_seconds) + + +def clear_rego_cache() -> None: + """Reset all module-level state. Intended for tests.""" + global _rego_evaluator, _rego_prefetch_event, _rego_refresh_started + with _rego_prefetch_lock: + _rego_evaluator = None + _rego_prefetch_event = None + with _rego_refresh_lock: + _rego_refresh_started = False + logger.debug("Rego loader cache cleared") diff --git a/src/uipath/runtime/governance/rego/models.py b/src/uipath/runtime/governance/rego/models.py new file mode 100644 index 0000000..5f3be6c --- /dev/null +++ b/src/uipath/runtime/governance/rego/models.py @@ -0,0 +1,18 @@ +"""Data models for the governance Rego/WASM bundle API response.""" +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class HookBundle: + """Metadata for a single per-hook WASM bundle returned by /all-policies.""" + hook_type: str + bundle_url: str + etag: str + + +@dataclass(frozen=True) +class AllPoliciesResponse: + """Parsed response from GET /all-policies/:tenantId.""" + hook_bundles: list[HookBundle] diff --git a/src/uipath/runtime/governance/runtime.py b/src/uipath/runtime/governance/runtime.py index c9cf699..aaf6c78 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,84 @@ 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 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 +334,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 +365,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: From 1917086fed5160427586582508899b96d51d28ef Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Thu, 9 Jul 2026 15:27:40 +0530 Subject: [PATCH 02/12] fix: remove native policy loading files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The /runtime/policy fetch (YAML policy loading, PolicyIndex construction) lives in uipath-python under platform, not here. Removes native/loader.py, native/policy_api_client.py, and native/_yaml_to_index.py that were incorrectly ported. native/backend_client.py stays — rego/api_client.py uses it for URL building and headers when fetching WASM bundles from /all-policies. Generated with Claude Code Co-Authored-By: Claude --- .../governance/native/_yaml_to_index.py | 459 ------------------ .../runtime/governance/native/loader.py | 284 ----------- .../governance/native/policy_api_client.py | 227 --------- 3 files changed, 970 deletions(-) delete mode 100644 src/uipath/runtime/governance/native/_yaml_to_index.py delete mode 100644 src/uipath/runtime/governance/native/loader.py delete mode 100644 src/uipath/runtime/governance/native/policy_api_client.py diff --git a/src/uipath/runtime/governance/native/_yaml_to_index.py b/src/uipath/runtime/governance/native/_yaml_to_index.py deleted file mode 100644 index 2deb463..0000000 --- a/src/uipath/runtime/governance/native/_yaml_to_index.py +++ /dev/null @@ -1,459 +0,0 @@ -"""Runtime YAML → PolicyIndex parser. - -Mirrors the shape produced by ``packs/compile_packs.py`` but builds the -PolicyIndex directly from parsed YAML data rather than generating Python -source. Used by :mod:`uipath.runtime.governance.native.loader` when policies are fetched -from the governance backend at startup. - -Accepts either a single YAML document (one pack) or a multi-document -stream (``---``-separated packs). Unknown check types and malformed -rules are skipped with a warning — partial packs are preferred over -failing the whole load. -""" - -from __future__ import annotations - -import logging -from typing import Any - -import yaml -from uipath.core.governance.models import Action, LifecycleHook - -from uipath.runtime.governance.native.models import ( - Check, - Condition, - PolicyIndex, - PolicyPack, - Rule, - Severity, -) - -logger = logging.getLogger(__name__) - - -_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, - "wrap_tool_call": LifecycleHook.TOOL_CALL, - "tool_call": LifecycleHook.TOOL_CALL, - "after_tool": LifecycleHook.AFTER_TOOL, -} - -_ACTION_MAP: dict[str, Action] = { - "block": Action.DENY, - "deny": Action.DENY, - "log": Action.AUDIT, - "audit": Action.AUDIT, - "allow": Action.ALLOW, - "require_approval": Action.ESCALATE, - "escalate": Action.ESCALATE, -} - -_SEVERITY_MAP: dict[str, Severity] = { - "low": Severity.LOW, - "medium": Severity.MEDIUM, - "high": Severity.HIGH, - "critical": Severity.CRITICAL, -} - - -def build_policy_index_from_yaml(yaml_text: str) -> PolicyIndex: - """Parse YAML policy packs into a PolicyIndex. - - Args: - yaml_text: YAML body, either a single document or ``---``-separated - multi-document stream. Each document is one pack. - - Returns: - PolicyIndex with all successfully parsed packs added. Empty when - the input has no parseable packs. - - Raises: - yaml.YAMLError: If the YAML itself is malformed. Callers are - expected to fall back to the compiled index on this error. - """ - index = PolicyIndex() - documents = list(yaml.safe_load_all(yaml_text)) - - for doc in documents: - if not isinstance(doc, dict): - continue - pack = _build_pack(doc) - if pack is not None and pack.rules: - index.add_pack(pack) - - logger.debug( - "Built PolicyIndex from YAML: packs=%s, rules=%d", - index.pack_names, - index.total_rules, - ) - return index - - -def _build_pack(data: dict[str, Any]) -> PolicyPack | None: - """Build a PolicyPack from one YAML document.""" - name = data.get("standard") or data.get("name") - if not name: - logger.warning("Skipping pack: missing 'standard'/'name' field") - return None - - default_action_str = data.get("default_action", "block") - default_action = _ACTION_MAP.get(default_action_str, Action.DENY) - - rules: list[Rule] = [] - for i, rule_data in enumerate(data.get("rules", []) or []): - if not isinstance(rule_data, dict): - continue - rule = _build_rule(rule_data, default_action, i) - if rule is not None: - rules.append(rule) - - return PolicyPack( - name=str(name), - version=str(data.get("version", "1.0.0")), - description=str(data.get("description", "")), - rules=rules, - ) - - -def _build_rule( - data: dict[str, Any], default_action: Action, index: int -) -> Rule | None: - """Build a single Rule from a YAML rule entry.""" - hook = _HOOK_MAP.get(data.get("hook", "before_model")) - if hook is None: - logger.warning( - "Skipping rule %s: unknown hook %r", data.get("id"), data.get("hook") - ) - return None - - action_str = data.get("action") - action = ( - _ACTION_MAP.get(action_str, default_action) if action_str else default_action - ) - - default_sev = "high" if action == Action.DENY else "medium" - severity = _SEVERITY_MAP.get(data.get("severity", default_sev), Severity.HIGH) - - checks = _build_checks( - data.get("checks", []) or [], - action, - mapped_to_uipath=bool(data.get("mapped_to_uipath", False)), - policy_enabled=bool(data.get("policy_enabled", True)), - ) - - # If checks were declared but none could be parsed (e.g. all unknown - # types), skip the rule. A rule with zero checks "always matches" in - # the evaluator, so keeping it would make it fire on every request. - declared = data.get("checks", []) or [] - if declared and not checks: - logger.warning( - "Skipping rule %s: none of its %d declared check(s) could be parsed", - data.get("id"), - len(declared), - ) - return None - - return Rule( - rule_id=str(data.get("id", f"RULE-{index}")), - name=str(data.get("name", data.get("id", f"RULE-{index}"))), - clause=str(data.get("clause", data.get("owasp_ref", ""))), - hook=hook, - action=action, - severity=severity, - checks=checks, - enabled=bool(data.get("enabled", True)), - description=str(data.get("description", "")), - ) - - -def _build_checks( - checks_data: list[dict[str, Any]], - default_action: Action, - *, - mapped_to_uipath: bool = False, - policy_enabled: bool = True, -) -> list[Check]: - """Build the checks list for a rule. - - ``mapped_to_uipath`` / ``policy_enabled`` are rule-level flags read - by ``guardrail_fallback`` checks so the per-check condition can - decide whether to fire the compensating governance call. - """ - checks: list[Check] = [] - for check_data in checks_data: - if not isinstance(check_data, dict): - continue - check = _build_check( - check_data, - default_action, - mapped_to_uipath=mapped_to_uipath, - policy_enabled=policy_enabled, - ) - if check is not None: - checks.append(check) - return checks - - -def _build_check( - data: dict[str, Any], - default_action: Action, - *, - mapped_to_uipath: bool = False, - policy_enabled: bool = True, -) -> Check | None: - """Build one Check from a YAML check entry. - - Supports the same check types as ``compile_packs.py``: explicit - conditions, regex, budget, tool_allowlist, parameter_validation, - rate_limit, field_regex, sentiment_concern, data_quality_score, - incident_taxonomy, commitment_extractor, plus ``guardrail_fallback`` - (reads the rule-level ``mapped_to_uipath`` / ``policy_enabled`` flags - threaded in from ``_build_rule``). - """ - conditions: list[Condition] = [] - message = "" - - raw_conditions = data.get("conditions") - has_explicit_conditions = ( - isinstance(raw_conditions, list) - and raw_conditions - and isinstance(raw_conditions[0], dict) - and "operator" in raw_conditions[0] - ) - - check_type = data.get("type", "regex") - - if has_explicit_conditions: - assert isinstance(raw_conditions, list) # narrowed by has_explicit_conditions - conditions.extend(_make_conditions(raw_conditions)) - message = str(data.get("message", "")) - - elif check_type == "regex": - patterns = data.get("patterns", []) or [] - scope = data.get("scope", ["human", "ai"]) - field = _field_for_scope(scope) - for pattern in patterns: - conditions.append(Condition(operator="regex", field=field, value=pattern)) - message = f"Pattern matched in {scope}" - - elif check_type == "budget": - if "max_tool_calls_per_session" in data: - conditions.append( - Condition( - operator="gt", - field="session_state.tool_calls", - value=data["max_tool_calls_per_session"], - ) - ) - if "max_tool_calls_per_minute" in data: - conditions.append( - Condition( - operator="gt", - field="session_state.tool_calls_per_minute", - value=data["max_tool_calls_per_minute"], - ) - ) - if "max_consecutive_tool_calls" in data: - conditions.append( - Condition( - operator="gt", - field="session_state.consecutive_tool_calls", - value=data["max_consecutive_tool_calls"], - ) - ) - message = "Tool budget exceeded" - - elif check_type == "tool_allowlist": - blocked_tools = data.get("blocked_tools", []) or [] - if blocked_tools: - conditions.append( - Condition(operator="in_list", field="tool_name", value=blocked_tools) - ) - message = "Tool not allowed" - - elif check_type == "parameter_validation": - for pattern in data.get("additional_patterns", []) or []: - conditions.append( - Condition(operator="regex", field="tool_args", value=pattern) - ) - message = "Suspicious pattern in tool parameters" - - elif check_type == "rate_limit": - if "max_llm_calls_per_session" in data: - conditions.append( - Condition( - operator="gt", - field="session_state.llm_calls", - value=data["max_llm_calls_per_session"], - ) - ) - if "max_llm_calls_per_minute" in data: - conditions.append( - Condition( - operator="gt", - field="session_state.llm_calls_per_minute", - value=data["max_llm_calls_per_minute"], - ) - ) - message = "Rate limit exceeded" - - elif check_type == "field_regex": - conditions.extend(_make_conditions(data.get("conditions", []) or [])) - message = str(data.get("message", "Field regex check failed")) - - elif check_type == "data_quality_score": - field = data.get("field", "tool_result") - if data.get("check_encoding", True): - conditions.append( - Condition( - operator="encoding_concern", - field=field, - value={ - "min_confidence": float(data.get("min_confidence", 0.5)), - "max_replacement_ratio": float( - data.get("max_replacement_ratio", 0.05) - ), - "min_corruption_events": int( - data.get("min_corruption_events", 2) - ), - }, - ) - ) - if data.get("check_entropy", True): - conditions.append( - Condition( - operator="entropy_concern", - field=field, - value={ - "min": float(data.get("entropy_min", 1.5)), - "max": float(data.get("entropy_max", 7.5)), - }, - ) - ) - message = str( - data.get("message", "A.7.4: Data quality signal (encoding or entropy)") - ) - - elif check_type == "incident_taxonomy": - field = data.get("field", "model_output") - categories = data.get("categories") - value: dict[str, Any] = {} - if categories: - value["categories"] = list(categories) - conditions.append( - Condition(operator="incident_concern", field=field, value=value) - ) - message = str(data.get("message", "A.8.4: Incident signal detected")) - - elif check_type == "commitment_extractor": - field = data.get("field", "model_output") - conditions.append( - Condition( - operator="commitment_concern", - field=field, - value={ - "require_amount": bool(data.get("require_amount", True)), - "require_deadline": bool(data.get("require_deadline", False)), - }, - ) - ) - message = str( - data.get("message", "A.10.4: Customer commitment language detected") - ) - - elif check_type == "sentiment_concern": - field = data.get("field", "model_input") - threshold = float(data.get("threshold", -0.3)) - conditions.append( - Condition( - operator="vader_concern", - field=field, - value={"threshold": threshold}, - ) - ) - message = str( - data.get( - "message", - f"Negative sentiment detected (VADER compound <= {threshold})", - ) - ) - - elif check_type == "guardrail_fallback": - # Centralized guardrail compensating control. The on/off state - # lives at the RULE level (mapped_to_uipath / policy_enabled), - # threaded in from ``_build_rule``; ``validator`` names which - # guardrail check the server should run on behalf of the agent. - # The condition matches only when the guardrail is mapped to - # UiPath but disabled — see the ``guardrail_fallback`` operator - # in :class:`GovernanceEvaluator`. - conditions.append( - Condition( - operator="guardrail_fallback", - field="", - value={ - "validator": str(data.get("validator", "")), - "mapped_to_uipath": mapped_to_uipath, - "policy_enabled": policy_enabled, - }, - ) - ) - message = str( - data.get("message", "Guardrail disabled — compensating check needed.") - ) - - else: - logger.debug("Skipping check: unknown type %r", check_type) - return None - - if not conditions: - return None - - action_str = data.get("action") - action = ( - _ACTION_MAP.get(action_str, default_action) if action_str else default_action - ) - - message = str(data.get("message", message)) - - # Multi-pattern regex/parameter_validation defaults to OR semantics - # (any pattern indicates a hit); explicit `logic` in YAML wins. - if check_type in ("parameter_validation", "regex") and len(conditions) > 1: - default_logic = "any" - else: - default_logic = "all" - logic = str(data.get("logic", default_logic)) - - return Check(conditions=conditions, action=action, message=message, logic=logic) - - -def _make_conditions(raw: list[dict[str, Any]]) -> list[Condition]: - """Translate a list of YAML condition dicts into Condition objects.""" - out: list[Condition] = [] - for cond in raw: - if not isinstance(cond, dict): - continue - out.append( - Condition( - operator=str(cond.get("operator", "regex")), - field=str(cond.get("field", "model_input")), - value=cond.get("value", ""), - negate=bool(cond.get("negate", False)), - ) - ) - return out - - -def _field_for_scope(scope: list[str] | str) -> str: - """Map a YAML `scope` value to the CheckContext field it targets.""" - if isinstance(scope, str): - scope = [scope] - if "system" in scope or "human" in scope: - return "model_input" - if "ai" in scope: - return "model_output" - if "tool_result" in scope: - return "tool_result" - return "model_input" diff --git a/src/uipath/runtime/governance/native/loader.py b/src/uipath/runtime/governance/native/loader.py deleted file mode 100644 index 4d8271f..0000000 --- a/src/uipath/runtime/governance/native/loader.py +++ /dev/null @@ -1,284 +0,0 @@ -"""Policy pack loader. - -Resolves the active PolicyIndex at startup. Policies are fetched -exclusively from the governance backend (``api/v1/policy``); there is -no local compiled fallback. When the backend is unavailable, the -access token is unset, or the fetch times out, the loader returns an -empty PolicyIndex and the agent runs without any rules. -""" - -from __future__ import annotations - -import logging -import os -import threading -import time -from collections import Counter - -import yaml -from uipath.core.governance import EnforcementMode - -from uipath.runtime.governance.config import set_enforcement_mode -from uipath.runtime.governance.native._yaml_to_index import build_policy_index_from_yaml -from uipath.runtime.governance.native.backend_client import ENV_ACCESS_TOKEN -from uipath.runtime.governance.native.models import PolicyIndex -from uipath.runtime.governance.native.policy_api_client import ( - ENV_ORGANIZATION_ID, - ENV_TENANT_ID, - POLICY_API_TIMEOUT_SECONDS, - fetch_policy_response, - resolve_organization_id, - resolve_tenant_id, -) - -logger = logging.getLogger(__name__) - -# Pack name aliases for backward compatibility -PACK_ALIASES: dict[str, str] = { - "owasp": "owasp_agentic", - "hipaa": "hipaa_runtime", - "soc2": "soc2_runtime", - "nist": "nist_ai_rmf_runtime", - "eu_ai": "eu_ai_act_runtime", - "iso": "iso42001_runtime", -} - - -# Module-level cache -_policy_index: PolicyIndex | None = None - -# Background-prefetch coordination. ``_prefetch_event`` is set once the -# background load_policy_index() call finishes (success OR failure); -# callers of ``get_policy_index()`` wait on it. ``_prefetch_lock`` -# protects the start-once semantics so concurrent ``prefetch`` calls -# don't kick off duplicate threads. -_prefetch_event: threading.Event | None = None -_prefetch_lock = threading.Lock() - -# Default wait when ``get_policy_index()`` blocks on an in-flight -# prefetch. Matched to the policy-API HTTP timeout so a stuck backend -# bounds the total time spent waiting at first hook fire to -# ~POLICY_API_TIMEOUT_SECONDS. If the wait expires we return an empty -# PolicyIndex — the agent runs without any policies rather than -# blocking further or retrying. -_PREFETCH_WAIT_SECONDS = POLICY_API_TIMEOUT_SECONDS - - -def prefetch_policy_index() -> None: - """Kick off a background load of the policy index. - - Non-blocking. Designed to be called as early as possible (at - ``GovernanceRuntime.__init__``) so the HTTP call to the governance - backend overlaps with the rest of agent setup. Idempotent. - """ - global _prefetch_event - - with _prefetch_lock: - if _policy_index is not None: - return # already loaded - if _prefetch_event is not None: - return # already in flight - event = threading.Event() - _prefetch_event = event - - def _worker() -> None: - global _policy_index - try: - loaded = load_policy_index() - except Exception as exc: # noqa: BLE001 - logger.warning("Policy prefetch failed: %s", exc) - else: - with _prefetch_lock: - _policy_index = loaded - finally: - event.set() - - threading.Thread( - target=_worker, - name="governance-policy-prefetch", - daemon=True, - ).start() - - -def get_policy_index() -> PolicyIndex: - """Get the cached policy index, loading if necessary. - - Resolution order on first call: - 1. If a prefetch (see :func:`prefetch_policy_index`) is in flight, - wait for it to complete (bounded by ``_PREFETCH_WAIT_SECONDS``). - 2. Governance backend at ``api/v1/policy`` (one HTTP GET, cached). - 3. Empty PolicyIndex when the backend is unavailable or times out. - - Result is cached for the process lifetime; per-hook evaluation never - touches the network. Call :func:`clear_policy_cache` to force a - refetch (mainly for tests). - """ - global _policy_index - - if _policy_index is not None: - return _policy_index - - event = _prefetch_event - if event is not None: - completed = event.wait(timeout=_PREFETCH_WAIT_SECONDS) - if completed and _policy_index is not None: - return _policy_index - if not completed: - logger.warning( - "Policy prefetch did not complete in %.1fs; " - "agent will run without any policies", - _PREFETCH_WAIT_SECONDS, - ) - else: - logger.warning( - "Policy prefetch completed but produced no PolicyIndex " - "(see prior WARN for the root cause); agent will run " - "without any policies" - ) - _policy_index = PolicyIndex() - return _policy_index - - # No prefetch was started (direct callers / tests). Sync load. - _policy_index = load_policy_index() - return _policy_index - - -def load_policy_index(pack_name: str | None = None) -> PolicyIndex: - """Load the active PolicyIndex from the governance backend. - - Args: - pack_name: Ignored. Pack selection is controlled entirely by the - backend. - - Returns: - PolicyIndex parsed from the backend response. Empty PolicyIndex - when the backend is unavailable, the token is unset, the YAML - is malformed, or the response yields zero rules. - """ - start = time.perf_counter() - - api_index = _load_from_api() - if api_index is not None: - _log_index_summary(api_index) - logger.info( - "Policy index ready: source=backend, total_ms=%.1f", - (time.perf_counter() - start) * 1000, - ) - return api_index - - reason = _empty_index_reason() - logger.info( - "Policy index ready: source=empty (%s), total_ms=%.1f", - reason, - (time.perf_counter() - start) * 1000, - ) - return PolicyIndex() - - -def _empty_index_reason() -> str: - """Diagnose why the policy fetch produced nothing.""" - if not resolve_organization_id(): - return ( - f"UiPathConfig.organization_id unavailable — set {ENV_ORGANIZATION_ID} " - "or install uipath-platform; backend API not contacted" - ) - if not resolve_tenant_id(): - return ( - f"UiPathConfig.tenant_id unavailable — set {ENV_TENANT_ID} " - "or install uipath-platform; backend API not contacted" - ) - if not os.environ.get(ENV_ACCESS_TOKEN): - return f"{ENV_ACCESS_TOKEN} unset — backend API not contacted" - return "backend returned no policies (timeout / error / empty body)" - - -def _apply_enforcement_mode(mode_str: str | None) -> None: - """Map a backend-supplied mode string onto :class:`EnforcementMode`.""" - if not mode_str: - return - try: - mode = EnforcementMode(mode_str.lower()) - except ValueError: - logger.warning( - "Backend returned unknown enforcement mode %r; keeping current mode", - mode_str, - ) - return - set_enforcement_mode(mode) - logger.info("Enforcement mode set from backend: %s", mode.value) - - -def _load_from_api() -> PolicyIndex | None: - """Fetch and parse the policy index from the governance backend.""" - start = time.perf_counter() - response = fetch_policy_response() - if response is None: - return None - - _apply_enforcement_mode(response.mode) - - if not response.policy: - logger.warning( - "Policy fetch returned empty policy field; " - "agent will run without any policies" - ) - return None - - try: - index = build_policy_index_from_yaml(response.policy) - except yaml.YAMLError as exc: - logger.warning("Policy YAML from backend was malformed: %s", exc) - return None - except Exception as exc: # noqa: BLE001 - logger.warning("Failed to build PolicyIndex from backend YAML: %s", exc) - return None - - if index.total_rules == 0: - logger.warning( - "Policy YAML from backend yielded zero rules; " - "agent will run without any policies" - ) - return None - - elapsed_ms = (time.perf_counter() - start) * 1000 - logger.info( - "Loaded policy index from backend: packs=%s, rules=%d, elapsed_ms=%.1f", - index.pack_names, - index.total_rules, - elapsed_ms, - ) - return index - - -def _log_index_summary(index: PolicyIndex) -> None: - """Log summary of loaded policy index.""" - hook_counts: Counter[str] = Counter() - for rule in index.all_rules: - hook_counts[rule.hook.value] += 1 - - logger.debug( - "Policy packs: %s, total rules: %d, by hook: %s", - index.pack_names, - index.total_rules, - dict(hook_counts), - ) - - -def get_available_packs() -> list[str]: - """Get list of pack names from the currently loaded policy index.""" - if _policy_index is None: - return [] - return _policy_index.pack_names - - -def clear_policy_cache() -> None: - """Clear the cached policy index and any in-flight prefetch state.""" - global _policy_index, _prefetch_event - with _prefetch_lock: - _policy_index = None - _prefetch_event = None - logger.debug("Policy index cache cleared") - - -# Backward compatibility alias -reset_policy_index = clear_policy_cache diff --git a/src/uipath/runtime/governance/native/policy_api_client.py b/src/uipath/runtime/governance/native/policy_api_client.py deleted file mode 100644 index 325b4e0..0000000 --- a/src/uipath/runtime/governance/native/policy_api_client.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Governance policy API client. - -Fetches the governance backend response so policies can be controlled -centrally without redeploying agents. Called once at process startup -from :mod:`uipath.runtime.governance.native.loader`; per-hook evaluation -stays in-process. - -Response shape (JSON):: - - { - "mode": "audit" | "enforce" | "disabled", - "policies": "" - } - -``mode`` is the platform-controlled enforcement mode for the tenant; -the loader applies it via -:func:`uipath.runtime.governance.config.set_enforcement_mode`. ``policies`` -is the YAML the evaluator compiles into a :class:`PolicyIndex`. - -Failure mode is fail-open: when the organization id is unknown, the -access token is missing, the backend errors, or the body can't be -parsed, the caller falls back to an empty PolicyIndex. The fetch is -single-shot (no retry by design — see :func:`_get_once`) so a slow -backend can't extend agent startup beyond -:data:`BACKEND_REQUEST_TIMEOUT_SECONDS`. Nothing in this module ever -raises to the caller. -""" - -from __future__ import annotations - -import json -import logging -import os -import urllib.error -import urllib.request -from dataclasses import dataclass -from urllib.parse import urlencode - -from uipath.runtime.governance.native.backend_client import ( - AGENT_TYPE_PARAM, - BACKEND_REQUEST_TIMEOUT_SECONDS, - ENV_ACCESS_TOKEN, - ENV_ORGANIZATION_ID, - ENV_TENANT_ID, - POLICY_API_PATH, - TENANT_HEADER, - agent_type_param, - build_governance_url, - governance_request_headers, - resolve_organization_id, - resolve_tenant_id, -) - -logger = logging.getLogger(__name__) - -# Re-exported alias kept for callers that imported the old name. -POLICY_API_TIMEOUT_SECONDS = BACKEND_REQUEST_TIMEOUT_SECONDS - - -@dataclass(frozen=True) -class PolicyResponse: - """Parsed governance backend response. - - Attributes: - mode: Enforcement mode string the backend returned - (``"audit"`` / ``"enforce"`` / ``"disabled"``), or ``None`` - when the backend omitted it. Loader applies this via - :func:`uipath.runtime.governance.config.set_enforcement_mode`. - policy: Policy pack YAML to compile into a ``PolicyIndex``. May - be an empty string if the backend returned no rules. - """ - - mode: str | None - policy: str - - -def build_policy_url(org_id: str) -> str: - """Build the policy endpoint URL for the given organization id. - - The tenant id is not part of the URL; it travels in the - ``x-uipath-internal-tenantid`` request header (see - :func:`fetch_policy_response`). - - When the hosted agent's type is known (see - :func:`uipath.runtime.governance.native.backend_client.set_agent_conversational`), - an ``agentType`` query param is appended so the server resolves the - conversational-vs-autonomous container key. Omitted when unknown — the - server then applies its default. - """ - url = build_governance_url(org_id, POLICY_API_PATH) - agent_type = agent_type_param() - if agent_type: - url = f"{url}?{urlencode({AGENT_TYPE_PARAM: agent_type})}" - return url - - -def fetch_policy_response() -> PolicyResponse | None: - """Fetch the governance backend's policy response. - - Single shot, no retry: a failed fetch (timeout / network error / - HTTP error / malformed body) returns ``None`` and the caller falls - back to an empty PolicyIndex. The agent must not spend time on a - second attempt — keeping governance off the critical path is more - important than maximising policy availability. - - Returns: - :class:`PolicyResponse` on success. ``None`` on any failure - path — caller falls back to an empty PolicyIndex. - - Never raises. - """ - try: - return _fetch_policy_response_inner() - except Exception as exc: # noqa: BLE001 - loader path must never raise - logger.warning("Policy fetch failed unexpectedly: %s", exc) - return None - - -def _fetch_policy_response_inner() -> PolicyResponse | None: - org_id = resolve_organization_id() - if not org_id: - logger.warning( - "Policy fetch skipped: UiPathConfig.organization_id is not " - "available (set %s in the environment, or ensure uipath-platform " - "is installed); governance will run with no policies. The " - "backend API was NOT contacted.", - ENV_ORGANIZATION_ID, - ) - return None - - tenant_id = resolve_tenant_id() - if not tenant_id: - logger.warning( - "Policy fetch skipped: UiPathConfig.tenant_id is not " - "available (set %s in the environment, or ensure uipath-platform " - "is installed); governance will run with no policies. The " - "backend API was NOT contacted.", - ENV_TENANT_ID, - ) - return None - - policy_url = build_policy_url(org_id) - - token = os.environ.get(ENV_ACCESS_TOKEN) - if not token: - logger.warning( - "Policy fetch skipped: %s is not set in the environment; " - "governance will run with no policies.", - ENV_ACCESS_TOKEN, - ) - return None - - # Policy fetch is a GET; ``json_body=False`` so ``Content-Type`` is - # omitted. Strict origin servers may 415 on unexpected Content-Type - # for GETs (see :func:`governance_request_headers` docstring). - headers = governance_request_headers(json_body=False) - headers[TENANT_HEADER] = tenant_id - logger.info("Policy fetch starting (org=%s, tenant=%s)", org_id, tenant_id) - - body = _get_once(policy_url, headers) - if body is None: - return None - return _parse_policy_body(body) - - -def _get_once(url: str, headers: dict[str, str]) -> bytes | None: - """GET ``url`` once. Returns body bytes, or ``None`` on any failure. - - No retry by design — see :func:`fetch_policy_response` for the - rationale. Every failure path logs a single WARNING and returns - ``None`` so the caller (the loader) falls back to an empty - PolicyIndex without delay. - """ - request = urllib.request.Request(url, headers=headers, method="GET") - try: - with urllib.request.urlopen( # noqa: S310 - URL is built from config - request, timeout=BACKEND_REQUEST_TIMEOUT_SECONDS - ) as response: - return response.read() - except urllib.error.HTTPError as exc: - logger.warning("Policy fetch returned HTTP %d: %s", exc.code, exc) - except (urllib.error.URLError, TimeoutError, OSError) as exc: - logger.warning("Policy fetch failed: %s", exc) - return None - - -def _parse_policy_body(body: bytes) -> PolicyResponse | None: - """Parse the JSON envelope into a :class:`PolicyResponse`.""" - if not body: - logger.warning("Policy fetch returned empty body") - return None - - try: - payload = json.loads(body.decode("utf-8")) - except UnicodeDecodeError as exc: - logger.warning("Policy fetch returned non-UTF8 body: %s", exc) - return None - except json.JSONDecodeError as exc: - logger.warning( - "Policy fetch returned malformed JSON " - "(server may have returned an HTML error page): %s", - exc, - ) - return None - - if not isinstance(payload, dict): - logger.warning( - "Policy fetch returned unexpected JSON shape (expected object, got %s)", - type(payload).__name__, - ) - return None - - raw_mode = payload.get("mode") - mode = raw_mode if isinstance(raw_mode, str) and raw_mode else None - - raw_policy = payload.get("policies", "") - if not isinstance(raw_policy, str): - logger.warning( - "Policy fetch returned non-string 'policies' field (got %s)", - type(raw_policy).__name__, - ) - return None - - logger.info( - "Policy fetch ok: mode=%s, policy_bytes=%d", mode, len(raw_policy) - ) - return PolicyResponse(mode=mode, policy=raw_policy) From 35299f00568ca22668eb5f9e83fc926cfbe44857 Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Fri, 10 Jul 2026 11:53:57 +0530 Subject: [PATCH 03/12] feat: add build_rego_evaluator_async to rego loader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds build_rego_evaluator_async(service) — a single async call that fetches /all-policies metadata via the injected platform service, downloads changed bundles to disk, and returns a RegoEvaluator or None. Mirrors the native evaluator bootstrap pattern (no background threads at startup). Bumps version to 0.12.3. Generated with Claude Code Co-Authored-By: Claude --- pyproject.toml | 2 +- .../runtime/governance/rego/__init__.py | 8 +- src/uipath/runtime/governance/rego/loader.py | 80 ++++++++++++++++++- 3 files changed, 87 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b0b413a..c1415a6 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" diff --git a/src/uipath/runtime/governance/rego/__init__.py b/src/uipath/runtime/governance/rego/__init__.py index 3d0b6a1..1d6c466 100644 --- a/src/uipath/runtime/governance/rego/__init__.py +++ b/src/uipath/runtime/governance/rego/__init__.py @@ -2,10 +2,16 @@ from __future__ import annotations from .evaluator import RegoEvaluator -from .loader import clear_rego_cache, get_rego_evaluator, prefetch_rego_bundles +from .loader import ( + build_rego_evaluator_async, + clear_rego_cache, + get_rego_evaluator, + prefetch_rego_bundles, +) __all__ = [ "RegoEvaluator", + "build_rego_evaluator_async", "clear_rego_cache", "get_rego_evaluator", "prefetch_rego_bundles", diff --git a/src/uipath/runtime/governance/rego/loader.py b/src/uipath/runtime/governance/rego/loader.py index e74e5fc..bd28988 100644 --- a/src/uipath/runtime/governance/rego/loader.py +++ b/src/uipath/runtime/governance/rego/loader.py @@ -1,4 +1,4 @@ -"""Rego bundle loader: startup prefetch, ETag-driven download, background refresh.""" +"""Rego bundle loader: startup build, ETag-driven download, background refresh.""" from __future__ import annotations import logging @@ -224,6 +224,84 @@ def _refresh_loop() -> None: logger.debug("Rego bundle background refresh started (interval=%ds)", refresh_seconds) +async def build_rego_evaluator_async(service: object) -> 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``. + + 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) + 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) -> object: + from pathlib import Path + + 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: + 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) + + def clear_rego_cache() -> None: """Reset all module-level state. Intended for tests.""" global _rego_evaluator, _rego_prefetch_event, _rego_refresh_started From 3b824d64ede721fefc372c23b80b392afbb452bb Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Sun, 12 Jul 2026 23:25:41 +0530 Subject: [PATCH 04/12] Rego: fix trace_id bug, remove dead code, add tests - Remove api_client.py and models.py (dead code since build_rego_evaluator_async uses the platform service layer; api_client also used wrong URL pattern) - Remove prefetch/background-refresh/get_rego_evaluator from loader.py - Fix AuditRecord construction: trace_id is not a field (reliability bug) - Add 40 tests covering bundle_cache, loader, and evaluator Generated with Claude Code Co-Authored-By: Claude --- .../runtime/governance/rego/__init__.py | 10 +- .../runtime/governance/rego/api_client.py | 145 -------- .../runtime/governance/rego/evaluator.py | 4 - src/uipath/runtime/governance/rego/loader.py | 216 +---------- src/uipath/runtime/governance/rego/models.py | 18 - tests/test_rego_bundle_cache.py | 109 ++++++ tests/test_rego_evaluator.py | 341 ++++++++++++++++++ tests/test_rego_loader.py | 209 +++++++++++ 8 files changed, 665 insertions(+), 387 deletions(-) delete mode 100644 src/uipath/runtime/governance/rego/api_client.py delete mode 100644 src/uipath/runtime/governance/rego/models.py create mode 100644 tests/test_rego_bundle_cache.py create mode 100644 tests/test_rego_evaluator.py create mode 100644 tests/test_rego_loader.py diff --git a/src/uipath/runtime/governance/rego/__init__.py b/src/uipath/runtime/governance/rego/__init__.py index 1d6c466..fddbb41 100644 --- a/src/uipath/runtime/governance/rego/__init__.py +++ b/src/uipath/runtime/governance/rego/__init__.py @@ -2,17 +2,9 @@ from __future__ import annotations from .evaluator import RegoEvaluator -from .loader import ( - build_rego_evaluator_async, - clear_rego_cache, - get_rego_evaluator, - prefetch_rego_bundles, -) +from .loader import build_rego_evaluator_async __all__ = [ "RegoEvaluator", "build_rego_evaluator_async", - "clear_rego_cache", - "get_rego_evaluator", - "prefetch_rego_bundles", ] diff --git a/src/uipath/runtime/governance/rego/api_client.py b/src/uipath/runtime/governance/rego/api_client.py deleted file mode 100644 index 0b82176..0000000 --- a/src/uipath/runtime/governance/rego/api_client.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Governance Rego bundle API client. - -Fetches per-hook WASM bundle metadata from ``GET /all-policies/:tenantId`` -and downloads individual bundles from their CDN URLs. - -Failure mode is fail-open throughout: every public function returns -``None`` / empty on any error so callers fall back to stale cache or -run without Rego evaluation. -""" -from __future__ import annotations - -import json -import logging -import os -import urllib.error -import urllib.request - -from uipath.runtime.governance.native.backend_client import ( - ALL_POLICIES_API_PATH, - BACKEND_REQUEST_TIMEOUT_SECONDS, - ENV_ACCESS_TOKEN, - TENANT_HEADER, - build_governance_url, - governance_request_headers, - resolve_organization_id, - resolve_tenant_id, -) -from uipath.runtime.governance.rego.models import AllPoliciesResponse, HookBundle - -logger = logging.getLogger(__name__) - - -def build_all_policies_url(org_id: str, tenant_id: str) -> str: - """Build the ``/all-policies/:tenantId`` endpoint URL.""" - return build_governance_url(org_id, f"{ALL_POLICIES_API_PATH}/{tenant_id}") - - -def fetch_all_policies() -> AllPoliciesResponse | None: - """Fetch per-hook bundle metadata from the governance backend. - - Returns AllPoliciesResponse on success, None on any failure. Never raises. - """ - try: - return _fetch_all_policies_inner() - except Exception as exc: # noqa: BLE001 - logger.warning("fetch_all_policies failed unexpectedly: %s", exc) - return None - - -def _fetch_all_policies_inner() -> AllPoliciesResponse | None: - org_id = resolve_organization_id() - if not org_id: - logger.warning( - "Rego bundle fetch skipped: organization_id unavailable; " - "agent will run without custom Rego policies." - ) - return None - - tenant_id = resolve_tenant_id() - if not tenant_id: - logger.warning( - "Rego bundle fetch skipped: tenant_id unavailable; " - "agent will run without custom Rego policies." - ) - return None - - token = os.environ.get(ENV_ACCESS_TOKEN) - if not token: - logger.warning( - "Rego bundle fetch skipped: %s not set; " - "agent will run without custom Rego policies.", - ENV_ACCESS_TOKEN, - ) - return None - - url = build_all_policies_url(org_id, tenant_id) - headers = governance_request_headers(json_body=False) - headers[TENANT_HEADER] = tenant_id - - body = _get_once(url, headers) - if body is None: - return None - return _parse_response(body) - - -def _get_once(url: str, headers: dict[str, str]) -> bytes | None: - request = urllib.request.Request(url, headers=headers, method="GET") - try: - with urllib.request.urlopen( # noqa: S310 - request, timeout=BACKEND_REQUEST_TIMEOUT_SECONDS - ) as response: - return response.read() - except urllib.error.HTTPError as exc: - logger.warning("Rego bundle fetch returned HTTP %d: %s", exc.code, exc) - except (urllib.error.URLError, TimeoutError, OSError) as exc: - logger.warning("Rego bundle fetch failed: %s", exc) - return None - - -def _parse_response(body: bytes) -> AllPoliciesResponse | None: - if not body: - logger.warning("Rego bundle fetch: empty response body") - return None - try: - payload = json.loads(body.decode("utf-8")) - except (json.JSONDecodeError, UnicodeDecodeError) as exc: - logger.warning("Rego bundle fetch: malformed response: %s", exc) - return None - - if not isinstance(payload, dict): - logger.warning("Rego bundle fetch: unexpected JSON shape") - return None - - raw_bundles = payload.get("hookBundles", []) - if not isinstance(raw_bundles, list): - logger.warning("Rego bundle fetch: hookBundles is not a list") - return None - - bundles: list[HookBundle] = [] - for item in raw_bundles: - if not isinstance(item, dict): - continue - hook_type = item.get("hookType", "") - url_str = item.get("bundleUrl", "") - etag = item.get("etag", "") - if hook_type and url_str: - bundles.append(HookBundle(hook_type=hook_type, bundle_url=url_str, etag=etag)) - - return AllPoliciesResponse(hook_bundles=bundles) - - -def download_bundle(url: str) -> bytes | None: - """Download a WASM bundle from a pre-signed CDN URL. - - No auth header is sent — the URL is pre-signed. Returns raw bytes or None. Never raises. - """ - try: - request = urllib.request.Request(url, method="GET") - with urllib.request.urlopen( # noqa: S310 - request, timeout=BACKEND_REQUEST_TIMEOUT_SECONDS - ) as response: - return response.read() - except Exception as exc: # noqa: BLE001 - logger.warning("Bundle download from %s failed: %s", url, exc) - return None diff --git a/src/uipath/runtime/governance/rego/evaluator.py b/src/uipath/runtime/governance/rego/evaluator.py index e9925bc..a29e154 100644 --- a/src/uipath/runtime/governance/rego/evaluator.py +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -201,7 +201,6 @@ def evaluate(self, context: CheckContext) -> AuditRecord: timestamp=datetime.now(timezone.utc), agent_name=context.agent_name, runtime_id=context.runtime_id, - trace_id="", hook=context.hook, evaluations=[], final_action=Action.ALLOW, @@ -214,7 +213,6 @@ def evaluate(self, context: CheckContext) -> AuditRecord: timestamp=datetime.now(timezone.utc), agent_name=context.agent_name, runtime_id=context.runtime_id, - trace_id="", hook=context.hook, evaluations=[], final_action=Action.ALLOW, @@ -230,7 +228,6 @@ def evaluate(self, context: CheckContext) -> AuditRecord: timestamp=datetime.now(timezone.utc), agent_name=context.agent_name, runtime_id=context.runtime_id, - trace_id="", hook=context.hook, evaluations=[], final_action=Action.ALLOW, @@ -280,7 +277,6 @@ def evaluate(self, context: CheckContext) -> AuditRecord: timestamp=datetime.now(timezone.utc), agent_name=context.agent_name, runtime_id=context.runtime_id, - trace_id="", hook=context.hook, evaluations=evaluations, final_action=final_action, diff --git a/src/uipath/runtime/governance/rego/loader.py b/src/uipath/runtime/governance/rego/loader.py index bd28988..711da1c 100644 --- a/src/uipath/runtime/governance/rego/loader.py +++ b/src/uipath/runtime/governance/rego/loader.py @@ -1,22 +1,10 @@ -"""Rego bundle loader: startup build, ETag-driven download, background refresh.""" +"""Rego bundle loader: async bootstrap builder.""" from __future__ import annotations import logging -import os -import threading -import time from uipath.core.governance.models import LifecycleHook -from uipath.runtime.governance.native.backend_client import ( - BACKEND_REQUEST_TIMEOUT_SECONDS, - resolve_organization_id, - resolve_tenant_id, -) -from uipath.runtime.governance.rego.api_client import ( - download_bundle, - fetch_all_policies, -) from uipath.runtime.governance.rego.bundle_cache import ( get_cache_dir, get_cached_bundle_path, @@ -26,15 +14,6 @@ logger = logging.getLogger(__name__) -# Module-level state -_rego_evaluator = None # RegoEvaluator | None -_rego_prefetch_event: threading.Event | None = None -_rego_prefetch_lock = threading.Lock() -_rego_refresh_started = False -_rego_refresh_lock = threading.Lock() - -_DEFAULT_REFRESH_SECONDS = 30 - # Hook type string → LifecycleHook mapping _HOOK_MAP: dict[str, LifecycleHook] = { "before_agent": LifecycleHook.BEFORE_AGENT, @@ -46,184 +25,6 @@ } -def prefetch_rego_bundles() -> None: - """Kick off a background download of changed WASM bundles. Non-blocking. Idempotent.""" - global _rego_prefetch_event - - with _rego_prefetch_lock: - if _rego_evaluator is not None: - return - if _rego_prefetch_event is not None: - return - event = threading.Event() - _rego_prefetch_event = event - - def _worker() -> None: - try: - _sync_bundles_to_disk() - except Exception as exc: # noqa: BLE001 - logger.warning("Rego bundle prefetch failed: %s", exc) - finally: - event.set() - - threading.Thread( - target=_worker, - name="governance-rego-prefetch", - daemon=True, - ).start() - - _start_background_refresh() - - -def _sync_bundles_to_disk() -> None: - """Fetch /all-policies and download any bundle whose ETag has changed.""" - org_id = resolve_organization_id() - tenant_id = resolve_tenant_id() - if not org_id or not tenant_id: - logger.warning( - "Rego bundle sync skipped: org_id or tenant_id unavailable; " - "agent will run without custom Rego policies." - ) - return - - cache_dir = get_cache_dir(org_id, tenant_id) - response = fetch_all_policies() - if response is None: - logger.warning( - "Rego bundle sync: /all-policies fetch returned nothing; " - "using cached bundles if available." - ) - return - - for bundle in response.hook_bundles: - hook_type = bundle.hook_type - cached_etag = get_cached_etag(cache_dir, hook_type) - - if cached_etag == bundle.etag: - logger.debug("Rego bundle for hook=%s is up to date (etag=%s)", hook_type, bundle.etag) - continue - - logger.info( - "Rego bundle for hook=%s changed (cached=%s new=%s); downloading", - hook_type, cached_etag, bundle.etag, - ) - data = download_bundle(bundle.bundle_url) - if data is None: - stale = get_cached_bundle_path(cache_dir, hook_type) - if stale: - logger.warning( - "Rego bundle download failed for hook=%s; using stale cache.", hook_type - ) - else: - logger.warning( - "Rego bundle download failed for hook=%s and no cache exists; " - "hook will pass through without Rego evaluation.", - hook_type, - ) - continue - - save_bundle(cache_dir, hook_type, data, bundle.etag) - logger.info("Rego bundle saved for hook=%s (etag=%s)", hook_type, bundle.etag) - # Invalidate the in-memory evaluator so the next run picks up the new bundle. - global _rego_evaluator - _rego_evaluator = None - - -def get_rego_evaluator(): - """Return the cached RegoEvaluator, building it from disk if needed. - - Waits for the startup prefetch (bounded by BACKEND_REQUEST_TIMEOUT_SECONDS). - Returns None when no bundles are available (fail-open). - """ - global _rego_evaluator - - if _rego_evaluator is not None: - return _rego_evaluator - - event = _rego_prefetch_event - if event is not None: - completed = event.wait(timeout=BACKEND_REQUEST_TIMEOUT_SECONDS) - if not completed: - logger.warning( - "Rego bundle prefetch timed out after %.1fs; " - "using whatever is on disk (may be empty).", - BACKEND_REQUEST_TIMEOUT_SECONDS, - ) - - return _build_evaluator_from_disk() - - -def _build_evaluator_from_disk(): - """Read the disk cache and construct a RegoEvaluator. Returns None if no bundles exist.""" - from pathlib import Path - - 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: - return None - - cache_dir = get_cache_dir(org_id, tenant_id) - hook_wasm_paths: dict[LifecycleHook, Path] = {} - hook_data: dict[LifecycleHook, dict] = {} - - for hook_type_str, lifecycle_hook in _HOOK_MAP.items(): - path = get_cached_bundle_path(cache_dir, hook_type_str) - if path is not None: - hook_wasm_paths[lifecycle_hook] = path - data = _extract_data_json_from_bundle(path) - if data is not None: - hook_data[lifecycle_hook] = data - - if not hook_wasm_paths: - logger.warning( - "No Rego WASM bundles found on disk; " - "agent will run without custom Rego evaluation." - ) - return None - - logger.info( - "Building RegoEvaluator from disk: hooks=%s", - [h.value for h in hook_wasm_paths], - ) - return RegoEvaluator(hook_wasm_paths, hook_data=hook_data or None) - - -def _start_background_refresh() -> None: - """Start the background refresh daemon (once per process).""" - global _rego_refresh_started - - refresh_seconds = int( - os.environ.get("UIPATH_GOVERNANCE_BUNDLE_REFRESH_SECONDS", _DEFAULT_REFRESH_SECONDS) - ) - if refresh_seconds <= 0: - return - - with _rego_refresh_lock: - if _rego_refresh_started: - return - _rego_refresh_started = True - - def _refresh_loop() -> None: - while True: - time.sleep(refresh_seconds) - try: - _sync_bundles_to_disk() - except Exception as exc: # noqa: BLE001 - logger.warning("Rego bundle background refresh failed: %s", exc) - - threading.Thread( - target=_refresh_loop, - name="governance-rego-refresh", - daemon=True, - ).start() - logger.debug("Rego bundle background refresh started (interval=%ds)", refresh_seconds) - - async def build_rego_evaluator_async(service: object) -> object: """Fetch bundles via platform service and build a RegoEvaluator. @@ -249,6 +50,10 @@ async def _build_rego_evaluator_async_inner(service: object) -> object: RegoEvaluator, _extract_data_json_from_bundle, ) + from uipath.runtime.governance.native.backend_client import ( + resolve_organization_id, + resolve_tenant_id, + ) org_id = resolve_organization_id() tenant_id = resolve_tenant_id() @@ -300,14 +105,3 @@ async def _build_rego_evaluator_async_inner(service: object) -> object: 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) - - -def clear_rego_cache() -> None: - """Reset all module-level state. Intended for tests.""" - global _rego_evaluator, _rego_prefetch_event, _rego_refresh_started - with _rego_prefetch_lock: - _rego_evaluator = None - _rego_prefetch_event = None - with _rego_refresh_lock: - _rego_refresh_started = False - logger.debug("Rego loader cache cleared") diff --git a/src/uipath/runtime/governance/rego/models.py b/src/uipath/runtime/governance/rego/models.py deleted file mode 100644 index 5f3be6c..0000000 --- a/src/uipath/runtime/governance/rego/models.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Data models for the governance Rego/WASM bundle API response.""" -from __future__ import annotations - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class HookBundle: - """Metadata for a single per-hook WASM bundle returned by /all-policies.""" - hook_type: str - bundle_url: str - etag: str - - -@dataclass(frozen=True) -class AllPoliciesResponse: - """Parsed response from GET /all-policies/:tenantId.""" - hook_bundles: list[HookBundle] 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..4334bb3 --- /dev/null +++ b/tests/test_rego_evaluator.py @@ -0,0 +1,341 @@ +"""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} + + +# --------------------------------------------------------------------------- +# _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 From 0aea9339bd96c6d2fb65c318a9e674c14cc9f8fc Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Mon, 13 Jul 2026 23:00:59 +0530 Subject: [PATCH 05/12] fix: handle nested OPA bundle format in WASM extraction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The backend produces bundles where the outer bundle.tar.gz contains a 'policy.wasm' that is itself a gzip-compressed inner tar bundle. The actual WASM binary lives at /policy.wasm inside that inner tar. _extract_wasm_from_bundle now handles three levels: 1. Direct WASM in outer bundle (rare) 2. Gzip-compressed WASM in outer bundle 3. Gzip-compressed inner tar → /policy.wasm (current backend format) Also adds warning log when all-policies returns no bundles and adds the governance-rego optional dependency group to pyproject.toml. Generated with Claude Code Co-Authored-By: Claude --- pyproject.toml | 3 + .../runtime/governance/rego/evaluator.py | 95 +++++++++++++++++-- src/uipath/runtime/governance/rego/loader.py | 6 ++ uv.lock | 2 +- 4 files changed, 95 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c1415a6..8791f4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,9 @@ dependencies = [ "vaderSentiment>=3.3.2, <4.0", "chardet>=5.2.0, <8.0", ] + +[project.optional-dependencies] +governance-rego = ["opa-wasmtime>=0.1.1"] classifiers = [ "Intended Audience :: Developers", "Topic :: Software Development :: Build Tools", diff --git a/src/uipath/runtime/governance/rego/evaluator.py b/src/uipath/runtime/governance/rego/evaluator.py index a29e154..58fe6f0 100644 --- a/src/uipath/runtime/governance/rego/evaluator.py +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -64,19 +64,94 @@ def context_to_input( } -def _extract_wasm_from_bundle(bundle_path: Path) -> bytes: - """Extract ``policy.wasm`` bytes from an OPA ``.tar.gz`` bundle.""" - with open(bundle_path, "rb") as f: - bundle_bytes = f.read() - with tarfile.open(fileobj=io.BytesIO(bundle_bytes), mode="r:gz") as tf: +_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: - member = tf.getmember("/policy.wasm") + fobj = tf.extractfile(tf.getmember(candidate)) except KeyError: - member = tf.getmember("policy.wasm") - fobj = tf.extractfile(member) + continue if fobj is None: - raise ValueError(f"policy.wasm is not a regular file in {bundle_path}") - return fobj.read() + 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: + """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: diff --git a/src/uipath/runtime/governance/rego/loader.py b/src/uipath/runtime/governance/rego/loader.py index 711da1c..f9a4325 100644 --- a/src/uipath/runtime/governance/rego/loader.py +++ b/src/uipath/runtime/governance/rego/loader.py @@ -66,6 +66,12 @@ async def _build_rego_evaluator_async_inner(service: object) -> object: 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) 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" }, From 8cbe51c8e16f311c15df9bfe567ca3e93fd27086 Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Mon, 13 Jul 2026 23:33:46 +0530 Subject: [PATCH 06/12] fix: move classifiers and maintainers back into [project] table [project.optional-dependencies] section boundary caused classifiers and maintainers to be parsed as optional dependency lists. Generated with Claude Code Co-Authored-By: Claude --- pyproject.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8791f4c..503bda8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,8 +10,6 @@ dependencies = [ "chardet>=5.2.0, <8.0", ] -[project.optional-dependencies] -governance-rego = ["opa-wasmtime>=0.1.1"] classifiers = [ "Intended Audience :: Developers", "Topic :: Software Development :: Build Tools", @@ -23,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" From e359acb8ea7382c7b84cccb49663fd671e74f718 Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Tue, 14 Jul 2026 11:13:56 +0530 Subject: [PATCH 07/12] feat: thread AuditManager into build_rego_evaluator_async Rego WASM violations were evaluated but never emitted as trace spans because the loader constructed RegoEvaluator without an AuditManager. Adds an optional audit_manager param to build_rego_evaluator_async and _build_rego_evaluator_async_inner, passing it through to RegoEvaluator so callers can wire the shared AuditManager and get Rego rule violations surfaced alongside native checker events in LLMOps traces. Generated with Claude Code Co-Authored-By: Claude --- src/uipath/runtime/governance/rego/loader.py | 34 +++++++++++++++----- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/src/uipath/runtime/governance/rego/loader.py b/src/uipath/runtime/governance/rego/loader.py index f9a4325..6d8a19f 100644 --- a/src/uipath/runtime/governance/rego/loader.py +++ b/src/uipath/runtime/governance/rego/loader.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging +from typing import TYPE_CHECKING from uipath.core.governance.models import LifecycleHook @@ -12,6 +13,9 @@ save_bundle, ) +if TYPE_CHECKING: + from uipath.runtime.governance._audit.base import AuditManager + logger = logging.getLogger(__name__) # Hook type string → LifecycleHook mapping @@ -25,7 +29,10 @@ } -async def build_rego_evaluator_async(service: object) -> object: +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: @@ -33,27 +40,38 @@ async def build_rego_evaluator_async(service: object) -> object: 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) + 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) -> object: +async def _build_rego_evaluator_async_inner( + service: object, + audit_manager: "AuditManager | None" = None, +) -> object: from pathlib import Path - from uipath.runtime.governance.rego.evaluator import ( - RegoEvaluator, - _extract_data_json_from_bundle, - ) 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() @@ -110,4 +128,4 @@ async def _build_rego_evaluator_async_inner(service: object) -> object: 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) + return RegoEvaluator(hook_wasm_paths, hook_data=hook_data or None, audit_manager=audit_manager) From 687a5aca39c3e29e3cacff9ccfb19c36c63be612 Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Tue, 14 Jul 2026 11:16:55 +0530 Subject: [PATCH 08/12] fix: RegoEvaluator creates default AuditManager when none provided MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Without an AuditManager, Rego WASM violations were silently dropped — evaluated correctly but never emitted as trace spans. Now the evaluator self-provisions a default AuditManager (traces + track_events sinks) when the caller passes None, so Rego rule violations surface in LLMOps traces alongside native checker events without requiring any changes in caller repos (e.g. uipath-langchain). Generated with Claude Code Co-Authored-By: Claude --- src/uipath/runtime/governance/rego/evaluator.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/uipath/runtime/governance/rego/evaluator.py b/src/uipath/runtime/governance/rego/evaluator.py index 58fe6f0..dc12df5 100644 --- a/src/uipath/runtime/governance/rego/evaluator.py +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -95,12 +95,12 @@ def _find_wasm_in_tar(tf: Any) -> bytes | None: def _extract_wasm_from_bundle(bundle_path: Path) -> bytes: - """Extract ``policy.wasm`` bytes from an OPA ``.tar.gz`` bundle. + 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``) + └── /policy.wasm ← the actual WASM binary (``\x00asm``) """ import gzip @@ -220,6 +220,8 @@ def __init__( """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 {} From 5272e537c00ca28b1d3b2d481e6ea4bfc9e437e9 Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Tue, 14 Jul 2026 11:48:32 +0530 Subject: [PATCH 09/12] fix: wire Rego evaluator into GovernanceEvaluator for per-step hooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The langchain adapter holds a reference to GovernanceEvaluator and calls evaluate_before_model / after_model / tool_call / after_tool on it for per-step lifecycle hooks. The Rego evaluator was only wired to UiPathGovernedRuntime._fire_before_agent / _fire_after_agent, so WASM rules scoped to before_model (e.g. SSN detection) never ran during actual agent execution and produced no trace spans. Fix: add set_rego_evaluator() to GovernanceEvaluator. Each per-step hook method now calls rego_evaluator.evaluate_*() after native evaluation when a Rego evaluator is wired. UiPathGovernedRuntime.__init__ calls evaluator.set_rego_evaluator(rego_evaluator) when both are provided, which mutates the shared object in place so the callback handler already holding the evaluator reference automatically picks up Rego evaluation without any changes to the framework adapter. BEFORE_AGENT / AFTER_AGENT are intentionally excluded — those remain owned by UiPathGovernedRuntime._fire_before_agent / _fire_after_agent to avoid double-firing. Generated with Claude Code Co-Authored-By: Claude --- .../runtime/governance/native/evaluator.py | 82 +++++++++++++++++-- src/uipath/runtime/governance/runtime.py | 10 +++ 2 files changed, 87 insertions(+), 5 deletions(-) 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/runtime.py b/src/uipath/runtime/governance/runtime.py index aaf6c78..92e9fd0 100644 --- a/src/uipath/runtime/governance/runtime.py +++ b/src/uipath/runtime/governance/runtime.py @@ -241,6 +241,16 @@ def __init__( 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 through native then Rego evaluators. From 6f86bdb3d8c0bd382912b371c90e0abd215a226f Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Tue, 14 Jul 2026 12:22:07 +0530 Subject: [PATCH 10/12] fix: normalize agent_input for BEFORE_MODEL Rego evaluation WASM rules like block_ssn_in_prompts check input.agent_input for user prompt text. The BEFORE_MODEL callback path only populates model_input, leaving agent_input empty and causing SSN/PII rules to silently miss violations. Fix: context_to_input copies model_input to agent_input when agent_input is empty and hook is BEFORE_MODEL. Both fields represent the same user content so this normalization is safe. Also adds INFO logging of raw WASM result for cloud diagnostics. Generated with Claude Code Co-Authored-By: Claude --- .../runtime/governance/rego/evaluator.py | 20 ++++++++- tests/test_rego_evaluator.py | 41 +++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/src/uipath/runtime/governance/rego/evaluator.py b/src/uipath/runtime/governance/rego/evaluator.py index dc12df5..72ac4f5 100644 --- a/src/uipath/runtime/governance/rego/evaluator.py +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -43,9 +43,17 @@ def context_to_input( """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 return { "hook": context.hook.value, - "agent_input": context.agent_input, + "agent_input": agent_input, "agent_output": context.agent_output, "model_input": context.model_input, "model_output": context.model_output, @@ -298,7 +306,15 @@ def evaluate(self, context: CheckContext) -> AuditRecord: try: plan = self._feature_plans.get(context.hook, []) - raw = engine.evaluate(context_to_input(context, feature_plan=plan)) + wasm_input = context_to_input(context, feature_plan=plan) + raw = engine.evaluate(wasm_input) + logger.info( + "Rego WASM hook=%s agent_input=%r model_input=%r raw=%r", + context.hook.value, + wasm_input.get("agent_input", "")[:80], + wasm_input.get("model_input", "")[:80], + raw, + ) except Exception as exc: logger.warning("Rego WASM evaluation failed for hook=%s: %s", context.hook.value, exc) return AuditRecord( diff --git a/tests/test_rego_evaluator.py b/tests/test_rego_evaluator.py index 4334bb3..d74dda8 100644 --- a/tests/test_rego_evaluator.py +++ b/tests/test_rego_evaluator.py @@ -113,6 +113,47 @@ def test_context_to_input_session_state_defaults() -> None: 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"] == "" + + # --------------------------------------------------------------------------- # _extract_wasm_from_bundle / _extract_data_json_from_bundle # --------------------------------------------------------------------------- From 514564687f0cd7d142b50d299eb3dfba5ebf2f2d Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Tue, 14 Jul 2026 12:24:39 +0530 Subject: [PATCH 11/12] fix: redact user content from WASM diagnostic log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The INFO log was including agent_input/model_input text which can contain PII (SSNs, etc.). Replace with field lengths — enough to confirm data was populated without exposing content. Generated with Claude Code Co-Authored-By: Claude --- src/uipath/runtime/governance/rego/evaluator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/uipath/runtime/governance/rego/evaluator.py b/src/uipath/runtime/governance/rego/evaluator.py index 72ac4f5..93bdf75 100644 --- a/src/uipath/runtime/governance/rego/evaluator.py +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -309,10 +309,10 @@ def evaluate(self, context: CheckContext) -> AuditRecord: wasm_input = context_to_input(context, feature_plan=plan) raw = engine.evaluate(wasm_input) logger.info( - "Rego WASM hook=%s agent_input=%r model_input=%r raw=%r", + "Rego WASM hook=%s agent_input_len=%d model_input_len=%d raw=%r", context.hook.value, - wasm_input.get("agent_input", "")[:80], - wasm_input.get("model_input", "")[:80], + len(wasm_input.get("agent_input", "")), + len(wasm_input.get("model_input", "")), raw, ) except Exception as exc: From 4ae60e8e3028d7d4e4e85c98474b50e024982d2d Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Tue, 14 Jul 2026 12:37:50 +0530 Subject: [PATCH 12/12] fix: synthesize messages for BEFORE_MODEL when messages is empty MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Production WASM rules (e.g. block_ssn_in_prompts) check input.messages[].content (plain string) rather than the flat model_input field. The callback-handler path never passes messages, so context.messages is always [] and rules silently miss violations. Fix: context_to_input now builds a synthetic [{"role": "user", "content": model_input}] entry when messages is empty for BEFORE_MODEL. Verified against the production bundle — the SSN rule now fires: fired_deny=["Message contains a value matching the SSN pattern (###-##-####)"], deny=true. Generated with Claude Code Co-Authored-By: Claude --- .../runtime/governance/rego/evaluator.py | 10 ++++- tests/test_rego_evaluator.py | 41 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/src/uipath/runtime/governance/rego/evaluator.py b/src/uipath/runtime/governance/rego/evaluator.py index 93bdf75..0c58a76 100644 --- a/src/uipath/runtime/governance/rego/evaluator.py +++ b/src/uipath/runtime/governance/rego/evaluator.py @@ -51,6 +51,14 @@ def context_to_input( 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, @@ -67,7 +75,7 @@ def context_to_input( "llm_calls": session.get("llm_calls", 0), }, "ring": context.ring, - "messages": context.messages, + "messages": messages, "features": features, } diff --git a/tests/test_rego_evaluator.py b/tests/test_rego_evaluator.py index d74dda8..5dd3117 100644 --- a/tests/test_rego_evaluator.py +++ b/tests/test_rego_evaluator.py @@ -154,6 +154,47 @@ def test_context_to_input_other_hooks_no_normalization() -> None: 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 # ---------------------------------------------------------------------------