Track 0: refresh/ask reliability + opt-in multi-provider LLM failover

engine/.env.example

@@ -36,6 +36,22 @@
 # Initial backoff delay in seconds (doubles each attempt).
 # AG_REFRESH_RETRY_DELAY=1.0
 
+# -------------------------------------------------------------------
+# Ask retry policy (optional)
+# Same-provider retry for transient ask-time failures (timeouts, 5xx,
+# litellm ServiceUnavailableError). Backoff doubles each attempt.
+# AG_ASK_RETRY_COUNT=3
+# AG_ASK_RETRY_DELAY=5.0
+
+# -------------------------------------------------------------------
+# Multi-provider failover (optional)
+# Ordered JSON array of backup LLM endpoints. When the primary OPENAI_*
+# endpoint keeps failing with a transient/provider error, ag-ask switches
+# to the next provider and re-runs the answer. Each entry may set
+# base_url / api_key / model / label; a missing api_key or model inherits
+# the primary value. Leave unset to disable (default behaviour unchanged).
+# AG_LLM_FALLBACKS=[{"base_url":"https://api.openai.com/v1","api_key":"sk-...","model":"gpt-4o","label":"openai"}]
+
 # --------------------------------------------------------------------
 # Sandbox Configuration (optional)
 # These settings control how code execution is sandboxed.

engine/antigravity_engine/hub/_providers.py

@@ -0,0 +1,191 @@
+"""Multi-provider LLM failover for the Knowledge Hub.
+
+The hub talks to a single OpenAI-compatible endpoint by default
+(``OPENAI_BASE_URL`` / ``OPENAI_API_KEY`` / ``OPENAI_MODEL``).  When that
+endpoint suffers a *sustained* outage, same-provider retries cannot help —
+every retry hits the same dead host.  This module adds an opt-in ordered
+list of backup providers (``AG_LLM_FALLBACKS``) plus a wrapper that re-runs
+an operation against the next provider when the active one keeps failing
+with a transient/provider error.
+
+Behaviour is unchanged when ``AG_LLM_FALLBACKS`` is unset: the chain holds
+exactly one provider and the wrapper is a pass-through with no environment
+mutation.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import sys
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Awaitable, Callable, TypeVar
+
+if TYPE_CHECKING:
+    from antigravity_engine.config import Settings
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+# Substrings (lower-cased) that mark an error as a transient provider failure.
+# Kept as the single source of truth for both the retry path and failover.
+_RETRYABLE_KEYWORDS = (
+    "timeout",
+    "gateway time-out",
+    "504",
+    "connection",
+    "network",
+    "unreachable",
+    "refused",
+    "rate limit",
+    "ratelimit",
+    "429",
+    "502",
+    "503",
+    "500",
+    "serviceunavailable",
+    "service unavailable",
+    "service temporarily unavailable",
+    "temporarily unavailable",
+    "bad gateway",
+    "internalservererror",
+    "internal server error",
+)
+
+
+def is_retryable_provider_error(exc: Exception) -> bool:
+    """Return True for transient model/provider failures worth retrying.
+
+    LiteLLM often wraps a provider 503 in an exception whose text preserves
+    the wording ("Service temporarily unavailable") but not the numeric
+    code, so the classifier matches on both.
+    """
+    if isinstance(exc, (TimeoutError, asyncio.TimeoutError)):
+        return True
+    msg = f"{type(exc).__module__}.{type(exc).__name__}: {exc}".lower()
+    return any(keyword in msg for keyword in _RETRYABLE_KEYWORDS)
+
+
+@dataclass(frozen=True)
+class ProviderConfig:
+    """One LLM endpoint: an OpenAI-compatible base URL, key, and model."""
+
+    model: str
+    base_url: str = ""
+    api_key: str = ""
+    label: str = "primary"
+
+
+def get_provider_chain(settings: "Settings") -> list[ProviderConfig]:
+    """Build the ordered provider list: primary first, then fallbacks.
+
+    The primary comes from the ``OPENAI_*`` settings.  Fallbacks are parsed
+    from the ``AG_LLM_FALLBACKS`` env var — a JSON array of objects with
+    optional ``base_url`` / ``api_key`` / ``model`` / ``label`` keys; a
+    missing ``api_key`` or ``model`` inherits the primary's value.  A
+    malformed value degrades to the primary alone and never breaks the
+    default path.
+    """
+    primary = ProviderConfig(
+        model=settings.OPENAI_MODEL,
+        base_url=settings.OPENAI_BASE_URL,
+        api_key=settings.OPENAI_API_KEY,
+        label="primary",
+    )
+    chain = [primary]
+
+    raw = os.environ.get("AG_LLM_FALLBACKS", "").strip()
+    if not raw:
+        return chain
+
+    try:
+        entries = json.loads(raw)
+    except (ValueError, TypeError) as exc:
+        logger.warning("Ignoring invalid AG_LLM_FALLBACKS (not JSON): %s", exc)
+        return chain
+    if not isinstance(entries, list):
+        logger.warning("Ignoring AG_LLM_FALLBACKS: expected a JSON array")
+        return chain
+
+    for idx, entry in enumerate(entries):
+        if not isinstance(entry, dict):
+            logger.warning("Skipping non-object AG_LLM_FALLBACKS entry #%d", idx)
+            continue
+        model = str(entry.get("model") or primary.model).strip()
+        if not model:
+            logger.warning("Skipping AG_LLM_FALLBACKS entry #%d: empty model", idx)
+            continue
+        chain.append(
+            ProviderConfig(
+                model=model,
+                base_url=str(entry.get("base_url") or "").strip(),
+                api_key=str(entry.get("api_key") or primary.api_key),
+                label=str(entry.get("label") or f"fallback{idx + 1}"),
+            )
+        )
+    return chain
+
+
+def activate_provider(provider: ProviderConfig) -> None:
+    """Make ``provider`` the active LLM endpoint for subsequent agent calls.
+
+    Sets the ``OPENAI_*`` environment variables and resets the cached
+    settings so the next ``get_settings()`` / ``create_model()`` resolves to
+    this provider.  ``litellm`` reads these at request time, so even
+    already-built agents pick up the change on their next call.
+    """
+    from antigravity_engine.config import reset_settings
+
+    os.environ["OPENAI_BASE_URL"] = provider.base_url or ""
+    if provider.api_key:
+        os.environ["OPENAI_API_KEY"] = provider.api_key
+    os.environ["OPENAI_MODEL"] = provider.model
+    reset_settings()
+
+
+async def run_with_provider_failover(
+    operation: Callable[[], Awaitable[T]],
+    *,
+    providers: list[ProviderConfig],
+    is_retryable: Callable[[Exception], bool] | None = None,
+    label: str = "operation",
+) -> T:
+    """Run ``operation`` against each provider until one succeeds.
+
+    With a single provider the behaviour is unchanged: the operation runs
+    once and the environment is left untouched.  With fallbacks configured, a
+    transient/provider failure (per ``is_retryable``) on the active provider
+    triggers a switch to the next provider and a full re-run.  A
+    non-retryable error is raised immediately without failing over.
+    """
+    if is_retryable is None:
+        is_retryable = is_retryable_provider_error
+
+    # No fallback configured: preserve the exact default behaviour.
+    if len(providers) <= 1:
+        return await operation()
+
+    last_exc: Exception | None = None
+    for idx, provider in enumerate(providers):
+        activate_provider(provider)
+        try:
+            return await operation()
+        except Exception as exc:  # noqa: BLE001 — re-raised below unless we fail over
+            last_exc = exc
+            is_last = idx >= len(providers) - 1
+            if is_last or not is_retryable(exc):
+                raise
+            raw_msg = str(exc).replace("\n", " ").replace("\r", "")[:150]
+            next_label = providers[idx + 1].label
+            print(
+                f"  ⚠ Provider '{provider.label}' failed for {label}: "
+                f"{raw_msg or type(exc).__name__}. "
+                f"Failing over to '{next_label}'...",
+                file=sys.stderr,
+            )
+
+    # Unreachable: the loop returns on success or raises on the last provider.
+    assert last_exc is not None
+    raise last_exc

engine/antigravity_engine/hub/ask_pipeline.py

@@ -33,6 +33,31 @@
 logger = logging.getLogger(__name__)
 
 
+def _get_ask_retry_config() -> tuple[int, float]:
+    """Return retry settings for transient ask-time model failures."""
+    try:
+        max_retries = max(0, int(os.environ.get("AG_ASK_RETRY_COUNT", "3")))
+    except (TypeError, ValueError):
+        max_retries = 3
+    try:
+        base_delay = max(0.0, float(os.environ.get("AG_ASK_RETRY_DELAY", "5.0")))
+    except (TypeError, ValueError):
+        base_delay = 5.0
+    return max_retries, base_delay
+
+
+def _is_retryable_ask_error(exc: Exception) -> bool:
+    """Return true for transient model/provider failures.
+
+    Delegates to the shared classifier in ``_providers`` so the same-provider
+    retry path and the cross-provider failover path agree on what counts as
+    transient.
+    """
+    from antigravity_engine.hub._providers import is_retryable_provider_error
+
+    return is_retryable_provider_error(exc)
+
+
 async def _run_with_optional_stream(
     agent: "Agent",
     prompt: str,
@@ -44,32 +69,48 @@ async def _run_with_optional_stream(
     """Execute agent with optional streaming support."""
     from agents import Runner
 
-    if not stream_enabled:
-        # Non-streaming: use existing pattern
-        if timeout and timeout > 0:
-            result = await asyncio.wait_for(
-                Runner.run(agent, prompt, max_turns=max_turns),
-                timeout=timeout,
-            )
-        else:
-            result = await Runner.run(agent, prompt, max_turns=max_turns)
-        return str(result.final_output)
+    async def _run_once() -> str:
+        if not stream_enabled:
+            if timeout and timeout > 0:
+                result = await asyncio.wait_for(
+                    Runner.run(agent, prompt, max_turns=max_turns),
+                    timeout=timeout,
+                )
+            else:
+                result = await Runner.run(agent, prompt, max_turns=max_turns)
+            return str(result.final_output)
 
-    # Streaming mode
-    stream_result = Runner.run_streamed(agent, prompt, max_turns=max_turns)
+        stream_result = Runner.run_streamed(agent, prompt, max_turns=max_turns)
+        try:
+            if timeout and timeout > 0:
+                return await asyncio.wait_for(
+                    _consume_stream_events(stream_result, progress_label),
+                    timeout=timeout,
+                )
+            return await _consume_stream_events(stream_result, progress_label)
+        except Exception:
+            stream_result.cancel()
+            raise
 
-    # For timeout with streaming, wrap the event consumption
-    try:
-        if timeout and timeout > 0:
-            return await asyncio.wait_for(
-                _consume_stream_events(stream_result, progress_label),
-                timeout=timeout,
+    max_retries, base_delay = _get_ask_retry_config()
+    for attempt in range(max_retries + 1):
+        try:
+            return await _run_once()
+        except Exception as exc:
+            if attempt >= max_retries or not _is_retryable_ask_error(exc):
+                raise
+            delay = base_delay * (2 ** attempt)
+            label = f" ({progress_label})" if progress_label else ""
+            raw_msg = str(exc).replace("\n", " ").replace("\r", "")[:150]
+            error_msg = raw_msg or type(exc).__name__
+            print(
+                f"  ⚠ Ask attempt {attempt + 1} failed{label}: "
+                f"{error_msg}. Retrying in {delay}s...",
+                file=sys.stderr,
             )
-        else:
-            return await _consume_stream_events(stream_result, progress_label)
-    except TimeoutError:
-        stream_result.cancel()
-        raise
+            await asyncio.sleep(delay)
+
+    raise RuntimeError("unreachable ask retry state")
 
 
 async def _consume_stream_events(
@@ -138,7 +179,33 @@ async def ask_pipeline(workspace: Path, question: str) -> str:
     Notes:
         MCP servers are only auto-connected when both ``MCP_ENABLED=true`` and
         ``AG_ALLOW_MCP=true`` are set in the runtime environment.
+
+        When ``AG_LLM_FALLBACKS`` configures backup providers, a sustained
+        provider outage on the active endpoint transparently fails over to
+        the next provider and re-runs the answer. With no fallbacks the call
+        is unchanged.
     """
+    from antigravity_engine.config import get_settings
+    from antigravity_engine.hub._providers import (
+        get_provider_chain,
+        run_with_provider_failover,
+    )
+
+    providers = get_provider_chain(get_settings())
+
+    async def _once() -> str:
+        return await _ask_pipeline_once(workspace, question)
+
+    return await run_with_provider_failover(
+        _once,
+        providers=providers,
+        is_retryable=_is_retryable_ask_error,
+        label="ask",
+    )
+
+
+async def _ask_pipeline_once(workspace: Path, question: str) -> str:
+    """Run one ask attempt: structured path first, then legacy swarm."""
     from agents import set_tracing_disabled
 
     set_tracing_disabled(True)

engine/antigravity_engine/hub/module_grouping.py

@@ -374,12 +374,14 @@ def group_files(
     test_files = [f for f in files if f.category == "test"]
     non_test_files = [f for f in files if f.category != "test"]
 
-    # If everything fits in one group, don't split
+    # If non-test files fit in one group, keep them together, but still route
+    # tests through chunking. Large all-test modules can otherwise produce a
+    # single agent instruction string that exceeds provider limits.
     total_eff = sum(f.effective_tokens for f in non_test_files)
     if total_eff <= token_budget and len(non_test_files) <= MAX_FILES_PER_GROUP:
-        groups = [_make_group("main", non_test_files)]
+        groups = [_make_group("main", non_test_files)] if non_test_files else []
         if test_files:
-            groups.append(_make_group("tests", test_files))
+            groups.extend(_chunk_files("tests", test_files, token_budget))
         return groups
 
     # Signal 1: Import graph connected components
@@ -592,7 +594,10 @@ def _chunk_files(
         would_exceed_chars = (
             current_raw_chars + raw_chars > _MAX_RAW_CHARS_PER_GROUP
         )
-        if current.files and (would_exceed_budget or would_exceed_chars):
+        would_exceed_files = len(current.files) >= MAX_FILES_PER_GROUP
+        if current.files and (
+            would_exceed_budget or would_exceed_chars or would_exceed_files
+        ):
             groups.append(current)
             idx += 1
             current = FileGroup(name=f"{base_name}_{idx}", files=[])