From c895b3728734d3a9c6f44627ac14b045f3fce62b Mon Sep 17 00:00:00 2001 From: Evan Bijoy Date: Thu, 9 Jul 2026 15:20:36 +0530 Subject: [PATCH 1/3] 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 2/3] 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 3/3] 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