Aegis

"""
Aegis: Agent Loop Prevention & Oscillation Breaker
Detects infinite tool-call loops via sliding-window n-gram trajectory hashing.
"""
import hashlib
from typing import List, Dict, Any, Optional
from pydantic import BaseModel, Field

class ToolCallEvent(BaseModel):
    tool_name: str
    arguments_hash: str
    iteration: int

class TrajectoryEvaluation(BaseModel):
    is_oscillating: bool
    detected_cycle_length: int
    confidence_score: float
    mitigation_action: str

class AegisLoopGuard:
    def __init__(self, window_size: int = 6, max_cycle_repetitions: int = 2):
        self.window_size = window_size
        self.max_cycle_repetitions = max_cycle_repetitions
        self.history: List[str] = []

    def record_and_evaluate(self, tool_name: str, args: Dict[str, Any], iteration: int) -> TrajectoryEvaluation:
        # Create deterministic fingerprint of the tool call
        canonical_str = f"{tool_name}:{sorted(args.items())}"
        step_hash = hashlib.sha256(canonical_str.encode()).hexdigest()[:12]
        self.history.append(step_hash)

        # Inspect sliding window for repeating cycles (periods 1 to window/2)
        history_len = len(self.history)
        for period in range(1, (self.window_size // 2) + 1):
            if history_len >= period * self.max_cycle_repetitions:
                pattern = self.history[-period:]
                repeats = True
                for r in range(1, self.max_cycle_repetitions):
                    start_idx = -(period * (r + 1))
                    end_idx = -(period * r)
                    if self.history[start_idx:end_idx] != pattern:
                        repeats = False
                        break
                if repeats:
                    return TrajectoryEvaluation(
                        is_oscillating=True,
                        detected_cycle_length=period,
                        confidence_score=0.98,
                        mitigation_action="HALT_TOOL_EXECUTION_AND_INJECT_SYSTEM_CORRECTION"
                    )

        return TrajectoryEvaluation(
            is_oscillating=False,
            detected_cycle_length=0,
            confidence_score=0.0,
            mitigation_action="PROCEED"
        )