diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 7f0045c1d93..0269070dc44 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 887716a383a..cb1b0fe8091 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 c83ba142011..9afd2666e75 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 @@ class TestStripTotalTokensFeatureFlag(unittest.TestCase): 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 1de3ed6e56d..8b5953e3c34 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