feat(governance): policy backend client, YAML compiler, loader

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: aditik0303

src/uipath/runtime/governance/native/_yaml_to_index.py

@@ -0,0 +1,340 @@
+"""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.config import is_governance_enabled
+
+from uipath.runtime.governance.config import EnforcementMode, 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. The result lands in
+    the same module cache that ``get_policy_index()`` reads from;
+    ``get_policy_index()`` waits on this prefetch when it's in flight.
+
+    Idempotent: subsequent calls while the first is running are no-ops,
+    and calls after completion are no-ops. Skipped entirely when the
+    governance feature flag is OFF so no network call is made.
+    """
+    global _prefetch_event
+
+    if not is_governance_enabled():
+        return
+
+    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 - logged; first hook will retry sync
+            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 the governance feature flag is OFF, return an empty
+         PolicyIndex (cached). No network call.
+      2. If a prefetch (see :func:`prefetch_policy_index`) is in flight,
+         wait for it to complete (bounded by ``_PREFETCH_WAIT_SECONDS``).
+      3. Governance backend at ``api/v1/policy`` (one HTTP GET, cached).
+      4. 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
+
+    if not is_governance_enabled():
+        logger.info(
+            "Governance feature flag is OFF; returning empty PolicyIndex. "
+            "No rules will fire. Set EnablePythonGovernanceChecker=True to enable."
+        )
+        _policy_index = PolicyIndex()
+        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:
+            # Distinguish from the timeout path so production triage
+            # can tell "prefetch hung" from "prefetch returned empty"
+            # (auth failure, server error, parse failure).
+            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 — bounded
+    # by the HTTP timeout in the API client.
+    _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`.
+
+    Unknown values log a warning and leave the existing mode untouched.
+    """
+    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.
+
+    Applies the backend-supplied enforcement mode as a side effect.
+    Returns ``None`` when the backend skips/errors, when the YAML is
+    malformed, or when the resulting index has no rules — caller returns
+    an empty PolicyIndex in those cases.
+    """
+    start = time.perf_counter()
+    response = fetch_policy_response()
+    if response is None:
+        return None
+
+    # Apply the platform-controlled enforcement mode before building the
+    # index, so anything that reads ``get_enforcement_mode()`` during
+    # index compilation already sees the right value.
+    _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 - never let load break agent startup
+        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 _backend_base_url() -> str:
+    """Return the backend base URL for logging; imported lazily to avoid cycles."""
+    try:
+        from uipath.runtime.governance.native.backend_client import (
+            get_backend_base_url,
+        )
+
+        return get_backend_base_url()
+    except Exception:  # noqa: BLE001
+        return "backend"
+
+
+def _log_index_summary(index: PolicyIndex) -> None:
+    """Log summary of loaded policy index."""
+    # Count rules by hook
+    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.
+
+    Returns whatever the backend supplied on the most recent load.
+    Empty list if no index has been loaded yet or the backend yielded
+    no packs.
+    """
+    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.
+
+    Next call to ``get_policy_index()`` will refetch from the backend.
+    """
+    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