agent_hypervisor 3.1.0__py3-none-any.whl

hypervisor/saga/schema.py

@@ -0,0 +1,244 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# Public Preview — basic implementation
+"""
+JSON Schema validation for Saga DSL definitions.
+
+Validates saga definitions at parse time with clear error messages
+for missing fields, invalid types, and constraint violations.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import jsonschema
+
+# Valid action type prefixes for step action_ids
+VALID_ACTION_PREFIXES = (
+    "model.", "data.", "deploy.", "validate.", "notify.",
+    "infra.", "security.", "monitor.", "config.", "test.",
+)
+
+SAGA_STEP_SCHEMA: dict[str, Any] = {
+    "type": "object",
+    "required": ["id", "action_id", "agent"],
+    "properties": {
+        "id": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Unique step identifier",
+        },
+        "action_id": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Action type (e.g. 'model.validate', 'deploy.k8s')",
+        },
+        "agent": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Agent DID or identifier",
+        },
+        "execute_api": {
+            "type": "string",
+            "description": "API endpoint for execution",
+        },
+        "undo_api": {
+            "type": ["string", "null"],
+            "description": "API endpoint for compensation/rollback",
+        },
+        "timeout": {
+            "type": "integer",
+            "minimum": 1,
+            "maximum": 86400,
+            "description": "Timeout in seconds (1–86400)",
+        },
+        "retries": {
+            "type": "integer",
+            "minimum": 0,
+            "maximum": 10,
+            "description": "Max retries (0–10)",
+        },
+        "checkpoint_goal": {
+            "type": ["string", "null"],
+            "description": "Semantic checkpoint goal",
+        },
+        "depends_on": {
+            "type": "array",
+            "items": {"type": "string", "minLength": 1},
+            "uniqueItems": True,
+            "description": "Step IDs this step depends on",
+        },
+    },
+    "additionalProperties": False,
+}
+
+SAGA_DEFINITION_SCHEMA: dict[str, Any] = {
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "title": "SagaDefinition",
+    "description": "Schema for saga DSL definitions",
+    "type": "object",
+    "required": ["name", "session_id", "steps"],
+    "properties": {
+        "name": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Saga name",
+        },
+        "session_id": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Session identifier",
+        },
+        "saga_id": {
+            "type": "string",
+            "description": "Optional saga identifier",
+        },
+        "steps": {
+            "type": "array",
+            "minItems": 1,
+            "items": SAGA_STEP_SCHEMA,
+            "description": "Ordered list of saga steps",
+        },
+        "fan_out": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "policy": {"type": "string"},
+                    "branches": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                    },
+                },
+            },
+            "description": "Fan-out groups (Public Preview: ignored)",
+        },
+        "metadata": {
+            "type": "object",
+            "description": "Arbitrary metadata",
+        },
+    },
+    "additionalProperties": False,
+}
+
+
+class SagaSchemaValidator:
+    """Validates saga definitions against JSON schema and semantic rules."""
+
+    def __init__(self) -> None:
+        self._validator = jsonschema.Draft202012Validator(SAGA_DEFINITION_SCHEMA)
+
+    def validate(self, definition: dict[str, Any]) -> list[str]:
+        """Validate definition and return list of error messages (empty = valid).
+
+        Performs both JSON schema validation and semantic checks:
+        - Required fields and types
+        - Step structure constraints
+        - Unique step IDs
+        - Valid action type prefixes
+        - Timeout and retry ranges
+        - Compensation requirements
+        - Step dependency references
+        """
+        errors: list[str] = []
+
+        # JSON schema validation
+        for error in sorted(self._validator.iter_errors(definition), key=lambda e: list(e.path)):
+            path = ".".join(str(p) for p in error.absolute_path) or "(root)"
+            errors.append(f"[{path}] {error.message}")
+
+        # Semantic checks only if basic structure is valid
+        if not errors and isinstance(definition.get("steps"), list):
+            errors.extend(self._check_semantic_rules(definition))
+
+        return errors
+
+    def validate_or_raise(self, definition: dict[str, Any]) -> None:
+        """Validate and raise SagaSchemaError if invalid."""
+        errors = self.validate(definition)
+        if errors:
+            raise SagaSchemaError(
+                f"Saga definition has {len(errors)} validation error(s):\n"
+                + "\n".join(f"  - {e}" for e in errors),
+                errors=errors,
+            )
+
+    def _check_semantic_rules(self, definition: dict[str, Any]) -> list[str]:
+        errors: list[str] = []
+        steps = definition["steps"]
+        step_ids: set[str] = set()
+
+        for i, step in enumerate(steps):
+            sid = step.get("id", f"<index {i}>")
+
+            # Duplicate step IDs
+            if sid in step_ids:
+                errors.append(f"Duplicate step ID: '{sid}'")
+            step_ids.add(sid)
+
+            # Action type prefix validation
+            action_id = step.get("action_id", "")
+            if action_id and not any(action_id.startswith(p) for p in VALID_ACTION_PREFIXES):
+                errors.append(
+                    f"Step '{sid}': action_id '{action_id}' does not start with a "
+                    f"valid prefix ({', '.join(VALID_ACTION_PREFIXES)})"
+                )
+
+            # Compensation requirement: every step should have undo_api
+            if step.get("undo_api") is None:
+                errors.append(
+                    f"Step '{sid}': missing 'undo_api' — every action should have a compensation endpoint"
+                )
+
+        # Dependency validation
+        for step in steps:
+            for dep in step.get("depends_on", []):
+                if dep not in step_ids:
+                    errors.append(
+                        f"Step '{step['id']}': depends_on references unknown step '{dep}'"
+                    )
+
+        # Circular dependency detection
+        errors.extend(self._check_circular_deps(steps))
+
+        return errors
+
+    def _check_circular_deps(self, steps: list[dict[str, Any]]) -> list[str]:
+        """Detect circular dependencies via DFS."""
+        graph: dict[str, list[str]] = {}
+        for step in steps:
+            sid = step.get("id", "")
+            graph[sid] = step.get("depends_on", [])
+
+        visited: set[str] = set()
+        in_stack: set[str] = set()
+        errors: list[str] = []
+
+        def dfs(node: str) -> bool:
+            if node in in_stack:
+                errors.append(f"Circular dependency detected involving step '{node}'")
+                return True
+            if node in visited:
+                return False
+            visited.add(node)
+            in_stack.add(node)
+            for dep in graph.get(node, []):
+                if dfs(dep):
+                    return True
+            in_stack.discard(node)
+            return False
+
+        for sid in graph:
+            if sid not in visited:
+                dfs(sid)
+
+        return errors
+
+
+class SagaSchemaError(Exception):
+    """Raised when a saga definition fails schema validation."""
+
+    def __init__(self, message: str, errors: list[str] | None = None) -> None:
+        super().__init__(message)
+        self.errors = errors or []

hypervisor/saga/state_machine.py

@@ -0,0 +1,157 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Saga State Machine
+
+Formal state tracking for individual saga steps and overall saga lifecycle,
+with persistence support for crash recovery.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from enum import Enum
+from typing import Any
+
+
+class StepState(str, Enum):
+    """State of an individual saga step."""
+
+    PENDING = "pending"
+    EXECUTING = "executing"
+    COMMITTED = "committed"
+    COMPENSATING = "compensating"
+    COMPENSATED = "compensated"
+    COMPENSATION_FAILED = "compensation_failed"
+    FAILED = "failed"
+
+
+class SagaState(str, Enum):
+    """State of the overall saga."""
+
+    RUNNING = "running"
+    COMPENSATING = "compensating"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    ESCALATED = "escalated"
+
+
+# Valid state transitions
+STEP_TRANSITIONS: dict[StepState, set[StepState]] = {
+    StepState.PENDING: {StepState.EXECUTING},
+    StepState.EXECUTING: {StepState.COMMITTED, StepState.FAILED},
+    StepState.COMMITTED: {StepState.COMPENSATING},
+    StepState.COMPENSATING: {StepState.COMPENSATED, StepState.COMPENSATION_FAILED},
+    StepState.COMPENSATED: set(),
+    StepState.COMPENSATION_FAILED: set(),
+    StepState.FAILED: set(),
+}
+
+SAGA_TRANSITIONS: dict[SagaState, set[SagaState]] = {
+    SagaState.RUNNING: {SagaState.COMPENSATING, SagaState.COMPLETED, SagaState.FAILED},
+    SagaState.COMPENSATING: {SagaState.COMPLETED, SagaState.FAILED, SagaState.ESCALATED},
+    SagaState.COMPLETED: set(),
+    SagaState.FAILED: set(),
+    SagaState.ESCALATED: set(),
+}
+
+
+@dataclass
+class SagaStep:
+    """A single step in a saga."""
+
+    step_id: str
+    action_id: str
+    agent_did: str
+    execute_api: str
+    undo_api: str | None = None
+    state: StepState = StepState.PENDING
+    execute_result: Any | None = None
+    compensation_result: Any | None = None
+    error: str | None = None
+    started_at: datetime | None = None
+    completed_at: datetime | None = None
+    timeout_seconds: int = 300
+    max_retries: int = 0
+    retry_count: int = 0
+
+    def transition(self, new_state: StepState) -> None:
+        """Transition to a new state, enforcing valid transitions."""
+        allowed = STEP_TRANSITIONS.get(self.state, set())
+        if new_state not in allowed:
+            raise SagaStateError(
+                f"Invalid step transition: {self.state.value} → {new_state.value}. "
+                f"Allowed: {[s.value for s in allowed]}"
+            )
+        self.state = new_state
+        now = datetime.now(UTC)
+        if new_state == StepState.EXECUTING:
+            self.started_at = now
+        elif new_state in (
+            StepState.COMMITTED,
+            StepState.COMPENSATED,
+            StepState.COMPENSATION_FAILED,
+            StepState.FAILED,
+        ):
+            self.completed_at = now
+
+
+@dataclass
+class Saga:
+    """A saga consisting of ordered steps."""
+
+    saga_id: str
+    session_id: str
+    steps: list[SagaStep] = field(default_factory=list)
+    state: SagaState = SagaState.RUNNING
+    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
+    completed_at: datetime | None = None
+    error: str | None = None
+
+    def transition(self, new_state: SagaState) -> None:
+        """Transition the saga to a new state."""
+        allowed = SAGA_TRANSITIONS.get(self.state, set())
+        if new_state not in allowed:
+            raise SagaStateError(
+                f"Invalid saga transition: {self.state.value} → {new_state.value}. "
+                f"Allowed: {[s.value for s in allowed]}"
+            )
+        self.state = new_state
+        if new_state in (SagaState.COMPLETED, SagaState.FAILED, SagaState.ESCALATED):
+            self.completed_at = datetime.now(UTC)
+
+    @property
+    def committed_steps(self) -> list[SagaStep]:
+        """Steps that completed execution (need compensation on rollback)."""
+        return [s for s in self.steps if s.state == StepState.COMMITTED]
+
+    @property
+    def committed_steps_reversed(self) -> list[SagaStep]:
+        """Committed steps in reverse order for rollback."""
+        return list(reversed(self.committed_steps))
+
+    def to_dict(self) -> dict:
+        """Serialize for VFS persistence."""
+        return {
+            "saga_id": self.saga_id,
+            "session_id": self.session_id,
+            "state": self.state.value,
+            "created_at": self.created_at.isoformat(),
+            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
+            "error": self.error,
+            "steps": [
+                {
+                    "step_id": s.step_id,
+                    "action_id": s.action_id,
+                    "agent_did": s.agent_did,
+                    "state": s.state.value,
+                    "error": s.error,
+                }
+                for s in self.steps
+            ],
+        }
+
+
+class SagaStateError(Exception):
+    """Raised for invalid saga state transitions."""

hypervisor/security/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Security module — rate limiting, kill switch, and agent protection."""
+
+from hypervisor.security.kill_switch import KillResult, KillSwitch
+from hypervisor.security.rate_limiter import AgentRateLimiter, RateLimitExceeded
+
+__all__ = [
+    "AgentRateLimiter",
+    "RateLimitExceeded",
+    "KillSwitch",
+    "KillResult",
+]

hypervisor/security/kill_switch.py

@@ -0,0 +1,200 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Kill Switch — agent termination with optional handoff.
+
+Terminates agent processes via registered callbacks and hands off
+in-flight saga steps to a substitute agent when one is available.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from enum import Enum
+
+_logger = logging.getLogger(__name__)
+
+
+class KillReason(str, Enum):
+    """Why an agent was killed."""
+
+    BEHAVIORAL_DRIFT = "behavioral_drift"
+    RATE_LIMIT = "rate_limit"
+    RING_BREACH = "ring_breach"
+    MANUAL = "manual"
+    QUARANTINE_TIMEOUT = "quarantine_timeout"
+    SESSION_TIMEOUT = "session_timeout"
+
+
+class HandoffStatus(str, Enum):
+    """Status of a saga step handoff."""
+
+    PENDING = "pending"
+    HANDED_OFF = "handed_off"
+    FAILED = "failed"
+    COMPENSATED = "compensated"
+
+
+@dataclass
+class StepHandoff:
+    """A saga step being handed off to a substitute or compensated."""
+
+    step_id: str
+    saga_id: str
+    from_agent: str
+    to_agent: str | None = None
+    status: HandoffStatus = HandoffStatus.COMPENSATED
+
+
+@dataclass
+class KillResult:
+    """Result of a kill switch operation."""
+
+    kill_id: str = field(default_factory=lambda: f"kill:{uuid.uuid4().hex[:8]}")
+    agent_did: str = ""
+    session_id: str = ""
+    reason: KillReason = KillReason.MANUAL
+    timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
+    handoffs: list[StepHandoff] = field(default_factory=list)
+    handoff_success_count: int = 0
+    compensation_triggered: bool = False
+    terminated: bool = False
+    details: str = ""
+
+
+class KillSwitch:
+    """
+    Kill switch with agent process registry and handoff support.
+
+    Agents register termination callbacks via ``register_agent``.  When
+    ``kill`` is called the switch hands in-flight saga steps to a
+    registered substitute (if any) and then invokes the termination
+    callback to stop the agent process.
+    """
+
+    def __init__(self) -> None:
+        self._kill_history: list[KillResult] = []
+        self._substitutes: dict[str, list[str]] = {}
+        self._agents: dict[str, Callable[[], None]] = {}
+
+    # ── Agent process registry ─────────────────────────────────────
+
+    def register_agent(
+        self, agent_did: str, process_handle: Callable[[], None]
+    ) -> None:
+        """Register an agent with its termination callback."""
+        self._agents[agent_did] = process_handle
+
+    def unregister_agent(self, agent_did: str) -> None:
+        """Remove an agent from the process registry."""
+        self._agents.pop(agent_did, None)
+
+    # ── Substitute management ──────────────────────────────────────
+
+    def register_substitute(
+        self, session_id: str, agent_did: str
+    ) -> None:
+        """Register a substitute agent for a session."""
+        self._substitutes.setdefault(session_id, []).append(agent_did)
+
+    def unregister_substitute(
+        self, session_id: str, agent_did: str
+    ) -> None:
+        subs = self._substitutes.get(session_id, [])
+        if agent_did in subs:
+            subs.remove(agent_did)
+
+    # ── Kill ───────────────────────────────────────────────────────
+
+    def kill(
+        self,
+        agent_did: str,
+        session_id: str,
+        reason: KillReason,
+        in_flight_steps: list[dict] | None = None,
+        details: str = "",
+    ) -> KillResult:
+        """Kill an agent, handing off in-flight steps to a substitute if available."""
+        in_flight = in_flight_steps or []
+
+        # Attempt to find a substitute for handoff
+        substitute = self._find_substitute(session_id, agent_did)
+
+        handoffs: list[StepHandoff] = []
+        handoff_success_count = 0
+        for step_info in in_flight:
+            if substitute is not None:
+                handoffs.append(
+                    StepHandoff(
+                        step_id=step_info.get("step_id", ""),
+                        saga_id=step_info.get("saga_id", ""),
+                        from_agent=agent_did,
+                        to_agent=substitute,
+                        status=HandoffStatus.HANDED_OFF,
+                    )
+                )
+                handoff_success_count += 1
+            else:
+                handoffs.append(
+                    StepHandoff(
+                        step_id=step_info.get("step_id", ""),
+                        saga_id=step_info.get("saga_id", ""),
+                        from_agent=agent_did,
+                        status=HandoffStatus.COMPENSATED,
+                    )
+                )
+
+        # Terminate the agent process
+        terminated = False
+        callback = self._agents.get(agent_did)
+        if callback is not None:
+            callback()
+            terminated = True
+        else:
+            _logger.warning(
+                "No termination callback registered for agent %s",
+                agent_did,
+            )
+
+        result = KillResult(
+            agent_did=agent_did,
+            session_id=session_id,
+            reason=reason,
+            handoffs=handoffs,
+            handoff_success_count=handoff_success_count,
+            compensation_triggered=any(
+                h.status == HandoffStatus.COMPENSATED for h in handoffs
+            ),
+            terminated=terminated,
+            details=details,
+        )
+        self._kill_history.append(result)
+        self.unregister_substitute(session_id, agent_did)
+        self.unregister_agent(agent_did)
+        return result
+
+    def _find_substitute(
+        self, session_id: str, exclude_did: str
+    ) -> str | None:
+        """Find a registered substitute for the session, excluding the given agent."""
+        subs = self._substitutes.get(session_id, [])
+        for sub in subs:
+            if sub != exclude_did:
+                return sub
+        return None
+
+    @property
+    def kill_history(self) -> list[KillResult]:
+        return list(self._kill_history)
+
+    @property
+    def total_kills(self) -> int:
+        return len(self._kill_history)
+
+    @property
+    def total_handoffs(self) -> int:
+        return sum(r.handoff_success_count for r in self._kill_history)