fix(oidc): answer 503 on token-count routes during a metadata outage and fall back on a bad GCE_METADATA_TIMEOUT

litellm/proxy/proxy_server.py

@@ -14741,6 +14741,13 @@ async def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]
             param="model",
             code=401,
         ) from e
+    except litellm.ServiceUnavailableError as e:
+        raise ProxyException(
+            message=e.message,
+            type="service_unavailable",
+            param="model",
+            code=503,
+        ) from e
 
 
 async def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None:

litellm/secret_managers/main.py

@@ -14,6 +14,7 @@ from pydantic import TypeAdapter, ValidationError
 import litellm
 from litellm._logging import verbose_logger
 from litellm.caching.caching import DualCache
+from litellm.litellm_core_utils.env_utils import get_env_float
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
 from litellm.secret_managers.get_azure_ad_token_provider import (
     get_azure_ad_token_provider,
@@ -170,7 +171,7 @@ LOCAL_OIDC_SECRET_PREFIXES: Final = ("oidc/file/", "oidc/env/")
 GOOGLE_OIDC_SECRET_PREFIX: Final = "oidc/google/"
 API_KEY_OIDC_SECRET_PREFIXES: Final = (*LOCAL_OIDC_SECRET_PREFIXES, GOOGLE_OIDC_SECRET_PREFIX)
 _GOOGLE_METADATA_DEFAULT_HOST: Final = "metadata.google.internal"
-_GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS: Final = "3"
+_GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS: Final = 3.0
 _GOOGLE_TOKEN_UNAVAILABLE: Final = "the metadata server did not return an ID token"
 
 
@@ -180,7 +181,8 @@ def _google_metadata_identity_url() -> str:
 
 
 def _google_metadata_timeout() -> float:
-    return float(os.getenv("GCE_METADATA_TIMEOUT") or _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS)
+    seconds: Final = get_env_float("GCE_METADATA_TIMEOUT", _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS)
+    return seconds if seconds > 0 else _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS
 
 
 class GoogleMetadataStatusError(ValueError):

tests/integration/security/test_api_key_secret_reference.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import base64
 import json
+import threading
 import time
 import uuid
 from collections.abc import Callable, Iterator, Mapping, Sequence
@@ -58,6 +59,7 @@ _GOOGLE_AUDIENCES: Final = MappingProxyType(
     }
 )
 _UNAVAILABLE_AUDIENCE: Final = "https://vllm-unavailable.example.test"
+_HUNG_AUDIENCE: Final = "https://vllm-hung.example.test"
 _IDENTITY_PATH: Final = "/computeMetadata/v1/instance/service-accounts/default/identity"
 _RUN_SECRETS_TOKEN: Final = "/run/secrets/wif-absent-token"
 _SOURCES: Final = ("config_credential", "config_inline", "db_credential", "db_inline")
@@ -152,6 +154,8 @@ def _metadata_server(google: _Google) -> Callable[[Request], Reply]:
             return Reply(status=404, body=b"not found", content_type="text/plain")
         if audience == _UNAVAILABLE_AUDIENCE:
             return Reply(status=503, body=b"metadata server unavailable", content_type="text/plain")
+        if audience == _HUNG_AUDIENCE:
+            return Reply(chunks=(b"never sent",), gate_before_first=threading.Event(), content_type="text/plain")
         return Reply(body=google.token(audience).encode(), content_type="text/plain")
 
     return respond
@@ -286,6 +290,14 @@ def _config(directory: Path, secrets: _Secrets, provider: str) -> Path:
                 "wif-google-unavailable",
                 {"model": _GOOGLE_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider},
             ),
+            _deployment(
+                "wif-google-unavailable-count",
+                {"model": _CHAT_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider},
+            ),
+            _deployment(
+                "wif-google-hung",
+                {"model": _GOOGLE_MODEL, "api_key": _google_reference(_HUNG_AUDIENCE), "api_base": provider},
+            ),
             _deployment(
                 "wif-mistral-passthrough",
                 {
@@ -420,6 +432,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
                 **secrets.overrides(),
                 "MISTRAL_API_BASE": provider_wire.url,
                 "GCE_METADATA_HOST": httpx.URL(metadata_wire.url).netloc.decode(),
+                "GCE_METADATA_TIMEOUT": "1",
             },
             config=_config(directory, secrets, provider_base),
             workers=2,
@@ -847,6 +860,28 @@ def test_google_reference_fails_closed_when_the_metadata_server_fails(rig: _Rig)
     assert _UNAVAILABLE_AUDIENCE in _asked_audiences(rig)
 
 
+def test_google_reference_answers_503_when_the_metadata_server_never_finishes_its_reply(rig: _Rig) -> None:
+    marker: Final = _marker()
+    response: Final = _chat(rig.owned.gateway, "wif-google-hung", marker)
+    reached: Final = _carrying(rig.provider.requests(), marker)
+    assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
+    assert "the metadata server did not return an ID token" in response.text, response.text
+    assert _HUNG_AUDIENCE in _asked_audiences(rig)
+
+
+def test_input_token_count_answers_503_when_the_metadata_server_fails(rig: _Rig) -> None:
+    marker: Final = _marker()
+    response: Final = rig.owned.gateway.request(
+        "POST",
+        "/v1/responses/input_tokens",
+        {"model": "wif-google-unavailable-count", "input": f"say {marker}"},
+        headers=_CLOSE,
+    )
+    reached: Final = _carrying(rig.provider.requests(), marker)
+    assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
+    assert "oidc/google secret reference could not be resolved" in response.text, response.text
+
+
 def _side_call(rig: _Rig, seam: str, marker: str) -> tuple[httpx.Response, Callable[[Request], bool]]:
     gateway: Final = rig.owned.gateway
     match seam:

tests/unit/proxy/test_proxy_token_counter.py

@@ -279,6 +279,39 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid
     assert str(tmp_path) not in excinfo.value.message
 
 
+@pytest.mark.asyncio
+async def test_token_counter_answers_503_when_the_metadata_server_behind_an_oidc_google_api_key_fails(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    audience: Final = f"https://count.test/{uuid.uuid4().hex}"
+    monkeypatch.setattr(
+        litellm.proxy.proxy_server,
+        "llm_router",
+        Router(
+            model_list=[
+                {
+                    "model_name": "claude",
+                    "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": f"oidc/google/{audience}"},
+                }
+            ]
+        ),
+    )
+
+    with respx.mock(assert_all_called=True) as upstream:
+        upstream.get(url__regex=r"/computeMetadata/v1/instance/service-accounts/default/identity").respond(
+            500, text="metadata server unavailable"
+        )
+        with pytest.raises(ProxyException) as excinfo:
+            await token_counter(
+                request=TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}]),
+                call_endpoint=True,
+            )
+
+    assert excinfo.value.code == "503"
+    assert "oidc/google secret reference could not be resolved" in excinfo.value.message
+    assert audience.removeprefix("https://") not in excinfo.value.message
+
+
 @pytest.mark.asyncio
 async def test_native_openai_responses_count_sends_the_token_an_oidc_api_key_points_to(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch

tests/unit/test_utils.py

@@ -10,7 +10,6 @@ import json
 import logging
 import os
 import queue
-import socket
 import threading
 import time
 import uuid
@@ -7576,32 +7575,68 @@ async def test_an_oidc_google_api_key_token_cached_when_the_call_starts_is_sent_
 
 
 @pytest.mark.asyncio
-async def test_a_hung_metadata_server_fails_an_oidc_google_api_key_call_with_a_503_and_leaves_the_default_executor_free(
+async def test_a_pending_oidc_google_metadata_fetch_leaves_the_default_executor_free(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
     _, reference = _google_reference()
     monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)])
-    monkeypatch.setenv("GCE_METADATA_TIMEOUT", "1")
     loop: Final = asyncio.get_running_loop()
     loop.set_default_executor(ThreadPoolExecutor(max_workers=1))
+    fetch_started: Final[queue.SimpleQueue[int]] = queue.SimpleQueue()
+    release: Final = threading.Event()
 
-    with socket.create_server(("127.0.0.1", 0)) as never_answering, respx.mock(assert_all_called=False) as upstream:
-        monkeypatch.setenv("GCE_METADATA_HOST", f"127.0.0.1:{never_answering.getsockname()[1]}")
-        upstream.route(host="127.0.0.1").pass_through()
+    def pending_metadata(_request: httpx.Request) -> httpx.Response:
+        fetch_started.put(threading.get_ident())
+        release.wait(timeout=30)
+        return httpx.Response(500, text="metadata server unavailable")
+
+    with ThreadPoolExecutor(max_workers=1) as observer, respx.mock(assert_all_called=False) as upstream:
+        upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=pending_metadata)
         chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY)
         call: Final = asyncio.ensure_future(_ping_vllm_through("async", "vllm-sa"))
-        await asyncio.sleep(0.2)
-        other_work: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "done"), timeout=0.5)
+        try:
+            await asyncio.wait_for(loop.run_in_executor(observer, fetch_started.get), timeout=30)
+            other_work: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "done"), timeout=30)
+        finally:
+            release.set()
         with pytest.raises(litellm.ServiceUnavailableError) as excinfo:
-            await asyncio.wait_for(call, timeout=10)
+            await call
 
     assert other_work == "done"
     assert excinfo.value.status_code == 503
-    assert "the metadata server did not return an ID token" in excinfo.value.message
     assert chat.call_count == 0
 
 
+@pytest.mark.parametrize(
+    ("configured", "bound"),
+    [(None, 3.0), ("1.5", 1.5), ("three seconds", 3.0), ("0", 3.0), ("-2", 3.0)],
+    ids=["unset", "configured", "malformed", "zero", "negative"],
+)
+def test_oidc_google_api_key_metadata_fetch_waits_at_most_gce_metadata_timeout_seconds(
+    monkeypatch: pytest.MonkeyPatch, configured: str | None, bound: float
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    if configured is None:
+        monkeypatch.delenv("GCE_METADATA_TIMEOUT", raising=False)
+    else:
+        monkeypatch.setenv("GCE_METADATA_TIMEOUT", configured)
+    monkeypatch.setattr(time, "time", lambda: _CLOCK_START)
+    audience, reference = _google_reference()
+    monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)])
+    token: Final = _google_id_token(audience, _CLOCK_START)
+
+    with respx.mock(assert_all_called=True) as upstream:
+        metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond(200, text=token)
+        chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY)
+        litellm.completion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name="vllm-sa")
+        sent: Final = chat.calls.last.request.headers["Authorization"]
+        waits: Final = metadata.calls.last.request.extensions["timeout"]
+
+    assert sent == f"Bearer {token}"
+    assert waits == {"connect": bound, "read": bound, "write": bound, "pool": bound}
+
+
 def test_an_oidc_google_api_key_sent_with_the_call_goes_upstream_as_written_without_a_metadata_fetch() -> None:
     _, reference = _google_reference()