fix(litellm): retry transient upstream 5xx at the proxy

src/benchflow/providers/litellm_config.py

@@ -425,6 +425,14 @@ def litellm_proxy_config(
             # storms or deployment cooldown fast-fails (#830).
             "num_retries": 0,
             "disable_cooldowns": True,
+            # ...with ONE narrow exception honoring #830's rationale: transient
+            # upstream 5xx (InternalServerError) are not deterministic rejects,
+            # and some agents silently swallow a passed-through 500 into a
+            # zero-event turn (observed: MiMo Code ending turns with no error on
+            # a deepseek 500, repro x3). Absorb up to 2 transient 500s at the
+            # proxy; every deterministic class (4xx, auth, content) still fails
+            # fast with zero retries.
+            "retry_policy": {"InternalServerErrorRetries": 2},
         },
         "litellm_settings": {
             "callbacks": [f"{callback_module}.proxy_handler_instance"],

src/benchflow/providers/litellm_retry_patch.py

@@ -0,0 +1,102 @@
+"""LiteLLM startup patch for transient upstream 5xx retry policy.
+
+LiteLLM 1.89.0's ``RetryPolicy`` model accepts
+``InternalServerErrorRetries``, but the Router helper that reads retry-policy
+fields never checks that attribute. BenchFlow still keeps the proxy-wide
+``num_retries=0`` fail-fast default; this patch makes the existing
+``InternalServerErrorRetries`` policy work for the one transient 5xx class.
+
+This module is copied into the per-run LiteLLM runtime and imported by
+``sitecustomize``. Keep it standalone: sandbox proxy processes may not be able
+to import the BenchFlow package tree.
+"""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+
+def _effective_retry_policy(
+    *,
+    retry_policy: Any,
+    model_group: str | None,
+    model_group_retry_policy: dict[str, Any] | None,
+) -> Any:
+    if (
+        model_group_retry_policy is not None
+        and model_group is not None
+        and model_group in model_group_retry_policy
+    ):
+        return model_group_retry_policy.get(model_group)
+    return retry_policy
+
+
+def _internal_server_error_retries(policy: Any) -> int | None:
+    if policy is None:
+        return None
+    if isinstance(policy, dict):
+        value = policy.get("InternalServerErrorRetries")
+    else:
+        value = getattr(policy, "InternalServerErrorRetries", None)
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def _patch_internal_server_error_retry_policy() -> None:
+    try:
+        import litellm
+        import litellm.router as router_mod
+        from litellm.router_utils import get_retry_from_policy as policy_mod
+    except Exception:
+        return
+
+    original = getattr(router_mod, "_get_num_retries_from_retry_policy", None)
+    if original is None:
+        original = getattr(policy_mod, "get_num_retries_from_retry_policy", None)
+    if original is None or getattr(original, "__benchflow_retry_patch__", False):
+        return
+
+    internal_server_error = getattr(litellm, "InternalServerError", None)
+    if internal_server_error is None:
+        return
+
+    def patched_get_num_retries_from_retry_policy(
+        exception: Exception,
+        retry_policy: Any = None,
+        model_group: str | None = None,
+        model_group_retry_policy: dict[str, Any] | None = None,
+    ) -> int | None:
+        retries = original(
+            exception=exception,
+            retry_policy=retry_policy,
+            model_group=model_group,
+            model_group_retry_policy=model_group_retry_policy,
+        )
+        if retries is not None:
+            return retries
+        if not isinstance(exception, internal_server_error):
+            return None
+        policy = _effective_retry_policy(
+            retry_policy=retry_policy,
+            model_group=model_group,
+            model_group_retry_policy=model_group_retry_policy,
+        )
+        return _internal_server_error_retries(policy)
+
+    patched_any = cast(Any, patched_get_num_retries_from_retry_policy)
+    patched_any.__benchflow_retry_patch__ = True
+    policy_module_any = cast(Any, policy_mod)
+    router_module_any = cast(Any, router_mod)
+    policy_module_any.get_num_retries_from_retry_policy = (
+        patched_get_num_retries_from_retry_policy
+    )
+    router_module_any._get_num_retries_from_retry_policy = (
+        patched_get_num_retries_from_retry_policy
+    )
+
+
+_patch_internal_server_error_retry_policy()

src/benchflow/providers/litellm_retry_preflight.py

@@ -0,0 +1,123 @@
+"""Fail-closed checks for BenchFlow's LiteLLM 5xx retry patch."""
+
+from __future__ import annotations
+
+import shlex
+import subprocess
+from typing import Any
+
+from benchflow.providers.litellm_bedrock_preflight import (
+    _exec_details,
+    _host_python_for_litellm,
+)
+
+
+class RetryPatchPreflightError(RuntimeError):
+    """Raised when the transient-5xx retry patch cannot be proven active."""
+
+
+RETRY_PATCH_PREFLIGHT_SOURCE = """\
+import sys
+
+failures = []
+try:
+    import litellm
+    import litellm.router as router_mod
+
+    helper = getattr(router_mod, "_get_num_retries_from_retry_policy", None)
+    if not getattr(helper, "__benchflow_retry_patch__", False):
+        failures.append("transient-5xx retry helper is unpatched")
+    else:
+        policy = {"InternalServerErrorRetries": 2}
+        transient = litellm.InternalServerError(
+            "upstream 500",
+            llm_provider="openai",
+            model="openai/test",
+        )
+        retries = helper(exception=transient, retry_policy=policy)
+        if retries != 2:
+            failures.append(
+                f"InternalServerError retry probe returned {retries!r}, expected 2"
+            )
+
+        permanent = litellm.BadRequestError(
+            "bad request",
+            llm_provider="openai",
+            model="openai/test",
+        )
+        permanent_retries = helper(exception=permanent, retry_policy=policy)
+        if permanent_retries is not None:
+            failures.append(
+                "BadRequestError retry probe returned "
+                f"{permanent_retries!r}, expected fail-fast None"
+            )
+except Exception as exc:  # noqa: BLE001 - report any LiteLLM shape drift
+    failures.append(f"retry preflight probe failed: {exc}")
+
+if failures:
+    print("; ".join(failures))
+    sys.exit(1)
+"""
+
+
+def _failure_message(runtime: str, detail: str) -> str:
+    return (
+        "BenchFlow LiteLLM transient-5xx retry patch is NOT active in the "
+        f"{runtime} LiteLLM runtime: {detail[:2000]}. Failing closed before "
+        "agent launch - a silent fallback would expose transient upstream "
+        "500s to agents instead of retrying them at the proxy."
+    )
+
+
+def preflight_host_retry_patch(
+    *,
+    env: dict[str, str],
+    litellm_executable: str,
+) -> None:
+    """Fail closed if the retry patch is not active for the host proxy."""
+    try:
+        result = subprocess.run(
+            [
+                _host_python_for_litellm(litellm_executable, env=env),
+                "-c",
+                RETRY_PATCH_PREFLIGHT_SOURCE,
+            ],
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=120,
+        )
+    except Exception as exc:
+        raise RetryPatchPreflightError(
+            _failure_message("host", f"preflight execution failed: {exc}")
+        ) from exc
+    if result.returncode != 0:
+        detail = (result.stdout or "").strip() or (result.stderr or "").strip()
+        raise RetryPatchPreflightError(_failure_message("host", detail))
+
+
+async def preflight_sandbox_retry_patch(
+    sandbox: Any,
+    *,
+    python: str,
+    runtime_dir: str,
+    preflight_path: str,
+) -> None:
+    """Fail closed if the retry patch is not active in a sandbox proxy."""
+    command = (
+        f"PYTHONPATH={shlex.quote(runtime_dir)} "
+        f"{shlex.quote(python)} {shlex.quote(preflight_path)}"
+    )
+    try:
+        result = await sandbox.exec(command, timeout_sec=120)
+    except Exception as exc:
+        raise RetryPatchPreflightError(
+            _failure_message("sandbox", f"preflight execution failed: {exc}")
+        ) from exc
+    if result.return_code != 0:
+        raise RetryPatchPreflightError(
+            _failure_message(
+                "sandbox",
+                _exec_details("retry patch preflight", result),
+            )
+        )

src/benchflow/providers/litellm_runtime.py

@@ -48,6 +48,11 @@
     extract_usage_from_trajectory,
     trajectory_from_litellm_callback_log,
 )
+from benchflow.providers.litellm_retry_preflight import (
+    RETRY_PATCH_PREFLIGHT_SOURCE,
+    preflight_host_retry_patch,
+    preflight_sandbox_retry_patch,
+)
 from benchflow.sandbox.providers import OFF_BOX_MODEL_PROVIDERS
 from benchflow.trajectories.types import Trajectory
 from benchflow.usage_tracking import UsageTrackingConfig, usage_unavailable
@@ -57,7 +62,8 @@
 LITELLM_VERSION_SPEC = "litellm[proxy]==1.89.0"
 LITELLM_SANDBOX_ROOT = "/tmp/benchflow-litellm"
 _CALLBACK_MODULE = "benchflow_litellm_callback"
-_PATCH_MODULE = "benchflow_litellm_bedrock_patch"
+_BEDROCK_PATCH_MODULE = "benchflow_litellm_bedrock_patch"
+_RETRY_PATCH_MODULE = "benchflow_litellm_retry_patch"
 
 # The proxy is an internal single-route gateway — it must never register the
 # FastAPI Swagger docs route. litellm's `_get_docs_url()` honours an inherited
@@ -409,13 +415,18 @@ def _write_runtime_files(
 ) -> tuple[Path, Path, Path]:
     runtime_dir.mkdir(parents=True, exist_ok=True)
     callback_path = runtime_dir / f"{_CALLBACK_MODULE}.py"
-    patch_path = runtime_dir / f"{_PATCH_MODULE}.py"
+    patch_path = runtime_dir / f"{_BEDROCK_PATCH_MODULE}.py"
+    retry_patch_path = runtime_dir / f"{_RETRY_PATCH_MODULE}.py"
     sitecustomize_path = runtime_dir / "sitecustomize.py"
     config_path = runtime_dir / "config.yaml"
     callback_path.write_text(callback_module_source())
     patch_source = Path(__file__).with_name("litellm_bedrock_patch.py").read_text()
     patch_path.write_text(patch_source)
-    sitecustomize_path.write_text(f"import {_PATCH_MODULE}\n")
+    retry_patch_source = Path(__file__).with_name("litellm_retry_patch.py").read_text()
+    retry_patch_path.write_text(retry_patch_source)
+    sitecustomize_path.write_text(
+        f"import {_BEDROCK_PATCH_MODULE}\nimport {_RETRY_PATCH_MODULE}\n"
+    )
     config_path.write_text(yaml.safe_dump(config, sort_keys=False))
     return config_path, callback_path, patch_path
 
@@ -536,6 +547,11 @@ async def _start_host_litellm(
     )
     try:
         await _poll_host_health(runner)
+        await asyncio.to_thread(
+            preflight_host_retry_patch,
+            env=env,
+            litellm_executable=litellm_executable,
+        )
         if route_requires_bedrock_patch(route):
             # Fail closed before the agent launches when the Bedrock 4.8+
             # thinking patch did not activate (#602).
@@ -626,7 +642,8 @@ async def _upload_runtime_files_to_sandbox(
     paths = {
         "config": f"{runtime_dir}/config.yaml",
         "callback": f"{runtime_dir}/{_CALLBACK_MODULE}.py",
-        "patch": f"{runtime_dir}/{_PATCH_MODULE}.py",
+        "patch": f"{runtime_dir}/{_BEDROCK_PATCH_MODULE}.py",
+        "retry_patch": f"{runtime_dir}/{_RETRY_PATCH_MODULE}.py",
         "sitecustomize": f"{runtime_dir}/sitecustomize.py",
         "launcher": f"{runtime_dir}/launcher.py",
         "stdout": f"{runtime_dir}/stdout.log",
@@ -637,6 +654,7 @@ async def _upload_runtime_files_to_sandbox(
         "venv": f"{runtime_dir}/venv",
         "launch_config": f"{runtime_dir}/launch_config.json",
         "preflight": f"{runtime_dir}/bedrock_patch_preflight.py",
+        "retry_preflight": f"{runtime_dir}/retry_patch_preflight.py",
     }
     result = await sandbox.exec(f"mkdir -p {shlex.quote(runtime_dir)}", timeout_sec=20)
     if result.return_code != 0:
@@ -652,12 +670,24 @@ async def _upload_runtime_files_to_sandbox(
         ".py",
     )
     await _upload_text(
-        sandbox, f"import {_PATCH_MODULE}\n", paths["sitecustomize"], ".py"
+        sandbox,
+        Path(__file__).with_name("litellm_retry_patch.py").read_text(),
+        paths["retry_patch"],
+        ".py",
+    )
+    await _upload_text(
+        sandbox,
+        f"import {_BEDROCK_PATCH_MODULE}\nimport {_RETRY_PATCH_MODULE}\n",
+        paths["sitecustomize"],
+        ".py",
     )
     await _upload_text(sandbox, _sandbox_launcher_source(), paths["launcher"], ".py")
     await _upload_text(
         sandbox, BEDROCK_PATCH_PREFLIGHT_SOURCE, paths["preflight"], ".py"
     )
+    await _upload_text(
+        sandbox, RETRY_PATCH_PREFLIGHT_SOURCE, paths["retry_preflight"], ".py"
+    )
     return paths
 
 
@@ -838,6 +868,12 @@ async def _start_sandbox_litellm(
             port=port,
             stderr_path=paths["stderr"],
         )
+        await preflight_sandbox_retry_patch(
+            sandbox,
+            python=python,
+            runtime_dir=runtime_dir,
+            preflight_path=paths["retry_preflight"],
+        )
         if route_requires_bedrock_patch(route):
             # Fail closed before the agent launches when the Bedrock 4.8+
             # thinking patch did not activate (#602). On Daytona this proxy is

tests/test_litellm_config.py

@@ -181,6 +181,10 @@ def test_proxy_config_registers_plain_and_openai_aliases():
     assert config["router_settings"] == {
         "num_retries": 0,
         "disable_cooldowns": True,
+        # transient-5xx-only exception to the no-retry rule (#830 honored:
+        # deterministic classes keep zero retries) — absorbs upstream 500s that
+        # some agents (MiMo) silently swallow into zero-event turns.
+        "retry_policy": {"InternalServerErrorRetries": 2},
     }