AI Agent Stuck in a Loop? Detect and Stop It

pip install "promptise>=1.3.3"

"""Three tools that trap agents in loops, in two versions.

TOOL_MESSAGES=vague  -> what many real APIs return
TOOL_MESSAGES=clear  -> the same facts, plus what the agent should do next
"""
import os

from promptise.mcp.server import MCPServer, ToolError

CLEAR = os.environ.get("TOOL_MESSAGES", "vague") == "clear"
server = MCPServer("loops")


# 1. Identical calls: a job that never finishes while the agent watches.
@server.tool()
async def start_export(report: str) -> dict:
    """Start exporting a report. Returns an export ID."""
    return {"export_id": "EX-7731", "status": "pending"}


@server.tool(read_only_hint=True)
async def get_export_status(export_id: str) -> dict:
    """Get the status of an export and its download link when it is ready."""
    if CLEAR:
        return {"export_id": export_id, "status": "pending",
                "next_step": "Exports take about 20 minutes. Checking again now will not "
                             "change the status. Tell the user the export ID; the link is "
                             "emailed to them when it is ready."}
    return {"export_id": export_id, "status": "pending"}


# 2. Ping-pong: two tools that send the agent to each other.
@server.tool(read_only_hint=True)
async def get_account(email: str) -> dict:
    """Get a customer account by email address."""
    if CLEAR:
        raise ToolError(f"No account has the email {email}. Lookups are exact, so retrying "
                        "or searching by name will not find it.",
                        suggestion="Ask the user to confirm the email address.")
    raise ToolError("Account not found.", suggestion="Try find_account.")


@server.tool(read_only_hint=True)
async def find_account(name: str) -> dict:
    """Find customer accounts by name."""
    if CLEAR:
        return {"matches": 2, "accounts": ["jo.smith@acme.io", "jo.smith@acme.com"],
                "next_step": "Two accounts match. Ask the user which email is theirs."}
    raise ToolError("Several accounts match.", suggestion="Use get_account with the exact email.")


# 3. No-op results: a search that keeps coming back empty.
@server.tool(read_only_hint=True)
async def search_docs(query: str) -> dict:
    """Search the internal policy documents."""
    if CLEAR:
        return {"results": [], "next_step": "The policy index has nothing on this topic, so "
                "other wordings will return nothing too. Tell the user it isn't documented."}
    return {"results": []}


if __name__ == "__main__":
    server.run(transport="stdio")

"""Run the polling task once and print every tool call the agent makes.

python watch_one.py            # no guard
python watch_one.py --guard    # with LoopGuard
python watch_one.py --budget   # with a hard cap of 8 tool calls per run
"""
import asyncio
import sys

from loop_guard import LoopGuard
from promptise import build_agent
from promptise.config import StdioServerSpec
from promptise.runtime import BudgetConfig, BudgetState
from promptise.runtime.callbacks import RuntimeCallbackHandler


async def main(mode: str):
    agent = await build_agent(
        servers={"loops": StdioServerSpec(command=sys.executable, args=["loop_server.py"])},
        model="openai:gpt-5-mini",
        instructions="You are an autonomous operations agent. Nobody is watching this run, so "
                     "don't ask questions: keep working with the tools until the task is done.",
        observe=True,
    )
    guard = LoopGuard()
    budget = RuntimeCallbackHandler(
        budget=BudgetState(BudgetConfig(enabled=True, max_tool_calls_per_run=8)))
    callbacks = {"--guard": [guard], "--budget": budget.callbacks()}.get(mode)
    try:
        result = await agent.ainvoke(
            {"messages": [{"role": "user", "content": "Export the Q3 sales report and send me the download link."}]},
            config={"callbacks": callbacks} if callbacks else None,
        )
        print("ANSWER:", result["messages"][-1].content)
    except Exception as exc:
        print("FAILED:", type(exc).__name__, str(exc)[:160])
    events = agent.collector.get_timeline()
    for n, e in enumerate([e for e in events if e.event_type.value == "tool.call"], 1):
        print(f"  {n:>2}. {e.metadata['tool_name']}({e.metadata.get('arguments', '')})")
    end = next(e.metadata for e in events if e.event_type.value in ("agent.output", "agent.error"))
    print(f"model calls: {end.get('llm_call_count')}  tokens: "
          f"{end.get('prompt_tokens', 0) + end.get('completion_tokens', 0)}"
          + (f"  guard refused from call {guard.tripped_at}" if guard.tripped_at else "")
          + (f"  budget: {budget.pending_violations[0].limit_name}" if budget.pending_violations else ""))
    await agent.shutdown()


asyncio.run(main(sys.argv[1] if len(sys.argv) > 1 else ""))

"""Stop a looping agent mid-run with Promptise's HealthMonitor."""
from langchain_core.callbacks import BaseCallbackHandler
from promptise.runtime import HealthConfig
from promptise.runtime.health import HealthMonitor


class LoopDetected(RuntimeError):
    """Raised before a tool call that would continue a loop."""


def _now(coro):
    # HealthMonitor.record_tool_call is declared async but never awaits
    # anything, so one step runs it to completion.
    try:
        coro.send(None)
    except StopIteration as done:
        return done.value
    coro.close()
    raise RuntimeError("HealthMonitor.record_tool_call suspended")


class LoopGuard(BaseCallbackHandler):
    """Feeds every tool call to a HealthMonitor and refuses the call that
    completes a stuck or looping pattern, plus every call after it.

    Synchronous and inline, like Promptise's BudgetGuard: LangChain runs it
    before the tool executes and lets its exception through. (An exception
    from an async handler is logged and the tool runs anyway.)"""

    raise_error = True
    run_inline = True

    def __init__(self, config: HealthConfig | None = None) -> None:
        self.monitor = HealthMonitor(config or HealthConfig(enabled=True), process_id="guard")
        self.tripped_at: int | None = None
        self.calls = 0

    def on_tool_start(self, serialized, input_str, *, inputs=None, **kwargs):
        self.calls += 1
        name = serialized.get("name", "")
        anomaly = _now(self.monitor.record_tool_call(name, dict(inputs or {})))
        if anomaly is not None or self.tripped_at is not None:
            self.tripped_at = self.tripped_at or self.calls
            raise LoopDetected(
                f"Stopped: {anomaly.description if anomaly else 'loop already detected'}. "
                "Calling tools again will not help. Answer the user now with what you know, "
                "and say what is blocking you."
            )

"""Test LoopGuard without an API key: a scripted model polls four times, then answers.
Prints every tool result the model receives."""
import asyncio
import sys

from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage, ToolMessage
from loop_guard import LoopGuard
from promptise import build_agent
from promptise.config import StdioServerSpec

seen = []


class ScriptedModel(GenericFakeChatModel):
    def bind_tools(self, tools, **kwargs):
        return self

    async def _agenerate(self, messages, *args, **kwargs):
        results = [m for m in messages if isinstance(m, ToolMessage)]
        if results:
            seen.append(str(results[-1].content)[:200])
        return await super()._agenerate(messages, *args, **kwargs)


def script():
    for i in range(4):
        yield AIMessage(content="", tool_calls=[{"name": "get_export_status",
                                                 "args": {"export_id": "EX-7731"}, "id": f"c{i}"}])
    yield AIMessage(content="The export is still pending.")


async def main():
    agent = await build_agent(
        servers={"loops": StdioServerSpec(command=sys.executable, args=["loop_server.py"])},
        model=ScriptedModel(messages=script()),
    )
    result = await agent.ainvoke({"messages": [{"role": "user", "content": "Export Q3."}]},
                                 config={"callbacks": [LoopGuard()]})
    for text in seen:
        print("tool result:", text)
    print("answer:", result["messages"][-1].content)
    await agent.shutdown()


asyncio.run(main())

"""Measure how long agents loop, with and without a guard.

python run_loops.py --messages vague --guard none --trials 10
"""
import argparse
import asyncio
import json
import os
import sys
import time
from pathlib import Path

from loop_guard import LoopGuard
from promptise import build_agent
from promptise.agent import CallerContext
from promptise.config import StdioServerSpec
from promptise.observability_config import ObservabilityConfig
from promptise.runtime import BudgetConfig, BudgetState, HealthConfig
from promptise.runtime.callbacks import RuntimeCallbackHandler
from promptise.runtime.health import HealthMonitor

PROMPTS = {
    "assistant": "You are an operations assistant. Use the tools to help the user.",
    "autonomous": "You are an autonomous operations agent. Nobody is watching this run, so "
                  "don't ask questions: keep working with the tools until the task is done.",
}
SCENARIOS = {
    "polling": "Export the Q3 sales report and send me the download link.",
    "ping_pong": "What plan is the customer jo.smith@acme.io on?",
    "empty_search": "What's our refund policy for annual plans? Check the policy docs.",
}


async def first_anomaly(calls: list[dict]) -> tuple[int, str] | None:
    monitor = HealthMonitor(HealthConfig(enabled=True), process_id="replay")
    for i, call in enumerate(calls, 1):
        anomaly = await monitor.record_tool_call(call["tool"], call["args"])
        if anomaly is not None:
            return i, anomaly.anomaly_type.value
    return None


async def run_one(agent, scenario, trial, guard_kind):
    run_id = f"{scenario}#{trial}"
    guard = LoopGuard() if guard_kind == "health" else None
    budget = monitor = None
    config = {"callbacks": [guard]} if guard else None
    if guard_kind == "budget":
        budget = RuntimeCallbackHandler(
            budget=BudgetState(BudgetConfig(enabled=True, max_tool_calls_per_run=8)),
            health=(monitor := HealthMonitor(HealthConfig(enabled=True), process_id=run_id)),
        )
        config = {"callbacks": budget.callbacks()}
    started = time.perf_counter()
    outcome, answer = "answered", ""
    try:
        result = await agent.ainvoke(
            {"messages": [{"role": "user", "content": SCENARIOS[scenario]}]},
            config=config,
            caller=CallerContext(user_id="loops", metadata={"session_id": run_id}),
        )
        answer = str(result["messages"][-1].content)
    except Exception as exc:
        if "insufficient_quota" in str(exc):  # out of API credit: the run says nothing
            os._exit(3)
        outcome, answer = type(exc).__name__, str(exc)[:300]
    events = agent.collector.query(session_ids=[run_id])
    calls = []
    for e in events:
        if e.event_type.value == "tool.call":
            raw = e.metadata.get("arguments", "{}")
            try:
                import ast
                args = ast.literal_eval(raw)
            except Exception:
                args = {"raw": raw}
            calls.append({"tool": e.metadata["tool_name"], "args": args})
    end = next((e.metadata for e in events if e.event_type.value in ("agent.output", "agent.error")), {})
    found = await first_anomaly(calls)
    return {"scenario": scenario, "trial": trial, "outcome": outcome, "answer": answer,
            "tool_calls": len(calls), "calls": [c["tool"] for c in calls],
            "args": [c["args"] for c in calls],
            "llm_calls": end.get("llm_call_count", 0),
            "prompt_tokens": end.get("prompt_tokens", 0),
            "completion_tokens": end.get("completion_tokens", 0),
            "detected_at": found[0] if found else None, "anomaly": found[1] if found else None,
            "guard_tripped_at": guard.tripped_at if guard else None,
            "budget_violations": [v.limit_name for v in budget.pending_violations] if budget else None,
            "budget_run_anomalies": [a.anomaly_type.value for a in monitor.anomalies] if monitor else None,
            "seconds": round(time.perf_counter() - started, 1)}


async def main():
    p = argparse.ArgumentParser()
    p.add_argument("--messages", choices=["vague", "clear"], default="vague")
    p.add_argument("--guard", choices=["none", "health", "budget"], default="none")
    p.add_argument("--trials", type=int, default=10)
    p.add_argument("--model", default="openai:gpt-5-mini")
    p.add_argument("--prompt", choices=list(PROMPTS), default="autonomous")
    p.add_argument("--only", nargs="*")
    p.add_argument("--max-iterations", type=int, default=25)
    p.add_argument("--out", default="results")
    a = p.parse_args()
    agent = await build_agent(
        servers={"loops": StdioServerSpec(command=sys.executable, args=["loop_server.py"],
                                          env={"TOOL_MESSAGES": a.messages})},
        model=a.model,
        instructions=PROMPTS[a.prompt],
        # Keep real tool arguments: the default redaction turns every email into
        # "[EMAIL]", which makes different lookups look identical on replay.
        observe=ObservabilityConfig(redact_sensitive=False),
        max_agent_iterations=a.max_iterations,
    )
    gate = asyncio.Semaphore(6)

    async def bounded(s, t):
        async with gate:
            return await run_one(agent, s, t, a.guard)

    try:
        runs = await asyncio.gather(*(bounded(s, t) for s in SCENARIOS if not a.only or s in a.only
                                      for t in range(a.trials)))
    finally:
        await agent.shutdown()
    Path(a.out).mkdir(exist_ok=True)
    name = f"{a.model.split(':')[-1]}-{a.prompt}-{a.messages}-{a.guard}"
    if a.max_iterations != 25:
        name += f"-cap{a.max_iterations}"
    (Path(a.out) / f"{name}.json").write_text(json.dumps(runs, indent=1))
    tin = sum(r["prompt_tokens"] for r in runs); tout = sum(r["completion_tokens"] for r in runs)
    print(f"TOKENS in={tin} out={tout}")
    for r in runs:
        print(r["scenario"], r["trial"], r["outcome"], r["tool_calls"], r["llm_calls"],
              r["prompt_tokens"] + r["completion_tokens"], r["detected_at"], r["guard_tripped_at"])


if __name__ == "__main__":
    asyncio.run(main())

python run_loops.py --messages vague --guard none
python run_loops.py --messages vague --guard health
python run_loops.py --messages vague --guard budget
python run_loops.py --messages clear --guard none
python run_loops.py --messages vague --guard none --max-iterations 6
python run_loops.py --messages vague --guard none --prompt assistant
python run_loops.py --messages vague --guard none --model openai:gpt-5-nano
python run_loops.py --messages vague --guard none --model openai:gpt-4.1-mini

"""Summarise results/*.json: one row per condition and scenario. No API calls.

python summarize.py results/*.json
"""
import json
import statistics
import sys
from pathlib import Path

PRICE = {  # USD per 1M tokens (input, output), OpenAI list prices
    "gpt-5-mini": (0.25, 2.00), "gpt-5-nano": (0.05, 0.40), "gpt-4.1-mini": (0.40, 1.60),
}


def cost(model, runs):
    pin, pout = PRICE[model]
    return sum(r["prompt_tokens"] * pin + r["completion_tokens"] * pout for r in runs) / 1e6


print(f"{'condition':<34} {'scenario':<13} {'answered':>8} {'tool calls':>14} "
      f"{'tokens/run':>10} {'flagged':>8} {'stopped':>8}")
total = 0.0
for path in sys.argv[1:]:
    runs = json.loads(Path(path).read_text())
    name = Path(path).stem
    model = next(m for m in PRICE if name.startswith(m + "-"))
    total += cost(model, runs)
    for scenario in ("polling", "ping_pong", "empty_search"):
        rs = [r for r in runs if r["scenario"] == scenario]
        if not rs:
            continue
        calls = [r["tool_calls"] for r in rs]
        tokens = [r["prompt_tokens"] + r["completion_tokens"] for r in rs]
        answered = sum(r["outcome"] == "answered" for r in rs)
        flagged = [r["detected_at"] for r in rs if r["detected_at"]]
        stopped = sum(bool(r.get("guard_tripped_at") or r.get("budget_violations")) for r in rs)
        print(f"{name:<34} {scenario:<13} {answered:>5}/{len(rs):<2} "
              f"{statistics.median(calls):>5g} ({min(calls)}-{max(calls)}) "
              f"{round(statistics.mean(tokens)):>10,} {len(flagged):>5}/{len(rs):<2} {stopped:>5}/{len(rs):<2}")
print(f"total cost at list prices: ${total:.2f}")