retry.py

"""
Self-repair retry policy (Phase 2).

A pure helper that decides, given a captured error / traceback / exit code,
whether the agent should retry the same step with a corrected code snippet,
and if so, packages up the feedback for the planner.

This file intentionally has no external state — call sites in agent.py drive
the loop. That keeps the existing Phase-1 ``stream_execute`` path untouched
when ``stream_agent_plan`` is not used.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import List, Optional


# ---------------------------------------------------------------------------
# Heuristics
# ---------------------------------------------------------------------------

# Errors that almost never get better by re-running the same code: surface
# them quickly instead of burning provider calls.
_FATAL_PATTERNS = [
    r"\bAPI_KEY\b.*\b(missing|not set)\b",
    r"\bPermissionError: \[Errno 13\]",
    r"\bquota\b.*\bexceeded\b",
    r"\bbilling\b",
]

_FATAL_RE = [re.compile(p, re.IGNORECASE) for p in _FATAL_PATTERNS]


_TRACEBACK_HINT_RE = re.compile(r"(Traceback|Exception|Error:|exit_code\s*[:=]\s*[1-9])",
                                re.IGNORECASE)


# ---------------------------------------------------------------------------
# API
# ---------------------------------------------------------------------------

@dataclass
class RetryDecision:
    should_retry: bool
    reason: str
    feedback: str = ""   # passed back to planner.code_for_step
    delay_seconds: float = 0.0


def build_feedback(stdout: str, stderr: str, error: Optional[str],
                   exit_code: Optional[int]) -> str:
    """Compose a short, deterministic feedback blob for the next LLM call."""
    parts: List[str] = []
    if exit_code is not None and exit_code != 0:
        parts.append(f"exit_code={exit_code}")
    if error:
        parts.append(f"ERROR: {error.strip()}")
    if stderr and stderr.strip():
        parts.append(f"STDERR:\n{stderr.strip()[-1500:]}")
    if stdout and stdout.strip() and not parts:
        # Only include stdout when there's no clearer signal — it can be huge.
        parts.append(f"STDOUT_TAIL:\n{stdout.strip()[-800:]}")
    return "\n\n".join(parts) or "Step did not produce a success marker."


def is_failure(stderr: str, error: Optional[str], exit_code: Optional[int],
               stdout: str = "") -> bool:
    """True when the step's output looks like a real failure."""
    if error:
        return True
    if exit_code is not None and exit_code != 0:
        return True
    if stderr and _TRACEBACK_HINT_RE.search(stderr):
        return True
    # As a last resort, check stdout — some scripts print "ERROR: ..." but
    # exit cleanly because they swallow exceptions.
    if stdout and re.search(r"^ERROR[: ]", stdout, re.MULTILINE | re.IGNORECASE):
        return True
    return False


def decide(*,
           attempt: int,
           max_attempts: int,
           stdout: str,
           stderr: str,
           error: Optional[str],
           exit_code: Optional[int]) -> RetryDecision:
    """Decide whether to retry the step.

    ``attempt`` is 1-based. Stops retrying once attempt >= max_attempts, or
    when a fatal-pattern error is detected.
    """
    if not is_failure(stderr, error, exit_code, stdout):
        return RetryDecision(False, "step succeeded")

    blob = f"{error or ''}\n{stderr or ''}"
    for rx in _FATAL_RE:
        if rx.search(blob):
            return RetryDecision(False, f"fatal pattern matched: {rx.pattern}")

    if attempt >= max_attempts:
        return RetryDecision(False, f"max attempts reached ({attempt}/{max_attempts})")

    feedback = build_feedback(stdout, stderr, error, exit_code)
    # Light exponential back-off keeps us under rate limits when the failure
    # is upstream rather than logic.
    delay = min(1.5 * attempt, 6.0)
    return RetryDecision(True, f"attempt {attempt + 1}/{max_attempts}", feedback, delay)