agent_os_kernel 3.1.0__py3-none-any.whl

iatp/policy_engine.py

@@ -0,0 +1,337 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+IATP Policy Engine - Protocol Layer Policy Validation.
+
+This module provides policy validation for IATP capability manifests
+using a standalone, protocol-native implementation.
+
+Design Note: IATP is Layer 2 (Infrastructure/Protocol). It defines the
+message format, handshake protocols, and trust scores. Higher layers
+(like agent-control-plane) USE this protocol; this protocol does NOT
+depend on them.
+"""
+from typing import Any, Dict, List, Optional, Protocol, Tuple, runtime_checkable
+
+from iatp.models import CapabilityManifest, ReversibilityLevel
+
+
+@runtime_checkable
+class PolicyEvaluator(Protocol):
+    """
+    Protocol class for policy evaluation.
+
+    This allows duck typing for any policy evaluator implementation
+    without requiring a specific import.
+    """
+
+    def evaluate(self, context: Dict[str, Any]) -> str:
+        """Evaluate context and return action: 'allow', 'warn', or 'deny'."""
+        ...
+
+
+class PolicyRule:
+    """
+    A single policy rule for manifest validation.
+
+    Rules are evaluated in order and the first matching rule determines
+    the action. Actions can be 'allow', 'warn', or 'deny'.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        action: str,
+        conditions: Dict[str, List[Any]],
+        description: str = ""
+    ):
+        """
+        Initialize a policy rule.
+
+        Args:
+            name: Rule identifier
+            action: Action to take ('allow', 'warn', 'deny')
+            conditions: Dictionary mapping field names to allowed values
+            description: Human-readable description
+        """
+        self.name = name
+        self.action = action
+        self.conditions = conditions
+        self.description = description
+
+    def matches(self, context: Dict[str, Any]) -> bool:
+        """
+        Check if this rule matches the given context.
+
+        Args:
+            context: Policy evaluation context
+
+        Returns:
+            True if any condition matches
+        """
+        for key, values in self.conditions.items():
+            if key in context and context[key] in values:
+                return True
+        return False
+
+
+class IATPPolicyEngine:
+    """
+    IATP Policy Engine for capability manifest validation.
+
+    This is a protocol-native policy engine that validates incoming
+    agent manifests against configurable policy rules. It uses the
+    IATP trust scoring algorithm and provides warn/block decisions.
+
+    The engine is designed to be:
+    - Self-contained (no external dependencies beyond IATP models)
+    - Extensible (custom rules can be added)
+    - Protocol-compliant (follows IATP trust semantics)
+    """
+
+    def __init__(self):
+        """Initialize the IATP Policy Engine."""
+        self.rules: List[PolicyRule] = []
+        self._setup_default_policies()
+
+    def _setup_default_policies(self):
+        """Setup default security policies for IATP."""
+        # Default IATP policy rules
+        # These augment (not replace) the SecurityValidator checks
+        self.rules = [
+            PolicyRule(
+                name="WarnUntrustedAgent",
+                description="Warn when agents are marked as untrusted",
+                action="warn",
+                conditions={"trust_level": ["untrusted"]}
+            ),
+            PolicyRule(
+                name="RequireReversibility",
+                description="Warn when agents don't support reversibility",
+                action="warn",
+                conditions={"reversibility": ["none"]}
+            ),
+            PolicyRule(
+                name="AllowEphemeral",
+                description="Allow agents with ephemeral data retention",
+                action="allow",
+                conditions={"retention_policy": ["ephemeral"]}
+            ),
+        ]
+
+    def add_custom_rule(self, rule: Dict[str, Any]):
+        """
+        Add a custom policy rule.
+
+        Args:
+            rule: Dictionary defining the policy rule with keys:
+                - name: Rule name
+                - description: Rule description (optional)
+                - action: "allow", "warn", or "deny"
+                - conditions: Dictionary of conditions to match
+        """
+        policy_rule = PolicyRule(
+            name=rule.get("name", "CustomRule"),
+            action=rule.get("action", "warn"),
+            conditions=rule.get("conditions", {}),
+            description=rule.get("description", "")
+        )
+        # Insert at beginning so custom rules take precedence
+        self.rules.insert(0, policy_rule)
+
+    def validate_manifest(
+        self,
+        manifest: CapabilityManifest
+    ) -> Tuple[bool, Optional[str], Optional[str]]:
+        """
+        Validate a capability manifest against policies.
+
+        This is the main integration point that validates incoming
+        agent manifests against the configured policy rules.
+
+        Args:
+            manifest: The capability manifest to validate
+
+        Returns:
+            Tuple of (is_allowed, error_message, warning_message)
+            - is_allowed: True if request should proceed
+            - error_message: Blocking error if is_allowed is False
+            - warning_message: Warning for user if there are concerns
+        """
+        # Convert manifest to policy context
+        context = self._manifest_to_context(manifest)
+
+        # Evaluate against policy rules
+        action = self._evaluate_rules(context)
+
+        # Generate appropriate response
+        if action == "deny":
+            return False, self._generate_deny_message(manifest, context), None
+        elif action == "warn":
+            return True, None, self._generate_warn_message(manifest, context)
+        else:  # allow
+            return True, None, None
+
+    def _evaluate_rules(self, context: Dict[str, Any]) -> str:
+        """
+        Evaluate policy rules against context.
+
+        Args:
+            context: Policy evaluation context
+
+        Returns:
+            Action string: "allow", "warn", or "deny"
+        """
+        for rule in self.rules:
+            if rule.matches(context):
+                return rule.action
+
+        # Default to allow if no rules match
+        return "allow"
+
+    def _manifest_to_context(self, manifest: CapabilityManifest) -> Dict[str, Any]:
+        """
+        Convert a capability manifest to a policy evaluation context.
+
+        Args:
+            manifest: The capability manifest
+
+        Returns:
+            Dictionary context for policy evaluation
+        """
+        return {
+            "agent_id": manifest.agent_id,
+            "trust_level": manifest.trust_level.value,
+            "retention_policy": manifest.privacy_contract.retention.value,
+            "reversibility": manifest.capabilities.reversibility.value,
+            "idempotency": manifest.capabilities.idempotency,
+            "human_review": manifest.privacy_contract.human_review,
+            "encryption_at_rest": manifest.privacy_contract.encryption_at_rest,
+            "encryption_in_transit": manifest.privacy_contract.encryption_in_transit,
+            "scopes": manifest.scopes,
+        }
+
+    def _generate_deny_message(
+        self,
+        manifest: CapabilityManifest,
+        context: Dict[str, Any]
+    ) -> str:
+        """
+        Generate a denial message for blocked requests.
+
+        Args:
+            manifest: The capability manifest
+            context: Policy context
+
+        Returns:
+            User-friendly error message
+        """
+        reasons = []
+
+        if context["retention_policy"] in ["permanent", "forever"]:
+            reasons.append(
+                f"Agent '{manifest.agent_id}' stores data permanently, "
+                "which violates privacy policies"
+            )
+
+        if context["trust_level"] == "untrusted":
+            reasons.append(
+                f"Agent '{manifest.agent_id}' is marked as untrusted"
+            )
+
+        if reasons:
+            return "Policy Violation: " + "; ".join(reasons)
+
+        return f"Policy Violation: Agent '{manifest.agent_id}' failed policy validation"
+
+    def _generate_warn_message(
+        self,
+        manifest: CapabilityManifest,
+        context: Dict[str, Any]
+    ) -> str:
+        """
+        Generate a warning message for risky requests.
+
+        Args:
+            manifest: The capability manifest
+            context: Policy context
+
+        Returns:
+            User-friendly warning message
+        """
+        warnings = []
+
+        if context["reversibility"] == "none":
+            warnings.append(
+                f"Agent '{manifest.agent_id}' does not support transaction reversal"
+            )
+
+        if not context["idempotency"]:
+            warnings.append(
+                f"Agent '{manifest.agent_id}' may not handle duplicate requests safely"
+            )
+
+        if context["human_review"]:
+            warnings.append(
+                f"Agent '{manifest.agent_id}' may have humans review your data"
+            )
+
+        if warnings:
+            return "⚠️  Policy Warning:\n" + "\n".join(f"  • {w}" for w in warnings)
+
+        return None
+
+    def validate_handshake(
+        self,
+        manifest: CapabilityManifest,
+        required_capabilities: Optional[List[str]] = None,
+        required_scopes: Optional[List[str]] = None
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Validate handshake compatibility between agents.
+
+        This checks if the remote agent's capabilities meet the
+        local agent's requirements.
+
+        Args:
+            manifest: Remote agent's capability manifest
+            required_capabilities: List of required capability keys
+            required_scopes: List of required RBAC scopes (e.g., ['repo:write'])
+
+        Returns:
+            Tuple of (is_compatible, error_message)
+        """
+        if not required_capabilities:
+            required_capabilities = []
+        if not required_scopes:
+            required_scopes = []
+
+        # Always validate against base policies first
+        is_allowed, error_msg, _ = self.validate_manifest(manifest)
+        if not is_allowed:
+            return False, error_msg
+
+        # Check specific capability requirements
+        missing = []
+        for capability in required_capabilities:
+            if capability == "reversibility" and \
+               manifest.capabilities.reversibility == ReversibilityLevel.NONE:
+                missing.append("reversibility support")
+            elif capability == "idempotency" and \
+                 not manifest.capabilities.idempotency:
+                missing.append("idempotency support")
+
+        if missing:
+            return False, f"Agent missing required capabilities: {', '.join(missing)}"
+
+        # Check RBAC scope requirements
+        missing_scopes = []
+        agent_scopes = set(manifest.scopes)
+        for scope in required_scopes:
+            if scope not in agent_scopes:
+                missing_scopes.append(scope)
+
+        if missing_scopes:
+            return False, f"Agent missing required scopes: {', '.join(missing_scopes)}"
+
+        return True, None

iatp/py.typed

@@ -0,0 +1,2 @@
+# PEP 561 marker file - indicates this package supports type checking
+# https://peps.python.org/pep-0561/

iatp/recovery.py

@@ -0,0 +1,321 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Recovery Engine Integration with SCAK (Self-Correcting Agent Kernel).
+
+This module wraps the agent_kernel (scak) to provide failure recovery
+and rollback capabilities for IATP.
+"""
+import asyncio
+import logging
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Callable, Dict, Optional
+
+from agent_primitives import (
+    AgentFailure,
+    FailureType,
+    FailureSeverity,
+)
+
+from iatp.models import CapabilityManifest
+
+logger = logging.getLogger(__name__)
+
+
+class RecoveryStrategy(str, Enum):
+    """Recovery strategies for IATP."""
+    ROLLBACK = "rollback"  # Execute compensation transaction
+    RETRY = "retry"  # Retry the operation
+    GIVE_UP = "give_up"  # No recovery possible
+
+
+class IATPRecoveryEngine:
+    """
+    Wrapper around SCAK for IATP failure recovery.
+
+    This integrates the Self-Correcting Agent Kernel (scak) to provide
+    automatic failure detection, triage, and recovery for agent interactions.
+    """
+
+    def __init__(self):
+        """Initialize the IATP Recovery Engine."""
+        # Note: SelfCorrectingAgentKernel requires specific initialization
+        # We'll use the triage and diagnosis functions directly
+        self.recovery_history: Dict[str, Any] = {}
+
+    async def handle_failure(
+        self,
+        trace_id: str,
+        error: Exception,
+        manifest: CapabilityManifest,
+        payload: Dict[str, Any],
+        compensation_callback: Optional[Callable] = None
+    ) -> Dict[str, Any]:
+        """
+        Handle a failure in agent communication.
+
+        This is the main entry point for failure recovery. It:
+        1. Diagnoses the failure using scak
+        2. Determines the appropriate recovery strategy
+        3. Executes compensation if available
+
+        Args:
+            trace_id: Unique trace ID for this request
+            error: The exception that occurred
+            manifest: The agent's capability manifest
+            payload: The original request payload
+            compensation_callback: Optional callback for rollback
+
+        Returns:
+            Dictionary with recovery result and actions taken
+        """
+        logger.info(f"[Recovery] Handling failure for trace {trace_id}")
+
+        # Determine failure type from exception
+        error_name = type(error).__name__
+        if "Timeout" in error_name:
+            failure_type = FailureType.TIMEOUT
+        elif "Resource" in error_name or "Limit" in error_name:
+            failure_type = FailureType.RESOURCE_EXHAUSTED
+        elif "Invalid" in error_name or "Validation" in error_name:
+            failure_type = FailureType.INVALID_ACTION
+        else:
+            failure_type = FailureType.UNKNOWN
+
+        # Create failure record using scak's AgentFailure model
+        failure = AgentFailure(
+            agent_id=manifest.agent_id,
+            failure_type=failure_type,
+            error_message=str(error),
+            context={
+                "trace_id": trace_id,
+                "payload": payload,
+                "reversibility": manifest.capabilities.reversibility.value,
+            },
+            timestamp=datetime.now(timezone.utc)
+        )
+
+        # Determine recovery strategy based on manifest and failure type
+        # Note: scak's diagnose_failure requires complex tool traces,
+        # so we use a simplified diagnosis based on the failure record
+        diagnosis = f"{failure_type.value}: {error_name}"
+
+        strategy = self._determine_strategy(
+            diagnosis,
+            manifest,
+            compensation_callback is not None
+        )
+
+        # Execute recovery
+        recovery_result = await self._execute_recovery(
+            strategy=strategy,
+            trace_id=trace_id,
+            failure=failure,
+            manifest=manifest,
+            compensation_callback=compensation_callback
+        )
+
+        # Record in history
+        self.recovery_history[trace_id] = {
+            "failure": failure,
+            "diagnosis": diagnosis,
+            "strategy": strategy,
+            "result": recovery_result,
+            "timestamp": datetime.now(timezone.utc).isoformat()
+        }
+
+        return recovery_result
+
+    def _determine_strategy(
+        self,
+        diagnosis: str,
+        manifest: CapabilityManifest,
+        has_compensation: bool
+    ) -> RecoveryStrategy:
+        """
+        Determine the appropriate recovery strategy.
+
+        Args:
+            diagnosis: Failure diagnosis from scak
+            manifest: Agent capability manifest
+            has_compensation: Whether compensation callback is available
+
+        Returns:
+            RecoveryStrategy to use for recovery
+        """
+        # Check if agent supports reversibility
+        reversibility = manifest.capabilities.reversibility.value
+
+        if reversibility in ["full", "partial"] and has_compensation:
+            # Agent supports rollback and we have compensation logic
+            return RecoveryStrategy.ROLLBACK
+        elif reversibility == "partial":
+            # Partial rollback - log and warn
+            return RecoveryStrategy.RETRY
+        elif diagnosis.lower().find("timeout") >= 0:
+            # Timeout errors might be transient
+            return RecoveryStrategy.RETRY
+        else:
+            # No recovery possible
+            return RecoveryStrategy.GIVE_UP
+
+    async def _execute_recovery(
+        self,
+        strategy: RecoveryStrategy,
+        trace_id: str,
+        failure: AgentFailure,
+        manifest: CapabilityManifest,
+        compensation_callback: Optional[Callable] = None
+    ) -> Dict[str, Any]:
+        """
+        Execute the determined recovery strategy.
+
+        Args:
+            strategy: Recovery strategy to execute
+            trace_id: Trace ID
+            failure: Failure details
+            manifest: Agent manifest
+            compensation_callback: Optional compensation callback
+
+        Returns:
+            Dictionary with recovery results
+        """
+        result = {
+            "strategy": strategy.value,
+            "success": False,
+            "actions_taken": [],
+            "trace_id": trace_id
+        }
+
+        if strategy == RecoveryStrategy.ROLLBACK:
+            logger.info(f"[Recovery] Executing rollback for {trace_id}")
+            result["actions_taken"].append("initiated_rollback")
+
+            if compensation_callback:
+                try:
+                    # Execute compensation transaction
+                    if asyncio.iscoroutinefunction(compensation_callback):
+                        await compensation_callback()
+                    else:
+                        compensation_callback()
+
+                    result["success"] = True
+                    result["actions_taken"].append("compensation_executed")
+                    logger.info(f"[Recovery] Rollback successful for {trace_id}")
+                except Exception as e:
+                    result["actions_taken"].append(f"compensation_failed: {str(e)}")
+                    logger.error(f"[Recovery] Rollback failed for {trace_id}: {e}")
+            else:
+                result["actions_taken"].append("no_compensation_available")
+                logger.warning(f"[Recovery] No compensation callback for {trace_id}")
+
+        elif strategy == RecoveryStrategy.RETRY:
+            logger.info(f"[Recovery] Retry recommended for {trace_id}")
+            result["actions_taken"].append("retry_recommended")
+            result["retry_possible"] = True
+            # Note: Actual retry logic would be implemented by the caller
+
+        else:  # GIVE_UP
+            logger.info(f"[Recovery] No recovery possible for {trace_id}")
+            result["actions_taken"].append("recovery_not_possible")
+            result["message"] = (
+                f"Agent '{manifest.agent_id}' does not support rollback and "
+                "error is not recoverable. Transaction may be in inconsistent state."
+            )
+
+        return result
+
+    def get_recovery_history(self, trace_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Get recovery history for a trace ID.
+
+        Args:
+            trace_id: Trace ID to look up
+
+        Returns:
+            Recovery history or None if not found
+        """
+        return self.recovery_history.get(trace_id)
+
+    def should_attempt_recovery(
+        self,
+        error: Exception,
+        manifest: CapabilityManifest
+    ) -> bool:
+        """
+        Determine if recovery should be attempted for this error.
+
+        Args:
+            error: The exception that occurred
+            manifest: Agent capability manifest
+
+        Returns:
+            True if recovery should be attempted
+        """
+        # Always attempt recovery if agent supports reversibility
+        if manifest.capabilities.reversibility.value in ["full", "partial"]:
+            return True
+
+        # Attempt recovery for certain error types
+        error_name = type(error).__name__
+        recoverable_errors = [
+            "TimeoutError",
+            "ConnectionError",
+            "HTTPError",
+            "ServiceUnavailable"
+        ]
+
+        return any(err in error_name for err in recoverable_errors)
+
+    async def execute_compensation_transaction(
+        self,
+        trace_id: str,
+        manifest: CapabilityManifest,
+        compensation_endpoint: str,
+        compensation_payload: Dict[str, Any]
+    ) -> bool:
+        """
+        Execute a compensation transaction using the agent's compensation endpoint.
+
+        This is used when the agent provides a specific compensation/rollback
+        endpoint as specified in the handshake.
+
+        Args:
+            trace_id: Trace ID
+            manifest: Agent manifest with undo_window info
+            compensation_endpoint: URL for compensation
+            compensation_payload: Payload for compensation request
+
+        Returns:
+            True if compensation succeeded
+        """
+        import httpx
+
+        logger.info(f"[Recovery] Executing compensation transaction for {trace_id}")
+
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    compensation_endpoint,
+                    json=compensation_payload,
+                    headers={
+                        "X-Agent-Trace-ID": trace_id,
+                        "X-Compensation-Request": "true"
+                    },
+                    timeout=30.0
+                )
+
+                if 200 <= response.status_code < 300:
+                    logger.info(f"[Recovery] Compensation successful for {trace_id}")
+                    return True
+                else:
+                    logger.error(
+                        f"[Recovery] Compensation failed for {trace_id}: "
+                        f"status {response.status_code}"
+                    )
+                    return False
+
+        except Exception as e:
+            logger.error(f"[Recovery] Compensation exception for {trace_id}: {e}")
+            return False