icdev 0.0.3__py3-none-any.whl

tools/core/circuit_breaker.py

@@ -0,0 +1,353 @@
+# CUI // SP-CTI
+"""Circuit breaker implementation for ICDEV agent and LLM provider calls.
+
+Reads configuration from ``args/resilience_config.yaml`` (circuit_breaker section)
+and provides a thread-safe circuit breaker that can be used as a decorator or
+context manager.
+
+States:
+  CLOSED   - Normal operation; failures are counted.
+  OPEN     - Tripped; all calls are rejected immediately.
+  HALF_OPEN - Recovery probe; limited calls allowed to test recovery.
+
+Usage::
+
+    from tools.core.circuit_breaker import CircuitBreaker, CircuitOpenError
+
+    breaker = CircuitBreaker("bedrock_llm")
+
+    # As decorator
+    @breaker
+    def call_bedrock(prompt):
+        ...
+
+    # As context manager
+    with breaker:
+        response = call_bedrock(prompt)
+
+    # Manual
+    if breaker.allow_request():
+        try:
+            result = call_bedrock(prompt)
+            breaker.record_success()
+        except Exception:
+            breaker.record_failure()
+            raise
+
+Configuration (args/resilience_config.yaml)::
+
+    circuit_breaker:
+      enabled: true
+      failure_threshold: 5
+      reset_timeout_seconds: 60
+      half_open_success_threshold: 2
+      window_size: 20
+      min_calls_before_trip: 5
+      overrides:
+        bedrock_llm:
+          failure_threshold: 3
+          reset_timeout_seconds: 120
+
+Classification: CUI // SP-CTI
+"""
+
+from __future__ import annotations
+
+import functools
+import logging
+import threading
+import time
+from collections import deque
+from enum import Enum
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger("icdev.core.circuit_breaker")
+
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
+
+
+# ---------------------------------------------------------------------------
+# Exceptions
+# ---------------------------------------------------------------------------
+
+class CircuitOpenError(Exception):
+    """Raised when the circuit breaker is OPEN and rejects a call."""
+
+    def __init__(self, breaker_name: str, retry_after_seconds: float):
+        self.breaker_name = breaker_name
+        self.retry_after_seconds = retry_after_seconds
+        super().__init__(
+            f"Circuit breaker '{breaker_name}' is OPEN; retry after "
+            f"{retry_after_seconds:.1f}s"
+        )
+
+
+# ---------------------------------------------------------------------------
+# State enum
+# ---------------------------------------------------------------------------
+
+class CircuitState(Enum):
+    CLOSED = "closed"
+    OPEN = "open"
+    HALF_OPEN = "half_open"
+
+
+# ---------------------------------------------------------------------------
+# Configuration loader
+# ---------------------------------------------------------------------------
+
+def _load_config() -> Dict[str, Any]:
+    """Load circuit breaker config from resilience_config.yaml."""
+    try:
+        import yaml
+    except ImportError:
+        return {}
+    config_path = _PROJECT_ROOT / "args" / "resilience_config.yaml"
+    if not config_path.exists():
+        return {}
+    try:
+        with open(config_path, "r", encoding="utf-8") as f:
+            data = yaml.safe_load(f) or {}
+        return data.get("circuit_breaker", {})
+    except Exception as exc:
+        logger.warning("Failed to load resilience config: %s", exc)
+        return {}
+
+
+_CONFIG: Optional[Dict[str, Any]] = None
+_CONFIG_LOCK = threading.Lock()
+
+
+def _get_config() -> Dict[str, Any]:
+    """Lazily load and cache config."""
+    global _CONFIG
+    if _CONFIG is None:
+        with _CONFIG_LOCK:
+            if _CONFIG is None:
+                _CONFIG = _load_config()
+    return _CONFIG
+
+
+# ---------------------------------------------------------------------------
+# CircuitBreaker
+# ---------------------------------------------------------------------------
+
+class CircuitBreaker:
+    """Thread-safe circuit breaker with sliding window failure tracking.
+
+    Parameters
+    ----------
+    name : str
+        Identifier for the protected subsystem (matches overrides key in YAML).
+    failure_threshold : int | None
+        Override; defaults to config or 5.
+    reset_timeout_seconds : float | None
+        Override; defaults to config or 60.
+    half_open_success_threshold : int | None
+        Override; defaults to config or 2.
+    window_size : int | None
+        Sliding window of recent outcomes; defaults to config or 20.
+    min_calls_before_trip : int | None
+        Minimum calls in window before breaker can trip; defaults to config or 5.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        failure_threshold: Optional[int] = None,
+        reset_timeout_seconds: Optional[float] = None,
+        half_open_success_threshold: Optional[int] = None,
+        window_size: Optional[int] = None,
+        min_calls_before_trip: Optional[int] = None,
+    ):
+        self.name = name
+        cfg = _get_config()
+        overrides = cfg.get("overrides", {}).get(name, {})
+
+        def _resolve(explicit, override_key, default):
+            if explicit is not None:
+                return explicit
+            return overrides.get(override_key, cfg.get(override_key, default))
+
+        self.failure_threshold: int = _resolve(failure_threshold, "failure_threshold", 5)
+        self.reset_timeout: float = _resolve(
+            reset_timeout_seconds, "reset_timeout_seconds", 60.0
+        )
+        self.half_open_success_threshold: int = _resolve(
+            half_open_success_threshold, "half_open_success_threshold", 2
+        )
+        self.window_size: int = _resolve(window_size, "window_size", 20)
+        self.min_calls_before_trip: int = _resolve(
+            min_calls_before_trip, "min_calls_before_trip", 5
+        )
+
+        self._lock = threading.Lock()
+        self._state = CircuitState.CLOSED
+        self._outcomes: deque = deque(maxlen=self.window_size)  # True=success, False=failure
+        self._last_failure_time: float = 0.0
+        self._half_open_successes: int = 0
+        self._half_open_calls: int = 0
+
+        self._enabled: bool = cfg.get("enabled", True)
+
+    # -- State inspection ---------------------------------------------------
+
+    @property
+    def state(self) -> CircuitState:
+        """Current state, accounting for reset timeout transition."""
+        with self._lock:
+            return self._compute_state()
+
+    def _compute_state(self) -> CircuitState:
+        """Must be called under lock."""
+        if self._state == CircuitState.OPEN:
+            elapsed = time.monotonic() - self._last_failure_time
+            if elapsed >= self.reset_timeout:
+                self._state = CircuitState.HALF_OPEN
+                self._half_open_successes = 0
+                self._half_open_calls = 0
+                logger.info(
+                    "Circuit breaker '%s' transitioned OPEN -> HALF_OPEN", self.name
+                )
+        return self._state
+
+    # -- Request gating -----------------------------------------------------
+
+    def allow_request(self) -> bool:
+        """Check whether a call should proceed."""
+        if not self._enabled:
+            return True
+        with self._lock:
+            state = self._compute_state()
+            if state == CircuitState.CLOSED:
+                return True
+            if state == CircuitState.HALF_OPEN:
+                if self._half_open_calls < self.half_open_success_threshold:
+                    self._half_open_calls += 1
+                    return True
+                return False
+            return False  # OPEN
+
+    def _retry_after(self) -> float:
+        """Seconds until OPEN -> HALF_OPEN transition."""
+        elapsed = time.monotonic() - self._last_failure_time
+        return max(0.0, self.reset_timeout - elapsed)
+
+    # -- Outcome recording --------------------------------------------------
+
+    def record_success(self) -> None:
+        """Record a successful call."""
+        with self._lock:
+            self._outcomes.append(True)
+            if self._state == CircuitState.HALF_OPEN:
+                self._half_open_successes += 1
+                if self._half_open_successes >= self.half_open_success_threshold:
+                    self._state = CircuitState.CLOSED
+                    self._outcomes.clear()
+                    logger.info(
+                        "Circuit breaker '%s' CLOSED after successful probes",
+                        self.name,
+                    )
+
+    def record_failure(self) -> None:
+        """Record a failed call; may trip the breaker."""
+        with self._lock:
+            self._outcomes.append(False)
+            self._last_failure_time = time.monotonic()
+
+            if self._state == CircuitState.HALF_OPEN:
+                # Any failure in half-open re-opens
+                self._state = CircuitState.OPEN
+                logger.warning(
+                    "Circuit breaker '%s' re-OPENED from HALF_OPEN", self.name
+                )
+                return
+
+            # Check sliding window
+            if len(self._outcomes) >= self.min_calls_before_trip:
+                failure_count = sum(1 for o in self._outcomes if not o)
+                if failure_count >= self.failure_threshold:
+                    self._state = CircuitState.OPEN
+                    logger.warning(
+                        "Circuit breaker '%s' OPENED (%d failures in last %d calls)",
+                        self.name,
+                        failure_count,
+                        len(self._outcomes),
+                    )
+
+    # -- Reset (for testing) ------------------------------------------------
+
+    def reset(self) -> None:
+        """Force-reset to CLOSED (testing only)."""
+        with self._lock:
+            self._state = CircuitState.CLOSED
+            self._outcomes.clear()
+            self._half_open_successes = 0
+            self._half_open_calls = 0
+
+    # -- Context manager / decorator ----------------------------------------
+
+    def __enter__(self):
+        if not self.allow_request():
+            raise CircuitOpenError(self.name, self._retry_after())
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is None:
+            self.record_success()
+        else:
+            self.record_failure()
+        return False  # Don't suppress exceptions
+
+    def __call__(self, func: Callable) -> Callable:
+        """Use as a decorator."""
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            with self:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    # -- Serialisation ------------------------------------------------------
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Snapshot for health endpoints and monitoring."""
+        with self._lock:
+            state = self._compute_state()
+            total = len(self._outcomes)
+            failures = sum(1 for o in self._outcomes if not o)
+            return {
+                "name": self.name,
+                "state": state.value,
+                "failure_count": failures,
+                "total_calls_in_window": total,
+                "failure_threshold": self.failure_threshold,
+                "reset_timeout_seconds": self.reset_timeout,
+                "retry_after_seconds": (
+                    self._retry_after() if state == CircuitState.OPEN else 0.0
+                ),
+            }
+
+
+# ---------------------------------------------------------------------------
+# Registry — singleton breakers per subsystem name
+# ---------------------------------------------------------------------------
+
+_REGISTRY: Dict[str, CircuitBreaker] = {}
+_REGISTRY_LOCK = threading.Lock()
+
+
+def get_breaker(name: str, **kwargs: Any) -> CircuitBreaker:
+    """Get or create a circuit breaker for the named subsystem."""
+    if name not in _REGISTRY:
+        with _REGISTRY_LOCK:
+            if name not in _REGISTRY:
+                _REGISTRY[name] = CircuitBreaker(name, **kwargs)
+    return _REGISTRY[name]
+
+
+def all_breakers() -> Dict[str, Dict[str, Any]]:
+    """Snapshot of all registered breakers (for /health endpoints)."""
+    return {name: cb.to_dict() for name, cb in _REGISTRY.items()}

tools/core/compliance_sidecar.py

@@ -0,0 +1,344 @@
+# CUI // SP-CTI
+"""Compliance sidecar decorator — ADR D-ARCH-3.
+
+Wraps tool entry points to enforce CUI marking, audit trail logging,
+and classification validation.  Eliminates duplicated compliance
+boilerplate across 255+ tools.
+
+Usage::
+
+    from tools.core.compliance_sidecar import compliance_sidecar
+
+    @compliance_sidecar()
+    def generate_ssp(project_id: str) -> dict:
+        return {"ssp": "...", "controls": [...]}
+
+    @compliance_sidecar(classification="SECRET", audit_event="ssp_generated")
+    def generate_classified_report() -> dict: ...
+
+Classification: CUI // SP-CTI
+"""
+from __future__ import annotations
+
+import argparse
+import datetime
+import functools
+import json
+import logging
+import sys
+import threading
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger("icdev.core.compliance_sidecar")
+
+CLASSIFICATION_HIERARCHY: List[str] = [
+    "UNCLASSIFIED", "CUI", "CUI // SP-CTI", "SECRET", "TOP SECRET", "TS//SCI",
+]
+DEFAULT_CLASSIFICATION = "CUI // SP-CTI"
+DEFAULT_MAX_LEVEL = "CUI // SP-CTI"
+
+# Module-level stats (thread-safe)
+_stats_lock = threading.Lock()
+_stats: Dict[str, Any] = {
+    "total_wrapped": 0, "total_calls": 0, "total_errors": 0,
+    "calls_by_function": {}, "durations_by_function": {},
+}
+
+# Configuration cache
+_config_cache: Optional[Dict[str, Any]] = None
+_config_lock = threading.Lock()
+
+
+def _load_config() -> Dict[str, Any]:
+    """Load sidecar config from args/compliance_config.yaml (cached)."""
+    global _config_cache
+    with _config_lock:
+        if _config_cache is not None:
+            return _config_cache
+        defaults: Dict[str, Any] = {
+            "sidecar": {
+                "default_classification": DEFAULT_CLASSIFICATION,
+                "enforce_marking": True,
+                "audit_all_calls": True,
+                "classification_hierarchy": list(CLASSIFICATION_HIERARCHY),
+            }
+        }
+        config_path = Path(__file__).resolve().parents[2] / "args" / "compliance_config.yaml"
+        if config_path.exists():
+            try:
+                import yaml  # type: ignore[import-untyped]
+                with open(config_path, "r", encoding="utf-8") as fh:
+                    loaded = yaml.safe_load(fh)
+                if isinstance(loaded, dict):
+                    _config_cache = loaded
+                    return _config_cache
+            except Exception as exc:  # noqa: BLE001
+                logger.debug("Failed to load compliance config: %s", exc)
+        _config_cache = defaults
+        return _config_cache
+
+
+def _get_sidecar_cfg() -> Dict[str, Any]:
+    """Return the ``sidecar:`` block from config."""
+    return _load_config().get("sidecar", {})
+
+
+def _get_hierarchy() -> List[str]:
+    """Return classification hierarchy list."""
+    return _get_sidecar_cfg().get("classification_hierarchy", list(CLASSIFICATION_HIERARCHY))
+
+
+# -- Audit trail helpers ---------------------------------------------------
+
+def _log_audit(
+    event_type: str, action: str, details: Optional[str] = None,
+    classification: str = DEFAULT_CLASSIFICATION, project_id: Optional[str] = None,
+) -> None:
+    """Append an entry to audit_trail (best-effort, never raises)."""
+    try:
+        from tools.db.storage import get_connection  # type: ignore[import-untyped]
+    except ImportError:
+        logger.debug("storage module unavailable — skipping audit log")
+        return
+    try:
+        with get_connection() as conn:
+            conn.execute(
+                "INSERT INTO audit_trail "
+                "(project_id, event_type, actor, action, details, classification, created_at) "
+                "VALUES (?, ?, ?, ?, ?, ?, ?)",
+                (project_id, event_type, "compliance_sidecar", action,
+                 details, classification, datetime.datetime.utcnow().isoformat()),
+            )
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("Audit log write failed: %s", exc)
+
+
+# -- Public helpers --------------------------------------------------------
+
+def _classification_rank(label: str) -> int:
+    """Ordinal rank of *label* in the hierarchy (-1 if unknown)."""
+    hierarchy = _get_hierarchy()
+    upper = label.strip().upper()
+    for idx, level in enumerate(hierarchy):
+        if level.upper() == upper:
+            return idx
+    return -1
+
+
+def validate_classification(text: str, max_level: str = DEFAULT_MAX_LEVEL) -> Dict[str, Any]:
+    """Check *text* for classification markings that exceed *max_level*.
+
+    Returns dict with ``valid`` (bool), ``violations`` (list), ``classification``.
+    """
+    max_rank = _classification_rank(max_level)
+    if max_rank < 0:
+        return {"valid": False, "violations": [f"Unknown max_level: {max_level}"],
+                "classification": DEFAULT_CLASSIFICATION}
+    violations: List[str] = []
+    text_upper = text.upper()
+    for level in reversed(_get_hierarchy()):
+        rank = _classification_rank(level)
+        if rank > max_rank and level.upper() in text_upper:
+            violations.append(
+                f"Found marking '{level}' which exceeds maximum allowed level '{max_level}'"
+            )
+    return {"valid": len(violations) == 0, "violations": violations,
+            "classification": DEFAULT_CLASSIFICATION}
+
+
+def ensure_marking(response_dict: Dict[str, Any], classification: str = DEFAULT_CLASSIFICATION) -> Dict[str, Any]:
+    """Inject ``classification`` key into *response_dict* if missing.
+
+    Recurses into nested dicts and list items.
+    """
+    if not isinstance(response_dict, dict):
+        return response_dict
+    if "classification" not in response_dict:
+        response_dict["classification"] = classification
+    for key, value in response_dict.items():
+        if isinstance(value, dict) and key not in ("payload", "params", "context"):
+            ensure_marking(value, classification)
+        elif isinstance(value, list):
+            for item in value:
+                if isinstance(item, dict):
+                    ensure_marking(item, classification)
+    return response_dict
+
+
+def get_sidecar_stats() -> Dict[str, Any]:
+    """Return decorator call/error statistics.
+
+    Keys: total_wrapped, total_calls, total_errors, calls_by_function,
+    avg_duration_ms, classification.
+    """
+    with _stats_lock:
+        avg_duration: Dict[str, float] = {}
+        for fn_name, durations in _stats["durations_by_function"].items():
+            if durations:
+                avg_duration[fn_name] = round(sum(durations) / len(durations) * 1000.0, 2)
+        return {
+            "total_wrapped": _stats["total_wrapped"],
+            "total_calls": _stats["total_calls"],
+            "total_errors": _stats["total_errors"],
+            "calls_by_function": dict(_stats["calls_by_function"]),
+            "avg_duration_ms": avg_duration,
+            "classification": DEFAULT_CLASSIFICATION,
+        }
+
+
+# -- Core decorator --------------------------------------------------------
+
+def compliance_sidecar(
+    func: Optional[Callable] = None, *,
+    classification: str = DEFAULT_CLASSIFICATION,
+    audit_event: Optional[str] = None,
+) -> Any:
+    """Decorator factory enforcing CUI marking and audit trail on tool functions.
+
+    Parameters
+    ----------
+    classification : str
+        CUI marking to inject into dict responses.
+    audit_event : str | None
+        Audit event_type (defaults to ``"compliance_check"``).
+
+    Supports ``@compliance_sidecar`` and ``@compliance_sidecar(...)`` forms.
+    """
+    effective_event = audit_event or "compliance_check"
+    cfg = _get_sidecar_cfg()
+    enforce = cfg.get("enforce_marking", True)
+    audit_all = cfg.get("audit_all_calls", True)
+
+    def decorator(fn: Callable) -> Callable:
+        with _stats_lock:
+            _stats["total_wrapped"] += 1
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            fn_name = fn.__qualname__
+
+            # Pre-execution: validate input classification if present
+            input_cls = kwargs.get("classification")
+            if input_cls is None and args:
+                first = args[0]
+                if isinstance(first, dict):
+                    input_cls = first.get("classification")
+            if input_cls is not None:
+                v = validate_classification(str(input_cls), max_level=classification)
+                if not v["valid"]:
+                    logger.warning("Input classification violation in %s: %s",
+                                   fn_name, "; ".join(v["violations"]))
+
+            if audit_all:
+                _log_audit(event_type=effective_event,
+                           action=f"sidecar_start:{fn_name}",
+                           details=json.dumps({"function": fn_name, "phase": "start"}),
+                           classification=classification)
+
+            start_time = time.monotonic()
+            try:
+                result = fn(*args, **kwargs)
+            except Exception as exc:
+                elapsed = time.monotonic() - start_time
+                with _stats_lock:
+                    _stats["total_calls"] += 1
+                    _stats["total_errors"] += 1
+                    _stats["calls_by_function"].setdefault(fn_name, 0)
+                    _stats["calls_by_function"][fn_name] += 1
+                    _stats["durations_by_function"].setdefault(fn_name, [])
+                    _stats["durations_by_function"][fn_name].append(elapsed)
+                if audit_all:
+                    _log_audit(event_type=effective_event,
+                               action=f"sidecar_error:{fn_name}",
+                               details=json.dumps({"function": fn_name, "phase": "error",
+                                                    "error": str(exc)[:500],
+                                                    "duration_ms": round(elapsed * 1000.0, 2)}),
+                               classification=classification)
+                raise type(exc)(f"[{classification}] {fn_name}: {exc}") from exc
+
+            elapsed = time.monotonic() - start_time
+            with _stats_lock:
+                _stats["total_calls"] += 1
+                _stats["calls_by_function"].setdefault(fn_name, 0)
+                _stats["calls_by_function"][fn_name] += 1
+                _stats["durations_by_function"].setdefault(fn_name, [])
+                _stats["durations_by_function"][fn_name].append(elapsed)
+
+            # Post-execution: enforce marking on dict/list responses
+            if enforce:
+                if isinstance(result, dict):
+                    ensure_marking(result, classification)
+                elif isinstance(result, list):
+                    for item in result:
+                        if isinstance(item, dict):
+                            ensure_marking(item, classification)
+
+            if audit_all:
+                _log_audit(event_type=effective_event,
+                           action=f"sidecar_complete:{fn_name}",
+                           details=json.dumps({"function": fn_name, "phase": "complete",
+                                                "duration_ms": round(elapsed * 1000.0, 2),
+                                                "result_type": type(result).__name__}),
+                           classification=classification)
+            return result
+
+        return wrapper
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+# -- CLI -------------------------------------------------------------------
+
+def main() -> None:
+    """CLI entry point: --stats, --validate --text <text>, --json."""
+    parser = argparse.ArgumentParser(
+        description="Compliance sidecar — CUI marking enforcement and audit trail decorator.")
+    parser.add_argument("--stats", action="store_true", help="Show call/error statistics.")
+    parser.add_argument("--validate", action="store_true", help="Validate classification markings.")
+    parser.add_argument("--text", type=str, default="", help="Text to validate (with --validate).")
+    parser.add_argument("--max-level", type=str, default=DEFAULT_MAX_LEVEL,
+                        help="Maximum allowed classification level.")
+    parser.add_argument("--json", action="store_true", help="JSON output.")
+    args = parser.parse_args()
+
+    if args.stats:
+        result = get_sidecar_stats()
+        if args.json:
+            print(json.dumps(result, indent=2))
+        else:
+            print("Compliance Sidecar Statistics")
+            print("=" * 40)
+            print(f"  Functions wrapped : {result['total_wrapped']}")
+            print(f"  Total calls       : {result['total_calls']}")
+            print(f"  Total errors      : {result['total_errors']}")
+            if result["calls_by_function"]:
+                print("  Calls by function:")
+                for fn, count in sorted(result["calls_by_function"].items()):
+                    avg = result["avg_duration_ms"].get(fn, 0.0)
+                    print(f"    {fn}: {count} calls, avg {avg:.2f} ms")
+            print(f"  Classification    : {result['classification']}")
+        return
+
+    if args.validate:
+        if not args.text:
+            parser.error("--validate requires --text")
+        result = validate_classification(args.text, max_level=args.max_level)
+        if args.json:
+            print(json.dumps(result, indent=2))
+        else:
+            status = "VALID" if result["valid"] else "VIOLATIONS FOUND"
+            print(f"Classification Validation: {status}")
+            for v in result["violations"]:
+                print(f"  - {v}")
+            print(f"  Classification: {result['classification']}")
+        return
+
+    parser.print_help()
+
+
+if __name__ == "__main__":
+    main()