From e2d78b69cd8b8e51569c9334eda6dd2aa491b002 Mon Sep 17 00:00:00 2001 From: m6atlas <349022726+m6atlas@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:55:25 -0700 Subject: [PATCH] test(proxy): audit #46815 (#48768) * test(proxy): audit #46815 Add integration tests for every count route #46815 changed, and count a Claude subscription caller's tokens with their forwarded token when the deployment has no configured key, which #46815 turned into a 401 * test(proxy): accept the forwarded subscription token in the count endpoint mocks --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic/count_tokens/token_counter.py | 10 +- .../proxy/anthropic_endpoints/endpoints.py | 6 +- litellm/proxy/litellm_pre_call_utils.py | 39 +++-- litellm/proxy/proxy_server.py | 13 +- .../test_count_tokens_error_shapes_wire.py | 143 +++++++++++++++++- .../llms/anthropic/test_count_tokens_oauth.py | 39 +++++ .../unit/proxy/test_litellm_pre_call_utils.py | 20 +++ tests/unit/proxy/test_proxy_token_counter.py | 5 + 8 files changed, 250 insertions(+), 25 deletions(-) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 0764fb41941..368bb929eeb 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -27,6 +27,9 @@ anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() class AnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Anthropic provider using the CountTokens API.""" + def __init__(self, client_oauth_token: str | None = None) -> None: + self.client_oauth_token: Final = client_oauth_token + def should_use_token_counting_api( self, custom_llm_provider: str | None = None, @@ -95,12 +98,17 @@ class AnthropicTokenCounter(BaseTokenCounter): api_base: Final = litellm_params.get("api_base") try: - auth_header: Final = await AnthropicModelInfo.aget_auth_header( + configured_auth_header: Final = await AnthropicModelInfo.aget_auth_header( api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key, api_base=api_base, litellm_params=litellm_params, allow_workload_identity=True, ) + auth_header: Final = ( + configured_auth_header + if configured_auth_header is not None or self.client_oauth_token is None + else AnthropicModelInfo.get_auth_header(api_key=self.client_oauth_token, api_base=api_base) + ) if auth_header is None: return TokenCountResponse( total_tokens=0, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 13a1844a6de..1d2045eb743 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -50,7 +50,10 @@ from litellm.proxy.common_utils.openai_error_payload import ( with_litellm_call_id, ) from litellm.proxy.common_utils.token_count_settings import count_source_fields -from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli +from litellm.proxy.litellm_pre_call_utils import ( + forwarded_anthropic_oauth_token, + should_auto_drop_params_for_agentic_cli, +) from litellm.router import Router router: Final = APIRouter() @@ -412,6 +415,7 @@ async def count_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", + client_oauth_token=forwarded_anthropic_oauth_token(request.headers), **_count_tokens_body_options(data, auto_drop_params), ) except UnsupportedContentTypeError as e: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0f8ac58861b..314f14b6279 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1415,12 +1415,34 @@ def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None: ) +def litellm_auth_header_name(headers: Headers) -> str: + if "x-litellm-api-key" in headers: + return "x-litellm-api-key" + if "authorization" in headers: + return "authorization" + return "x-api-key" + + +def forwarded_anthropic_oauth_token(headers: Headers) -> str | None: + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + forwarded: Final = clean_headers(headers, authenticated_with_header=litellm_auth_header_name(headers)) + return next( + ( + value.removeprefix("Bearer ") + for header, value in forwarded.items() + if header.lower() == "authorization" and is_anthropic_oauth_key(value) + ), + None, + ) + + def clean_headers( headers: Headers, litellm_key_header_name: str | None = None, forward_llm_provider_auth_headers: bool = False, authenticated_with_header: str | None = None, -) -> dict: +) -> dict[str, str]: """ Removes litellm api key from headers @@ -2513,26 +2535,13 @@ async def add_litellm_data_to_request( forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) if not forward_llm_auth: forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False) - # Determine which header was used for authentication - # This enables forwarding provider keys (e.g., x-api-key) when they weren't used for LiteLLM auth - authenticated_with_header = None - if "x-litellm-api-key" in request.headers: - # If x-litellm-api-key is present, it was used for auth - authenticated_with_header = "x-litellm-api-key" - elif "authorization" in request.headers: - # Authorization header was used for auth - authenticated_with_header = "authorization" - else: - # x-api-key or another header was used for auth - authenticated_with_header = "x-api-key" - _headers: Final[dict[str, str]] = clean_headers( request.headers, litellm_key_header_name=( general_settings.get("litellm_key_header_name") if general_settings is not None else None ), forward_llm_provider_auth_headers=forward_llm_auth, - authenticated_with_header=authenticated_with_header, + authenticated_with_header=litellm_auth_header_name(request.headers), ) _strip_untrusted_request_header_controls( _headers, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b36ea5afce..04f84f1f7b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14628,6 +14628,7 @@ async def run_thread( # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy +from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient @@ -14649,6 +14650,14 @@ from litellm.responses.litellm_completion_transformation.count_request import ( ) +def _with_client_oauth_token( + counter: BaseTokenCounter | None, client_oauth_token: str | None +) -> BaseTokenCounter | None: + if client_oauth_token is None or not isinstance(counter, AnthropicTokenCounter): + return counter + return AnthropicTokenCounter(client_oauth_token=client_oauth_token) + + def _get_provider_token_counter( deployment: dict, model_to_use: str ) -> tuple[BaseTokenCounter | None, str | None, str | None]: @@ -15264,6 +15273,7 @@ async def count_request_tokens( context_management: JsonValue = None, request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), responses_request: Mapping[str, object] | None = None, + client_oauth_token: str | None = None, ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -15321,7 +15331,8 @@ async def count_request_tokens( custom_llm_provider: str | None = None if call_endpoint is True and deployment is not None: # Auto-route to the correct provider based on model - provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) + routed_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) + provider_counter = _with_client_oauth_token(routed_counter, client_oauth_token) if _model is not None: model_to_use = _model diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py index a6c3a6dc62f..59114e4aced 100644 --- a/tests/integration/providers/test_count_tokens_error_shapes_wire.py +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -1,5 +1,6 @@ import json -from collections.abc import Iterator +from collections.abc import Callable, Iterator +from types import MappingProxyType from typing import Final import anthropic @@ -8,6 +9,7 @@ import pytest from integration._support.client import Gateway, gateway_from_environment from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from litellm.types.llms.anthropic import ANTHROPIC_OAUTH_BETA_HEADER from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -26,6 +28,14 @@ _PROVIDER_KEY_VARIABLES: Final = ( "AZURE_TENANT_ID", "OPENAI_API_KEY", ) +_ENVIRONMENT_KEYS: Final = MappingProxyType( + { + "ANTHROPIC_API_KEY": "synthetic-anthropic-env-key", + "AZURE_API_KEY": "synthetic-azure-env-key", + "OPENAI_API_KEY": "synthetic-openai-env-key", + } +) +_SUBSCRIPTION_TOKEN: Final = "sk-ant-oat01-synthetic-subscription-token" @pytest.fixture(scope="module") @@ -44,14 +54,14 @@ def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedPro @pytest.fixture(scope="module") -def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: +def environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: with ( gateway_from_environment() as gateway, owned_proxy_process( gateway, - tmp_path_factory.mktemp("count-tokens-azure-environment-key"), - {"AZURE_API_KEY": "synthetic-azure-env-key"}, - remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + tmp_path_factory.mktemp("count-tokens-environment-keys"), + _ENVIRONMENT_KEYS, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name not in _ENVIRONMENT_KEYS), workers=1, ) as owned, ): @@ -147,9 +157,9 @@ def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( - azure_environment_key_proxy: OwnedProxy, + environment_key_proxy: OwnedProxy, ) -> None: - gateway: Final = azure_environment_key_proxy.gateway + gateway: Final = environment_key_proxy.gateway def count(request: Request) -> Reply: assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers @@ -257,3 +267,122 @@ def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: Owned assert _JSON_OBJECT.validate_json(malformed.content) == { "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} } + + +def _counted(tokens: int) -> Reply: + return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": tokens}).encode()) + + +def _anthropic_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + + +def _responses_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return "/v1/responses/input_tokens", {"model": model, "input": "Count this synthetic request."} + + +@pytest.mark.parametrize( + ("provider_model", "base_path", "count_request", "header", "credential"), + ( + ("anthropic/claude-sonnet-5-5", "", _anthropic_count, "x-api-key", "synthetic-anthropic-env-key"), + ("openai/gpt-4o-mini", "/v1", _responses_count, "authorization", "Bearer synthetic-openai-env-key"), + ), + ids=("anthropic", "openai"), +) +def test_count_uses_a_provider_key_set_only_in_the_proxy_environment( + environment_key_proxy: OwnedProxy, + provider_model: str, + base_path: str, + count_request: Callable[[str], tuple[str, dict[str, JsonValue]]], + header: str, + credential: str, +) -> None: + gateway: Final = environment_key_proxy.gateway + with wire_server(lambda _request: _counted(47)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=provider_model, api_base=f"{wire.url}{base_path}", api_key=None, num_retries=0 + ) + path, body = count_request(model) + counted: Final = gateway.request("POST", path, body) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == 47, counted.text + received: Final = wire.drain() + assert [request.headers.get(header) for request in received] == [credential], received + + +def _google_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return f"/v1beta/models/{model}:countTokens", { + "contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}] + } + + +def _utils_count(model: str) -> tuple[str, dict[str, JsonValue]]: + return "/utils/token_counter?call_endpoint=true", {"model": model, "messages": _MESSAGES} + + +@pytest.mark.parametrize( + ("count_request", "google_status"), + ((_google_count, "UNAUTHENTICATED"), (_utils_count, None)), + ids=("google", "utils"), +) +def test_count_entrypoints_report_a_missing_provider_key_in_their_own_shape( + keyless_proxy: OwnedProxy, + count_request: Callable[[str], tuple[str, dict[str, JsonValue]]], + google_status: str | None, +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key=None, num_retries=0 + ) + path, body = count_request(model) + rejected: Final = gateway.request("POST", path, body) + assert rejected.status_code == 401, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict), rejected.text + assert f"No Anthropic API key is set for model {model}" in str(error["message"]), rejected.text + if google_status is not None: + assert (error["code"], error["status"]) == (401, google_status), rejected.text + assert wire.drain() == () + + +def test_anthropic_count_uses_the_callers_subscription_token_when_the_deployment_has_no_key( + keyless_proxy: OwnedProxy, +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(lambda _request: _counted(41)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key=None, num_retries=0 + ) + counted: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": _MESSAGES}, + headers={"Authorization": f"Bearer {_SUBSCRIPTION_TOKEN}", "x-litellm-api-key": f"Bearer {gateway.key}"}, + ) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == 41, counted.text + (received,) = wire.drain() + assert received.headers.get("authorization") == f"Bearer {_SUBSCRIPTION_TOKEN}", received.headers + assert "x-api-key" not in received.headers, received.headers + assert ANTHROPIC_OAUTH_BETA_HEADER in received.headers.get("anthropic-beta", ""), received.headers + + +def test_anthropic_count_keeps_the_deployments_key_over_the_callers_subscription_token( + keyless_proxy: OwnedProxy, +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(lambda _request: _counted(43)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-5-5", api_base=wire.url, api_key="synthetic-deployment-key", num_retries=0 + ) + counted: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": _MESSAGES}, + headers={"Authorization": f"Bearer {_SUBSCRIPTION_TOKEN}", "x-litellm-api-key": f"Bearer {gateway.key}"}, + ) + assert counted.status_code == 200, counted.text + (received,) = wire.drain() + assert received.headers.get("x-api-key") == "synthetic-deployment-key", received.headers + assert _SUBSCRIPTION_TOKEN not in json.dumps(dict(received.headers)), received.headers diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index 16ec513916c..b8a7c001b37 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -221,6 +221,45 @@ class TestCountTokensUsesWorkloadIdentity: assert result.total_tokens == 13 assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("deployment_key", "expected_credentials"), + [ + (None, {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}", "x-api-key": None}), + (FAKE_REGULAR_KEY, {"authorization": None, "x-api-key": FAKE_REGULAR_KEY}), + ], + ids=["keyless-deployment-uses-the-callers-token", "deployment-key-wins-over-the-callers-token"], + ) + async def test_a_callers_subscription_token_counts_only_when_no_key_is_configured( + self, + monkeypatch: pytest.MonkeyPatch, + httpx_transport_clients: None, + deployment_key: str | None, + expected_credentials: dict[str, str | None], + ) -> None: + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", None) + monkeypatch.setattr(litellm, "api_key", None) + + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 17}) + ) + result = await AnthropicTokenCounter(client_oauth_token=FAKE_OAUTH_TOKEN).count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": deployment_key}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 17 + sent = route.calls.last.request.headers + assert {name: sent.get(name) for name in expected_credentials} == expected_credentials + assert (ANTHROPIC_OAUTH_BETA_HEADER in sent["anthropic-beta"].split(",")) is (deployment_key is None) + @pytest.mark.asyncio async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch): from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index f743a14c7a2..a32c40e5a43 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -53,6 +53,7 @@ from litellm.proxy.litellm_pre_call_utils import ( add_provider_specific_headers_to_request, check_if_token_is_service_account, clean_headers, + forwarded_anthropic_oauth_token, get_dynamic_logging_metadata, get_metadata_variable_name, move_guardrails_to_metadata, @@ -5966,6 +5967,25 @@ def test_clean_headers_strips_x_api_key_when_byok_disabled(): assert "x-api-key" not in result +@pytest.mark.parametrize( + ("headers", "expected_token"), + [ + ( + {"authorization": "Bearer sk-ant-oat01-caller", "x-litellm-api-key": "Bearer sk-proxy"}, + "sk-ant-oat01-caller", + ), + ({"authorization": "Bearer sk-ant-oat01-caller"}, None), + ({"authorization": "Bearer sk-ant-api03-caller", "x-litellm-api-key": "Bearer sk-proxy"}, None), + ({"x-litellm-api-key": "Bearer sk-proxy"}, None), + ], + ids=["oauth-beside-proxy-key", "oauth-used-as-proxy-key", "non-oauth-authorization", "no-authorization"], +) +def test_forwarded_anthropic_oauth_token_follows_chats_forwarding_rule( + headers: dict[str, str], expected_token: str | None +) -> None: + assert forwarded_anthropic_oauth_token(Headers(headers)) == expected_token + + def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth_header(): """ Anti-replay regression: even when forward_llm_provider_auth_headers=True, diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index fa0c9eda02d..ebe3c7407d8 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -364,6 +364,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" @@ -444,6 +445,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" @@ -1262,6 +1264,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', @@ -1319,6 +1322,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): raise ProxyException( message='{"Message":"Bearer Token has expired"}', @@ -1376,6 +1380,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): output_format=None, context_management=None, request_litellm_params=None, + client_oauth_token=None, ): raise ProxyException( message="Rate limit exceeded",