diff --git a/litellm/batches/main.py b/litellm/batches/main.py index a84aa4832d7..0a73a4b8afc 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -483,6 +483,7 @@ def _handle_retrieve_batch_providers_without_provider_config( ) api_key = optional_params.api_key or litellm.api_key or litellm.azure_key or get_secret_str("ANTHROPIC_API_KEY") + batch_params: Final = dict(litellm_params) # mutable-ok: handler contract, copied not shared response = anthropic_batches_instance.retrieve_batch( _is_async=_is_async, batch_id=batch_id, @@ -490,7 +491,7 @@ def _handle_retrieve_batch_providers_without_provider_config( api_key=api_key, timeout=timeout, max_retries=optional_params.max_retries, - litellm_params=dict(litellm_params), + litellm_params=batch_params, ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/llms/anthropic/batches/handler.py b/litellm/llms/anthropic/batches/handler.py index da3079769ac..3f3716e0de6 100644 --- a/litellm/llms/anthropic/batches/handler.py +++ b/litellm/llms/anthropic/batches/handler.py @@ -42,7 +42,7 @@ class AnthropicBatchesHandler: timeout: float | httpx.Timeout, max_retries: int | None, logging_obj: LiteLLMLoggingObj | None = None, - litellm_params: dict | None = None, + litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment ) -> LiteLLMBatch: """ Async: Retrieve a batch from Anthropic. @@ -87,8 +87,10 @@ class AnthropicBatchesHandler: litellm_params=resolved_litellm_params, ) - # Validate environment and get headers - headers: Final = self.provider_config.validate_environment( + # Validate environment and get headers. Offloaded to a worker thread: a WIF token + # exchange here would otherwise block the event loop. + headers: Final = await asyncio.to_thread( + self.provider_config.validate_environment, headers={}, model="", messages=[], @@ -129,7 +131,7 @@ class AnthropicBatchesHandler: timeout: float | httpx.Timeout, max_retries: int | None, logging_obj: LiteLLMLoggingObj | None = None, - litellm_params: dict | None = None, + litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment ) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: """ Retrieve a batch from Anthropic. diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 1f0d7a9355d..f122a75bbcc 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -687,6 +687,33 @@ class AnthropicModelInfo(BaseLLMModelInfo): return {"authorization": value} return {"x-api-key": api_key} + def _credential_headers( + self, + *, + api_key: str | None, + auth_token: str | None, + api_base: str | None, + use_bearer_for_custom_base: bool, + wif_minted: bool, + betas: set[str], # mutable-ok: the caller's beta accumulator, appended to by the oauth tier + ) -> Mapping[str, str]: + """The credential tier walk: a consumer OAuth token, then ANTHROPIC_AUTH_TOKEN, then an api key. + + A server-minted federation token takes the same Bearer shape as a consumer OAuth token but is + not browser-forwarded, so it does not get the direct-browser-access header. + """ + if api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX): + betas.add(ANTHROPIC_OAUTH_BETA_HEADER) + oauth_headers: Final = {"authorization": f"Bearer {api_key}"} + if wif_minted: + return oauth_headers + return {**oauth_headers, "anthropic-dangerous-direct-browser-access": "true"} + if auth_token and not api_key: + return {"authorization": f"Bearer {auth_token}"} + if api_key: + return self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base) + return {} + def get_anthropic_headers( self, api_key: str | None = None, @@ -744,21 +771,21 @@ class AnthropicModelInfo(BaseLLMModelInfo): if container_with_skills_used: betas.add("skills-2025-10-02") - _is_oauth: Final = api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) headers: Final = { "anthropic-version": anthropic_version or "2023-06-01", "accept": "application/json", "content-type": "application/json", } - if _is_oauth: - headers["authorization"] = f"Bearer {api_key}" - if not wif_minted: - headers["anthropic-dangerous-direct-browser-access"] = "true" - betas.add(ANTHROPIC_OAUTH_BETA_HEADER) - elif auth_token and not api_key: - headers["authorization"] = f"Bearer {auth_token}" - elif api_key: - headers.update(self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base)) + headers.update( + self._credential_headers( + api_key=api_key, + auth_token=auth_token, + api_base=api_base, + use_bearer_for_custom_base=use_bearer_for_custom_base, + wif_minted=wif_minted, + betas=betas, + ) + ) if user_anthropic_beta_headers is not None: betas.update(user_anthropic_beta_headers) diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index 8ee58d6b5a3..f8ebccb7215 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -14,6 +14,7 @@ Anthropic Files API endpoints: import calendar import time +from collections.abc import Mapping from typing import Any, Final, cast import httpx @@ -94,12 +95,41 @@ class AnthropicFilesConfig(BaseFilesConfig): api_key: str | None = None, api_base: str | None = None, ) -> dict: + params_mapping, resolved_api_base = self._resolve_params(litellm_params, api_base) + auth_header: Final = AnthropicModelInfo.get_auth_header( + api_key, resolved_api_base, litellm_params=params_mapping, allow_workload_identity=True + ) + return self._finalize_headers(headers, auth_header) + + async def avalidate_environment( + self, + headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + model: str, + messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides + optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: # mutable-ok: mirrors the sync validate_environment contract this overrides + """Async counterpart of validate_environment: the WIF tier can block on a token + exchange POST, so async callers await it off the event loop.""" + params_mapping, resolved_api_base = self._resolve_params(litellm_params, api_base) + auth_header: Final = await AnthropicModelInfo.aget_auth_header( + api_key, resolved_api_base, litellm_params=params_mapping, allow_workload_identity=True + ) + return self._finalize_headers(headers, auth_header) + + @staticmethod + def _resolve_params( + litellm_params: dict, api_base: str | None + ) -> tuple[dict | None, str | None]: # mutable-ok: mirrors the sync validate_environment contract this overrides params_mapping: Final = litellm_params if isinstance(litellm_params, dict) else None if api_base is None and params_mapping is not None: api_base = params_mapping.get("api_base") - auth_header: Final = AnthropicModelInfo.get_auth_header( - api_key, api_base, litellm_params=params_mapping, allow_workload_identity=True - ) + return params_mapping, api_base + + @staticmethod + def _finalize_headers(headers: dict, auth_header: Mapping[str, str] | None) -> dict: # mutable-ok: out-param if auth_header is None: raise ValueError( "Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter." diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 0048b76a166..38128d45755 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,4 +1,5 @@ import asyncio +import inspect import json import os import ssl @@ -6,7 +7,20 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequenc from contextlib import asynccontextmanager from functools import lru_cache from types import MappingProxyType, ModuleType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + Optional, + Protocol, + TypedDict, + TypeVar, + Union, + cast, + get_type_hints, + runtime_checkable, +) from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx @@ -202,6 +216,55 @@ class _MediaUploadKwargs(TypedDict, total=False): timeout: float | httpx.Timeout +@runtime_checkable +class _AsyncFilesEnvironmentValidator(Protocol): + async def avalidate_environment( + self, + headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + model: str, + messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides + optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: ... # mutable-ok: mirrors the sync validate_environment contract this overrides + + +async def _avalidate_files_environment( + provider_config: BaseFilesConfig, + *, + headers: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + model: str, + messages: list, # mutable-ok: mirrors the sync validate_environment contract this overrides + optional_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + litellm_params: dict, # mutable-ok: mirrors the sync validate_environment contract this overrides + api_key: str | None, +) -> dict: # mutable-ok: mirrors the sync validate_environment contract this overrides + """Await the provider's async credential hook when it has one (e.g. Anthropic's workload + identity token exchange); otherwise offload the sync hook to a worker thread. Either way + the caller, an async file handler, never blocks the event loop on it.""" + if isinstance(provider_config, _AsyncFilesEnvironmentValidator) and inspect.iscoroutinefunction( + provider_config.avalidate_environment + ): + return await provider_config.avalidate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + ) + return await asyncio.to_thread( + provider_config.validate_environment, + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + ) + + def _google_genai_streaming_hidden_params( *, api_base: str, @@ -4648,7 +4711,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", @@ -4772,7 +4836,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", @@ -4896,7 +4961,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", @@ -5027,7 +5093,8 @@ class BaseLLMHTTPHandler: ) # Validate environment and get headers - headers = provider_config.validate_environment( + headers = await _avalidate_files_environment( + provider_config, api_key=litellm_params.get("api_key"), headers=headers, model="", diff --git a/tests/test_litellm/llms/anthropic/batches/test_handler.py b/tests/test_litellm/llms/anthropic/batches/test_handler.py index 6fde6350127..398e0aeea0f 100644 --- a/tests/test_litellm/llms/anthropic/batches/test_handler.py +++ b/tests/test_litellm/llms/anthropic/batches/test_handler.py @@ -14,6 +14,8 @@ asyncio.run) is exercised directly, mirroring the dispatch-contract discipline i tests/test_litellm/batches/test_main.py. """ +import asyncio +import threading from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -281,3 +283,103 @@ def test_retrieve_batch_sync_runs_to_result(handler, patched_client): assert isinstance(batch, LiteLLMBatch) assert batch.id == "msgbatch_abc" assert batch.status == "completed" + + +# =========================================================================== # +# aretrieve_batch must not block the event loop on a WIF token exchange +# =========================================================================== # + +_WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_batches_seam", + "ANTHROPIC_ORGANIZATION_ID": "org-batches-seam", + "ANTHROPIC_IDENTITY_TOKEN": "batches-seam-inline-jwt", +} + + +class _BlockingPoster: + """A token-endpoint poster that blocks until released, so the test can prove + the exchange ran off the event loop's own thread instead of freezing it.""" + + def __init__(self): + self.release = threading.Event() + self.thread_ids = [] + + def post(self, url, *, content, headers, timeout): + self.thread_ids.append(threading.get_ident()) + self.release.wait(timeout=5) + return httpx.Response( + 200, + json={ + "access_token": "sk-ant-oat01-batches-seam", + "token_type": "Bearer", + "expires_in": 3600, + }, + ) + + +@pytest.mark.asyncio +async def test_aretrieve_batch_wif_exchange_does_not_block_event_loop( + handler, patched_client, monkeypatch +): + """Regression: aretrieve_batch called the synchronous validate_environment + directly, so a cold WIF mint ran inline on the event loop and froze every + other concurrent coroutine until the exchange finished.""" + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + fake_client, _ = patched_client + for name in ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_API_BASE", + "ANTHROPIC_BASE_URL", + ): + monkeypatch.delenv(name, raising=False) + for name, value in _WIF_ENV.items(): + monkeypatch.setenv(name, value) + + poster = _BlockingPoster() + engine = JwtBearerTokenExchangeEngine(poster=poster) + + def routed_through_injected_engine(litellm_params, api_base, model): + return get_anthropic_wif_token(litellm_params, api_base, model, engine) + + monkeypatch.setattr( + anthropic_common_utils, "get_anthropic_wif_token", routed_through_injected_engine + ) + + ticks = [] + + async def ticker(): + for i in range(20): + await asyncio.sleep(0.005) + ticks.append(i) + + ticker_task = asyncio.create_task(ticker()) + await asyncio.sleep(0.02) + + retrieve_task = asyncio.create_task( + handler.aretrieve_batch( + batch_id="msgbatch_abc", + api_base="https://api.anthropic.com", + api_key=None, + timeout=60.0, + max_retries=0, + ) + ) + await asyncio.sleep(0.05) + # The ticker kept advancing while the token exchange was still blocked on + # poster.release, proving the exchange did not run on the event loop. + assert len(ticks) > 0 + assert not retrieve_task.done() + + poster.release.set() + batch = await retrieve_task + await ticker_task + + assert batch.id == "msgbatch_abc" + assert poster.thread_ids + assert poster.thread_ids[0] != threading.get_ident() + sent_headers = fake_client.get.call_args.kwargs["headers"] + assert sent_headers["authorization"] == "Bearer sk-ant-oat01-batches-seam" diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py index d9763f173a7..01aa4799c8d 100644 --- a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py +++ b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py @@ -5,7 +5,9 @@ Tests the AnthropicFilesConfig class which transforms between OpenAI-compatible file operations and Anthropic's Files API format. """ +import asyncio import io +import threading import time import httpx @@ -90,6 +92,38 @@ class TestAnthropicFilesConfig: api_key=None, ) + @pytest.mark.asyncio + async def test_avalidate_environment_sets_headers(self): + headers = {} + result = await self.config.avalidate_environment( + headers=headers, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key="sk-ant-test-key", + ) + assert result["x-api-key"] == "sk-ant-test-key" + assert result["anthropic-version"] == "2023-06-01" + assert result["anthropic-beta"] == ANTHROPIC_FILES_BETA_HEADER + + @pytest.mark.asyncio + @patch.dict("os.environ", {}, clear=True) + @patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=None, + ) + async def test_avalidate_environment_missing_api_key(self, mock_get_key): + with pytest.raises(ValueError, match="Anthropic API key is required"): + await self.config.avalidate_environment( + headers={}, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + def test_get_supported_openai_params(self): params = self.config.get_supported_openai_params(model="") assert "purpose" in params @@ -411,6 +445,108 @@ class TestAnthropicFilesConfig: assert error.message == "Not found" +_WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_files_seam", + "ANTHROPIC_ORGANIZATION_ID": "org-files-seam", + "ANTHROPIC_IDENTITY_TOKEN": "files-seam-inline-jwt", +} + + +class _BlockingPoster: + """A token-endpoint poster that blocks until released, so the test can prove + the exchange ran off the event loop's own thread instead of freezing it.""" + + def __init__(self): + self.release = threading.Event() + self.thread_ids = [] + + def post(self, url, *, content, headers, timeout): + self.thread_ids.append(threading.get_ident()) + self.release.wait(timeout=5) + return httpx.Response( + 200, + json={ + "access_token": "sk-ant-oat01-files-seam", + "token_type": "Bearer", + "expires_in": 3600, + }, + ) + + +class TestAnthropicFilesConfigWifAsyncSeam: + """Regression (Greptile P1): avalidate_environment must resolve workload identity + federation through the async token-exchange facade, never the blocking sync one, + so a cold WIF mint on async file retrieval doesn't freeze the event loop.""" + + def setup_method(self): + self.config = AnthropicFilesConfig() + + @pytest.mark.asyncio + async def test_avalidate_environment_wif_exchange_does_not_block_event_loop(self, monkeypatch): + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + for name in ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_API_BASE", + "ANTHROPIC_BASE_URL", + ): + monkeypatch.delenv(name, raising=False) + for name, value in _WIF_ENV.items(): + monkeypatch.setenv(name, value) + + poster = _BlockingPoster() + engine = JwtBearerTokenExchangeEngine(poster=poster) + sync_calls = [] + + def sync_shim(litellm_params, api_base, model): + sync_calls.append(model) + return get_anthropic_wif_token(litellm_params, api_base, model, engine) + + async def async_shim(litellm_params, api_base, model): + return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + + monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) + monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) + + ticks = [] + + async def ticker(): + for i in range(20): + await asyncio.sleep(0.005) + ticks.append(i) + + ticker_task = asyncio.create_task(ticker()) + await asyncio.sleep(0.02) + + validate_task = asyncio.create_task( + self.config.avalidate_environment( + headers={}, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + ) + await asyncio.sleep(0.05) + # The ticker kept advancing while the exchange was still blocked on + # poster.release, proving avalidate_environment did not run it inline. + assert len(ticks) > 0 + assert not validate_task.done() + + poster.release.set() + headers = await validate_task + await ticker_task + + assert headers["authorization"] == "Bearer sk-ant-oat01-files-seam" + assert sync_calls == [] + assert poster.thread_ids + assert poster.thread_ids[0] != threading.get_ident() + + class TestProviderConfigRegistration: """Test that AnthropicFilesConfig is properly registered.""" diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 24f0dea8c05..d3647c50902 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -19,7 +19,9 @@ from unittest.mock import patch import pytest -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) +) # Fake tokens for testing (not real secrets) FAKE_OAUTH_TOKEN = "sk-ant-oat01-fake-token-for-testing-123456789abcdef" @@ -915,9 +917,7 @@ class TestValidateEnvironmentAuthToken: config = AnthropicModelInfo() with mock_patch.dict("os.environ", {}, clear=True): - with pytest.raises( - Exception, match=r"ANTHROPIC_API_KEY.*ANTHROPIC_AUTH_TOKEN" - ): + with pytest.raises(Exception, match=r"ANTHROPIC_API_KEY.*ANTHROPIC_AUTH_TOKEN"): config.validate_environment( headers={}, model="claude-sonnet-4-5-20250929", @@ -2847,7 +2847,7 @@ class TestWifRespxEndToEnd: "anthropic_federation_rule_id": "fdrl_e2e", "anthropic_organization_id": "org-e2e", "anthropic_identity_token_file": str(token_file), - } + }, ) assert result == { diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py index 37c09302b71..948e30b2c2b 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py @@ -131,7 +131,9 @@ class TestAnthropicFilesHandler: ), ) - with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -186,7 +188,9 @@ class TestAnthropicFilesHandler: ), ) - with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -227,7 +231,9 @@ class TestAnthropicFilesHandler: ), ) - with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -272,7 +278,9 @@ class TestAnthropicFilesHandler: ), ) - with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -316,7 +324,9 @@ class TestAnthropicFilesHandler: ), ) - with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -399,7 +409,9 @@ class TestAnthropicFilesHandler: side_effect=httpx.HTTPStatusError("Not Found", request=mock_response.request, response=mock_response) ) - with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + with patch( + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client @@ -473,7 +485,9 @@ class TestAnthropicFilesHandler: ), ) - with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + with patch( # test-quality-ok: the proxy wiring under test is what this patches + "litellm.llms.anthropic.files.handler.get_async_httpx_client" + ) as mock_get_client: # test-quality-ok: the proxy wiring under test is what this patches mock_client = AsyncMock() mock_client.get = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index b5efac07df4..114543e5617 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import threading import time from unittest.mock import AsyncMock, Mock, patch @@ -1630,7 +1631,7 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks( ), patch.object(handler, "_call_agentic_completion_hooks", side_effect=fake_agentic_hooks), patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client"), - patch( + patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers", return_value=None, ), @@ -1966,6 +1967,105 @@ def test_sync_retrieve_file_content_raises_on_http_error(): assert exc_info.value.status_code == 404 +_FILE_CONTENT_WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_llm_http_handler_seam", + "ANTHROPIC_ORGANIZATION_ID": "org-llm-http-handler-seam", + "ANTHROPIC_IDENTITY_TOKEN": "llm-http-handler-seam-inline-jwt", +} + + +class _BlockingWifPoster: + """A token-endpoint poster that blocks until released, so the test can prove + the exchange ran off the event loop's own thread instead of freezing it.""" + + def __init__(self): + self.release = threading.Event() + self.thread_ids = [] + + def post(self, url, *, content, headers, timeout): + self.thread_ids.append(threading.get_ident()) + self.release.wait(timeout=5) + return httpx.Response( + 200, + json={ + "access_token": "sk-ant-oat01-llm-http-handler-seam", + "token_type": "Bearer", + "expires_in": 3600, + }, + ) + + +@pytest.mark.asyncio +async def test_async_retrieve_file_content_wif_exchange_does_not_block_event_loop(monkeypatch): + """Regression (Greptile P1): async_retrieve_file_content called the synchronous + validate_environment directly, so a cold WIF mint on this call site froze the + event loop until the exchange finished. It must resolve credentials through the + async facade instead.""" + from litellm.llms.anthropic import common_utils as anthropic_common_utils + from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig + from litellm.llms.anthropic.wif import aget_anthropic_wif_token, get_anthropic_wif_token + from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine + + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + for name, value in _FILE_CONTENT_WIF_ENV.items(): + monkeypatch.setenv(name, value) + + poster = _BlockingWifPoster() + engine = JwtBearerTokenExchangeEngine(poster=poster) + sync_calls = [] + + def sync_shim(litellm_params, api_base, model): + sync_calls.append(model) + return get_anthropic_wif_token(litellm_params, api_base, model, engine) + + async def async_shim(litellm_params, api_base, model): + return await aget_anthropic_wif_token(litellm_params, api_base, model, engine) + + monkeypatch.setattr(anthropic_common_utils, "get_anthropic_wif_token", sync_shim) + monkeypatch.setattr(anthropic_common_utils, "aget_anthropic_wif_token", async_shim) + + handler = BaseLLMHTTPHandler() + client = Mock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=httpx.Response(status_code=200, content=b"file bytes")) + + ticks = [] + + async def ticker(): + for i in range(20): + await asyncio.sleep(0.005) + ticks.append(i) + + ticker_task = asyncio.create_task(ticker()) + await asyncio.sleep(0.02) + + retrieve_task = asyncio.create_task( + handler.async_retrieve_file_content( + file_content_request={"file_id": "file-abc"}, + provider_config=AnthropicFilesConfig(), + litellm_params={}, + headers={}, + logging_obj=Mock(), + client=client, + ) + ) + await asyncio.sleep(0.05) + # The ticker kept advancing while the token exchange was still blocked on + # poster.release, proving the exchange did not run inline on the event loop. + assert len(ticks) > 0 + assert not retrieve_task.done() + + poster.release.set() + await retrieve_task + await ticker_task + + assert sync_calls == [] + assert poster.thread_ids + assert poster.thread_ids[0] != threading.get_ident() + sent_headers = client.get.call_args.kwargs["headers"] + assert sent_headers["authorization"] == "Bearer sk-ant-oat01-llm-http-handler-seam" + + _UPSTREAM_NOT_FOUND_BODY = { "error": { "message": "Response with id 'resp_abc' not found.",