claude-code-generator 0.1.0__py3-none-any.whl

code_generator/runner/retry.py

@@ -0,0 +1,165 @@
+"""Retry logic with exponential backoff and a circuit breaker.
+
+Backoff applies only to transient errors. RateLimitHit and OverageAbort
+bypass both backoff and the circuit breaker — they are handled by the
+wait-and-resume loop in rate_limit.py, not by blind retrying.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+from code_generator.runner.sdk_runner import OverageAbort, RateLimitHit
+
+T = TypeVar("T")
+
+# ---------------------------------------------------------------------------
+# Backoff schedule (seconds). Capped at 120.
+# ---------------------------------------------------------------------------
+
+_BACKOFF_SCHEDULE: list[int] = [10, 20, 40, 80, 120]
+
+
+# ---------------------------------------------------------------------------
+# Sentinel exception for transient runner errors
+# ---------------------------------------------------------------------------
+
+
+class TransientRunnerError(RuntimeError):
+    """Marker exception for transient runner failures that should be retried."""
+
+
+# ---------------------------------------------------------------------------
+# Transient-error types (retry-eligible)
+# ---------------------------------------------------------------------------
+
+_TRANSIENT_TYPES: tuple[type[BaseException], ...] = (
+    asyncio.TimeoutError,
+    ConnectionError,
+    OSError,
+    TransientRunnerError,
+)
+
+
+def _is_transient(exc: BaseException) -> bool:
+    """Return True for exceptions that warrant a retry via backoff."""
+    return isinstance(exc, _TRANSIENT_TYPES)
+
+
+def _is_rate_limit_or_overage(exc: BaseException) -> bool:
+    """Return True for exceptions that bypass both backoff and the breaker."""
+    return isinstance(exc, (RateLimitHit, OverageAbort))
+
+
+# ---------------------------------------------------------------------------
+# Circuit breaker
+# ---------------------------------------------------------------------------
+
+
+class CircuitOpen(RuntimeError):
+    """Raised when the circuit breaker trips due to consecutive failures."""
+
+
+class CircuitBreaker:
+    """Counts consecutive non-rate-limit failures and opens after max_failures.
+
+    Args:
+        max_failures: Number of consecutive failures before the circuit opens.
+    """
+
+    def __init__(self, max_failures: int = 3) -> None:
+        self._max_failures = max_failures
+        self._consecutive_failures = 0
+
+    def record_success(self) -> None:
+        """Reset the failure counter after a successful operation."""
+        self._consecutive_failures = 0
+
+    def record_failure(self, exc: BaseException) -> None:
+        """Increment the failure counter; raise CircuitOpen when threshold is hit.
+
+        Rate-limit and overage exceptions are excluded from the counter —
+        they are expected transients managed elsewhere.
+
+        Args:
+            exc: The exception that caused the failure.
+
+        Raises:
+            CircuitOpen: When consecutive failures reach max_failures.
+        """
+        if _is_rate_limit_or_overage(exc):
+            # Non-negotiable bypass — never count rate-limit/overage towards circuit.
+            return
+
+        self._consecutive_failures += 1
+        if self._consecutive_failures >= self._max_failures:
+            raise CircuitOpen(
+                f"Circuit opened after {self._consecutive_failures} consecutive failures."
+            )
+
+
+# ---------------------------------------------------------------------------
+# with_backoff
+# ---------------------------------------------------------------------------
+
+
+async def with_backoff(
+    coro_factory: Callable[[], Awaitable[T]],
+    *,
+    attempts: int = 5,
+    sleep_fn: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    breaker: CircuitBreaker | None = None,
+) -> T:
+    """Execute coro_factory with exponential backoff on transient errors.
+
+    Only retries on transient errors (asyncio.TimeoutError, ConnectionError,
+    OSError, TransientRunnerError). RateLimitHit and OverageAbort propagate
+    immediately without touching the circuit breaker.
+
+    Args:
+        coro_factory: Zero-argument callable returning an awaitable.
+        attempts: Maximum number of attempts (1 means no retries).
+        sleep_fn: Async sleep function (injectable for tests).
+        breaker: Optional CircuitBreaker instance.
+
+    Returns:
+        The value returned by coro_factory on success.
+
+    Raises:
+        CircuitOpen: When the breaker trips.
+        RateLimitHit: Immediately, bypassing backoff.
+        OverageAbort: Immediately, bypassing backoff.
+        The original exception: When all attempts are exhausted.
+    """
+    last_exc: BaseException | None = None
+
+    for attempt in range(attempts):
+        try:
+            result = await coro_factory()
+            if breaker is not None:
+                breaker.record_success()
+            return result
+        except (RateLimitHit, OverageAbort):
+            # Non-negotiable: bypass backoff and breaker for rate-limit/overage.
+            raise
+        except BaseException as exc:
+            if not _is_transient(exc):
+                # Non-transient error — propagate without retry or breaker.
+                raise
+
+            last_exc = exc
+
+            if breaker is not None:
+                # record_failure may raise CircuitOpen — let it propagate.
+                breaker.record_failure(exc)
+
+            delay = _BACKOFF_SCHEDULE[min(attempt, len(_BACKOFF_SCHEDULE) - 1)]
+            await sleep_fn(float(delay))
+
+    # All attempts exhausted.
+    assert last_exc is not None
+    raise last_exc

code_generator/runner/sdk_runner.py

@@ -0,0 +1,267 @@
+"""Primary runner: async wrapper around claude_agent_sdk.ClaudeSDKClient.
+
+Handles RateLimitEvent processing, overage abort, and session resume.
+Non-negotiables: #1 (strip env), #3 (bypassPermissions), #4 (overage abort).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from code_generator import env as _env
+from code_generator import state as _state
+
+if TYPE_CHECKING:
+    import logging
+    from pathlib import Path
+
+# Guard SDK import so the module can be loaded without the SDK installed
+# (used by get_runner() to detect availability).
+try:
+    from claude_agent_sdk import (
+        AssistantMessage,
+        ClaudeSDKClient,
+        ResultMessage,
+        TextBlock,
+    )
+    from claude_agent_sdk.types import RateLimitEvent
+
+    _SDK_AVAILABLE = True
+except ImportError:  # pragma: no cover
+    _SDK_AVAILABLE = False
+    # Provide sentinel names so type checkers and duck-typed tests still work.
+    ClaudeSDKClient = None  # type: ignore[assignment,misc]
+    AssistantMessage = None  # type: ignore[assignment]
+    ResultMessage = None  # type: ignore[assignment]
+    TextBlock = None  # type: ignore[assignment]
+    RateLimitEvent = None  # type: ignore[assignment]
+
+
+# ---------------------------------------------------------------------------
+# Public exception types
+# ---------------------------------------------------------------------------
+
+
+class OverageAbort(RuntimeError):
+    """Raised when overage billing is active — abort to avoid charges."""
+
+
+class RateLimitHit(Exception):
+    """Raised when the rate limit is rejected and the session is paused."""
+
+    def __init__(self, resets_at: int, rate_limit_type: str) -> None:
+        self.resets_at = resets_at
+        self.rate_limit_type = rate_limit_type
+        super().__init__(f"Rate limit hit; resets at {resets_at}, type={rate_limit_type}")
+
+
+# ---------------------------------------------------------------------------
+# Result type
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class RunResult:
+    """Captured output from a single SDK run."""
+
+    text: str
+    session_id: str | None
+    tokens_in: int | None
+    tokens_out: int | None
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _extract_text(msg: Any) -> str:
+    """Extract concatenated text from an AssistantMessage's content blocks."""
+    parts: list[str] = []
+    for block in getattr(msg, "content", []):
+        text = getattr(block, "text", None)
+        if text is not None:
+            parts.append(text)
+    return "".join(parts)
+
+
+def _extract_usage(msg: Any) -> tuple[int | None, int | None]:
+    """Return (tokens_in, tokens_out) from a ResultMessage, defensively."""
+    usage = getattr(msg, "usage", None)
+    if usage is None:
+        return None, None
+    if isinstance(usage, dict):
+        return usage.get("input_tokens"), usage.get("output_tokens")
+    return getattr(usage, "input_tokens", None), getattr(usage, "output_tokens", None)
+
+
+def _handle_rate_limit_event(
+    msg: Any,
+    logger: logging.Logger,
+    state_path: Path,
+) -> None:
+    """Process a RateLimitEvent; raise OverageAbort or RateLimitHit as needed.
+
+    Args:
+        msg: The RateLimitEvent message (duck-typed).
+        logger: Phase logger for warnings.
+        state_path: Path to state.json for persistence.
+
+    Raises:
+        OverageAbort: When overage billing is active.
+        RateLimitHit: When the rate limit is rejected and resets_at is known.
+    """
+    info = getattr(msg, "rate_limit_info", None)
+    if info is None:
+        logger.warning("RateLimitEvent received but has no rate_limit_info — skipping.")
+        return
+
+    # Non-negotiable #4: overage check — abort before incurring charges.
+    # SDK values: "allowed" / "allowed_warning" mean overage billing IS active
+    # (unsafe); "rejected" means overage is off and requests are blocked (safe);
+    # None / "disabled" also safe.
+    overage_status = getattr(info, "overage_status", None)
+    if overage_status in ("allowed", "allowed_warning"):
+        raise OverageAbort(
+            f"Overage API billing active ({overage_status}) — aborting to avoid charges"
+        )
+
+    status = getattr(info, "status", None)
+
+    if status == "rejected":
+        resets_at = getattr(info, "resets_at", None)
+        rate_limit_type = getattr(info, "rate_limit_type", None) or "unknown"
+        session_id = getattr(msg, "session_id", None)
+
+        if resets_at is not None:
+            st = _state.load_state(state_path)
+            _state.mark_paused(
+                st,
+                resets_at=resets_at,
+                rate_limit_type=str(rate_limit_type),
+                session_id=session_id,
+            )
+            _state.save_state(state_path, st)
+            raise RateLimitHit(resets_at, str(rate_limit_type))
+
+        logger.warning("Rate limit rejected but no resets_at provided — cannot pause.")
+
+    elif status == "allowed_warning":
+        utilization = getattr(info, "utilization", None)
+        logger.warning(
+            "Rate limit allowed_warning: approaching limit (utilization=%s).", utilization
+        )
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+async def run(
+    prompt: str,
+    options: Any,
+    *,
+    logger: logging.Logger,
+    state_path: Path,
+) -> RunResult:
+    """Run a prompt via ClaudeSDKClient and return the accumulated result.
+
+    Non-negotiable #1: strips dangerous env vars at the start of every call.
+    Non-negotiable #3: forces permission_mode to bypassPermissions.
+
+    Args:
+        prompt: The prompt text to send.
+        options: ClaudeAgentOptions (or duck-typed equivalent).
+        logger: Phase logger for rate-limit warnings and debug output.
+        state_path: Path to state.json used for pause persistence.
+
+    Returns:
+        RunResult with accumulated text, session_id, and token counts.
+
+    Raises:
+        OverageAbort: When overage billing is detected.
+        RateLimitHit: When the rate limit is rejected.
+    """
+    # Non-negotiable #1 — strip dangerous env vars on every invocation.
+    _env.strip_dangerous_env()
+
+    # Non-negotiable #3 — always force bypassPermissions; never trust caller.
+    options.permission_mode = "bypassPermissions"
+
+    text_parts: list[str] = []
+    session_id: str | None = None
+    tokens_in: int | None = None
+    tokens_out: int | None = None
+
+    logger.info("SDK session starting (model=%s).", getattr(options, "model", "?"))
+    msg_count = 0
+    async with ClaudeSDKClient(options=options) as client:
+        await client.query(prompt)
+        async for msg in client.receive_messages():
+            msg_count += 1
+            # Detect message type by duck-typing on attributes rather than
+            # isinstance() so fake test objects work without SDK installed.
+            if _is_rate_limit_event(msg):
+                logger.info("msg #%d: RateLimitEvent", msg_count)
+                _handle_rate_limit_event(msg, logger, state_path)
+                continue
+
+            if _is_assistant_message(msg):
+                chunk = _extract_text(msg)
+                text_parts.append(chunk)
+                if chunk.strip():
+                    logger.info(
+                        "msg #%d: assistant (%d chars)\n%s",
+                        msg_count, len(chunk), chunk.rstrip(),
+                    )
+                else:
+                    logger.debug("msg #%d: assistant (empty)", msg_count)
+                # Some SDK versions attach session_id to AssistantMessage.
+                if session_id is None:
+                    session_id = getattr(msg, "session_id", None)
+
+            elif _is_result_message(msg):
+                session_id = getattr(msg, "session_id", session_id)
+                tokens_in, tokens_out = _extract_usage(msg)
+                logger.info(
+                    "msg #%d: ResultMessage (tokens_in=%s tokens_out=%s)",
+                    msg_count, tokens_in, tokens_out,
+                )
+            else:
+                logger.debug("msg #%d: %s", msg_count, type(msg).__name__)
+
+    logger.info("SDK session complete (%d messages).", msg_count)
+
+    return RunResult(
+        text="".join(text_parts),
+        session_id=session_id,
+        tokens_in=tokens_in,
+        tokens_out=tokens_out,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Duck-type discriminators
+# ---------------------------------------------------------------------------
+
+
+def _is_rate_limit_event(msg: Any) -> bool:
+    """Return True when msg looks like a RateLimitEvent."""
+    return hasattr(msg, "rate_limit_info")
+
+
+def _is_assistant_message(msg: Any) -> bool:
+    """Return True when msg looks like an AssistantMessage."""
+    return hasattr(msg, "content") and hasattr(msg, "model")
+
+
+def _is_result_message(msg: Any) -> bool:
+    """Return True when msg looks like a ResultMessage."""
+    return (
+        hasattr(msg, "session_id")
+        and hasattr(msg, "is_error")
+        and not hasattr(msg, "rate_limit_info")
+        and not hasattr(msg, "model")
+    )

code_generator/runner/subprocess_runner.py

@@ -0,0 +1,200 @@
+"""Subprocess fallback runner: shells out to the `claude` CLI.
+
+Used when claude_agent_sdk is not importable. Spawns:
+  claude -p <prompt> --output-format stream-json --verbose
+         --dangerously-skip-permissions --model <model>
+         --max-turns <n> --allowedTools <csv>
+
+Non-negotiables: #1 (build_agent_env), #2 (never --bare).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import subprocess
+import time
+from typing import TYPE_CHECKING, Any
+
+from code_generator import env as _env
+from code_generator import state as _state
+
+# Re-export the shared exception types so callers can import from one place.
+from code_generator.runner.sdk_runner import RateLimitHit, RunResult
+
+if TYPE_CHECKING:
+    import logging
+    from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Model name translation map
+# SDK uses full model IDs; the `claude` CLI may use short aliases.
+# ---------------------------------------------------------------------------
+
+_MODEL_ALIASES: dict[str, str] = {
+    "claude-opus-4-6": "claude-opus-4-6",
+    "claude-sonnet-4-6": "claude-sonnet-4-6",
+    "claude-haiku-4-5": "claude-haiku-4-5",
+}
+
+
+def _translate_model(model: str | None) -> str:
+    """Return the CLI model alias for a given SDK model name."""
+    if model is None:
+        return "claude-sonnet-4-6"
+    return _MODEL_ALIASES.get(model, model)
+
+
+# ---------------------------------------------------------------------------
+# NDJSON line handlers
+# ---------------------------------------------------------------------------
+
+
+def _handle_line(
+    line: str,
+    text_parts: list[str],
+    result_holder: dict[str, Any],
+    logger: logging.Logger,
+    state_path: Path,
+    proc: Any,
+) -> None:
+    """Parse one NDJSON line and mutate text_parts / result_holder in place.
+
+    Args:
+        line: A single JSON line from the CLI stdout.
+        text_parts: Accumulator for assistant text fragments.
+        result_holder: Dict capturing session_id and token counts.
+        logger: Phase logger.
+        state_path: Path to state.json for pause persistence.
+        proc: The Popen process (so it can be killed on hard rate-limit).
+
+    Raises:
+        RateLimitHit: When retry_delay_ms exceeds 120_000 ms.
+    """
+    try:
+        data = json.loads(line)
+    except json.JSONDecodeError:
+        logger.debug("Ignoring malformed NDJSON line: %r", line[:200])
+        return
+
+    msg_type = data.get("type")
+
+    if msg_type == "assistant":
+        message = data.get("message", {})
+        for block in message.get("content", []):
+            if block.get("type") == "text":
+                text_parts.append(block.get("text", ""))
+
+    elif msg_type == "result":
+        result_holder["session_id"] = data.get("session_id")
+        usage = data.get("usage") or {}
+        result_holder["tokens_in"] = usage.get("input_tokens")
+        result_holder["tokens_out"] = usage.get("output_tokens")
+
+    elif msg_type == "system" and data.get("subtype") == "api_retry":
+        if data.get("error") == "rate_limit":
+            delay_ms: int = data.get("retry_delay_ms", 0)
+            session_id: str | None = data.get("session_id")
+            if delay_ms > 120_000:
+                paused_until = int(time.time() + delay_ms / 1000)
+                st = _state.load_state(state_path)
+                _state.mark_paused(
+                    st,
+                    resets_at=paused_until,
+                    rate_limit_type="five_hour",
+                    session_id=session_id,
+                )
+                _state.save_state(state_path, st)
+                proc.kill()
+                raise RateLimitHit(paused_until, "five_hour")
+            else:
+                logger.info(
+                    "CLI rate-limit retry in %dms — letting CLI handle internally.", delay_ms
+                )
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+async def run(
+    prompt: str,
+    options: Any,
+    *,
+    logger: logging.Logger,
+    state_path: Path,
+) -> RunResult:
+    """Spawn the Claude CLI and parse its stream-json output.
+
+    Non-negotiable #1: uses build_agent_env() — no dangerous vars reach the CLI.
+    Non-negotiable #2: --bare is never included; an assertion guards this.
+
+    Args:
+        prompt: The prompt text to pass via -p.
+        options: Duck-typed options object with model, max_turns, allowed_tools, cwd.
+        logger: Phase logger for debug/info/warning messages.
+        state_path: Path to state.json for pause persistence.
+
+    Returns:
+        RunResult with accumulated text, session_id, and token counts.
+
+    Raises:
+        RateLimitHit: When retry_delay_ms > 120_000 in the stream.
+    """
+    model = _translate_model(getattr(options, "model", None))
+    max_turns = getattr(options, "max_turns", 50)
+    allowed_tools: list[str] = getattr(options, "allowed_tools", [])
+    cwd: str | None = getattr(options, "cwd", None)
+
+    argv: list[str] = [
+        "claude",
+        "-p",
+        prompt,
+        "--output-format",
+        "stream-json",
+        "--verbose",
+        "--dangerously-skip-permissions",
+        "--model",
+        model,
+        "--max-turns",
+        str(max_turns),
+    ]
+    if allowed_tools:
+        argv += ["--allowedTools", ",".join(allowed_tools)]
+
+    # Non-negotiable #2 — never pass --bare.
+    assert "--bare" not in argv, "BUG: --bare must never appear in the claude CLI argv."
+
+    safe_env = _env.build_agent_env()
+
+    text_parts: list[str] = []
+    result_holder: dict[str, Any] = {
+        "session_id": None,
+        "tokens_in": None,
+        "tokens_out": None,
+    }
+
+    # Run subprocess in a thread to keep the coroutine non-blocking.
+    def _run_subprocess() -> None:
+        with subprocess.Popen(
+            argv,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            env=safe_env,
+            cwd=cwd,
+        ) as proc:
+            for raw_line in proc.stdout:  # type: ignore[union-attr]
+                line = raw_line.decode(errors="replace").rstrip()
+                if not line:
+                    continue
+                _handle_line(line, text_parts, result_holder, logger, state_path, proc)
+
+    await asyncio.to_thread(_run_subprocess)
+
+    return RunResult(
+        text="".join(text_parts),
+        session_id=result_holder["session_id"],
+        tokens_in=result_holder["tokens_in"],
+        tokens_out=result_holder["tokens_out"],
+    )