From 26c7bf59671c9ec2c97749a6b03bd2ce5908bd98 Mon Sep 17 00:00:00 2001 From: JingWen Fan <106414602+study8677@users.noreply.github.com> Date: Sun, 24 May 2026 18:17:41 +0800 Subject: [PATCH 1/2] fix(hub): chunk all-test modules, fix underscore module paths, retry transient ask failures - module_grouping: route all-test modules through chunking and cap files per group so large test suites don't overflow the provider instruction limit (FastAPI tests module hit 2.4M chars > 1.05M). - scanner: resolve_module_path tries the longest existing parent prefix first, fixing underscored parents like docs_src_additional_responses. - ask_pipeline: retry transient provider failures (503/timeout/litellm ServiceUnavailableError) with exponential backoff on the ask path. Co-Authored-By: Claude Opus 4.7 (1M context) --- engine/antigravity_engine/hub/ask_pipeline.py | 105 ++++++++++++++---- .../antigravity_engine/hub/module_grouping.py | 13 ++- engine/antigravity_engine/hub/scanner.py | 23 ++-- engine/tests/test_hub_module_grouping.py | 39 ++++++- engine/tests/test_hub_pipeline.py | 15 +++ 5 files changed, 159 insertions(+), 36 deletions(-) diff --git a/engine/antigravity_engine/hub/ask_pipeline.py b/engine/antigravity_engine/hub/ask_pipeline.py index 1bb720f3e..f272b4b0a 100644 --- a/engine/antigravity_engine/hub/ask_pipeline.py +++ b/engine/antigravity_engine/hub/ask_pipeline.py @@ -33,6 +33,49 @@ 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.""" + if isinstance(exc, (TimeoutError, asyncio.TimeoutError)): + return True + msg = f"{type(exc).__module__}.{type(exc).__name__}: {exc}".lower() + 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", + ) + return any(keyword in msg for keyword in retryable_keywords) + + async def _run_with_optional_stream( agent: "Agent", prompt: str, @@ -44,32 +87,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( diff --git a/engine/antigravity_engine/hub/module_grouping.py b/engine/antigravity_engine/hub/module_grouping.py index 43cfd3eb5..26f33fd16 100644 --- a/engine/antigravity_engine/hub/module_grouping.py +++ b/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=[]) diff --git a/engine/antigravity_engine/hub/scanner.py b/engine/antigravity_engine/hub/scanner.py index 60a7614ec..13f69df48 100644 --- a/engine/antigravity_engine/hub/scanner.py +++ b/engine/antigravity_engine/hub/scanner.py @@ -779,21 +779,28 @@ def resolve_module_path(root: Path, module_id: str) -> Path: return direct # Two-level case: "parent_child" → root/parent//child - # OR direct auto-split: "parent_child" → root/parent/child + # OR direct auto-split: "parent_child" → root/parent/child. + # Parent directory names can themselves contain underscores + # (e.g. ``docs_src_additional_responses`` → ``docs_src/additional_responses``), + # so try the longest existing parent prefix first instead of splitting once. if "_" in module_id: - parts = module_id.split("_", 1) - parent_dir = root / parts[0] - if parent_dir.is_dir(): - venv_dirs = _find_venv_dirs(root) - skip = _MODULE_SKIP_DIRS | venv_dirs + parts = module_id.split("_") + venv_dirs = _find_venv_dirs(root) + skip = _MODULE_SKIP_DIRS | venv_dirs + for split_at in range(len(parts) - 1, 0, -1): + parent_name = "_".join(parts[:split_at]) + child_name = "_".join(parts[split_at:]) + parent_dir = root / parent_name + if not parent_dir.is_dir(): + continue # First try: single code-bearing inner dir inner = _find_single_code_child(parent_dir, venv_dirs, skip) if inner is not None: - child_dir = inner / parts[1] + child_dir = inner / child_name if child_dir.is_dir(): return child_dir # Second try: direct child (extensions_slack → extensions/slack) - child_dir = parent_dir / parts[1] + child_dir = parent_dir / child_name if child_dir.is_dir(): return child_dir diff --git a/engine/tests/test_hub_module_grouping.py b/engine/tests/test_hub_module_grouping.py index a8784139f..7a1b18c5a 100644 --- a/engine/tests/test_hub_module_grouping.py +++ b/engine/tests/test_hub_module_grouping.py @@ -4,7 +4,12 @@ import pytest from antigravity_engine.hub._constants import WORKSPACE_ROOT_MODULE_ID -from antigravity_engine.hub.module_grouping import group_files, load_module_files +from antigravity_engine.hub.module_grouping import ( + MAX_FILES_PER_GROUP, + format_group_context, + group_files, + load_module_files, +) from antigravity_engine.hub.scanner import detect_modules, resolve_module_path @@ -122,3 +127,35 @@ def _write_text(path: Path, content: str) -> None: """Write a text fixture file, creating parent directories as needed.""" path.parent.mkdir(parents=True, exist_ok=True) path.write_text(content, encoding="utf-8") + + +def test_group_files_chunks_all_test_modules(tmp_path: Path) -> None: + """All-test modules should not bypass group chunking.""" + tests_dir = tmp_path / "tests" + tests_dir.mkdir() + for idx in range(MAX_FILES_PER_GROUP + 5): + (tests_dir / f"test_feature_{idx}.py").write_text( + "def test_feature():\n" + " assert True\n" + + ("# padding\n" * 250), + encoding="utf-8", + ) + + loaded = load_module_files(tests_dir, tmp_path) + groups = group_files(loaded, tmp_path, token_budget=100_000) + + assert len(groups) == 2 + assert all(group.name.startswith("tests") for group in groups) + assert all(0 < len(group.files) <= MAX_FILES_PER_GROUP for group in groups) + assert all(len(format_group_context(group)) < 1_048_576 for group in groups) + + +def test_resolve_module_path_handles_underscore_parent_names(tmp_path: Path) -> None: + """Auto-split module ids should resolve when the parent has underscores.""" + target = tmp_path / "docs_src" / "additional_responses" + target.mkdir(parents=True) + (target / "tutorial.py").write_text("def example():\n return None\n", encoding="utf-8") + + resolved = resolve_module_path(tmp_path, "docs_src_additional_responses") + + assert resolved == target diff --git a/engine/tests/test_hub_pipeline.py b/engine/tests/test_hub_pipeline.py index b1f4ce616..14fb5f622 100644 --- a/engine/tests/test_hub_pipeline.py +++ b/engine/tests/test_hub_pipeline.py @@ -261,6 +261,21 @@ def test_load_project_context_respects_total_budget(tmp_path: Path) -> None: assert "REGISTRY_MARKER" in section +def test_ask_retry_classifier_handles_litellm_service_unavailable() -> None: + """LiteLLM wraps provider 503s without always preserving the numeric code.""" + from antigravity_engine.hub.ask_pipeline import _is_retryable_ask_error + + class ServiceUnavailableError(Exception): + pass + + exc = ServiceUnavailableError( + "litellm.ServiceUnavailableError: OpenAIException - " + "Service temporarily unavailable" + ) + + assert _is_retryable_ask_error(exc) + + # --------------------------------------------------------------------------- # Phase 1: config/entry/git in _format_scan_report # --------------------------------------------------------------------------- From 995655c2aaf3d2ab6a17fb062446e6464c0c0b58 Mon Sep 17 00:00:00 2001 From: JingWen Fan <106414602+study8677@users.noreply.github.com> Date: Sun, 24 May 2026 18:25:39 +0800 Subject: [PATCH 2/2] feat(hub): add opt-in multi-provider LLM failover for the ask path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sustained provider outages (e.g. the shared-proxy 503s that invalidated several benchmark runs) cannot be recovered by same-provider retries — every retry hits the same dead host. Add an ordered provider chain and wrap the ask pipeline so a transient/provider error on the active endpoint fails over to the next provider and re-runs the answer. - _providers.py: ProviderConfig, get_provider_chain (parses AG_LLM_FALLBACKS, degrades to primary on bad input), activate_provider, and the run_with_provider_failover wrapper, plus one shared retryable-error classifier (is_retryable_provider_error). - ask_pipeline: split into a failover wrapper + _ask_pipeline_once; _is_retryable_ask_error now delegates to the shared classifier. - Unchanged when AG_LLM_FALLBACKS is unset (chain length 1 => pass-through, no env mutation). - Tests for chain parsing, inheritance, bad-JSON degradation, failover rotation, and non-transient pass-through. - Document AG_ASK_RETRY_* and AG_LLM_FALLBACKS in .env.example. Co-Authored-By: Claude Opus 4.7 (1M context) --- engine/.env.example | 16 ++ engine/antigravity_engine/hub/_providers.py | 191 ++++++++++++++++++ engine/antigravity_engine/hub/ask_pipeline.py | 62 +++--- engine/tests/test_hub_providers.py | 162 +++++++++++++++ 4 files changed, 404 insertions(+), 27 deletions(-) create mode 100644 engine/antigravity_engine/hub/_providers.py create mode 100644 engine/tests/test_hub_providers.py diff --git a/engine/.env.example b/engine/.env.example index 514f149ad..b3f4fdc1c 100644 --- a/engine/.env.example +++ b/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. diff --git a/engine/antigravity_engine/hub/_providers.py b/engine/antigravity_engine/hub/_providers.py new file mode 100644 index 000000000..0bcd08d66 --- /dev/null +++ b/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 diff --git a/engine/antigravity_engine/hub/ask_pipeline.py b/engine/antigravity_engine/hub/ask_pipeline.py index f272b4b0a..665e70157 100644 --- a/engine/antigravity_engine/hub/ask_pipeline.py +++ b/engine/antigravity_engine/hub/ask_pipeline.py @@ -47,33 +47,15 @@ def _get_ask_retry_config() -> tuple[int, float]: def _is_retryable_ask_error(exc: Exception) -> bool: - """Return true for transient model/provider failures.""" - if isinstance(exc, (TimeoutError, asyncio.TimeoutError)): - return True - msg = f"{type(exc).__module__}.{type(exc).__name__}: {exc}".lower() - 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", - ) - return any(keyword in msg for keyword in retryable_keywords) + """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( @@ -197,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) diff --git a/engine/tests/test_hub_providers.py b/engine/tests/test_hub_providers.py new file mode 100644 index 000000000..4c72edad3 --- /dev/null +++ b/engine/tests/test_hub_providers.py @@ -0,0 +1,162 @@ +"""Tests for hub._providers multi-provider LLM failover.""" +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from antigravity_engine.hub import _providers +from antigravity_engine.hub._providers import ( + ProviderConfig, + get_provider_chain, + is_retryable_provider_error, + run_with_provider_failover, +) + + +def _settings() -> SimpleNamespace: + return SimpleNamespace( + OPENAI_MODEL="primary-model", + OPENAI_BASE_URL="https://primary/v1", + OPENAI_API_KEY="primary-key", + ) + + +# --- classifier ----------------------------------------------------------- + + +@pytest.mark.parametrize( + "exc", + [ + Exception("litellm.ServiceUnavailableError: Service temporarily unavailable"), + Exception("HTTP 503 from upstream"), + Exception("Connection refused"), + TimeoutError("deadline exceeded"), + asyncio.TimeoutError(), + ], +) +def test_is_retryable_provider_error_true(exc: Exception) -> None: + assert is_retryable_provider_error(exc) is True + + +@pytest.mark.parametrize( + "exc", + [ValueError("bad question"), KeyError("missing"), RuntimeError("logic bug")], +) +def test_is_retryable_provider_error_false(exc: Exception) -> None: + assert is_retryable_provider_error(exc) is False + + +# --- provider chain parsing ---------------------------------------------- + + +def test_get_provider_chain_without_fallbacks_is_single(monkeypatch) -> None: + monkeypatch.delenv("AG_LLM_FALLBACKS", raising=False) + chain = get_provider_chain(_settings()) + assert len(chain) == 1 + assert chain[0].label == "primary" + assert chain[0].model == "primary-model" + assert chain[0].base_url == "https://primary/v1" + + +def test_get_provider_chain_parses_and_inherits(monkeypatch) -> None: + monkeypatch.setenv( + "AG_LLM_FALLBACKS", + json.dumps( + [ + { + "base_url": "https://backup/v1", + "api_key": "bk", + "model": "gpt-x", + "label": "backup", + }, + {"model": "only-model"}, # inherits api_key, default label + ] + ), + ) + chain = get_provider_chain(_settings()) + assert [p.label for p in chain] == ["primary", "backup", "fallback2"] + assert chain[1] == ProviderConfig( + model="gpt-x", base_url="https://backup/v1", api_key="bk", label="backup" + ) + # Second fallback inherits the primary key and has no base_url. + assert chain[2].api_key == "primary-key" + assert chain[2].base_url == "" + assert chain[2].model == "only-model" + + +def test_get_provider_chain_degrades_on_bad_json(monkeypatch) -> None: + monkeypatch.setenv("AG_LLM_FALLBACKS", "not json {{") + chain = get_provider_chain(_settings()) + assert len(chain) == 1 + + +def test_get_provider_chain_degrades_on_non_array(monkeypatch) -> None: + monkeypatch.setenv("AG_LLM_FALLBACKS", json.dumps({"model": "x"})) + chain = get_provider_chain(_settings()) + assert len(chain) == 1 + + +# --- failover wrapper ----------------------------------------------------- + + +def test_failover_single_provider_is_passthrough(monkeypatch) -> None: + """With one provider the operation runs once and env is left untouched.""" + activated: list[str] = [] + monkeypatch.setattr( + _providers, "activate_provider", lambda p: activated.append(p.label) + ) + + async def op() -> str: + return "ok" + + result = asyncio.run( + run_with_provider_failover( + op, providers=[ProviderConfig(model="m", label="primary")], label="ask" + ) + ) + assert result == "ok" + assert activated == [] # never touched the environment + + +def test_failover_switches_provider_on_transient_error(monkeypatch) -> None: + activated: list[str] = [] + monkeypatch.setattr( + _providers, "activate_provider", lambda p: activated.append(p.label) + ) + calls = {"n": 0} + + async def op() -> str: + calls["n"] += 1 + if calls["n"] == 1: + raise Exception("OpenAIException - Service temporarily unavailable") + return "answered-on-backup" + + providers = [ + ProviderConfig(model="m1", label="primary"), + ProviderConfig(model="m2", label="backup"), + ] + result = asyncio.run( + run_with_provider_failover(op, providers=providers, label="ask") + ) + assert result == "answered-on-backup" + assert activated == ["primary", "backup"] + assert calls["n"] == 2 + + +def test_failover_does_not_retry_non_transient_error(monkeypatch) -> None: + activated: list[str] = [] + monkeypatch.setattr( + _providers, "activate_provider", lambda p: activated.append(p.label) + ) + + async def op() -> str: + raise ValueError("genuine logic error") + + providers = [ + ProviderConfig(model="m1", label="primary"), + ProviderConfig(model="m2", label="backup"), + ] + with pytest.raises(ValueError): + asyncio.run(run_with_provider_failover(op, providers=providers, label="ask")) + assert activated == ["primary"] # no failover on a non-transient error