scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

test(proxy): audit #46815 (#48768)

Browse source 
* 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>
This commit is contained in:
m6atlas 2026-10-28 03:55:25 -07:00 • committed by GitHub
parent 6ba77b7013
commit e2d78b69cd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 250 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/llms/anthropic/count_tokens/token_counter.py
View file

@ -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,

6
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -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:

39
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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,

13
litellm/proxy/proxy_server.py
View file

@ -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

143
tests/integration/providers/test_count_tokens_error_shapes_wire.py
View file

@ -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

39
tests/unit/llms/anthropic/test_count_tokens_oauth.py
View file

@ -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

20
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -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,

5
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -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",