feat(providers): discover OAuth model catalogs online

2026-09-04 02:01:48 +03:00 · 2026-08-29 21:22:20 +08:00
parent 2389ab1f5a
commit bc4de246a4
12 changed files with 920 additions and 283 deletions
@@ -0,0 +1,257 @@
+from __future__ import annotations
+
+import base64
+import json
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import httpx
+import pytest
+
+from nanobot.providers.oauth_model_catalog import (
+    DEFAULT_XAI_GROK_MODELS_URL,
+    OAuthModelCatalog,
+    OAuthModelInfo,
+    get_oauth_model_catalog,
+    invalidate_oauth_model_catalog,
+)
+from nanobot.providers.xai_oauth import XAIToken
+
+
+@pytest.fixture(autouse=True)
+def _clear_xai_catalog() -> None:
+    invalidate_oauth_model_catalog("xai_grok")
+    yield
+    invalidate_oauth_model_catalog("xai_grok")
+
+
+def _fallback_model() -> OAuthModelInfo:
+    return OAuthModelInfo(id="provider/fallback", label="Fallback")
+
+
+def test_xai_catalog_fetches_remote_models_and_reuses_capability_metadata(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    original_client = httpx.Client
+    captured: dict[str, object] = {}
+    payload = base64.urlsafe_b64encode(
+        json.dumps({"sub": "user-42", "email": "user@example.com"}).encode()
+    ).decode().rstrip("=")
+    token = XAIToken(
+        access=f"header.{payload}.signature",
+        refresh="refresh-token",
+        expires=int(time.time() * 1000) + 3_600_000,
+        account_id="user@example.com",
+    )
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        captured["request"] = request
+        return httpx.Response(
+            200,
+            json={
+                "data": [
+                    {
+                        "id": "grok-4.6",
+                        "name": "Grok 4.6",
+                        "description": "Latest frontier model",
+                        "owned_by": "xAI",
+                        "context_window": 500_000,
+                        "supports_backend_search": True,
+                        "reasoning_efforts": [
+                            {"value": "xhigh"},
+                            {"value": "high"},
+                            {"value": "low"},
+                        ],
+                    },
+                    {
+                        "id": "grok-next",
+                        "_meta": {
+                            "name": "Grok Next",
+                            "context_window": 750_000,
+                            "reasoning_efforts": ["high", "low"],
+                        },
+                    },
+                ]
+            },
+            request=request,
+        )
+
+    def fake_client(**kwargs: object) -> httpx.Client:
+        captured["kwargs"] = kwargs
+        return original_client(
+            transport=httpx.MockTransport(handler),
+            timeout=kwargs["timeout"],
+            follow_redirects=kwargs["follow_redirects"],
+        )
+
+    monkeypatch.setattr(
+        "nanobot.providers.oauth_model_catalog._xai_oauth_storage_path",
+        lambda: tmp_path / "auth" / "xai.json",
+    )
+    monkeypatch.setattr(
+        "nanobot.providers.oauth_model_catalog._xai_oauth_token",
+        lambda _proxy: token,
+    )
+    monkeypatch.setattr("nanobot.providers.oauth_model_catalog.httpx.Client", fake_client)
+
+    catalog = get_oauth_model_catalog("xai_grok")
+
+    assert catalog.source == "remote"
+    assert [model.id for model in catalog.models] == [
+        "xai-grok/grok-4.6",
+        "xai-grok/grok-next",
+    ]
+    grok = catalog.find("grok-4.6")
+    assert grok is not None
+    assert grok.description == "Latest frontier model"
+    assert grok.context_window == 500_000
+    assert grok.reasoning_efforts == ("xhigh", "high", "low")
+    assert grok.supports_backend_search is True
+    next_model = catalog.find("xai-grok/grok-next")
+    assert next_model is not None
+    assert next_model.label == "Grok Next"
+    assert next_model.context_window == 750_000
+    assert next_model.reasoning_efforts == ("high", "low")
+
+    request = captured["request"]
+    assert isinstance(request, httpx.Request)
+    assert str(request.url) == DEFAULT_XAI_GROK_MODELS_URL
+    assert request.headers["Authorization"] == f"Bearer {token.access}"
+    assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
+    assert request.headers["x-userid"] == "user-42"
+    assert request.headers["x-email"] == "user@example.com"
+    assert captured["kwargs"] == {"timeout": 10.0, "follow_redirects": False}
+    assert get_oauth_model_catalog("xai_grok").source == "cache"
+
+
+def test_catalog_single_flights_concurrent_refreshes() -> None:
+    calls = 0
+    calls_lock = threading.Lock()
+    barrier = threading.Barrier(8)
+
+    def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
+        nonlocal calls
+        with calls_lock:
+            calls += 1
+        time.sleep(0.05)
+        return (OAuthModelInfo(id="provider/remote", label="Remote"),)
+
+    catalog = OAuthModelCatalog(fallback_models=(_fallback_model(),), fetch=fetch)
+
+    def get_catalog(_index: int):
+        barrier.wait()
+        return catalog.get(cache_key="shared")
+
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        results = list(pool.map(get_catalog, range(8)))
+
+    assert calls == 1
+    assert {result.models[0].id for result in results} == {"provider/remote"}
+    assert [result.source for result in results].count("remote") == 1
+    assert [result.source for result in results].count("cache") == 7
+
+
+def test_catalog_returns_stale_then_negative_caches_refresh_failure() -> None:
+    now = [0.0]
+    calls = 0
+
+    def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
+        nonlocal calls
+        calls += 1
+        if calls > 1:
+            raise httpx.ConnectError("offline")
+        return (OAuthModelInfo(id="provider/remote", label="Remote"),)
+
+    catalog = OAuthModelCatalog(
+        fallback_models=(_fallback_model(),),
+        fetch=fetch,
+        fresh_ttl_s=10,
+        stale_ttl_s=100,
+        failure_ttl_s=30,
+        monotonic=lambda: now[0],
+        wall_clock=lambda: 123.0,
+    )
+
+    assert catalog.get(cache_key="one").source == "remote"
+    now[0] = 11
+    stale = catalog.get(cache_key="one")
+    assert stale.source == "stale"
+    assert stale.models[0].id == "provider/remote"
+    assert catalog.get(cache_key="one").source == "stale"
+    assert calls == 2
+
+    now[0] = 101
+    fallback = catalog.get(cache_key="one")
+    assert fallback.source == "fallback"
+    assert fallback.models[0].id == "provider/fallback"
+    assert calls == 3
+
+
+@pytest.mark.parametrize(
+    "failure",
+    [
+        httpx.ConnectError("offline"),
+        ValueError("invalid JSON"),
+        httpx.HTTPStatusError(
+            "unauthorized",
+            request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
+            response=httpx.Response(401),
+        ),
+        httpx.HTTPStatusError(
+            "rate limited",
+            request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
+            response=httpx.Response(429),
+        ),
+        httpx.HTTPStatusError(
+            "upstream failure",
+            request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
+            response=httpx.Response(503),
+        ),
+    ],
+)
+def test_catalog_falls_back_for_remote_failures(failure: Exception) -> None:
+    calls = 0
+
+    def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
+        nonlocal calls
+        calls += 1
+        raise failure
+
+    catalog = OAuthModelCatalog(
+        fallback_models=(_fallback_model(),),
+        fetch=fetch,
+        failure_ttl_s=30,
+    )
+
+    first = catalog.get(cache_key="one")
+    second = catalog.get(cache_key="one")
+
+    assert first.source == "fallback"
+    assert second.source == "fallback"
+    assert first.models == (_fallback_model(),)
+    assert calls == 1
+
+
+def test_catalog_treats_empty_remote_list_as_failure_and_can_be_invalidated() -> None:
+    calls = 0
+
+    def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
+        nonlocal calls
+        calls += 1
+        return () if calls == 1 else (OAuthModelInfo(id="provider/new", label="New"),)
+
+    catalog = OAuthModelCatalog(
+        fallback_models=(_fallback_model(),),
+        fetch=fetch,
+        failure_ttl_s=30,
+    )
+
+    assert catalog.get(cache_key="one").source == "fallback"
+    catalog.invalidate()
+    refreshed = catalog.get(cache_key="one")
+    assert refreshed.source == "remote"
+    assert refreshed.models[0].id == "provider/new"
+    assert calls == 2