From 719c44db3f2ec5a13d51fbf29524e001720306be Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 13 Oct 2026 12:08:01 -0700 Subject: [PATCH] fix(proxy): raise on provider count_tokens errors instead of falling back to tiktoken (#46121) Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yuneng-jiang --- litellm/proxy/proxy_server.py | 17 ++--- ...anthropic_count_tokens_system_lift_wire.py | 31 ++++---- .../test_bedrock_mantle_count_tokens_wire.py | 71 ++++++++++--------- .../test_vertex_partner_count_tokens_wire.py | 49 +++++++------ tests/unit/proxy/test_proxy_token_counter.py | 55 +++++++------- 5 files changed, 110 insertions(+), 113 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c2ca045236c..60b9ade82c0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14116,19 +14116,12 @@ 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: - raise ProxyException( - message=result.error_message or "Token counting failed", - type="token_counting_error", - param="model", - code=result.status_code or 500, - ) - verbose_proxy_logger.warning( - "Provider token counting failed (%s): %s. Falling back to local tokenizer.", - result.status_code, - result.error_message, + raise ProxyException( + message=result.error_message or "Token counting failed", + type="token_counting_error", + param="model", + code=result.status_code or 500, ) - return None return result diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index 13091ccd14a..a4b5c15ad54 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -398,24 +398,24 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place( ) -def test_messages_count_tokens_falls_back_locally_when_every_message_is_system( +def test_messages_count_tokens_surfaces_the_peer_error_when_every_message_is_system( deployment: _Deployment, gateway: Gateway ) -> None: with _listening(deployment) as peer: body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [LEADING]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [], "system": [LIFTED]},) -def test_messages_count_tokens_leaves_the_request_untouched_for_a_non_text_system( +def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system( deployment: _Deployment, gateway: Gateway ) -> None: with _listening(deployment) as peer: - local: Final = _local_count(gateway, count_request(deployment.model, STRING_CASE)) response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "system": 5}) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [LEADING, USER], "system": 5},) @@ -428,14 +428,14 @@ def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Dep assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2 -def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_peer_call( +def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content( deployment: _Deployment, gateway: Gateway ) -> None: with _listening(deployment) as peer: body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [{"role": "system"}, USER]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert peer.drain() == () follow_up: Final = _count(gateway, count_request(deployment.model, STRING_CASE)) assert _payload(follow_up) == {"input_tokens": _COUNT}, follow_up.text @@ -458,14 +458,14 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value( @pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) -def test_messages_count_tokens_falls_back_locally_when_the_peer_rejects_the_lifted_body( +def test_messages_count_tokens_surfaces_the_peer_status_when_the_peer_rejects_the_lifted_body( deployment: _Deployment, gateway: Gateway, status: int ) -> None: with _listening(deployment, _rejecting(status)) as peer: body: Final = count_request(deployment.model, STRING_CASE) - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == status, response.text + assert "input_tokens" not in response.text, response.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) @@ -478,14 +478,13 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer( assert peer.drain() == () -def test_peer_outage_between_concurrent_waves_falls_back_then_recovers( +def test_peer_outage_between_concurrent_waves_errors_then_recovers( deployment: _Deployment, gateway: Gateway ) -> None: with ExitStack() as stack: clients: Final = _clients(stack, str(gateway.client.base_url), 8) pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) body: Final = count_request(deployment.model, STRING_CASE) - local: Final = _local_count(gateway, body) def count(client: httpx.Client) -> tuple[int, JsonValue]: return _counted_on(client, gateway.key, body) @@ -493,7 +492,9 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers( with _listening(deployment) as peer: assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients) assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * len(clients) - assert tuple(pool.map(count, clients)) == ((200, local),) * len(clients) + outage: Final = tuple(pool.map(count, clients)) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with _listening(deployment) as revived: assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients) assert _bodies(revived, deployment.provider) == (_STRING_BODY,) * len(clients) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 7f59b818cba..a5ab085c492 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -522,7 +522,7 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx assert mantle.drain() == () -def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_mantle( +def test_non_claude_model_the_runtime_cannot_count_surfaces_the_error_without_mantle( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -533,7 +533,8 @@ def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_ma ): model: Final = _deployment(scenario, runtime.url, _NOVA) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert response.status_code == 400, response.text + assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text assert _runtime_count_targets(runtime) == (f"/model/{_NOVA}/count-tokens",) assert mantle.drain() == () @@ -551,17 +552,19 @@ def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: Owne "messages": [{"role": "user", "content": "status=403 Count this message"}], } response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _local_count(gateway, body)}, response.text + assert response.status_code == 403, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert mantle.drain() == () -def test_unreachable_runtime_falls_back_locally_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: +def test_unreachable_runtime_surfaces_the_error_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario: model: Final = _deployment(scenario, _closed_port_url()) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert mantle.drain() == () @@ -770,7 +773,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant ) -def test_messages_count_tokens_falls_back_locally_when_every_message_is_system_for_mantle( +def test_messages_count_tokens_surfaces_the_mantle_error_when_every_message_is_system( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -781,14 +784,14 @@ def test_messages_count_tokens_falls_back_locally_when_every_message_is_system_f ): model: Final = _deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [LEADING]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(messages=[], system=[LIFTED]),) -def test_messages_count_tokens_leaves_a_leading_system_in_place_beside_a_non_text_system_for_mantle( +def test_messages_count_tokens_surfaces_the_mantle_error_beside_a_non_text_system( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -798,9 +801,9 @@ def test_messages_count_tokens_leaves_a_leading_system_in_place_beside_a_non_tex gateway.scenario() as scenario, ): model: Final = _deployment(scenario, runtime.url) - local: Final = _local_count(gateway, count_request(model, STRING_CASE)) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "system": 5}) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(messages=[LEADING, USER], system=5),) @@ -821,7 +824,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2 -def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_mantle_call( +def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -832,9 +835,9 @@ def test_messages_count_tokens_answers_a_leading_system_without_content_before_a ): model: Final = _deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "system"}, USER]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert mantle.drain() == () follow_up: Final = _count(gateway, count_request(model, STRING_CASE)) assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT}, follow_up.text @@ -889,7 +892,7 @@ def test_disabled_token_counter_counts_a_leading_system_through_mantle( assert len(_runtime_count_targets(runtime)) == 1 -def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( +def test_mantle_outage_between_concurrent_lifted_waves_errors_then_recovers( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -900,7 +903,6 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( scenario: Final = stack.enter_context(gateway.scenario()) model: Final = _deployment(scenario, runtime.url) body: Final = count_request(model, STRING_CASE) - local: Final = _local_count(gateway, body) def count(client: httpx.Client) -> tuple[int, JsonValue]: return _counted_on(client, gateway.key, body) @@ -908,7 +910,9 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( with wire_server(_mantle(_strict), port=mantle_port) as mantle: assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * len(clients) - assert tuple(pool.map(count, clients)) == ((200, local),) * len(clients) + outage: Final = tuple(pool.map(count, clients)) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with wire_server(_mantle(_strict), port=mantle_port) as revived: assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(revived) == (_MANTLE_LIFTED,) * len(clients) @@ -916,7 +920,7 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( @pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) -def test_messages_count_tokens_falls_back_locally_when_mantle_errors( +def test_messages_count_tokens_surfaces_the_mantle_status_when_mantle_errors( counting_proxy: OwnedProxy, mantle_port: int, status: int ) -> None: gateway: Final = counting_proxy.gateway @@ -927,7 +931,8 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_errors( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert response.status_code == status, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_FULL,) @@ -937,7 +942,7 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_errors( [b'{"inputTokens": 7}', b"not json at all", b"{}"], ids=["runtime_key", "not_json", "empty_object"], ) -def test_messages_count_tokens_falls_back_locally_when_mantle_answers_without_input_tokens( +def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens( counting_proxy: OwnedProxy, mantle_port: int, body: bytes ) -> None: gateway: Final = counting_proxy.gateway @@ -948,21 +953,22 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_answers_without_in ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert 500 <= response.status_code < 600, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) -def test_messages_count_tokens_falls_back_locally_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None: +def test_messages_count_tokens_errors_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_runtime) as runtime, gateway.scenario() as scenario: model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 -def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers( +def test_messages_count_tokens_errors_when_mantle_never_answers( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -983,18 +989,18 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers( scenario: Final = stack.enter_context(gateway.scenario()) model: Final = _deployment(scenario, runtime.url) body: Final = _full(model) - local: Final = _local_count(gateway, body) future: Final = pool.submit(_counted_on, client, gateway.key, body) eventually(held.qsize, lambda size: size == 1, seconds=30) assert gateway.request("GET", "/health/liveliness").status_code == 200 - assert _local_count(gateway, body) == local + assert _local_count(gateway, body) > 0 assert not future.done() - assert future.result(timeout=120) == (200, local) + status, tokens = future.result(timeout=120) + assert 500 <= status < 600 and tokens is None, (status, tokens) assert len(_runtime_count_targets(runtime)) == 1 assert len(mantle.drain()) == 1 -def test_messages_count_tokens_falls_back_locally_when_mantle_rejects_a_non_text_system( +def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -1005,7 +1011,8 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_rejects_a_non_text ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": 5}) - assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=5),) @@ -1174,7 +1181,7 @@ def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_lo assert len(_runtime_count_targets(runtime)) == 4 -def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( +def test_mantle_outage_between_concurrent_waves_errors_then_recovers( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -1185,7 +1192,6 @@ def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( scenario: Final = stack.enter_context(gateway.scenario()) model: Final = _deployment(scenario, runtime.url) body: Final = _full(model) - local: Final = _local_count(gateway, body) def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: return _generated_then_counted(client, gateway.key, model, body) @@ -1197,7 +1203,8 @@ def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients) outage: Final = tuple(pool.map(count_only, clients)) - assert outage == ((200, local),) * len(clients) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with wire_server(_mantle(), port=mantle_port) as revived: assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(revived) == (_MANTLE_FULL,) * len(clients) diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 558ee3a004e..e181e38b1bf 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -344,15 +344,14 @@ def test_utils_token_counter_local_mode_never_calls_the_peer(gateway: Gateway) - assert wire.drain() == () -def test_utils_token_counter_falls_back_locally_when_peer_rejects_openai_tools(gateway: Gateway) -> None: +def test_utils_token_counter_surfaces_the_peer_error_for_openai_tools(gateway: Gateway) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) body: Final = {**_bare(model), "tools": _OPENAI_TOOLS} response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"}) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "total_tokens" not in response.text, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": _OPENAI_TOOLS},) - assert payload["total_tokens"] == _local_count(gateway, body), response.text - assert payload["tokenizer_type"] != "vertex_ai_partner_models", response.text def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway) -> None: @@ -365,7 +364,7 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway assert _counted_bodies(wire) == (_PEER_BARE,) -def test_responses_input_tokens_falls_back_locally_when_peer_rejects(gateway: Gateway) -> None: +def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = gateway.request( @@ -373,11 +372,10 @@ def test_responses_input_tokens_falls_back_locally_when_peer_rejects(gateway: Ga "/v1/responses/input_tokens", {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS}, ) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text (sent,) = _counted_bodies(wire) assert sent["tools"] == _RESPONSES_TOOLS, sent - local: Final = _local_count(gateway, {"model": model, "messages": sent["messages"], "tools": _RESPONSES_TOOLS}) - assert payload == {"object": "response.input_tokens", "input_tokens": local}, response.text def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None: @@ -388,26 +386,26 @@ def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},) -def test_gemini_count_tokens_route_falls_back_locally_when_peer_rejects(gateway: Gateway) -> None: +def test_gemini_count_tokens_route_surfaces_the_peer_error(gateway: Gateway) -> None: with wire_server(_peer(_rejecting(400))) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "totalTokens" not in response.text, response.text assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},) - assert payload["totalTokens"] == _local_count(gateway, {"model": model, "messages": _GEMINI_MESSAGES}) @pytest.mark.parametrize("status", [400, 500, 503]) -def test_messages_count_tokens_falls_back_locally_when_peer_errors(gateway: Gateway, status: int) -> None: +def test_messages_count_tokens_surfaces_the_peer_error_when_peer_errors(gateway: Gateway, status: int) -> None: with wire_server(_peer(_rejecting(status))) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, _full(model)) - payload: Final = _payload(response) + assert response.status_code == status, response.text + assert "input_tokens" not in response.text, response.text assert _counted_bodies(wire) == (_PEER_FULL,) - assert payload == {"input_tokens": _local_count(gateway, _full(model))}, response.text -def test_messages_count_tokens_falls_back_locally_when_token_endpoint_rejects(gateway: Gateway) -> None: +def test_messages_count_tokens_surfaces_the_error_when_token_endpoint_rejects(gateway: Gateway) -> None: def respond(request: Request) -> Reply: if request.target == "/_oauth/token": return Reply( @@ -421,17 +419,18 @@ def test_messages_count_tokens_falls_back_locally_when_token_endpoint_rejects(ga gateway, scenario, wire.url, vertex_credentials=service_account_json(_PROJECT, wire.url) ) response: Final = _count(gateway, _full(model)) - payload: Final = _payload(response) + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text targets: Final = frozenset(request.target for request in wire.drain()) assert targets == {"/_oauth/token"}, targets - assert payload == {"input_tokens": _local_count(gateway, _full(model))}, response.text -def test_messages_count_tokens_falls_back_locally_when_peer_is_unreachable(gateway: Gateway) -> None: +def test_messages_count_tokens_errors_when_peer_is_unreachable(gateway: Gateway) -> None: with gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, _closed_port_url()) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text @pytest.mark.parametrize( @@ -476,13 +475,13 @@ def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway assert _counted_bodies(wire) == (_PEER_BARE,) -def test_messages_count_tokens_falls_back_locally_when_peer_rejects_a_non_text_system(gateway: Gateway) -> None: +def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system(gateway: Gateway) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": 5}) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": 5},) - assert payload == {"input_tokens": _local_count(gateway, _bare(model))}, response.text def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway) -> None: @@ -578,7 +577,7 @@ def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_l assert len(_counted_bodies(wire)) == 2 -def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway: Gateway) -> None: +def test_peer_outage_between_concurrent_waves_errors_then_recovers(gateway: Gateway) -> None: with ExitStack() as stack: clients: Final = _clients(stack, str(gateway.client.base_url), 8) pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) @@ -587,7 +586,6 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway: port: Final = int(wire.url.rsplit(":", 1)[1]) model: Final = _deployment(gateway, scenario, wire.url) body: Final = _full(model) - local: Final = _local_count(gateway, body) def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: return _generated_then_counted(client, gateway.key, model, body) @@ -595,7 +593,8 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway: assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients) assert _counted_bodies(wire) == (_PEER_FULL,) * len(clients) outage: Final = tuple(pool.map(lambda client: _counted_on(client, gateway.key, body), clients)) - assert outage == ((200, local),) * len(clients) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with wire_server(_peer(), port=port) as revived: assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients) assert _counted_bodies(revived) == (_PEER_FULL,) * len(clients) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index e4c4ddce64f..c93c199cfa2 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -310,20 +310,16 @@ 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, - ) - - print("Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "claude-test" - assert response.model_used == "claude-3-sonnet-20240229" - assert response.total_tokens > 0 + # The provider counter is engaged: its API error surfaces instead of a local count + with pytest.raises(ProxyException) as exc_info: + await token_counter( + request=TokenCountRequest( + model="claude-test", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + assert exc_info.value.type == "token_counting_error" # Test with non-Anthropic provider llm_router = Router( @@ -1018,10 +1014,10 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): @pytest.mark.asyncio -async def test_proxy_token_counter_error_falls_back_when_enabled(): +async def test_proxy_token_counter_error_raises_when_enabled(): """ - Test that proxy token_counter falls back to local tokenizer when disable_token_counter=False - and provider returns an error response. + Test that proxy token_counter raises ProxyException when disable_token_counter=False + and the provider returns an error response, never falling back to the local tokenizer. """ # Create error response error_response = TokenCountResponse( @@ -1069,19 +1065,20 @@ async def test_proxy_token_counter_error_falls_back_when_enabled(): mock_get_provider_token_counter ) - # Should not raise, should fall back to local tokenizer - result = await token_counter( - request=TokenCountRequest( - model="claude-bedrock", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) + with patch("litellm.token_counter") as mock_local_counter: + with pytest.raises(ProxyException) as exc_info: + await token_counter( + request=TokenCountRequest( + model="claude-bedrock", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) - # Should have used the fallback tokenizer - assert result.error is False - assert result.total_tokens > 0 - assert result.tokenizer_type != "bedrock_api" + assert exc_info.value.code == "429" + assert "Rate limit exceeded" in exc_info.value.message + assert exc_info.value.type == "token_counting_error" + mock_local_counter.assert_not_called() finally: litellm.disable_token_counter = original_disable litellm.proxy.proxy_server._get_provider_token_counter = (