agent_os_kernel 3.1.0__py3-none-any.whl

agent_os/_circuit_breaker_impl.py

@@ -0,0 +1,232 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Standalone circuit breaker implementation.
+
+This module provides a self-contained circuit breaker that requires no
+external packages beyond the Python standard library.  It is used as a
+fallback by ``agent_os.circuit_breaker`` when ``agent_sre`` is not installed.
+"""
+
+from __future__ import annotations
+
+import inspect
+import threading
+import time
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+
+class CircuitState(str, Enum):
+    CLOSED = "CLOSED"
+    OPEN = "OPEN"
+    HALF_OPEN = "HALF_OPEN"
+
+
+class CircuitOpenError(Exception):
+    """Raised when a call is attempted on an open circuit."""
+
+    def __init__(self, agent_id: str, retry_after: float) -> None:
+        self.agent_id = agent_id
+        self.retry_after = retry_after
+        super().__init__(
+            f"Circuit breaker OPEN for agent '{agent_id}'. "
+            f"Retry after {retry_after:.1f}s."
+        )
+
+
+CircuitBreakerOpen = CircuitOpenError
+
+
+@dataclass(init=False)
+class CircuitBreakerConfig:
+    """Configuration for a circuit breaker instance."""
+
+    failure_threshold: int = 5
+    recovery_timeout_seconds: float = 30.0
+    half_open_max_calls: int = 1
+
+    def __init__(
+        self,
+        failure_threshold: int = 5,
+        recovery_timeout_seconds: float | None = None,
+        half_open_max_calls: int = 1,
+        reset_timeout_seconds: float | None = None,
+    ) -> None:
+        if (
+            recovery_timeout_seconds is not None
+            and reset_timeout_seconds is not None
+            and recovery_timeout_seconds != reset_timeout_seconds
+        ):
+            raise ValueError(
+                "recovery_timeout_seconds and reset_timeout_seconds must match"
+            )
+        timeout = recovery_timeout_seconds
+        if timeout is None:
+            timeout = reset_timeout_seconds
+        if timeout is None:
+            timeout = 30.0
+        self.failure_threshold = failure_threshold
+        self.recovery_timeout_seconds = timeout
+        self.half_open_max_calls = half_open_max_calls
+
+    @property
+    def reset_timeout_seconds(self) -> float:
+        return self.recovery_timeout_seconds
+
+    @reset_timeout_seconds.setter
+    def reset_timeout_seconds(self, value: float) -> None:
+        self.recovery_timeout_seconds = value
+
+
+class CircuitBreaker:
+    """Standalone circuit breaker (used when agent_sre is not installed)."""
+
+    def __init__(
+        self,
+        agent_id: str | CircuitBreakerConfig | None = None,
+        config: CircuitBreakerConfig | None = None,
+    ) -> None:
+        if isinstance(agent_id, CircuitBreakerConfig) and config is None:
+            config = agent_id
+            agent_id = None
+        self.agent_id = agent_id or "legacy"
+        self.config = config or CircuitBreakerConfig()
+        self._config = self.config
+        self._state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._success_count = 0
+        self._half_open_calls = 0
+        self._last_failure_time = 0.0
+        self._lock = threading.Lock()
+
+    @property
+    def state(self) -> str:
+        return self.get_state().value
+
+    @property
+    def failure_count(self) -> int:
+        return self._failure_count
+
+    def get_state(self) -> CircuitState:
+        with self._lock:
+            self._maybe_transition_to_half_open()
+            return self._state
+
+    def call(self, func: Any, *args: Any, fallback: Any = None, **kwargs: Any) -> Any:
+        retry_after = self._prepare_call()
+        if retry_after is not None:
+            if fallback is not None:
+                return fallback
+            raise CircuitOpenError(self.agent_id, retry_after)
+        try:
+            result = func(*args, **kwargs)
+        except Exception:
+            self.record_failure()
+            raise
+        if inspect.isawaitable(result):
+            async def _await_result() -> Any:
+                try:
+                    value = await result
+                except Exception:
+                    self.record_failure()
+                    raise
+                self.record_success()
+                return value
+            return _await_result()
+        self.record_success()
+        return result
+
+    def record_success(self) -> None:
+        with self._lock:
+            if self._state is CircuitState.HALF_OPEN:
+                self._transition(CircuitState.CLOSED)
+            self._failure_count = 0
+            self._success_count += 1
+            self._half_open_calls = 0
+
+    def record_failure(self) -> None:
+        with self._lock:
+            self._failure_count += 1
+            self._last_failure_time = time.monotonic()
+            if self._state is CircuitState.HALF_OPEN:
+                self._transition(CircuitState.OPEN)
+                self._half_open_calls = 0
+            elif self._failure_count >= self.config.failure_threshold:
+                self._transition(CircuitState.OPEN)
+
+    def reset(self) -> None:
+        with self._lock:
+            self._transition(CircuitState.CLOSED)
+            self._failure_count = 0
+            self._success_count = 0
+            self._half_open_calls = 0
+            self._last_failure_time = 0.0
+
+    def _prepare_call(self) -> float | None:
+        with self._lock:
+            self._maybe_transition_to_half_open()
+            if self._state is CircuitState.OPEN:
+                return self._time_until_recovery()
+            if self._state is CircuitState.HALF_OPEN:
+                if self._half_open_calls >= self.config.half_open_max_calls:
+                    return self._time_until_recovery()
+                self._half_open_calls += 1
+            return None
+
+    def _maybe_transition_to_half_open(self) -> None:
+        if self._state is CircuitState.OPEN:
+            elapsed = time.monotonic() - self._last_failure_time
+            if elapsed >= self.config.recovery_timeout_seconds:
+                self._transition(CircuitState.HALF_OPEN)
+
+    def _transition(self, new_state: CircuitState) -> None:
+        self._state = new_state
+        if new_state is CircuitState.HALF_OPEN:
+            self._half_open_calls = 0
+
+    def _time_until_recovery(self) -> float:
+        elapsed = time.monotonic() - self._last_failure_time
+        return max(0.0, self.config.recovery_timeout_seconds - elapsed)
+
+
+class CascadeDetector:
+    """Detects cascading failures across multiple agents."""
+
+    def __init__(
+        self,
+        agents: list[str],
+        cascade_threshold: int = 3,
+        config: CircuitBreakerConfig | None = None,
+    ) -> None:
+        self.cascade_threshold = cascade_threshold
+        self._breakers: dict[str, CircuitBreaker] = {
+            agent_id: CircuitBreaker(agent_id, config) for agent_id in agents
+        }
+
+    def get_breaker(self, agent_id: str) -> CircuitBreaker | None:
+        return self._breakers.get(agent_id)
+
+    def check_cascade(self) -> bool:
+        return len(self.get_affected_agents()) >= self.cascade_threshold
+
+    def get_affected_agents(self) -> list[str]:
+        return [
+            agent_id
+            for agent_id, breaker in self._breakers.items()
+            if breaker.state == CircuitState.OPEN.value
+        ]
+
+    def reset_all(self) -> None:
+        for breaker in self._breakers.values():
+            breaker.reset()
+
+
+__all__ = [
+    "CascadeDetector",
+    "CircuitBreaker",
+    "CircuitBreakerConfig",
+    "CircuitBreakerOpen",
+    "CircuitOpenError",
+    "CircuitState",
+]

agent_os/_mcp_metrics.py

@@ -0,0 +1,193 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""OpenTelemetry-friendly metrics helpers for MCP governance components."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Protocol
+
+logger = logging.getLogger(__name__)
+
+try:
+    from opentelemetry import metrics as _otel_metrics
+
+    _HAS_OTEL = True
+except ImportError:  # pragma: no cover
+    _otel_metrics = None  # type: ignore[assignment]
+    _HAS_OTEL = False
+
+
+class MCPMetricsRecorder(Protocol):
+    """Protocol for MCP governance metric emission."""
+
+    def record_decision(
+        self,
+        *,
+        allowed: bool,
+        agent_id: str,
+        tool_name: str,
+        stage: str,
+    ) -> None:
+        """Record an allow or deny decision."""
+
+    def record_threats_detected(
+        self,
+        count: int,
+        *,
+        tool_name: str,
+        server_name: str,
+    ) -> None:
+        """Record detected threats from an MCP scan."""
+
+    def record_rate_limit_hit(self, *, agent_id: str, tool_name: str) -> None:
+        """Record a rate limit rejection."""
+
+    def record_scan(
+        self,
+        *,
+        operation: str,
+        tool_name: str,
+        server_name: str,
+    ) -> None:
+        """Record an MCP scan invocation."""
+
+
+class NoOpMCPMetrics:
+    """No-op metrics recorder used when OTel is unavailable."""
+
+    def record_decision(
+        self,
+        *,
+        allowed: bool,
+        agent_id: str,
+        tool_name: str,
+        stage: str,
+    ) -> None:
+        return None
+
+    def record_threats_detected(
+        self,
+        count: int,
+        *,
+        tool_name: str,
+        server_name: str,
+    ) -> None:
+        return None
+
+    def record_rate_limit_hit(self, *, agent_id: str, tool_name: str) -> None:
+        return None
+
+    def record_scan(
+        self,
+        *,
+        operation: str,
+        tool_name: str,
+        server_name: str,
+    ) -> None:
+        return None
+
+
+class MCPMetrics(NoOpMCPMetrics):
+    """MCP governance counters backed by OpenTelemetry when available."""
+
+    def __init__(self, meter_provider: Any | None = None) -> None:
+        self._enabled = _HAS_OTEL
+        self._decisions = None
+        self._threats_detected = None
+        self._rate_limit_hits = None
+        self._scans = None
+        if not _HAS_OTEL:
+            return
+
+        try:
+            if meter_provider is not None:
+                meter = meter_provider.get_meter("agent_os.mcp", version="3.1.0")
+            else:
+                meter = _otel_metrics.get_meter("agent_os.mcp", version="3.1.0")
+
+            self._decisions = meter.create_counter(
+                "mcp_decisions",
+                description="Total MCP gateway allow and deny decisions.",
+            )
+            self._threats_detected = meter.create_counter(
+                "mcp_threats_detected",
+                description="Threats detected by MCP scanners.",
+            )
+            self._rate_limit_hits = meter.create_counter(
+                "mcp_rate_limit_hits",
+                description="MCP requests denied by rate limiting.",
+            )
+            self._scans = meter.create_counter(
+                "mcp_scans",
+                description="MCP scan operations performed.",
+            )
+        except Exception:  # pragma: no cover - defensive opt-in path
+            logger.debug("Failed to initialize MCP OpenTelemetry counters", exc_info=True)
+            self._enabled = False
+
+    def record_decision(
+        self,
+        *,
+        allowed: bool,
+        agent_id: str,
+        tool_name: str,
+        stage: str,
+    ) -> None:
+        if not self._enabled or self._decisions is None:
+            return
+        self._decisions.add(
+            1,
+            {
+                "agent_id": agent_id,
+                "tool_name": tool_name,
+                "decision": "allow" if allowed else "deny",
+                "stage": stage,
+            },
+        )
+
+    def record_threats_detected(
+        self,
+        count: int,
+        *,
+        tool_name: str,
+        server_name: str,
+    ) -> None:
+        if count <= 0 or not self._enabled or self._threats_detected is None:
+            return
+        self._threats_detected.add(
+            count,
+            {
+                "tool_name": tool_name,
+                "server_name": server_name,
+            },
+        )
+
+    def record_rate_limit_hit(self, *, agent_id: str, tool_name: str) -> None:
+        if not self._enabled or self._rate_limit_hits is None:
+            return
+        self._rate_limit_hits.add(
+            1,
+            {
+                "agent_id": agent_id,
+                "tool_name": tool_name,
+            },
+        )
+
+    def record_scan(
+        self,
+        *,
+        operation: str,
+        tool_name: str,
+        server_name: str,
+    ) -> None:
+        if not self._enabled or self._scans is None:
+            return
+        self._scans.add(
+            1,
+            {
+                "operation": operation,
+                "tool_name": tool_name,
+                "server_name": server_name,
+            },
+        )

agent_os/adversarial.py

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Adversarial Evaluation — backward-compatibility shim.
+
+Attempts to import the canonical implementation from
+``agent_sre.chaos.adversarial_policy``.  When ``agent_sre`` is not
+installed the standalone fallback in ``agent_os._adversarial_impl`` is
+re-exported so that ``agent_os`` continues to work without requiring the
+optional SRE package.
+
+.. deprecated::
+    Import from ``agent_sre.chaos.adversarial_policy`` instead.
+"""
+
+from __future__ import annotations
+
+try:
+    from agent_sre.chaos.adversarial_policy import *  # noqa: F401,F403
+except ImportError:
+    from agent_os._adversarial_impl import *  # noqa: F401,F403