""" 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), )