diff --git a/.github/scripts/desktop_backend_candidate_probe.py b/.github/scripts/desktop_backend_candidate_probe.py index d9d2053d86b..2c1d4df77d9 100644 --- a/.github/scripts/desktop_backend_candidate_probe.py +++ b/.github/scripts/desktop_backend_candidate_probe.py @@ -36,7 +36,13 @@ MAX_FIRST_EVENT_SECONDS = 20 SHA_PATTERN = re.compile(r"^[0-9a-f]{40}$") CONTRACT_PATTERN = re.compile(r"^[1-9][0-9]{0,5}$") -REAL_GEMINI_PROVIDER_ROUTES = frozenset({"vertex_ai", "ai_studio", "ai_studio_byok"}) +REAL_GEMINI_PROVIDER_ROUTES = frozenset({"vertex_ai", "ai_studio", "ai_studio_byok", "llm_gateway"}) +# `llm_gateway` is the route the desktop proxy stamps on X-Omi-Provider for +# company-paid Gemini traffic since #12337 routed it through the LLM gateway +# (backend/utils/llm/desktop_gemini_gateway.py proxy_company_paid_via_gateway). +# The gateway's desktop-vertex lanes pin the Vertex provider with no fallbacks +# (backend/llm_gateway/gateway/config_loader.py), so the hop is still real +# Gemini-on-Vertex; stub and unknown routes stay rejected fail-closed. class ProbeError(RuntimeError): diff --git a/.github/scripts/test_desktop_backend_candidate_probe.py b/.github/scripts/test_desktop_backend_candidate_probe.py index 7d519d1861f..f79055803a1 100644 --- a/.github/scripts/test_desktop_backend_candidate_probe.py +++ b/.github/scripts/test_desktop_backend_candidate_probe.py @@ -74,10 +74,33 @@ def read(self) -> bytes: PROBE._gemini_request("https://candidate.example", token="token") def test_gemini_probe_rejects_stub_or_unknown_provider_routes(self) -> None: - for provider in ("offline_stub", "unknown", ""): + for provider in ("offline_stub", "desktop_llm_stub", "unknown", ""): with self.assertRaisesRegex(PROBE.ProbeError, "admitted provider"): PROBE._require_real_gemini_provider(provider) - self.assertEqual(PROBE._require_real_gemini_provider("vertex_ai"), "vertex_ai") + for admitted in ("vertex_ai", "ai_studio", "ai_studio_byok"): + self.assertEqual(PROBE._require_real_gemini_provider(admitted), admitted) + + def test_gemini_probe_admits_post_gateway_llm_gateway_route(self) -> None: + # Since #12337 the desktop proxy serves company-paid Gemini traffic via + # the LLM gateway's Vertex-backed desktop-vertex lanes and stamps + # `llm_gateway` on X-Omi-Provider; the probe must admit that real route. + self.assertEqual(PROBE._require_real_gemini_provider("llm_gateway"), "llm_gateway") + + class GatewayResponse: + headers = {"x-omi-provider": "llm_gateway", "x-omi-request-id": "server-request-id"} + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def read(self) -> bytes: + return b'{"candidates":[{"content":{"parts":[{"text":"OK"}]}}]}' + + with mock.patch.object(PROBE.urllib.request, "urlopen", return_value=GatewayResponse()): + summary = PROBE._gemini_request("https://candidate.example", token="token") + self.assertEqual(summary["provider_route"], "llm_gateway") def test_gemini_request_cannot_pass_against_offline_stub(self) -> None: class StubResponse: