How to stop an AI agent infinite loop in Python (step caps, repeat checks)

import hashlib
import json
from collections import Counter, deque

class LoopStop(Exception):
    pass

TEXT_KEYS = {"query", "q", "question"}  # arguments where case and spacing do not matter

def normalize(args: dict) -> dict:
    """Treat "Python loops " and "python loops" as the same search; leave paths and code alone."""
    return {k: " ".join(v.lower().split()) if k in TEXT_KEYS and isinstance(v, str) else v
            for k, v in args.items()}

def call_key(name: str, args: dict, state: str = "") -> str:
    """state: a version of the world for tools whose answer depends on it (e.g. the git tree hash)."""
    return hashlib.sha256(json.dumps([name, normalize(args), state], sort_keys=True, default=str).encode()).hexdigest()

class LoopGuard:
    def __init__(self, max_steps: int = 25, max_repeats: int = 1, stall_steps: int = 5):
        self.max_steps, self.max_repeats, self.stall_steps = max_steps, max_repeats, stall_steps
        self.steps = 0
        self.calls: Counter[str] = Counter()
        self.recent: deque[str] = deque(maxlen=stall_steps + 1)

    def tick(self, fingerprint: str) -> None:
        """Call at the start of every model turn with a fingerprint of the useful state."""
        self.steps += 1
        if self.steps > self.max_steps:
            raise LoopStop(f"step cap of {self.max_steps} reached")
        self.recent.append(fingerprint)
        if len(self.recent) == self.recent.maxlen and len(set(self.recent)) == 1:
            raise LoopStop(f"no new useful result in the last {self.stall_steps} steps")

    def record_call(self, key: str) -> bool:
        """Records a tool call. Returns True if it repeats an earlier one; raises LoopStop
        when a call has been repeated more than max_repeats times."""
        self.calls[key] += 1
        if self.calls[key] > self.max_repeats + 1:
            raise LoopStop(f"the same call was made {self.calls[key]} times")
        return self.calls[key] > 1

def check_finish(args: dict) -> str | None:
    """None if the result meets the contract, else the reason it does not."""
    findings = args.get("findings")
    if not isinstance(findings, list) or len(findings) < 3:
        return "finish needs at least 3 findings"
    for i, f in enumerate(findings, 1):
        if not (isinstance(f, dict) and str(f.get("url", "")).startswith(("http://", "https://"))
                and f.get("quote")):
            return f"finding {i} needs a url and a quote"
    return None

def useful(entry: dict) -> bool:
    obs = entry["observation"]
    return bool(obs) and not (isinstance(obs, dict) and "error" in obs)

def fingerprint(results: dict) -> str:
    return hashlib.sha256(json.dumps(sorted(k for k, e in results.items() if useful(e))).encode()).hexdigest()

def run_agent(task: str, call_model, tools: dict, guard: LoopGuard | None = None) -> dict:
    """call_model(messages) returns {"tool": name, "args": {...}}; "finish" ends the run."""
    guard = guard or LoopGuard()  # one guard per run: it holds the run's counts
    messages = [{"role": "user", "content": task}]
    results: dict[str, dict] = {}  # per distinct call: tool, args, observation
    try:
        while True:
            guard.tick(fingerprint(results))
            action = call_model(messages)
            if not isinstance(action, dict):
                messages.append({"role": "user", "content": 'Reply with JSON: {"tool": ..., "args": {...}}'})
                continue
            name, args = action.get("tool"), action.get("args") or {}
            messages.append({"role": "assistant", "content": json.dumps(action)})
            if name == "finish":
                problem = check_finish(args)
                if problem is None:
                    return {"status": "done", "result": args, "steps": guard.steps}
                messages.append({"role": "user", "content": f"Not finished: {problem}"})
                continue
            if name not in tools:
                options = sorted([*tools, "finish"])
                messages.append({"role": "user", "content": f"Unknown tool {name}. Use one of {options}."})
                continue
            key = call_key(name, args)
            if guard.record_call(key):
                observation = {"note": "You already made this call. Same result as before; "
                                       "try a different approach or finish.",
                               "result": results[key]["observation"]}
            else:
                try:
                    observation = tools[name](**args)
                except Exception as e:  # errors become observations, not crashes
                    observation = {"error": f"{name} failed: {e}"}
                results[key] = {"tool": name, "args": args, "observation": observation}
            messages.append({"role": "user", "content": json.dumps(observation, default=str)})
    except LoopStop as e:
        return {"status": "stopped", "reason": str(e), "steps": guard.steps,
                "partial": [x for x in results.values() if useful(x)]}

def search(query: str) -> list:
    return []  # nothing found, every time

def stuck_model(messages):
    return {"tool": "search", "args": {"query": "agent loops"}}

print(run_agent("Find three sources on agent loops", stuck_model, {"search": search}))
# {'status': 'stopped', 'reason': 'the same call was made 3 times', 'steps': 3, 'partial': []}