mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(proxy): return provider auth errors from /v1/messages/count_tokens instead of masking or 500
This commit is contained in:
mateo-berri
parent
d44d281d1d
commit
cf9c1c94bc
4 changed files with 208 additions and 1 deletions
|
|
@ -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}"})
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue