coderouter-cli 2.0.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

coderouter/guards/self_healing.py

@@ -0,0 +1,413 @@
+"""Self-healing routing orchestrator (v2.0-J).
+
+When the L5 :class:`BackendHealthMonitor` transitions a provider to
+UNHEALTHY and the profile's ``backend_health_action`` is ``exclude``,
+this orchestrator:
+
+1. **Excludes** the provider from the chain (complete removal, not
+   just demotion to the back).
+2. **Attempts restart** if the provider declares a ``restart_command``
+   in ``providers.yaml``.
+3. **Schedules recovery probes** with exponential backoff (default
+   30 s → 60 s → 120 s → 300 s cap) to detect when the backend
+   comes back online.
+4. **Restores** the provider to its original chain position on the
+   first successful recovery probe.
+
+Architecture
+============
+
+::
+
+    BackendHealthMonitor
+      └─ transition → UNHEALTHY
+           └─ engine calls orchestrator.on_unhealthy(provider)
+                ├─ add to _excluded set
+                ├─ try restart_command (if configured)
+                └─ schedule recovery probe (async)
+
+    recovery_probe_loop:
+      while provider in _excluded:
+        sleep(backoff interval)
+        probe_one(provider)
+        if success:
+          remove from _excluded
+          record_attempt(success=True)  → snap to HEALTHY
+          log restore
+          break
+        else:
+          interval = min(interval * 2, max_interval)
+
+    _resolve_chain (engine):
+      Pass 4b: if action == "exclude":
+        filter out providers in orchestrator._excluded
+
+Design choices
+==============
+
+- **Thread-safe** via an internal ``RLock`` (same pattern as
+  ``BackendHealthMonitor``). The exclude set and restart lock
+  are guarded independently.
+- **No new dependency** — subprocess (stdlib) for restart commands,
+  asyncio for recovery probe scheduling.
+- **Restart is opt-in** — only providers with ``restart_command``
+  set get automatic restart. Others rely solely on recovery probes
+  (waiting for manual restart by the operator).
+- **Double-restart prevention** — a per-provider ``_restart_lock``
+  prevents concurrent restart attempts (e.g. two profiles both
+  hitting UNHEALTHY on the same provider).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import subprocess
+import threading
+import time
+from dataclasses import dataclass
+
+from coderouter.config.schemas import ProviderConfig
+from coderouter.logging import (
+    get_logger,
+    log_self_healing_exclude,
+    log_self_healing_recovery_probe,
+    log_self_healing_restart,
+    log_self_healing_restore,
+)
+
+logger = get_logger(__name__)
+
+
+@dataclass(slots=True)
+class _ExcludedProvider:
+    """Metadata for a provider currently excluded from the chain."""
+
+    provider: str
+    excluded_at: float
+    profile: str
+    consecutive_failures: int
+
+
+class SelfHealingOrchestrator:
+    """Manages provider exclusion, restart, and recovery probing.
+
+    Public API:
+
+    - :meth:`on_unhealthy(provider, ...)` — called by the engine when
+      a provider transitions to UNHEALTHY with action ``exclude``.
+    - :meth:`on_recovered(provider, ...)` — called when a recovery
+      probe succeeds or a regular request succeeds on a previously
+      excluded provider.
+    - :meth:`is_excluded(provider)` — True iff the provider is
+      currently excluded from the chain.
+    - :meth:`excluded_providers()` — set of currently excluded
+      provider names.
+    - :meth:`try_restart(provider_config, ...)` — attempt to restart
+      a provider's backend process.
+    - :meth:`reset()` — clear all state. Mainly for tests.
+    """
+
+    def __init__(self) -> None:
+        self._lock: threading.RLock = threading.RLock()
+        self._excluded: dict[str, _ExcludedProvider] = {}
+        # Per-provider lock to prevent concurrent restart attempts.
+        self._restart_locks: dict[str, threading.Lock] = {}
+
+    # ------------------------------------------------------------------
+    # Exclusion management
+    # ------------------------------------------------------------------
+
+    def on_unhealthy(
+        self,
+        provider: str,
+        *,
+        profile: str,
+        consecutive_failures: int,
+    ) -> bool:
+        """Mark a provider as excluded from the chain.
+
+        Returns True if the provider was newly excluded (not already
+        excluded). Returns False if it was already excluded (idempotent).
+        """
+        with self._lock:
+            if provider in self._excluded:
+                return False
+            self._excluded[provider] = _ExcludedProvider(
+                provider=provider,
+                excluded_at=time.monotonic(),
+                profile=profile,
+                consecutive_failures=consecutive_failures,
+            )
+        log_self_healing_exclude(
+            logger,
+            provider=provider,
+            profile=profile,
+            consecutive_failures=consecutive_failures,
+        )
+        return True
+
+    def on_recovered(
+        self,
+        provider: str,
+        *,
+        profile: str,
+    ) -> float | None:
+        """Restore a provider to the chain after recovery.
+
+        Returns the duration (seconds) the provider was excluded,
+        or None if the provider was not in the excluded set.
+        """
+        with self._lock:
+            entry = self._excluded.pop(provider, None)
+            if entry is None:
+                return None
+        duration = time.monotonic() - entry.excluded_at
+        log_self_healing_restore(
+            logger,
+            provider=provider,
+            profile=profile,
+            excluded_duration_s=duration,
+        )
+        return duration
+
+    def is_excluded(self, provider: str) -> bool:
+        """True iff the provider is currently excluded from the chain."""
+        with self._lock:
+            return provider in self._excluded
+
+    def excluded_providers(self) -> set[str]:
+        """Return a snapshot of currently excluded provider names."""
+        with self._lock:
+            return set(self._excluded.keys())
+
+    def reset(self) -> None:
+        """Drop all state. Mainly for tests."""
+        with self._lock:
+            self._excluded.clear()
+            self._restart_locks.clear()
+
+    # ------------------------------------------------------------------
+    # v2.0-K: Persistence
+    # ------------------------------------------------------------------
+
+    def save_state(self) -> dict[str, object]:
+        """Export the current excluded-provider set for persistence."""
+        with self._lock:
+            return {
+                name: {
+                    "profile": entry.profile,
+                    "consecutive_failures": entry.consecutive_failures,
+                }
+                for name, entry in self._excluded.items()
+            }
+
+    def load_state(self, state: dict[str, object]) -> None:
+        """Restore excluded providers from a previously saved dict.
+
+        Re-creates ``_ExcludedProvider`` entries with ``excluded_at``
+        set to the current time (the original exclude timestamp is
+        lost across restarts — the important thing is that the provider
+        *stays* excluded until a recovery probe succeeds).
+        """
+        if not isinstance(state, dict):
+            return
+
+        with self._lock:
+            for name, data in state.items():
+                if not isinstance(data, dict):
+                    continue
+                if name in self._excluded:
+                    continue  # already excluded
+                profile = data.get("profile", "")
+                failures = data.get("consecutive_failures", 0)
+                if not isinstance(failures, int):
+                    failures = 0
+                self._excluded[name] = _ExcludedProvider(
+                    provider=name,
+                    excluded_at=time.monotonic(),
+                    profile=str(profile),
+                    consecutive_failures=failures,
+                )
+
+    # ------------------------------------------------------------------
+    # Restart helper
+    # ------------------------------------------------------------------
+
+    def try_restart(
+        self,
+        provider_config: ProviderConfig,
+        *,
+        timeout_s: float = 30.0,
+    ) -> bool:
+        """Attempt to restart a provider's backend process.
+
+        Returns True if the restart command succeeded (exit code 0),
+        False otherwise (or if no restart_command is configured).
+
+        Thread-safe: only one restart per provider at a time. A
+        concurrent call returns False immediately without blocking.
+        """
+        command = provider_config.restart_command
+        if not command:
+            return False
+
+        provider = provider_config.name
+
+        # Get or create a per-provider lock.
+        with self._lock:
+            if provider not in self._restart_locks:
+                self._restart_locks[provider] = threading.Lock()
+            restart_lock = self._restart_locks[provider]
+
+        # Non-blocking acquire — if another thread is already
+        # restarting this provider, we skip silently.
+        if not restart_lock.acquire(blocking=False):
+            return False
+
+        try:
+            result = subprocess.run(
+                command,
+                shell=True,
+                capture_output=True,
+                timeout=timeout_s,
+                text=True,
+            )
+            success = result.returncode == 0
+            error = result.stderr.strip() if not success else None
+            log_self_healing_restart(
+                logger,
+                provider=provider,
+                command=command,
+                success=success,
+                error=error,
+            )
+            return success
+        except subprocess.TimeoutExpired:
+            log_self_healing_restart(
+                logger,
+                provider=provider,
+                command=command,
+                success=False,
+                error=f"timeout after {timeout_s}s",
+            )
+            return False
+        except OSError as exc:
+            log_self_healing_restart(
+                logger,
+                provider=provider,
+                command=command,
+                success=False,
+                error=str(exc),
+            )
+            return False
+        finally:
+            restart_lock.release()
+
+
+# ---------------------------------------------------------------------------
+# Recovery probe loop (async, runs as a background task)
+# ---------------------------------------------------------------------------
+
+
+async def recovery_probe_loop(
+    provider_config: ProviderConfig,
+    *,
+    orchestrator: SelfHealingOrchestrator,
+    record_fn: object | None = None,
+    health_threshold: int = 3,
+    initial_interval_s: float = 30.0,
+    max_interval_s: float = 300.0,
+    restart_timeout_s: float = 30.0,
+    probe_timeout_s: float = 10.0,
+    shutdown_event: asyncio.Event | None = None,
+    profile: str = "",
+) -> None:
+    """Probe an excluded provider with exponential backoff until recovery.
+
+    This function runs as a long-lived asyncio task, one per excluded
+    provider. It terminates when:
+
+    - The provider recovers (probe succeeds) → restores to chain.
+    - The shutdown event is set → graceful exit.
+    - The provider is no longer excluded (external recovery).
+
+    On first invocation, attempts a restart if configured, then waits
+    for the initial interval before the first probe.
+    """
+    from coderouter.guards.continuous_probe import probe_one
+
+    _shutdown = shutdown_event or asyncio.Event()
+    provider_name = provider_config.name
+    interval = initial_interval_s
+
+    # Step 1: attempt restart if configured.
+    if provider_config.restart_command:
+        # Run in executor to avoid blocking the event loop.
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(
+            None,
+            lambda: orchestrator.try_restart(
+                provider_config,
+                timeout_s=restart_timeout_s,
+            ),
+        )
+
+    # Step 2: exponential backoff recovery probes.
+    while not _shutdown.is_set():
+        # Check if still excluded (external code may have restored it).
+        if not orchestrator.is_excluded(provider_name):
+            return
+
+        # Wait for the current interval (or shutdown).
+        try:
+            await asyncio.wait_for(_shutdown.wait(), timeout=interval)
+            return  # shutdown signalled
+        except TimeoutError:
+            pass  # normal: interval elapsed
+
+        # Still excluded? Probe.
+        if not orchestrator.is_excluded(provider_name):
+            return
+
+        result = await probe_one(provider_config, timeout_s=probe_timeout_s)
+
+        if result.success:
+            # Provider is back! Restore it.
+            orchestrator.on_recovered(provider_name, profile=profile)
+
+            # Feed success into the backend health state machine
+            # so it snaps back to HEALTHY.
+            if record_fn is not None:
+                with contextlib.suppress(Exception):
+                    record_fn(  # type: ignore[operator]
+                        provider_name,
+                        success=True,
+                        threshold=health_threshold,
+                    )
+
+            log_self_healing_recovery_probe(
+                logger,
+                provider=provider_name,
+                success=True,
+                next_interval_s=0,
+                latency_ms=result.latency_ms,
+            )
+            return
+
+        # Failed — exponential backoff.
+        next_interval = min(interval * 2, max_interval_s)
+        log_self_healing_recovery_probe(
+            logger,
+            provider=provider_name,
+            success=False,
+            next_interval_s=next_interval,
+            latency_ms=result.latency_ms,
+        )
+        interval = next_interval
+
+
+__all__ = [
+    "SelfHealingOrchestrator",
+    "recovery_probe_loop",
+]

coderouter/guards/tool_loop.py

@@ -148,6 +148,44 @@ class ToolLoopDetection:
     """
 
 
+@dataclass(frozen=True)
+class ToolCountExceeded:
+    """The outcome of a total tool-call count check.
+
+    Returned by :func:`check_total_tool_count` when the conversation's
+    cumulative tool_use count exceeds the configured hard cap. This is
+    a safety valve against runaway agents that call many *different*
+    tools without looping (which L3's identical-streak detector misses).
+    """
+
+    total_count: int
+    """How many tool_use blocks the conversation currently contains."""
+    max_allowed: int
+    """The configured ceiling that was exceeded."""
+
+
+class ToolCountExceededError(CodeRouterError):
+    """Raised when total tool-call count exceeds the hard cap.
+
+    The ingress converts this into a structured ``400`` response with
+    ``error: "tool_count_exceeded"`` so the client sees a programmable
+    failure rather than a 5xx.
+    """
+
+    def __init__(
+        self,
+        exceeded: ToolCountExceeded,
+        profile: str,
+    ) -> None:
+        super().__init__(
+            f"tool count exceeded on profile={profile!r}: "
+            f"{exceeded.total_count} tool calls exceed the limit of "
+            f"{exceeded.max_allowed}."
+        )
+        self.exceeded = exceeded
+        self.profile = profile
+
+
 class ToolLoopBreakError(CodeRouterError):
     """Raised when a loop is detected and the configured action is ``break``.
 
@@ -337,3 +375,36 @@ def inject_loop_break_hint(
         new_system = [*list(system), {"type": "text", "text": hint}]
 
     return request.model_copy(update={"system": new_system})
+
+
+# ---------------------------------------------------------------------------
+# Total tool-call count hard cap (v2.2)
+# ---------------------------------------------------------------------------
+
+
+def check_total_tool_count(
+    request: AnthropicRequest,
+    *,
+    max_calls: int,
+) -> ToolCountExceeded | None:
+    """Return a detection if total tool_use count exceeds ``max_calls``.
+
+    Unlike :func:`detect_tool_loop` which catches *identical*
+    consecutive calls, this is a blunt hard cap on the cumulative
+    number of tool_use blocks across the entire conversation. It
+    catches runaway agents that cycle through many *different* tools
+    without ever repeating the same (name, args) pair — a pattern
+    that the streak-based L3 detector cannot see.
+
+    Default ceiling is 50 (configurable per-profile). This is
+    deliberately more permissive than Unsloth Studio's 25 — Claude
+    Code's long-running agent sessions routinely reach 25+ tool calls
+    in normal operation.
+
+    Returns ``None`` when the count is within limits.
+    """
+    history = _extract_tool_use_history(request)
+    count = len(history)
+    if count > max_calls:
+        return ToolCountExceeded(total_count=count, max_allowed=max_calls)
+    return None

coderouter/ingress/anthropic_routes.py

@@ -28,7 +28,7 @@ from typing import Any
 from fastapi import APIRouter, Header, HTTPException, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 
-from coderouter.guards.tool_loop import ToolLoopBreakError
+from coderouter.guards.tool_loop import ToolCountExceededError, ToolLoopBreakError
 from coderouter.logging import get_logger
 from coderouter.routing import (
     FallbackEngine,
@@ -49,6 +49,7 @@ _MODE_HEADER = "x-coderouter-mode"
 _ANTHROPIC_VERSION_HEADER = "anthropic-version"
 _ANTHROPIC_BETA_HEADER = "anthropic-beta"
 _CTX_BUDGET_HEADER = "X-CodeRouter-Context-Budget"
+_DRIFT_HEADER = "X-CodeRouter-Drift"
 
 
 @router.post("/messages", response_model=None)
@@ -144,6 +145,12 @@ async def messages(
         }
         if ctx_budget_status:
             stream_headers[_CTX_BUDGET_HEADER] = ctx_budget_status
+        # v2.0-G: drift header is set post-stream via a trailer-like
+        # mechanism — for streaming we cannot know the verdict before
+        # the first chunk ships. Instead, check pre-existing drift state.
+        drift_severity = engine.last_drift_severity
+        if drift_severity:
+            stream_headers[_DRIFT_HEADER] = drift_severity
         return StreamingResponse(
             _anthropic_sse_iterator(engine, anth_req),
             media_type="text/event-stream",
@@ -166,11 +173,31 @@ async def messages(
             status_code=400,
             detail=_tool_loop_break_detail(exc),
         ) from exc
+    except ToolCountExceededError as exc:
+        # v2.2: total tool-call count exceeded — surface as a 400.
+        raise HTTPException(
+            status_code=400,
+            detail={
+                "error": "tool_count_exceeded",
+                "message": str(exc),
+                "total_count": exc.exceeded.total_count,
+                "max_allowed": exc.exceeded.max_allowed,
+                "profile": exc.profile,
+            },
+        ) from exc
 
+    # v2.0-G: collect drift header after engine dispatch.
+    drift_severity = engine.last_drift_severity
+    resp_headers: dict[str, str] = {}
     if ctx_budget_status:
+        resp_headers[_CTX_BUDGET_HEADER] = ctx_budget_status
+    if drift_severity:
+        resp_headers[_DRIFT_HEADER] = drift_severity
+
+    if resp_headers:
         return JSONResponse(
             content=anth_resp.model_dump(exclude_none=True),
-            headers={_CTX_BUDGET_HEADER: ctx_budget_status},
+            headers=resp_headers,
         )
     return anth_resp.model_dump(exclude_none=True)
 
@@ -224,26 +251,93 @@ async def _anthropic_sse_iterator(
             },
         )
         yield _format_anthropic_sse(err_event)
-    except MidStreamError as exc:
-        # v0.3-B: a provider failed AFTER emitting at least one event. We
-        # cannot fall back (client already received partial content), so
-        # close the stream with an explicit error event. `api_error`
-        # distinguishes this from "no provider could start" (overloaded).
-        logger.warning(
-            "sse-midstream-error",
-            extra={"provider": exc.provider, "original": str(exc.original)},
-        )
+    except ToolCountExceededError as exc:
+        # v2.2: streaming counterpart of the tool-count-exceeded 400.
         err_event = AnthropicStreamEvent(
             type="error",
             data={
                 "type": "error",
                 "error": {
-                    "type": "api_error",
+                    "type": "invalid_request_error",
                     "message": str(exc),
+                    "tool_count": {
+                        "total_count": exc.exceeded.total_count,
+                        "max_allowed": exc.exceeded.max_allowed,
+                        "profile": exc.profile,
+                    },
                 },
             },
         )
         yield _format_anthropic_sse(err_event)
+    except MidStreamError as exc:
+        # v0.3-B: a provider failed AFTER emitting at least one event. We
+        # cannot fall back (client already received partial content), so
+        # close the stream with an explicit error event. `api_error`
+        # distinguishes this from "no provider could start" (overloaded).
+        logger.warning(
+            "sse-midstream-error",
+            extra={"provider": exc.provider, "original": str(exc.original)},
+        )
+
+        # v2.0-H (L6): partial stitch surface mode — synthesize a graceful
+        # stream termination that delivers accumulated text to the client.
+        profile_name = anth_req.profile or engine.config.default_profile
+        partial_action = "off"
+        try:
+            chain_cfg = engine.config.profile_by_name(profile_name)
+            partial_action = chain_cfg.partial_stitch_action
+        except (KeyError, ValueError):
+            pass
+
+        if partial_action == "surface" and exc.partial_content:
+            # Emit message_delta with accumulated usage (signals stream end).
+            yield _format_anthropic_sse(AnthropicStreamEvent(
+                type="message_delta",
+                data={
+                    "type": "message_delta",
+                    "delta": {"stop_reason": None, "stop_sequence": None},
+                    "usage": {"output_tokens": 0},
+                },
+            ))
+            # Emit message_stop so the client sees a complete stream.
+            yield _format_anthropic_sse(AnthropicStreamEvent(
+                type="message_stop",
+                data={"type": "message_stop"},
+            ))
+            # Emit coderouter_partial metadata event (client-optional).
+            yield _format_anthropic_sse(AnthropicStreamEvent(
+                type="coderouter_partial",
+                data={
+                    "type": "coderouter_partial",
+                    "partial_content": exc.partial_content,
+                    "provider": exc.provider,
+                    "reason": "mid_stream_failure",
+                    "original_error": str(exc.original)[:200],
+                },
+            ))
+            logger.info(
+                "partial-stitch-surfaced",
+                extra={
+                    "provider": exc.provider,
+                    "profile": profile_name,
+                    "text_blocks": len(exc.partial_content),
+                    "text_length": sum(
+                        len(b.get("text", "")) for b in exc.partial_content
+                    ),
+                },
+            )
+        else:
+            err_event = AnthropicStreamEvent(
+                type="error",
+                data={
+                    "type": "error",
+                    "error": {
+                        "type": "api_error",
+                        "message": str(exc),
+                    },
+                },
+            )
+            yield _format_anthropic_sse(err_event)
 
 
 def _format_anthropic_sse(ev: AnthropicStreamEvent) -> str: