Make budgeted GenPlan CLI runs reliable

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
    "mcp[cli]>=1.27,<2",
    "anthropic",
    "openai",
+   "tenacity",
 ]
 
 [project.optional-dependencies]

src/robocode/utils/apptainer_sandbox.py

@@ -63,6 +63,7 @@
 )
 from robocode.utils.docker_sandbox import (
     DOCKER_PYTHON,
+    GENPLAN_CONTAINER_TIMEOUT_S,
     _filtered_repo_mounts,
     _find_repo_root,
     _get_claude_oauth_token,
@@ -419,7 +420,7 @@ def run_genplan_in_apptainer(
     sandbox_dir: Path,
     completion_cfg: dict[str, Any],
     sif_path: Path = _DEFAULT_SIF,
-    timeout: float = 3600.0,
+    timeout: float = GENPLAN_CONTAINER_TIMEOUT_S,
     include_bilevel: bool = False,
 ) -> None:
     """Apptainer analog of :func:`docker_sandbox.run_genplan_in_docker`.

src/robocode/utils/docker_sandbox.py

@@ -99,6 +99,11 @@ def container_python(blackbox_strict: bool) -> str:
 # Default Docker image name.
 _DEFAULT_IMAGE: str = "robocode-sandbox"
 
+# A $20 Opus GenPlan run can take longer than an hour. Keep one shared limit for
+# Docker and Apptainer so the selected container backend does not change whether
+# the configured model budget can be exhausted.
+GENPLAN_CONTAINER_TIMEOUT_S: float = 10 * 60 * 60
+
 
 def _telemetry_docker(config: SandboxConfig) -> tuple[list[str], list[str]]:
     """(extra volumes, env args) enabling telemetry for a whitebox docker run.
@@ -397,7 +402,7 @@ def run_genplan_in_docker(
     sandbox_dir: Path,
     completion_cfg: dict[str, Any],
     image: str = _DEFAULT_IMAGE,
-    timeout: float = 3600.0,
+    timeout: float = GENPLAN_CONTAINER_TIMEOUT_S,
     include_bilevel: bool = False,
 ) -> None:
     """Run the whole LLM-GenPlan loop inside one sandbox container.
@@ -445,13 +450,25 @@ def run_genplan_in_docker(
             ss_pybullet=ss_pybullet,
         ) + [DOCKER_PYTHON, "-m", "robocode.approaches.genplan_driver"]
         logger.info("Starting genplan Docker container %s", container_name)
-        subprocess.run(
-            docker_cmd,
-            env={**os.environ, **auth_env},
-            stdin=subprocess.DEVNULL,
-            check=True,
-            timeout=timeout,
-        )
+        try:
+            subprocess.run(
+                docker_cmd,
+                env={**os.environ, **auth_env},
+                stdin=subprocess.DEVNULL,
+                check=True,
+                timeout=timeout,
+            )
+        except BaseException:
+            # A timeout kills the attached Docker client, not necessarily the named
+            # container. Remove it explicitly so an abandoned run cannot keep
+            # spending model budget or writing into the sandbox directory.
+            subprocess.run(
+                ["docker", "rm", "-f", container_name],
+                check=False,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            raise
 
 
 def _copy_src(

src/robocode/utils/llm/cli_client.py

@@ -1,6 +1,7 @@
 """Claude CLI driven as a plain LLM: the default completion backend.
 
-Runs ``claude -p`` single-shot with tools and system prompt stripped. Billing
+Runs ``claude -p`` with tools and system prompt stripped, resuming matching
+conversations by session ID to preserve message boundaries and cache reuse. Billing
 follows the authenticated CLI (no API key needed), and the same CLI serves the
 agentic backend, so completion and agentic runs share one auth path and one set
 of model ids.
@@ -13,17 +14,83 @@
 from __future__ import annotations
 
 import json
+import logging
 import os
 import subprocess
+from typing import Any
 
 from omegaconf import DictConfig
+from tenacity import (
+    Retrying,
+    before_sleep_log,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 
 from robocode.utils.backends.claude import anthropic_compatible_env, get_claude_cmd
 from robocode.utils.llm.base import LLMResponse
+from robocode.utils.rate_limit import wait_for_rate_limit_reset
+
+logger = logging.getLogger(__name__)
+
+
+class ClaudeCLIError(RuntimeError):
+    """A failed Claude CLI invocation, including its structured error output."""
+
+    def __init__(self, returncode: int, stdout: str | None, stderr: str | None) -> None:
+        self.returncode = returncode
+        self.stdout = stdout or ""
+        self.stderr = stderr or ""
+        super().__init__(
+            f"Claude exited with status {returncode}.\n"
+            f"stdout: {self.stdout or '(empty)'}\n"
+            f"stderr: {self.stderr or '(empty)'}"
+        )
+
+
+def _is_transient_claude_error(exc: BaseException) -> bool:
+    """Return whether Claude reported a retryable server-side API failure."""
+    data = _claude_error_data(exc)
+    if data is None:
+        return False
+    status = data.get("api_error_status")
+    # Only retry when the CLI explicitly confirms that the failed request was
+    # free. A missing cost is ambiguous and could duplicate a charged request.
+    return (
+        isinstance(status, int)
+        and 500 <= status < 600
+        and data.get("total_cost_usd") == 0
+    )
+
+
+def _claude_error_data(exc: BaseException) -> dict[str, Any] | None:
+    """Parse a Claude CLI error envelope when one is available."""
+    if not isinstance(exc, ClaudeCLIError):
+        return None
+    try:
+        data: Any = json.loads(exc.stdout)
+    except (json.JSONDecodeError, TypeError):
+        return None
+    return data if isinstance(data, dict) else None
+
+
+def _session_limit_reset(exc: BaseException) -> str | None:
+    """Return Claude's reset message for a safe-to-repeat usage-limit failure."""
+    data = _claude_error_data(exc)
+    if data is None or data.get("api_error_status") != 429:
+        return None
+    message = data.get("result")
+    if not isinstance(message, str) or data.get("total_cost_usd") != 0:
+        return None
+    normalized = message.lower()
+    if "limit" not in normalized or "reset" not in normalized:
+        return None
+    return message
 
 
 class ClaudeCLIClient:
-    """Single-shot completions via the Claude Code CLI (no tools)."""
+    """Text completions via the Claude Code CLI with per-client session reuse."""
 
     def __init__(self, cfg: DictConfig) -> None:
         self._model = cfg["model"]
@@ -32,10 +99,49 @@ def __init__(self, cfg: DictConfig) -> None:
         self._ollama_keep_alive = cfg.get("ollama_keep_alive", "")
         self._timeout_s = cfg.get("request_timeout_s", 1200.0)
         self._max_thinking_tokens = cfg.get("max_thinking_tokens", 0)
+        self._retry_attempts = cfg.get("retry_attempts", 3)
+        self._retry_wait_min_s = cfg.get("retry_wait_min_s", 2.0)
+        self._retry_wait_max_s = cfg.get("retry_wait_max_s", 30.0)
+        self._session_id: str | None = None
+        self._history: list[dict[str, str]] = []
 
     def complete(self, messages: list[dict[str, str]]) -> LLMResponse:
         """Return the model's reply to a message list."""
-        prompt = _flatten(messages)
+        retrying = Retrying(
+            retry=retry_if_exception(_is_transient_claude_error),
+            stop=stop_after_attempt(self._retry_attempts),
+            wait=wait_random_exponential(
+                min=self._retry_wait_min_s, max=self._retry_wait_max_s
+            ),
+            before_sleep=before_sleep_log(logger, logging.WARNING),
+            reraise=True,
+        )
+        while True:
+            try:
+                return retrying(self._complete_once, messages)
+            except ClaudeCLIError as exc:
+                reset_message = _session_limit_reset(exc)
+                if reset_message is None:
+                    raise
+                wait_for_rate_limit_reset(reset_message)
+
+    def _complete_once(self, messages: list[dict[str, str]]) -> LLMResponse:
+        """Make one attempt, retaining session state only for a safe reset retry."""
+        # Only resume an exact continuation. Callers may reuse this client for
+        # independent samples or edit history, which must start fresh sessions.
+        resume = (
+            self._session_id is not None
+            and len(messages) == len(self._history) + 1
+            and messages[:-1] == self._history
+            and messages[-1]["role"] == "user"
+        )
+        session_id = self._session_id if resume else None
+        previous_history = self._history
+        prompt = messages[-1]["content"] if resume else _flatten(messages)
+        # A failed request may have advanced the remote session. Do not resume
+        # that uncertain state if the caller retries.
+        self._session_id = None
+        self._history = []
         args = [
             get_claude_cmd(),
             "-p",
@@ -54,6 +160,8 @@ def complete(self, messages: list[dict[str, str]]) -> LLMResponse:
             "--max-thinking-tokens",
             str(self._max_thinking_tokens),
         ]
+        if session_id is not None:
+            args += ["--resume", session_id]
         env = {k: v for k, v in os.environ.items() if not k.startswith("CLAUDECODE")}
         env.update(
             anthropic_compatible_env(
@@ -72,15 +180,26 @@ def complete(self, messages: list[dict[str, str]]) -> LLMResponse:
                 timeout=self._timeout_s,
             )
         except subprocess.CalledProcessError as exc:
-            raise RuntimeError(
-                f"Claude exited with status {exc.returncode}.\n"
-                f"stdout: {exc.stdout or '(empty)'}\n"
-                f"stderr: {exc.stderr or '(empty)'}"
-            ) from exc
+            error = ClaudeCLIError(exc.returncode, exc.stdout, exc.stderr)
+            if _session_limit_reset(error) is not None:
+                self._session_id = session_id
+                self._history = previous_history
+            raise error from exc
         data = json.loads(result.stdout)
-        return LLMResponse(text=data["result"], cost_usd=data.get("total_cost_usd"))
+        if data.get("is_error"):
+            error = ClaudeCLIError(result.returncode, result.stdout, result.stderr)
+            if _session_limit_reset(error) is not None:
+                self._session_id = session_id
+                self._history = previous_history
+            raise error
+        response = LLMResponse(text=data["result"], cost_usd=data.get("total_cost_usd"))
+        self._session_id = data.get("session_id")
+        self._history = [dict(message) for message in messages] + [
+            {"role": "assistant", "content": response.text}
+        ]
+        return response
 
 
 def _flatten(messages: list[dict[str, str]]) -> str:
-    """Collapse the multi-turn history into one prompt for the stateless CLI."""
+    """Serialize initial or replaced history when starting a fresh CLI session."""
     return "\n\n".join(f"{m['role'].upper()}: {m['content']}" for m in messages)

src/robocode/utils/rate_limit.py

@@ -112,6 +112,23 @@ def seconds_until_reset(
     return min(wait, _MAX_WAIT_SECS)
 
 
+def wait_for_rate_limit_reset(reset_message: str) -> None:
+    """Sleep until Claude's reported usage reset, including a short grace period."""
+    reset_hour, reset_minute, is_utc = parse_reset_time(reset_message)
+    wait_secs = seconds_until_reset(reset_hour, is_utc, reset_minute)
+    logger.warning(
+        "Rate-limited (%s). Sleeping %.1f hours for reset at %02d:%02d %s "
+        "(plus 5-minute grace) ...",
+        reset_message,
+        wait_secs / 3600,
+        reset_hour,
+        reset_minute,
+        "UTC" if is_utc else "local",
+    )
+    time.sleep(wait_secs)
+    logger.info("Woke up after rate-limit sleep; retrying...")
+
+
 def _fold_retry_metrics(
     result: SandboxResult,
     aborted_tokens: int,
@@ -319,22 +336,7 @@ def budget_stop_reason(latest: SandboxResult) -> str | None:
 
         if rate_limited:
             assert result.rate_limit_reset is not None
-            reset_hour, reset_minute, is_utc = parse_reset_time(result.rate_limit_reset)
-            wait_secs = seconds_until_reset(reset_hour, is_utc, reset_minute)
-            hours = wait_secs / 3600
-            logger.warning(
-                "Rate-limited (%s). Sleeping %.1f hours for reset at %02d:%02d %s "
-                "(plus 5-minute grace) ...",
-                result.error,
-                hours,
-                reset_hour,
-                reset_minute,
-                "UTC" if is_utc else "local",
-            )
-            time.sleep(wait_secs)
-            logger.info(
-                "Woke up after rate-limit sleep, resuming with remaining budget..."
-            )
+            wait_for_rate_limit_reset(result.rate_limit_reset)
             resume_prompt = active.prompt
         else:
             output_token_retries += 1

tests/utils/test_apptainer_sandbox.py

@@ -16,7 +16,11 @@
     _build_apptainer_cmd,
     run_genplan_in_apptainer,
 )
-from robocode.utils.docker_sandbox import DOCKER_PYTHON, _find_repo_root
+from robocode.utils.docker_sandbox import (
+    DOCKER_PYTHON,
+    GENPLAN_CONTAINER_TIMEOUT_S,
+    _find_repo_root,
+)
 
 
 def test_apptainer_python_matches_docker_python() -> None:
@@ -182,8 +186,11 @@ def test_genplan_cmd_adds_containall(
     )
     calls: list[list[str]] = []
 
-    def fake_run(cmd: list[str], **_kwargs) -> None:
+    timeouts: list[float] = []
+
+    def fake_run(cmd: list[str], **kwargs) -> None:
         calls.append(cmd)
+        timeouts.append(kwargs["timeout"])
 
     monkeypatch.setattr("robocode.utils.apptainer_sandbox.subprocess.run", fake_run)
 
@@ -194,6 +201,7 @@ def fake_run(cmd: list[str], **_kwargs) -> None:
     )
 
     assert len(calls) == 1
+    assert timeouts == [GENPLAN_CONTAINER_TIMEOUT_S]
     assert calls[0][:3] == ["apptainer", "exec", "--containall"]
     assert "--pid" in calls[0]