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 27132c90e055..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: + 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/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 39ec4bb1887d..5a4ad1ed2721 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( # 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, + ): + 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 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 4ed6a4683719..386e1eb51b48 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12415,6 +12415,93 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin 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 +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(401, "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 == "401" + assert "API key is invalid." in exc_info.value.message + + +@pytest.mark.asyncio +@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. + 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")) + 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 + + def test_docs_redoc_openapi_are_reachable_by_default(): """ LIT-6745: the interactive/machine-readable docs surfaces are on by