orchestrata_ultimate.orchestrator.engine

Reference › orchestrata_ultimate › orchestrator › engine

UOxEngine is the governance decision point of the Orchestrator volk (U:Ox): it answers "may this volk run this tool with these parameters?" with allow, deny, escalate or require-approval. The Claude tool interceptor (core.adapters.claude.tool_interceptor) is its importer, so every governed tool call reaches it through evaluate_policy. See the orchestrator package narrative for how this governance substrate sits under the intent and pipeline flows.

A PolicyContext carries the session id, volk id, tool name, parameters, an optional pre-parsed intent and a timestamp; its constructor rejects empty identifiers, and to_cache_key derives the cache key. The engine is built from a rules directory, a RedisStateManager, a GraphAuditTrail (graph_audit), a PrologBridge and a cache TTL (300 seconds by default), creating any that are missing and raising TypeError or ValueError for bad arguments. initialize starts the graph audit and loads the Prolog rules when governance.pl exists; shutdown and health_check complete the lifecycle.

evaluate_policy works in order: look up a cached decision in Redis (a failed lookup is logged and ignored); use the supplied intent or have _parse_intent_from_context build one through intent_parser; apply _evaluate_intent; record the decision with _record_audit; cache the GovernanceResult with the TTL. _evaluate_intent denies protected paths, then queries the Prolog predicate governance:evaluate_policy/3 with a term cleaned by _sanitize_prolog_term (restricted character set, length limit). The engine is fail-closed: an unparsable intent, a Prolog error, an unknown decision value or an empty result all yield DENY with a reason. start_session and end_session record volk sessions in Redis, and facts are handled by fact_manager.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

GovernanceDecision (enum)

Bases: Enum · Line 58

Policy evaluation decision.

PolicyContext (dataclass)

Decorators: dataclass · Line 79

Context for policy evaluation.

GovernanceResult (dataclass)

Decorators: dataclass, traces(spec='000', story='US-007', reqs=['REQ-000-061']) · Line 144

Result of policy evaluation.

UOxEngine

Decorators: traces(spec='001', story='US-003', reqs=['REQ-001-US-003-002']) · Line 164

U:Ox governance engine.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/orchestrator/engine.py

# mypy: ignore-errors
"""U:Ox governance engine with Prolog-based policy evaluation.

Provides deterministic governance decisions using:
- Prolog rules for policy evaluation
- Redis for hot state and caching
- Graph for audit trail persistence
- Intent parsing for tool call analysis

Based on DESIGN_DECISIONS.md and INFRASTRUCTURE.md specifications.
"""

import hashlib
import json
import logging
import re
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from pathlib import Path
from typing import Any, cast

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.enums import FactScope
from orchestrata_ultimate.core.paths import get_prolog_dir
from orchestrata_ultimate.infra.prolog_bridge import PrologBridge
from orchestrata_ultimate.infra.redis_state import (
    RedisStateManager,
    SessionState,
)
from orchestrata_ultimate.orchestrator.fact_manager import FactManager
from orchestrata_ultimate.orchestrator.graph_audit import (
    DecisionOutcome,
    GraphAuditTrail,
)
from orchestrata_ultimate.orchestrator.graph_audit import (
    PolicyDecision as GraphPolicyDecision,
)
from orchestrata_ultimate.orchestrator.intent_parser import (
    Action,
    Intent,
    IntentParser,
    ResourceType,
)

# Spec: 001 - Orchestration Kernel


# Configuration constants
DEFAULT_CACHE_TTL_SECONDS: int = 300
MAX_PARAMETER_SIZE_BYTES: int = 1024 * 1024  # 1MB limit for parameters
MAX_PROLOG_TERM_LENGTH: int = (
    4096  # Truncation limit for sanitised Prolog terms; guards against DoS via oversized inputs
)


class GovernanceDecision(Enum):
    """Policy evaluation decision."""

    ALLOW = "allow"
    DENY = "deny"
    ESCALATE = "escalate"
    REQUIRE_APPROVAL = "require_approval"


def _utc_now() -> datetime:
    """Get current UTC time with timezone info."""
    return datetime.now(UTC)


@traces(spec="001", story="US-003", reqs=["REQ-001-US-003-005"])
def _sanitize_for_cache_key(value: str) -> str:
    """Sanitize string for use in cache key - only alphanumeric and limited chars."""
    return re.sub(r"[^a-zA-Z0-9_\-.]", "_", value)[:256]


@dataclass
class PolicyContext:
    """Context for policy evaluation.

    Contains all information needed to evaluate a governance decision:
    - Session and volk identification
    - Tool and parameters
    - Optional pre-parsed intent
    - Timestamp for audit trail
    """

    session_id: str
    volk_id: str
    tool: str
    parameters: dict[str, Any]
    intent: Intent | None = None
    timestamp: datetime = field(default_factory=_utc_now)

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-017"])
    def __post_init__(self) -> None:
        """Validate context after initialization."""
        if not self.session_id or not isinstance(self.session_id, str):
            raise ValueError("session_id must be a non-empty string")
        if not self.volk_id or not isinstance(self.volk_id, str):
            raise ValueError("volk_id must be a non-empty string")
        if not self.tool or not isinstance(self.tool, str):
            raise ValueError("tool must be a non-empty string")
        if not isinstance(self.parameters, dict):
            raise ValueError("parameters must be a dictionary")
        # Validate parameter size to prevent DoS
        try:
            params_json = json.dumps(self.parameters, default=str)
            if len(params_json.encode("utf-8")) > MAX_PARAMETER_SIZE_BYTES:
                raise ValueError(
                    f"parameters exceed maximum size of {MAX_PARAMETER_SIZE_BYTES} bytes"
                )
        except (TypeError, ValueError) as e:
            if "exceed maximum size" in str(e):
                raise
            # Allow non-JSON-serializable parameters but log warning

    @traces(spec="001", story="US-003", reqs=["REQ-001-US-003-005"])
    def to_cache_key(self) -> str:
        """Generate secure cache key for this context.

        Uses SHA256 hash to prevent cache key injection attacks.
        All inputs are sanitized before hashing.
        """
        # Sanitize inputs before hashing
        safe_volk = _sanitize_for_cache_key(self.volk_id)
        safe_tool = _sanitize_for_cache_key(self.tool)

        # Create deterministic key from context using secure hash
        try:
            params_str = json.dumps(self.parameters, sort_keys=True, default=str)
        except TypeError, ValueError:
            # Fallback to empty dict representation for non-serializable params
            params_str = "{}"

        # Use full SHA256 hash for collision resistance
        key_data = f"{safe_volk}:{safe_tool}:{params_str}"
        return hashlib.sha256(key_data.encode("utf-8")).hexdigest()


@dataclass
@traces(spec="000", story="US-007", reqs=["REQ-000-061"])
class GovernanceResult:
    """Result of policy evaluation.

    Contains:
    - Decision (ALLOW/DENY/ESCALATE)
    - Reason for decision
    - Parsed intent
    - Audit ID for traceability
    - Whether result was cached
    """

    decision: GovernanceDecision
    reason: str
    intent: Intent | None = None
    audit_id: str | None = None
    cached: bool = False
    metadata: dict[str, Any] = field(default_factory=dict)


@traces(spec="001", story="US-003", reqs=["REQ-001-US-003-002"])
class UOxEngine:
    """U:Ox governance engine.

    Coordinates policy evaluation using:
    - Prolog rules for deterministic decisions
    - Intent parsing for semantic analysis
    - Redis for hot state and caching
    - Graph for audit trail
    - Fact management for dynamic rules
    """

    # Protected paths that require special approval
    PROTECTED_PATHS = {
        "/etc/",
        "/sys/",
        "/proc/",
        "/dev/",
        "/boot/",
        "/root/",
    }

    @staticmethod
    def _validate_cache_ttl(cache_ttl_seconds: int) -> None:
        """Raise ValueError if cache_ttl_seconds is not a non-negative integer."""
        if not isinstance(cache_ttl_seconds, int) or cache_ttl_seconds < 0:
            raise ValueError("cache_ttl_seconds must be a non-negative integer")

    @staticmethod
    def _validate_dependencies(
        redis: RedisStateManager | None,
        graph: GraphAuditTrail | None,
        prolog: PrologBridge | None,
    ) -> None:
        """Type-check optional dependency objects via duck typing."""
        if redis is not None:
            if not hasattr(redis, "client") or not hasattr(redis, "create_session"):
                raise TypeError(
                    f"redis must be RedisStateManager-compatible, got {type(redis).__name__}"
                )
        if graph is not None:
            if not hasattr(graph, "start") or not hasattr(graph, "log_event"):
                raise TypeError(
                    f"graph must be GraphAuditTrail-compatible, got {type(graph).__name__}"
                )
        if prolog is not None:
            if not hasattr(prolog, "query") or not hasattr(prolog, "assert_fact"):
                raise TypeError(
                    f"prolog must be PrologBridge-compatible, got {type(prolog).__name__}"
                )

    @traces(spec="001", story="US-004", reqs=["REQ-001-US-004-011"])
    def _init_fact_manager(self) -> FactManager:
        """Build a FactManager wired to this engine's Prolog bridge."""
        return FactManager(
            assert_fn=lambda pred, args: self.prolog.assert_fact(
                pred, list(args), FactScope.SESSION
            ),
            retract_fn=self.prolog.retract_fact,
        )

    def __init__(
        self,
        rules_dir: Path | None = None,
        redis: RedisStateManager | None = None,
        graph: GraphAuditTrail | None = None,
        prolog: PrologBridge | None = None,
        cache_ttl_seconds: int = DEFAULT_CACHE_TTL_SECONDS,
    ) -> None:
        """Initialize governance engine.

        Args:
            rules_dir: Directory containing Prolog rules
            redis: Redis state manager (created if not provided)
            graph: Graph audit trail (created if not provided)
            prolog: Prolog bridge (created if not provided)
            cache_ttl_seconds: TTL for cached governance decisions

        Raises:
            TypeError: If provided dependencies are of wrong type
            ValueError: If cache_ttl_seconds is invalid
        """
        self._logger = logging.getLogger(__name__)

        self._validate_cache_ttl(cache_ttl_seconds)
        self._cache_ttl_seconds = cache_ttl_seconds

        self.rules_dir = rules_dir or get_prolog_dir()

        self._validate_dependencies(redis, graph, prolog)

        self.redis = redis or RedisStateManager()
        self.graph = graph or GraphAuditTrail()
        self.prolog = prolog or PrologBridge(rules_dir=self.rules_dir)

        self.intent_parser = IntentParser()
        self.fact_manager = self._init_fact_manager()

        self._initialized = False

    async def initialize(self) -> None:
        """Initialize the engine and load rules."""
        if self._initialized:
            return

        # Start graph audit trail - use interface method with async support
        try:
            await self.graph.start()
        except AttributeError:
            self._logger.warning("Graph audit trail does not have start method")
        except Exception as e:
            self._logger.error(f"Failed to start graph audit trail: {e}")

        # Load Prolog rules
        governance_rules = self.rules_dir / "governance.pl"
        if governance_rules.exists():
            self.prolog.load_all_rules()

        self._initialized = True

    async def shutdown(self) -> None:
        """Shutdown the engine and cleanup resources."""
        if not self._initialized:
            return

        # Close connections - use interface methods with proper error handling
        try:
            self.redis.close()
        except AttributeError:
            self._logger.debug("Redis state manager does not have close method")
        except Exception as e:
            self._logger.error(f"Error closing Redis connection: {e}")

        try:
            await self.graph.close()
        except AttributeError:
            self._logger.debug("Graph audit trail does not have close method")
        except Exception as e:
            self._logger.error(f"Error closing graph connection: {e}")

        self._initialized = False

    async def health_check(self) -> dict[str, Any]:
        """Check health of all engine components.

        Returns:
            Dictionary with health status of each component and error details.
        """
        health: dict[str, Any] = {
            "initialized": self._initialized,
            "redis": False,
            "redis_error": None,
            "graph": False,
            "graph_error": None,
            "prolog": False,
            "prolog_error": None,
        }

        # Check Redis
        try:
            self.redis.client.ping()
            health["redis"] = True
        except Exception as e:
            self._logger.error(f"Redis health check failed: {e}")
            health["redis_error"] = str(e)

        # Check graph - use async interface method
        try:
            repo = getattr(self.graph, "_repo", None)
            if repo is not None:
                await repo.verify_connectivity()
                health["graph"] = True
            else:
                health["graph_error"] = "No driver available"
        except Exception as e:
            self._logger.error(f"Graph health check failed: {e}")
            health["graph_error"] = str(e)

        # Check Prolog - verify it can execute a simple query
        try:
            if self.prolog is not None:
                # Actually verify Prolog is working with a simple query
                result = self.prolog.query_bool("true")
                health["prolog"] = result
                if not result:
                    health["prolog_error"] = (
                        "Prolog query returned false"  # pragma: no cover
                    )
        except Exception as e:
            self._logger.error(f"Prolog health check failed: {e}")
            health["prolog_error"] = str(e)

        return health

    @traces(
        spec="001", story="US-003", reqs=["REQ-001-US-003-005", "REQ-001-US-003-009"]
    )
    async def evaluate_policy(self, context: PolicyContext) -> GovernanceResult:
        """Evaluate policy for given context.

        Process:
        1. Check cache for previous decision
        2. Parse intent if not provided
        3. Evaluate against Prolog rules
        4. Apply built-in safety checks
        5. Cache result
        6. Record audit trail

        Args:
            context: Policy evaluation context

        Returns:
            GovernanceResult with decision and reasoning
        """
        # Check cache first
        cache_key = f"governance:{context.to_cache_key()}"
        try:
            # get() is typed as redis-py's shared ResponseT (Awaitable[Any] |
            # Any) because the same command mixin backs both sync and async
            # clients; this sync client always returns str or None.
            cached = cast("str | None", self.redis.client.get(cache_key))
            if cached:
                result = self._deserialize_result(cached)
                result.cached = True
                return result
        except Exception as e:
            # Cache miss is not fatal - continue with evaluation
            self._logger.warning(f"Cache lookup failed: {e}")

        # Parse intent if not provided
        intent = context.intent
        if intent is None:
            intent = self._parse_intent_from_context(context)

        # Handle case where intent parsing fails
        if intent is None:
            self._logger.error("Failed to parse intent from context")
            return GovernanceResult(
                decision=GovernanceDecision.DENY,
                reason="Failed to parse intent - denied for safety",
                intent=None,
                audit_id=None,
                cached=False,
            )

        # Evaluate policy
        decision, reason = self._evaluate_intent(intent, context)

        # Record audit trail with async support
        audit_id = await self._record_audit(context, decision, reason, intent)

        # Create result
        result = GovernanceResult(
            decision=decision,
            reason=reason,
            intent=intent,
            audit_id=audit_id,
            cached=False,
        )

        # Cache result with configurable TTL
        try:
            self.redis.client.setex(
                cache_key, self._cache_ttl_seconds, self._serialize_result(result)
            )
        except Exception as e:
            # Cache write failure is not fatal
            self._logger.warning(f"Cache write failed: {e}")

        return result

    @traces(spec="001", story="US-003", reqs=["REQ-001-US-003-017"])
    def _parse_intent_from_context(self, context: PolicyContext) -> Intent:
        """Parse intent from context parameters.

        Args:
            context: Policy context

        Returns:
            Parsed Intent object
        """
        # Extract action from tool name
        tool_to_action = {
            "Read": Action.READ,
            "Write": Action.WRITE,
            "Edit": Action.MODIFY,
            "Bash": Action.EXECUTE,
            "Delete": Action.DELETE,
        }

        action = tool_to_action.get(context.tool, Action.READ)

        # Extract resource path and type
        resource_path = "/"
        resource_type = ResourceType.SYSTEM

        if "file_path" in context.parameters:
            resource_path = context.parameters["file_path"]
            resource_type = ResourceType.FILE
        elif "path" in context.parameters:
            resource_path = context.parameters["path"]
            resource_type = ResourceType.DIRECTORY
        elif "command" in context.parameters:
            resource_path = context.parameters["command"]
            resource_type = ResourceType.COMMAND

        return Intent(
            action=action,
            resource_type=resource_type,
            resource_path=resource_path,
        )

    def _sanitize_prolog_term(self, term: str) -> str:
        """Sanitize a Prolog term to prevent injection attacks.

        Args:
            term: Raw Prolog term string

        Returns:
            Sanitized term safe for use in queries
        """
        # Remove potentially dangerous characters
        # Allow only alphanumeric, underscores, parentheses, brackets, commas, quotes, and spaces
        # Square brackets [] are needed for Prolog lists
        sanitized = re.sub(r'[^\w\s(),\[\]\'"\-./:]', "", term)
        # Limit length to prevent DoS
        return sanitized[:MAX_PROLOG_TERM_LENGTH]

    @traces(
        spec="001", story="US-003", reqs=["REQ-001-US-003-001", "REQ-001-US-003-008"]
    )
    def _evaluate_intent(
        self, intent: Intent, context: PolicyContext
    ) -> tuple[GovernanceDecision, str]:
        """Evaluate intent against governance rules.

        Args:
            intent: Parsed intent
            context: Policy context

        Returns:
            Tuple of (decision, reason)
        """
        # Check protected paths
        if self._is_protected_path(intent.resource_path):
            return (GovernanceDecision.DENY, f"Protected path: {intent.resource_path}")

        # Query Prolog rules - FAIL CLOSED (DENY) on any error
        try:
            # Sanitize the Prolog term to prevent injection
            prolog_term = self._sanitize_prolog_term(intent.to_prolog())
            # Use governance module prefix for proper module resolution
            query = f"governance:evaluate_policy({prolog_term}, Decision, Reason)"
            query_result = self.prolog.query(query)
            results = query_result.solutions  # QueryResult.solutions contains the list

            if results:
                result = results[0]
                decision_str = result.get("Decision", "deny")
                reason = result.get("Reason", "Policy evaluation succeeded")

                try:
                    decision = GovernanceDecision(decision_str)
                    return (decision, reason)
                except ValueError:
                    self._logger.error(f"Invalid decision from Prolog: {decision_str}")
                    return (
                        GovernanceDecision.DENY,
                        f"Invalid policy decision: {decision_str}",
                    )

            # No results from Prolog - DENY by default
            self._logger.warning("No policy evaluation results from Prolog")
            return (GovernanceDecision.DENY, "No policy matched - denied by default")

        except Exception as e:
            # CRITICAL: Policy evaluation failure = DENY (fail closed)
            self._logger.error(f"Policy evaluation failed: {e}")
            return (GovernanceDecision.DENY, f"Policy system error: {e}")

    @traces(spec="001", story="US-003", reqs=["REQ-001-US-003-013"])
    def _is_protected_path(self, path: str) -> bool:
        """Check if path is protected.

        Args:
            path: Path to check

        Returns:
            True if path is protected
        """
        return any(path.startswith(protected) for protected in self.PROTECTED_PATHS)

    @traces(spec="015", story="US-001", reqs=["REQ-015-115"])
    async def _record_audit(
        self,
        context: PolicyContext,
        decision: GovernanceDecision,
        reason: str,
        intent: Intent | None,
    ) -> str | None:
        """Record governance decision in audit trail.

        Args:
            context: Policy context
            decision: Governance decision
            reason: Decision reason
            intent: Parsed intent

        Returns:
            Audit ID if recorded, None otherwise
        """
        # Map GovernanceDecision to DecisionOutcome
        outcome_map = {
            GovernanceDecision.ALLOW: DecisionOutcome.ALLOWED,
            GovernanceDecision.DENY: DecisionOutcome.DENIED,
            GovernanceDecision.ESCALATE: DecisionOutcome.ESCALATED,
            GovernanceDecision.REQUIRE_APPROVAL: DecisionOutcome.PENDING,
        }

        decision_id = str(uuid.uuid4())

        policy_decision = GraphPolicyDecision(
            decision_id=decision_id,
            session_id=context.session_id,
            volk_id=context.volk_id,
            intent=intent.to_string() if intent else "",
            tool=context.tool,
            outcome=outcome_map[decision],
            reason=reason,
            timestamp=context.timestamp,
            cached=False,
        )

        try:
            # Record in graph with async support
            await self.graph.log_policy_decision(policy_decision)
            return decision_id
        except Exception as e:
            # Log the failure - audit is critical for compliance
            self._logger.critical(
                f"AUDIT FAILURE: Could not record decision {decision_id}: {e}"
            )

            # Queue for retry in Redis
            try:
                retry_key = f"audit:retry:{decision_id}"
                retry_data = {
                    "decision_id": decision_id,
                    "session_id": context.session_id,
                    "volk_id": context.volk_id,
                    "intent": intent.to_string() if intent else "",
                    "tool": context.tool,
                    "outcome": outcome_map[decision].value,
                    "reason": reason,
                    "timestamp": context.timestamp.isoformat(),
                    "retry_count": 0,
                    "error": str(e),
                }
                self.redis.client.setex(
                    retry_key, 86400, json.dumps(retry_data)
                )  # 24h TTL
                self._logger.warning(f"Queued audit {decision_id} for retry")
            except Exception as retry_err:
                self._logger.critical(
                    f"CRITICAL: Cannot queue audit for retry: {retry_err}"
                )

            return decision_id  # Return ID even if audit failed - decision was made

    @traces(
        spec="001",
        story="US-004",
        reqs=["REQ-001-US-004-001", "REQ-001-US-004-005"],
    )
    def start_session(self, session_id: str, volk_id: str, goals: list[str]) -> bool:
        """Start a new session.

        Args:
            session_id: Session identifier
            volk_id: Volk identifier
            goals: Session goals

        Returns:
            True if session started successfully

        Raises:
            RuntimeError: If session creation fails
        """
        try:
            session = SessionState(
                session_id=session_id,
                volk_id=volk_id,
                status="active",
                goals=goals,
                context={},
                started_at=_utc_now(),
                last_activity=_utc_now(),
            )

            result = self.redis.create_session(session)
            if not result:
                self._logger.error(f"Failed to create session {session_id}")
            return result
        except Exception as e:
            self._logger.error(f"Session start failed for {session_id}: {e}")
            raise RuntimeError(f"Failed to start session: {e}") from e

    @traces(
        spec="001",
        story="US-004",
        reqs=["REQ-001-US-004-023"],
    )
    def end_session(self, session_id: str, volk_id: str, outcome: str) -> bool:
        """End a session.

        Args:
            session_id: Session identifier
            volk_id: Volk identifier
            outcome: Session outcome

        Returns:
            True if session ended successfully

        Raises:
            RuntimeError: If session deletion fails
        """
        try:
            result = self.redis.delete_session(session_id)
            if not result:
                self._logger.warning(
                    f"Session {session_id} not found or already deleted"
                )
            return result
        except Exception as e:
            self._logger.error(f"Session end failed for {session_id}: {e}")
            raise RuntimeError(f"Failed to end session: {e}") from e

    @traces(
        spec="001",
        story="US-010",
        reqs=["REQ-001-US-003-005", "REQ-001-US-010-016"],
    )
    def _serialize_result(self, result: GovernanceResult) -> str:
        """Serialize result for caching.

        Args:
            result: Governance result

        Returns:
            JSON string
        """
        return json.dumps(
            {
                "decision": result.decision.value,
                "reason": result.reason,
                "intent": result.intent.to_string() if result.intent else None,
                "audit_id": result.audit_id,
                "metadata": result.metadata,
            }
        )

    @traces(
        spec="001",
        story="US-010",
        reqs=["REQ-001-US-003-005", "REQ-001-US-010-016"],
    )
    def _deserialize_result(self, data: str) -> GovernanceResult:
        """Deserialize cached result.

        Args:
            data: JSON string

        Returns:
            Governance result
        """
        obj = json.loads(data)

        # Parse intent if present
        intent = None
        if obj.get("intent"):
            intent = self.intent_parser.parse(obj["intent"])

        return GovernanceResult(
            decision=GovernanceDecision(obj["decision"]),
            reason=obj["reason"],
            intent=intent,
            audit_id=obj.get("audit_id"),
            metadata=obj.get("metadata", {}),
        )