opshield 0.2.0__py3-none-any.whl

opshield/__init__.py

@@ -0,0 +1,53 @@
+"""OpShield — rollback, cost control, and testing for AI agents."""
+
+from opshield.shield import OpShield
+from opshield.action import Action, ActionResult, RiskLevel
+from opshield.approval import ApprovalManager
+from opshield.async_shield import AsyncOpShield
+from opshield.capabilities import CapabilityManager
+from opshield.chaos import ChaosEngine, FaultType
+from opshield.circuitbreaker import CircuitBreaker, CircuitState
+from opshield.compliance import ComplianceReporter
+from opshield.forecast import CostForecaster
+from opshield.logger import ActionLogger
+from opshield.masking import DataMasker
+from opshield.multiagent import AgentPool
+from opshield.policy import load_policy, validate_policy
+from opshield.ratelimit import RateLimiter, ToolLimit
+from opshield.replay import ActionReplayer
+from opshield.retry import RetryPolicy
+from opshield.scoring import AgentScorer, PoolScorer
+from opshield.tracing import Tracer
+from opshield.webhooks import EventNotifier
+from opshield.config import rules_from_config
+
+__all__ = [
+    "OpShield",
+    "AsyncOpShield",
+    "AgentPool",
+    "AgentScorer",
+    "ApprovalManager",
+    "Action",
+    "ActionResult",
+    "CapabilityManager",
+    "ChaosEngine",
+    "CircuitBreaker",
+    "CircuitState",
+    "ComplianceReporter",
+    "CostForecaster",
+    "FaultType",
+    "PoolScorer",
+    "RiskLevel",
+    "ActionLogger",
+    "ActionReplayer",
+    "DataMasker",
+    "RateLimiter",
+    "ToolLimit",
+    "RetryPolicy",
+    "Tracer",
+    "EventNotifier",
+    "load_policy",
+    "validate_policy",
+    "rules_from_config",
+]
+__version__ = "0.2.0"

opshield/__main__.py

@@ -0,0 +1,3 @@
+from opshield.cli import main
+
+main()

opshield/action.py

@@ -0,0 +1,50 @@
+"""Action model — represents a single agent action with metadata."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class RiskLevel(Enum):
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+    CRITICAL = "critical"
+
+
+@dataclass
+class Action:
+    """A single action an agent wants to perform."""
+
+    tool_name: str
+    args: dict[str, Any] = field(default_factory=dict)
+    id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
+    timestamp: float = field(default_factory=time.time)
+    risk: RiskLevel = RiskLevel.LOW
+    cost: float = 0.0
+
+    @property
+    def description(self) -> str:
+        args_str = ", ".join(f"{k}={v!r}" for k, v in self.args.items())
+        return f"{self.tool_name}({args_str})"
+
+
+@dataclass
+class ActionResult:
+    """The result of executing an action."""
+
+    action: Action
+    output: Any = None
+    error: str | None = None
+    blocked: bool = False
+    block_reason: str | None = None
+    snapshot_id: str | None = None
+    cost: float = 0.0
+
+    @property
+    def success(self) -> bool:
+        return not self.blocked and self.error is None

opshield/approval.py

@@ -0,0 +1,275 @@
+"""Human-in-the-loop approval workflows for AI agent actions.
+
+Provides configurable approval gates that pause execution until a human
+approves or denies the action. Supports timeout with auto-deny, callback-
+based approval, and escalation chains.
+
+Usage:
+    approver = ApprovalManager(
+        default_timeout=30.0,
+        on_approval_request=my_callback,
+    )
+
+    # Register tool-specific policies
+    approver.require_approval("send_email", timeout=60.0)
+    approver.require_approval("delete_record", auto_deny=True)
+
+    # Check in shield flow
+    decision = approver.request_approval(action)
+    if decision.approved:
+        execute(action)
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class ApprovalDecision(Enum):
+    APPROVED = "approved"
+    DENIED = "denied"
+    TIMEOUT = "timeout"
+    ESCALATED = "escalated"
+    NOT_REQUIRED = "not_required"
+
+
+@dataclass
+class ApprovalResult:
+    decision: ApprovalDecision
+    tool_name: str
+    approver: str | None = None
+    reason: str | None = None
+    elapsed: float = 0.0
+
+    @property
+    def approved(self) -> bool:
+        return self.decision in (
+            ApprovalDecision.APPROVED,
+            ApprovalDecision.NOT_REQUIRED,
+        )
+
+
+@dataclass
+class ApprovalPolicy:
+    timeout: float = 30.0
+    auto_deny_on_timeout: bool = True
+    escalate_to: callable | None = None
+    required_for_risk: str | None = None
+
+
+class ApprovalManager:
+    """Manages human-in-the-loop approval for agent actions."""
+
+    def __init__(
+        self,
+        default_timeout: float = 30.0,
+        on_approval_request: callable | None = None,
+        auto_approve_low_risk: bool = True,
+    ) -> None:
+        self._default_timeout = default_timeout
+        self._on_request = on_approval_request
+        self._auto_approve_low = auto_approve_low_risk
+        self._policies: dict[str, ApprovalPolicy] = {}
+        self._pending: dict[str, threading.Event] = {}
+        self._decisions: dict[str, ApprovalResult] = {}
+        self._history: list[ApprovalResult] = []
+        self._lock = threading.Lock()
+
+    def require_approval(
+        self,
+        tool_name: str,
+        timeout: float | None = None,
+        auto_deny: bool = True,
+        escalate_to: callable | None = None,
+    ) -> None:
+        self._policies[tool_name] = ApprovalPolicy(
+            timeout=timeout or self._default_timeout,
+            auto_deny_on_timeout=auto_deny,
+            escalate_to=escalate_to,
+        )
+
+    def require_approval_for_risk(
+        self,
+        risk_level: str,
+        timeout: float | None = None,
+        auto_deny: bool = True,
+    ) -> None:
+        policy = ApprovalPolicy(
+            timeout=timeout or self._default_timeout,
+            auto_deny_on_timeout=auto_deny,
+            required_for_risk=risk_level.upper(),
+        )
+        self._policies[f"__risk__{risk_level.upper()}"] = policy
+
+    def needs_approval(self, tool_name: str, risk_level: str = "LOW") -> bool:
+        if tool_name in self._policies:
+            return True
+        risk_key = f"__risk__{risk_level.upper()}"
+        if risk_key in self._policies:
+            return True
+        return False
+
+    def request_approval(
+        self,
+        tool_name: str,
+        args: dict | None = None,
+        risk_level: str = "LOW",
+        action_id: str | None = None,
+    ) -> ApprovalResult:
+        if not self.needs_approval(tool_name, risk_level):
+            result = ApprovalResult(
+                decision=ApprovalDecision.NOT_REQUIRED,
+                tool_name=tool_name,
+            )
+            return result
+
+        policy = self._policies.get(tool_name)
+        if not policy:
+            risk_key = f"__risk__{risk_level.upper()}"
+            policy = self._policies.get(risk_key, ApprovalPolicy())
+
+        request_id = action_id or f"{tool_name}_{time.time()}"
+
+        event = threading.Event()
+        with self._lock:
+            self._pending[request_id] = event
+
+        if self._on_request:
+            self._on_request(
+                request_id=request_id,
+                tool_name=tool_name,
+                args=args or {},
+                risk_level=risk_level,
+            )
+
+        start = time.time()
+        approved = event.wait(timeout=policy.timeout)
+        elapsed = time.time() - start
+
+        with self._lock:
+            self._pending.pop(request_id, None)
+
+            if request_id in self._decisions:
+                result = self._decisions.pop(request_id)
+                result.elapsed = elapsed
+            elif approved:
+                result = ApprovalResult(
+                    decision=ApprovalDecision.APPROVED,
+                    tool_name=tool_name,
+                    elapsed=elapsed,
+                )
+            else:
+                if policy.escalate_to:
+                    try:
+                        escalation = policy.escalate_to(
+                            tool_name=tool_name,
+                            args=args,
+                            risk_level=risk_level,
+                        )
+                        if escalation:
+                            result = ApprovalResult(
+                                decision=ApprovalDecision.APPROVED,
+                                tool_name=tool_name,
+                                approver="escalation",
+                                elapsed=elapsed,
+                            )
+                        else:
+                            result = ApprovalResult(
+                                decision=ApprovalDecision.ESCALATED,
+                                tool_name=tool_name,
+                                reason="Escalation denied",
+                                elapsed=elapsed,
+                            )
+                    except Exception:
+                        result = ApprovalResult(
+                            decision=ApprovalDecision.DENIED,
+                            tool_name=tool_name,
+                            reason="Escalation failed",
+                            elapsed=elapsed,
+                        )
+                elif policy.auto_deny_on_timeout:
+                    result = ApprovalResult(
+                        decision=ApprovalDecision.TIMEOUT,
+                        tool_name=tool_name,
+                        reason=f"Approval timed out after {policy.timeout}s",
+                        elapsed=elapsed,
+                    )
+                else:
+                    result = ApprovalResult(
+                        decision=ApprovalDecision.DENIED,
+                        tool_name=tool_name,
+                        reason="No approval received",
+                        elapsed=elapsed,
+                    )
+
+        self._history.append(result)
+        return result
+
+    def approve(
+        self,
+        request_id: str,
+        approver: str | None = None,
+        reason: str | None = None,
+    ) -> bool:
+        with self._lock:
+            event = self._pending.get(request_id)
+            if not event:
+                return False
+            self._decisions[request_id] = ApprovalResult(
+                decision=ApprovalDecision.APPROVED,
+                tool_name="",
+                approver=approver,
+                reason=reason,
+            )
+            event.set()
+            return True
+
+    def deny(
+        self,
+        request_id: str,
+        approver: str | None = None,
+        reason: str | None = None,
+    ) -> bool:
+        with self._lock:
+            event = self._pending.get(request_id)
+            if not event:
+                return False
+            self._decisions[request_id] = ApprovalResult(
+                decision=ApprovalDecision.DENIED,
+                tool_name="",
+                approver=approver,
+                reason=reason,
+            )
+            event.set()
+            return True
+
+    @property
+    def pending_count(self) -> int:
+        with self._lock:
+            return len(self._pending)
+
+    @property
+    def history(self) -> list[ApprovalResult]:
+        return list(self._history)
+
+    def stats(self) -> dict:
+        total = len(self._history)
+        approved = sum(1 for r in self._history if r.approved)
+        denied = sum(
+            1 for r in self._history
+            if r.decision == ApprovalDecision.DENIED
+        )
+        timed_out = sum(
+            1 for r in self._history
+            if r.decision == ApprovalDecision.TIMEOUT
+        )
+        return {
+            "total_requests": total,
+            "approved": approved,
+            "denied": denied,
+            "timed_out": timed_out,
+            "pending": self.pending_count,
+        }

opshield/async_shield.py

@@ -0,0 +1,195 @@
+"""Async version of OpShield for use with asyncio-based agent frameworks."""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import inspect
+from typing import Any, Callable
+
+from opshield.action import Action, ActionResult, RiskLevel
+from opshield.shield import OpShield
+
+
+class AsyncOpShield(OpShield):
+    """OpShield with async tool support.
+
+    Usage:
+        shield = AsyncOpShield(budget=5.00)
+
+        @shield.async_tool(cost=0.03)
+        async def query_db(sql: str) -> str:
+            return await db.execute(sql)
+
+        result = await query_db(sql="SELECT * FROM users")
+    """
+
+    def async_tool(
+        self,
+        fn: Callable | None = None,
+        *,
+        cost: float = 0.0,
+        risk: RiskLevel | None = None,
+        snapshot_state: Callable[[], Any] | None = None,
+        restore_state: Callable[[Any], None] | None = None,
+    ) -> Callable:
+
+        def decorator(func: Callable) -> Callable:
+            tool_name = func.__name__
+
+            if snapshot_state and restore_state:
+                self._tool_states[tool_name] = {
+                    "capture": snapshot_state,
+                    "restore": restore_state,
+                }
+
+            @functools.wraps(func)
+            async def wrapper(*args: Any, **kwargs: Any) -> ActionResult:
+                action = Action(
+                    tool_name=tool_name,
+                    args=kwargs if kwargs else {"args": list(args)},
+                    cost=cost,
+                )
+                return await self._async_execute(action, func, args, kwargs, cost, risk)
+
+            wrapper._shield_tool = True
+            wrapper._tool_name = tool_name
+            return wrapper
+
+        if fn is not None:
+            return decorator(fn)
+        return decorator
+
+    async def _async_execute(
+        self,
+        action: Action,
+        func: Callable,
+        args: tuple,
+        kwargs: dict,
+        cost: float,
+        risk_override: RiskLevel | None,
+    ) -> ActionResult:
+        detected_risk, reasons = self._rules.evaluate(action)
+        action.risk = risk_override or detected_risk
+
+        if action.risk == RiskLevel.CRITICAL:
+            result = ActionResult(
+                action=action,
+                blocked=True,
+                block_reason=f"CRITICAL risk blocked: {'; '.join(reasons)}",
+            )
+            self._history.append(result)
+            if self._logger:
+                self._logger.log(result)
+            self._notifier.emit("action_blocked", {
+                "tool": action.tool_name, "reason": result.block_reason,
+            })
+            return result
+
+        if action.risk == RiskLevel.HIGH and self._on_block:
+            allowed = self._on_block(action, reasons)
+            if not allowed:
+                result = ActionResult(
+                    action=action,
+                    blocked=True,
+                    block_reason=f"Blocked by user: {'; '.join(reasons)}",
+                )
+                self._history.append(result)
+                if self._logger:
+                    self._logger.log(result)
+                self._notifier.emit("action_blocked", {
+                    "tool": action.tool_name, "reason": result.block_reason,
+                })
+                return result
+
+        if self._budget.would_exceed(cost):
+            result = ActionResult(
+                action=action,
+                blocked=True,
+                block_reason=f"Budget exceeded (spent: ${self._budget.spent:.2f}, "
+                f"budget: ${self._budget.budget:.2f}, action cost: ${cost:.2f})",
+            )
+            self._history.append(result)
+            if self._logger:
+                self._logger.log(result)
+            self._notifier.emit("budget_exceeded", {
+                "tool": action.tool_name, "spent": self._budget.spent,
+                "budget": self._budget.budget,
+            })
+            return result
+
+        exceeded, rate_reason = self._rate_limiter.would_exceed(action.tool_name)
+        if exceeded:
+            result = ActionResult(
+                action=action,
+                blocked=True,
+                block_reason=rate_reason,
+            )
+            self._history.append(result)
+            if self._logger:
+                self._logger.log(result)
+            self._notifier.emit("rate_limited", {
+                "tool": action.tool_name, "reason": rate_reason,
+            })
+            return result
+
+        snapshot_id = None
+        if self._auto_snapshot and action.tool_name in self._tool_states:
+            state = self._tool_states[action.tool_name]["capture"]()
+            snap = self._snapshots.capture(action.id, action.tool_name, state)
+            snapshot_id = snap.id
+
+        if self.is_test_mode:
+            result = ActionResult(
+                action=action,
+                output=f"[TEST MODE] Would execute: {action.description}",
+                snapshot_id=snapshot_id,
+                cost=cost,
+            )
+            self._budget.record(action.id, action.tool_name, cost)
+            self._rate_limiter.record(action.tool_name)
+            self._history.append(result)
+            if self._logger:
+                self._logger.log(result)
+            return result
+
+        try:
+            if inspect.iscoroutinefunction(func):
+                output = await func(*args, **kwargs)
+            else:
+                loop = asyncio.get_event_loop()
+                output = await loop.run_in_executor(None, lambda: func(*args, **kwargs))
+            result = ActionResult(
+                action=action,
+                output=output,
+                snapshot_id=snapshot_id,
+                cost=cost,
+            )
+        except Exception as e:
+            result = ActionResult(
+                action=action,
+                error=str(e),
+                snapshot_id=snapshot_id,
+                cost=cost,
+            )
+            self._notifier.emit("action_error", {
+                "tool": action.tool_name, "error": str(e),
+            })
+
+        self._budget.record(action.id, action.tool_name, cost)
+        self._rate_limiter.record(action.tool_name)
+        self._history.append(result)
+        if self._logger:
+            self._logger.log(result)
+
+        if result.success:
+            self._notifier.emit("action_completed", {
+                "tool": action.tool_name, "cost": cost,
+            })
+
+        if self._budget.warning:
+            self._notifier.emit("budget_warning", {
+                "spent": self._budget.spent, "budget": self._budget.budget,
+            })
+
+        return result

opshield/budget.py

@@ -0,0 +1,76 @@
+"""Budget tracker for agent cost control."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+
+
+@dataclass
+class CostEntry:
+    """A single cost record."""
+
+    action_id: str
+    tool_name: str
+    cost: float
+    timestamp: float
+
+
+class BudgetTracker:
+    """Tracks cumulative cost and enforces a budget limit."""
+
+    def __init__(self, budget: float | None = None, warn_at: float = 0.8) -> None:
+        self._budget = budget
+        self._warn_at = warn_at
+        self._entries: list[CostEntry] = []
+
+    @property
+    def budget(self) -> float | None:
+        return self._budget
+
+    @property
+    def spent(self) -> float:
+        return sum(e.cost for e in self._entries)
+
+    @property
+    def remaining(self) -> float | None:
+        if self._budget is None:
+            return None
+        return max(0.0, self._budget - self.spent)
+
+    @property
+    def over_budget(self) -> bool:
+        if self._budget is None:
+            return False
+        return self.spent >= self._budget
+
+    @property
+    def warning(self) -> bool:
+        if self._budget is None:
+            return False
+        return self.spent >= self._budget * self._warn_at
+
+    def would_exceed(self, cost: float) -> bool:
+        if self._budget is None:
+            return False
+        return self.spent + cost > self._budget
+
+    def record(self, action_id: str, tool_name: str, cost: float) -> CostEntry:
+        entry = CostEntry(
+            action_id=action_id,
+            tool_name=tool_name,
+            cost=cost,
+            timestamp=time.time(),
+        )
+        self._entries.append(entry)
+        return entry
+
+    def summary(self) -> dict:
+        return {
+            "budget": self._budget,
+            "spent": round(self.spent, 4),
+            "remaining": round(self.remaining, 4) if self.remaining is not None else None,
+            "entries": len(self._entries),
+            "over_budget": self.over_budget,
+            "warning": self.warning,
+        }