lambdagent/lambdagent_guard/core.py

"""
lambdagent_guard.core — Shared guard infrastructure (I10).

Provides:
  - GuardConfig: unified config for all framework guards
  - Compile-time checks: type safety, cost prediction, store independence
  - Runtime hooks: cost accumulation, loop detection, budget enforcement
"""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional


# ============================================================
# Exceptions
# ============================================================

class CostBudgetExceeded(RuntimeError):
    """Execution cost exceeded the configured budget ceiling."""
    def __init__(self, spent: float, budget: float, step: int = 0):
        self.spent = spent
        self.budget = budget
        self.step = step
        super().__init__(
            f"Cost budget exceeded at step {step}: "
            f"${spent:.4f} > ${budget:.4f}"
        )


class InfiniteLoopDetected(RuntimeError):
    """Repeated identical states suggest no forward progress."""
    def __init__(self, message: str, step: int = 0):
        self.step = step
        super().__init__(message)


class StoreConflictError(RuntimeError):
    """Parallel agents write to overlapping state keys."""
    pass


# ============================================================
# Configuration
# ============================================================

@dataclass
class GuardConfig:
    """Unified guard configuration for all framework wrappers."""
    cost_budget: float = float("inf")      # Max cost in USD
    type_check: bool = True                # T-Compose checking at wrap time
    loop_detection: bool = True            # Repeated-state detection
    loop_window: int = 5                   # Recent steps to check
    loop_threshold: int = 3                # Repeats before alert
    parallel_safety: bool = False          # Store independence check
    terminate_check: bool = True           # Verify termination exists
    cost_alert: Optional[Callable[[float], None]] = None  # Callback on cost
    empty_message_detection: bool = False  # AutoGen #108 blank-msg detection
    terminate_robustness: bool = False     # Don't rely on exact string match


@dataclass
class GuardedResult:
    """Result from a guarded execution."""
    result: Any
    total_cost_usd: float = 0.0
    total_tokens: int = 0
    steps: int = 0
    loop_detected: bool = False
    budget_remaining: float = float("inf")


# ============================================================
# Compile-Time Checks
# ============================================================

def run_compile_checks(config: dict, guard_cfg: GuardConfig) -> List[str]:
    """
    Run compile-time static checks on extracted config.
    Returns list of warning/error messages (empty = all clear).
    """
    warnings = []

    # Lint check
    try:
        from lambdagent.fromconfig.lint import lint_config
        results = lint_config(config)
        errors = [r for r in results if r.level == "ERROR"]
        if errors:
            for e in errors[:5]:
                warnings.append(f"[LINT {e.rule}] {e.message}")
    except Exception:
        pass

    # Cost prediction
    if guard_cfg.cost_budget < float("inf"):
        try:
            from lambdagent.cost_grade import estimate_cost
            from lambdagent.fromconfig import from_config
            import tempfile, yaml, os

            with tempfile.NamedTemporaryFile(mode="w", suffix=".yml", delete=False) as f:
                yaml.dump(config, f, allow_unicode=True)
                tmp = f.name
            try:
                term = from_config(tmp)
                grade = estimate_cost(term)
                if grade.money > guard_cfg.cost_budget:
                    warnings.append(
                        f"[COST] Predicted worst-case ${grade.money:.2f} exceeds "
                        f"budget ${guard_cfg.cost_budget:.2f}"
                    )
                if grade.probability < 0.05:
                    warnings.append(
                        f"[COST] Success probability {grade.probability:.1%} is critically low"
                    )
            finally:
                os.unlink(tmp)
        except Exception:
            pass

    # Terminate check
    if guard_cfg.terminate_check:
        tools = config.get("mcp", {}).get("localTools", [])
        agent_type = config.get("type", "simple")
        if agent_type == "react" and "terminate" not in tools:
            warnings.append(
                "[LINT L004] No terminate tool in ReAct loop — "
                "forced truncation at maxSteps, not clean exit"
            )

    return warnings


# ============================================================
# Runtime Monitor
# ============================================================

class RuntimeMonitor:
    """
    Tracks cost and state during execution.
    Injected as a hook into framework execution loops.
    """

    def __init__(self, guard_cfg: GuardConfig):
        self.cfg = guard_cfg
        self.total_cost = 0.0
        self.total_tokens = 0
        self.step_count = 0
        self._state_hashes: List[str] = []
        self._empty_msg_count = 0

    def on_step(self, step_info: Dict[str, Any]):
        """Called before/after each LLM/tool step."""
        self.step_count += 1

        # Cost tracking
        tokens = step_info.get("tokens", 0)
        cost = step_info.get("cost_usd", tokens * 0.000003)  # default: ~sonnet rate
        self.total_tokens += tokens
        self.total_cost += cost

        if self.cfg.cost_alert:
            self.cfg.cost_alert(self.total_cost)

        # Budget enforcement
        if self.total_cost > self.cfg.cost_budget:
            raise CostBudgetExceeded(self.total_cost, self.cfg.cost_budget, self.step_count)

        # Loop detection
        if self.cfg.loop_detection:
            state_str = str(step_info.get("output", ""))[:500]
            h = hashlib.md5(state_str.encode()).hexdigest()[:12]
            self._state_hashes.append(h)
            if len(self._state_hashes) > self.cfg.loop_window:
                self._state_hashes = self._state_hashes[-self.cfg.loop_window:]
            if self._state_hashes.count(h) >= self.cfg.loop_threshold:
                raise InfiniteLoopDetected(
                    f"Same output repeated {self._state_hashes.count(h)} times "
                    f"in last {self.cfg.loop_window} steps at step {self.step_count}",
                    self.step_count,
                )

        # Empty message detection (AutoGen #108)
        if self.cfg.empty_message_detection:
            output = str(step_info.get("output", "")).strip()
            if len(output) < 5:
                self._empty_msg_count += 1
                if self._empty_msg_count >= 3:
                    raise InfiniteLoopDetected(
                        f"Empty/near-empty messages detected {self._empty_msg_count} "
                        f"consecutive times (cf. AutoGen #108)",
                        self.step_count,
                    )
            else:
                self._empty_msg_count = 0

    def result(self, value: Any) -> GuardedResult:
        """Build final GuardedResult."""
        return GuardedResult(
            result=value,
            total_cost_usd=self.total_cost,
            total_tokens=self.total_tokens,
            steps=self.step_count,
            budget_remaining=max(0, self.cfg.cost_budget - self.total_cost),
        )