mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
fix(anthropic): mint the files and batches credential without blocking the event loop
The async file and batch handlers called the provider's sync validate_environment directly. For the workload identity tier that hook performs a blocking token exchange, so a mint stalled the whole loop. They now go through one facade that awaits the provider's async hook when it has one and offloads the sync hook to a worker thread otherwise, so every other surface keeps its existing behaviour Batch retrieval also copies the litellm_params it hands down rather than sharing the caller's dict, and get_anthropic_headers moves its credential selection into a helper so the four auth tiers read as four branches instead of one nested chain
This commit is contained in:
parent
2ee90c258c
commit
35626611cf
10 changed files with 515 additions and 36 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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="",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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 == {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue