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] 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 026fa9bd278..9d38f745dd8 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 5c2a307dedb..386e1eb51b4 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"))