From cf9c1c94bcbad7dec855dfb1305ae251055c3a0b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 10:03:48 -0700 Subject: [PATCH 1/4] fix(proxy): return provider auth errors from /v1/messages/count_tokens instead of masking or 500 --- .../proxy/anthropic_endpoints/endpoints.py | 9 ++ litellm/proxy/proxy_server.py | 2 +- .../anthropic_endpoints/test_endpoints.py | 113 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 85 +++++++++++++ 4 files changed, 208 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 7f0045c1d93d..0269070dc44f 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -310,6 +310,15 @@ async def count_tokens( detail=detail, ) except Exception as e: + exception_status: Final = getattr(e, "status_code", None) + if isinstance(exception_status, int) and 400 <= exception_status <= 599: + raise HTTPException( + status_code=exception_status, + detail=AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=exception_status, + raw_message=str(getattr(e, "message", None) or e), + ), + ) verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e) raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 887716a383a1..cb1b0fe80918 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12315,7 +12315,7 @@ async def _try_provider_token_count( code=status_code, ) if result is not None and result.error is True: - if litellm.disable_token_counter is True: + if litellm.disable_token_counter is True or result.status_code in (401, 403): raise ProxyException( message=result.error_message or "Token counting failed", type="token_counting_error", diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index c83ba1420115..9afd2666e75a 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -350,3 +350,116 @@ def test_flag_defaults_off(self): import litellm assert litellm.strip_anthropic_total_tokens is False + + +class TestCountTokensAuthErrorMapping: + """LIT-6507: /v1/messages/count_tokens must surface Anthropic auth failures + as the Anthropic error envelope with the provider's status, matching + /v1/messages, instead of masking them behind a 200 local count or a 500.""" + + def _count_tokens_body(self): + return { + "model": "claude-haiku-4-5", + "messages": [{"role": "user", "content": "count these tokens please"}], + } + + @pytest.mark.asyncio + async def test_provider_refused_credential_returns_401_envelope_not_local_count(self, monkeypatch): + """A real 401 from the provider's count-tokens API must reach the + client as a 401 authentication_error envelope, never a 200 with a + silently substituted local-tokenizer count.""" + import litellm + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from fastapi import HTTPException + from litellm.types.utils import TokenCountResponse + + provider_error = TokenCountResponse( + total_tokens=0, + request_model="claude-haiku-4-5", + model_used="claude-haiku-4-5", + tokenizer_type="anthropic_api", + error=True, + error_message="API key is invalid.", + status_code=401, + ) + + class _RefusedCredentialCounter: + def should_use_token_counting_api(self, custom_llm_provider=None): + return True + + async def count_tokens(self, **kwargs): + return provider_error + + mock_deployment = { + "litellm_params": {"model": "anthropic/claude-haiku-4-5"}, + "model_info": {}, + } + mock_router = MagicMock() + mock_router.async_get_available_deployment = AsyncMock(return_value=mock_deployment) + + monkeypatch.setattr(litellm, "disable_token_counter", False) + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value=self._count_tokens_body())), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "llm_router", mock_router), # test-quality-ok: module global read at call time; no injection seam + patch.object( # test-quality-ok: provider counter resolution is a module function; no injection seam + proxy_server, + "_get_provider_token_counter", + new=lambda deployment, model_to_use: (_RefusedCredentialCounter(), "claude-haiku-4-5", "anthropic"), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await ep.count_tokens(request=MagicMock(), user_api_key_dict=MagicMock()) + + assert exc_info.value.status_code == 401 + detail = exc_info.value.detail + assert detail["type"] == "error" + assert detail["error"]["type"] == "authentication_error" + assert "API key is invalid." in detail["error"]["message"] + + @pytest.mark.asyncio + async def test_status_carrying_exception_maps_to_its_status_not_500(self): + """A litellm.AuthenticationError escaping token counting (e.g. a + rejected workload-identity-federation token exchange) must map to its + status_code with the Anthropic envelope, not the blanket 500.""" + import litellm + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from fastapi import HTTPException + + auth_error = litellm.AuthenticationError( + message="Anthropic workload identity federation failed. The token endpoint returned HTTP 400", + llm_provider="anthropic", + model="claude-haiku-4-5", + ) + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value=self._count_tokens_body())), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=auth_error)), # test-quality-ok: endpoint imports the module attribute at call time; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await ep.count_tokens(request=MagicMock(), user_api_key_dict=MagicMock()) + + assert exc_info.value.status_code == 401 + detail = exc_info.value.detail + assert detail["type"] == "error" + assert detail["error"]["type"] == "authentication_error" + assert "workload identity federation failed" in detail["error"]["message"] + + @pytest.mark.asyncio + async def test_statusless_exception_stays_500(self): + """Exceptions without an HTTP status keep the internal-server-error + contract.""" + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from fastapi import HTTPException + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value=self._count_tokens_body())), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=ValueError("boom"))), # test-quality-ok: endpoint imports the module attribute at call time; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await ep.count_tokens(request=MagicMock(), user_api_key_dict=MagicMock()) + + assert exc_info.value.status_code == 500 + assert "Internal server error" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 1de3ed6e56d9..8b5953e3c34b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12165,3 +12165,88 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin router, _, _ = await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) assert router.fallback_access_check is router_fallback_access_check + + +class _StubProviderTokenCounter: + """LIT-6507: injected in place of a real provider counter so + _try_provider_token_count's error handling is exercised without network.""" + + def __init__(self, response): + self._response = response + + def should_use_token_counting_api(self, custom_llm_provider=None): + return True + + async def count_tokens( + self, + model_to_use, + messages, + contents, + deployment=None, + request_model="", + tools=None, + system=None, + ): + return self._response + + +def _provider_token_count_error(status_code, message): + from litellm.types.utils import TokenCountResponse + + return TokenCountResponse( + total_tokens=0, + request_model="claude-auth-test", + model_used="claude-haiku-4-5", + tokenizer_type="anthropic_api", + error=True, + error_message=message, + status_code=status_code, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_status", [401, 403]) +async def test_try_provider_token_count_raises_proxy_exception_on_provider_auth_error(auth_status, monkeypatch): + """LIT-6507: a provider-refused credential (401/403) must surface as a + ProxyException with that status, not silently fall back to the local + tokenizer and mask the auth failure behind a 200.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.proxy_server import _try_provider_token_count + + counter = _StubProviderTokenCounter(_provider_token_count_error(auth_status, "API key is invalid.")) + monkeypatch.setattr(litellm, "disable_token_counter", False) + with pytest.raises(ProxyException) as exc_info: + await _try_provider_token_count( + provider_counter=counter, + custom_llm_provider="anthropic", + model_to_use="claude-haiku-4-5", + messages=[{"role": "user", "content": "count these tokens please"}], + contents=None, + deployment=None, + request_model="claude-auth-test", + ) + + assert exc_info.value.code == str(auth_status) + assert "API key is invalid." in exc_info.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("non_auth_status", [429, 500]) +async def test_try_provider_token_count_falls_back_to_local_on_non_auth_error(non_auth_status, monkeypatch): + """Non-auth provider failures keep the deliberate silent fallback to the + local tokenizer (PR #34258): the caller gets None and counts locally.""" + from litellm.proxy.proxy_server import _try_provider_token_count + + counter = _StubProviderTokenCounter(_provider_token_count_error(non_auth_status, "provider unavailable")) + monkeypatch.setattr(litellm, "disable_token_counter", False) + result = await _try_provider_token_count( + provider_counter=counter, + custom_llm_provider="anthropic", + model_to_use="claude-haiku-4-5", + messages=[{"role": "user", "content": "count these tokens please"}], + contents=None, + deployment=None, + request_model="claude-auth-test", + ) + + assert result is None From 1e189df1748fa8ca2f4951a0cf17bfa8e698bcc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:28:20 -0700 Subject: [PATCH 2/4] test(proxy): stub the anthropic counter in the provider detection test --- .../test_proxy_token_counter.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 39ec4bb1887d..e97b4e02f15d 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -382,15 +382,24 @@ async def test_internal_token_counter_anthropic_provider_detection(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - # Test with is_direct_request=False (simulating call from Anthropic endpoint) - response = await token_counter( - request=TokenCountRequest( - model="claude-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, + mock_handler = MagicMock() + mock_handler.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 42} ) + # Test with is_direct_request=False (simulating call from Anthropic endpoint) + with patch( + "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", + mock_handler, + ): + response = await token_counter( + request=TokenCountRequest( + model="claude-test", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + print("Anthropic provider test response:", response) # Verify response structure From c79776d3da5c458990286a08499dd855aaf2ff67 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:03:21 -0700 Subject: [PATCH 3/4] test(proxy): give the anthropic counter patch its test-quality-ok reason --- tests/proxy_unit_tests/test_proxy_token_counter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index e97b4e02f15d..5a4ad1ed2721 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -388,7 +388,7 @@ async def test_internal_token_counter_anthropic_provider_detection(): ) # Test with is_direct_request=False (simulating call from Anthropic endpoint) - with patch( + with patch( # test-quality-ok: the module-global anthropic counter is the only seam; the test checks provider detection, not the counter "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", mock_handler, ): From 2d96db3bd9d1012ae693271879417b82008e2200 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:02:30 -0700 Subject: [PATCH 4/4] fix(proxy): only a provider 401 aborts token counting, 403 keeps the local fallback A Bedrock API key that can invoke models but lacks bedrock:CountTokens answers the count call with 403. Raising there turned every count-tokens surface into a 403 while sending kept working, so a 403 now falls back to the local estimate like every other non-auth failure. --- litellm/proxy/proxy_server.py | 2 +- tests/test_litellm/proxy/test_proxy_server.py | 16 +++++++++------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 026fa9bd2783..9d38f745dd8e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12428,7 +12428,7 @@ async def _try_provider_token_count( code=status_code, ) if result is not None and result.error is True: - if litellm.disable_token_counter is True or result.status_code in (401, 403): + if litellm.disable_token_counter is True or result.status_code == 401: raise ProxyException( message=result.error_message or "Token counting failed", type="token_counting_error", diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5c2a307dedb7..386e1eb51b48 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12453,15 +12453,14 @@ def _provider_token_count_error(status_code, message): @pytest.mark.asyncio -@pytest.mark.parametrize("auth_status", [401, 403]) -async def test_try_provider_token_count_raises_proxy_exception_on_provider_auth_error(auth_status, monkeypatch): - """LIT-6507: a provider-refused credential (401/403) must surface as a +async def test_try_provider_token_count_raises_proxy_exception_on_provider_auth_error(monkeypatch): + """LIT-6507: a provider-refused credential (401) must surface as a ProxyException with that status, not silently fall back to the local tokenizer and mask the auth failure behind a 200.""" from litellm.proxy._types import ProxyException from litellm.proxy.proxy_server import _try_provider_token_count - counter = _StubProviderTokenCounter(_provider_token_count_error(auth_status, "API key is invalid.")) + counter = _StubProviderTokenCounter(_provider_token_count_error(401, "API key is invalid.")) monkeypatch.setattr(litellm, "disable_token_counter", False) with pytest.raises(ProxyException) as exc_info: await _try_provider_token_count( @@ -12474,15 +12473,18 @@ async def test_try_provider_token_count_raises_proxy_exception_on_provider_auth_ request_model="claude-auth-test", ) - assert exc_info.value.code == str(auth_status) + assert exc_info.value.code == "401" assert "API key is invalid." in exc_info.value.message @pytest.mark.asyncio -@pytest.mark.parametrize("non_auth_status", [429, 500]) +@pytest.mark.parametrize("non_auth_status", [403, 429, 500]) async def test_try_provider_token_count_falls_back_to_local_on_non_auth_error(non_auth_status, monkeypatch): """Non-auth provider failures keep the deliberate silent fallback to the - local tokenizer (PR #34258): the caller gets None and counts locally.""" + local tokenizer (PR #34258): the caller gets None and counts locally. + 403 stays here on purpose: a credential that can invoke the model but is + denied the count action (a Bedrock API key without bedrock:CountTokens) + must keep its local estimate instead of losing every count-tokens surface.""" from litellm.proxy.proxy_server import _try_provider_token_count counter = _StubProviderTokenCounter(_provider_token_count_error(non_auth_status, "provider unavailable"))