fix(token_counter): stop large token counts from blocking the proxy event loop

litellm/constants.py

@@ -3,7 +3,7 @@
 from types import MappingProxyType
 from typing import Final, Literal
 
-from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none
+from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_range, get_env_int_or_none
 
 DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
 AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
@@ -323,6 +323,17 @@
 DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20))
 MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768))
 MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000))
+# tiktoken's BPE merge loop is quadratic in the length of a single regex piece, so a long run of one
+# repeated character (dot leaders, whitespace, zero-padded base64) can take minutes on a multi-MB payload.
+# Encoding in chunks makes the cost linear, at a drift of at most ~1 token per chunk boundary. The upper
+# bound keeps a misconfigured chunk size from restoring the quadratic cost this exists to remove.
+TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS: Final = 4096
+TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
+    "TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS",
+    default=1024,
+    minimum=1,
+    maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
+)
 MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
 MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
 OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))

litellm/litellm_core_utils/env_utils.py

@@ -2,6 +2,7 @@
 Utility helpers for reading and parsing environment variables.
 """
 
+import logging
 import os
 from typing import Final
 
@@ -22,6 +23,26 @@ def get_env_int(env_var: str, default: int) -> int:
         return default
 
 
+def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int) -> int:
+    """Parse an environment variable as an integer constrained to ``[minimum, maximum]``.
+
+    Values outside the range fall back to the default and warn, so a misconfigured knob can
+    neither crash the caller nor silently change the meaning of what it computes.
+    """
+    value: Final = get_env_int(env_var, default)
+    if minimum <= value <= maximum:
+        return value
+    logging.getLogger("LiteLLM").warning(
+        "%s=%s is outside the supported range [%s, %s]. Falling back to %s.",
+        env_var,
+        value,
+        minimum,
+        maximum,
+        default,
+    )
+    return default
+
+
 def get_env_int_or_none(env_var: str) -> int | None:
     """Parse an environment variable as an integer, returning None when it is unset or unusable.
 

litellm/litellm_core_utils/token_counter.py

@@ -19,6 +19,7 @@
     MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES,
     MAX_TILE_HEIGHT,
     MAX_TILE_WIDTH,
+    TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS,
 )
 from litellm.litellm_core_utils.default_encoding import encoding as default_encoding
 from litellm.litellm_core_utils.url_utils import safe_get
@@ -305,6 +306,16 @@ def calculate_img_tokens(
 """
 
 
+def _get_tiktoken_count_function(
+    encode_length: Callable[[str], int],
+    chunk_size: int = TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS,
+) -> TokenCounterFunction:
+    def count_tokens(text: str) -> int:
+        return sum(encode_length(text[start : start + chunk_size]) for start in range(0, len(text), chunk_size))
+
+    return count_tokens
+
+
 class _MessageCountParams:
     """
     A class to hold the parameters for counting tokens in messages.
@@ -531,6 +542,7 @@ def count_tokens(text: str) -> int:
                 enc: Final = tokenizer_json["tokenizer"].encode(text)
                 return len(enc.ids)
 
+            return count_tokens
         elif tokenizer_json["type"] == "openai_tokenizer":
             model_to_use: Final = _fix_model_name(model)
             try:
@@ -542,17 +554,18 @@ def count_tokens(text: str) -> int:
                 print_verbose("Warning: model not found. Using cl100k_base encoding.")
                 encoding = tiktoken.get_encoding("cl100k_base")
 
-            def count_tokens(text: str) -> int:
+            def encode_length(text: str) -> int:
                 return len(encoding.encode(text, disallowed_special=()))
 
+            return _get_tiktoken_count_function(encode_length)
         else:
             raise ValueError("Unsupported tokenizer type")
     else:
 
-        def count_tokens(text: str) -> int:
+        def encode_length(text: str) -> int:
             return len(default_encoding.encode(text, disallowed_special=()))
 
-    return count_tokens
+        return _get_tiktoken_count_function(encode_length)
 
 
 def _fix_model_name(model: str) -> str:

litellm/proxy/proxy_server.py

@@ -253,6 +253,7 @@ def generate_feedback_box():
 from litellm.integrations.custom_guardrail import ModifyResponseException
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
+from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.core_helpers import (
     _get_parent_otel_span_from_kwargs,
     get_litellm_metadata_from_kwargs,
@@ -11909,8 +11910,6 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
     Returns:
         TokenCountResponse
     """
-    from litellm import token_counter
-
     global llm_router
 
     prompt: Final = request.prompt
@@ -11994,7 +11993,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
     _tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer)
 
     tokenizer_used: Final = str(_tokenizer_used["type"])
-    total_tokens: Final = token_counter(
+    total_tokens: Final = await asyncify(litellm.token_counter)(
         model=model_to_use,
         text=prompt,
         messages=messages,

litellm/router.py

@@ -52,7 +52,7 @@
     SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
 )
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.asyncify import run_async_function
+from litellm.litellm_core_utils.asyncify import asyncify, run_async_function
 from litellm.litellm_core_utils.core_helpers import (
     _get_parent_otel_span_from_kwargs,
     coerce_token_limit,
@@ -10570,13 +10570,73 @@ def _count_pre_call_check_tokens(
             return litellm.token_counter(messages=cast(list, input_messages))  # cast-ok: transformed chat messages
         raise ValueError("Either messages or input must be provided to count tokens")
 
+    def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None:
+        """The deployment's declared context window, or None when it declares none or cannot be resolved."""
+        try:
+            model_info: Final = self.get_router_model_info(
+                deployment=cast(dict, deployment),  # cast-ok: router deployments are plain dicts
+                received_model_name=model,
+            )
+        except Exception as e:  # noqa: BLE001  # best-effort: an unmappable deployment must not hide the others
+            verbose_router_logger.debug(
+                "litellm.router.py::_deployment_max_input_tokens: skipping deployment. Got - %s", e
+            )
+            return None
+        max_input_tokens: Final = model_info.get("max_input_tokens")
+        return max_input_tokens if isinstance(max_input_tokens, int) else None
+
+    def _pre_call_checks_need_token_count(
+        self, model: str, healthy_deployments: Sequence[Mapping[str, object]]
+    ) -> bool:
+        """Whether any healthy deployment declares a context window that a token count could exceed.
+
+        Resolves each deployment the way ``_pre_call_checks`` does, so one unmappable deployment
+        cannot hide a later one that does declare a limit.
+        """
+        return any(
+            self._deployment_max_input_tokens(model, deployment) is not None for deployment in healthy_deployments
+        )
+
+    async def _acount_pre_call_check_tokens(
+        self,
+        model: str,
+        healthy_deployments: Sequence[Mapping[str, object]],
+        messages: Sequence[Mapping[str, str]] | None,
+        input: str | Sequence[object] | None,
+        request_kwargs: Mapping[str, object] | None,
+    ) -> int | None:
+        """Count input tokens off the event loop, so a multi-MB prompt cannot stall the proxy.
+
+        Returns None when no deployment limits its context window, and when counting fails. The
+        caller pairs this with ``skip_inline_token_count`` so neither case puts the count back on
+        the loop: a failed count leaves the deployments unfiltered, exactly as before.
+        """
+        if messages is None and input is None:
+            return None
+        raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None
+        try:
+            if not self._pre_call_checks_need_token_count(model, healthy_deployments):
+                return None
+            return await asyncify(self._count_pre_call_check_tokens)(
+                messages=cast(list[dict[str, str]] | None, messages),  # cast-ok: forwarded to the sync counter
+                input=cast(str | list | None, input),  # cast-ok: forwarded to the sync counter
+                instructions=raw_instructions if isinstance(raw_instructions, str) else None,
+            )
+        except Exception as e:  # noqa: BLE001  # best-effort: an uncountable prompt must not fail the request
+            verbose_router_logger.error(
+                "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e
+            )
+            return None
+
     def _pre_call_checks(
         self,
         model: str,
         healthy_deployments: list,
         messages: list[dict[str, str]] | None = None,
         input: str | list | None = None,
         request_kwargs: dict | None = None,
+        input_token_count: int | None = None,
+        skip_inline_token_count: bool = False,
     ):
         """
         Filter out model in model group, if:
@@ -10598,7 +10658,9 @@ def _pre_call_checks(
         # Token counting (tiktoken) is the dominant on-loop cost for large prompts.
         # Only count when a deployment actually declares max_input_tokens, and count
         # at most once; for model groups with no context-window limit it is skipped.
-        input_tokens: int | None = None
+        # Async callers pass the count in, already computed off the event loop, and set
+        # skip_inline_token_count so a failed off-loop count is not retried back on the loop.
+        input_tokens: int | None = input_token_count
 
         _context_window_error = False
         _potential_error_str = ""
@@ -10633,6 +10695,8 @@ def _pre_call_checks(
                 max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None
                 if isinstance(max_input_tokens, int) and has_countable_input:
                     if input_tokens is None:
+                        if skip_inline_token_count:
+                            return _returned_deployments
                         try:
                             input_tokens = self._count_pre_call_check_tokens(
                                 messages=messages, input=input, instructions=instructions
@@ -11115,12 +11179,21 @@ async def async_get_healthy_deployments(
         )
 
         if self.enable_pre_call_checks and (messages is not None or input is not None):
+            deployments_to_check: Final = cast(list[dict], healthy_deployments)
             healthy_deployments = self._pre_call_checks(
                 model=model,
-                healthy_deployments=cast(list[dict], healthy_deployments),
+                healthy_deployments=deployments_to_check,
                 messages=messages,
                 input=input,
                 request_kwargs=request_kwargs,
+                input_token_count=await self._acount_pre_call_check_tokens(
+                    model=model,
+                    healthy_deployments=deployments_to_check,
+                    messages=messages,
+                    input=input,
+                    request_kwargs=request_kwargs,
+                ),
+                skip_inline_token_count=True,
             )
         # check if user wants to do tag based routing
         healthy_deployments = await get_deployments_for_tag(

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -1,12 +1,14 @@
 #### What this tests ####
 #    This tests litellm.token_counter.token_counter() function
+import importlib
 import os
 import sys
 import time
 import traceback
 from unittest.mock import MagicMock
 
 import pytest
+import tiktoken
 
 sys.path.insert(
     0, os.path.abspath("../../..")
@@ -16,6 +18,8 @@
 import litellm
 from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens
 from litellm import token_counter as token_counter_old
+import litellm.constants
+from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function
 from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
 from tests.large_text import text
 from tests.test_litellm.litellm_core_utils.messages_with_counts import (
@@ -54,6 +58,73 @@ def test_token_counter_basic():
     )
 
 
+def test_token_counter_large_repeated_text_is_fast():
+    messages = [{"role": "user", "content": [{"type": "text", "text": "A" * 1024 * 1024}]}]
+
+    start_time = time.perf_counter()
+    tokens = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages)
+    elapsed = time.perf_counter() - start_time
+
+    assert elapsed < 2, f"Token counting took too long: {elapsed:.2f}s"
+    assert tokens > 0
+
+
+@pytest.mark.parametrize(
+    "text",
+    [
+        "Short text",
+        "This is a normal message with punctuation, numbers, and a few words.",
+    ],
+)
+def test_token_counter_short_text_matches_tiktoken(text):
+    encoding = tiktoken.get_encoding("cl100k_base")
+    expected = len(encoding.encode(text, disallowed_special=()))
+
+    assert token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text) == expected
+
+
+def test_token_counter_text_over_chunk_boundary_stays_close_to_tiktoken():
+    text = ("The quick brown fox jumps over the lazy dog. " * 30)[:1025]
+    encoding = tiktoken.get_encoding("cl100k_base")
+    expected = len(encoding.encode(text, disallowed_special=()))
+
+    actual = token_counter_new(model="us.anthropic.claude-sonnet-4-6", text=text)
+
+    assert abs(actual - expected) <= 4
+
+
+@pytest.mark.parametrize(
+    "configured",
+    ["0", "-1", "-1024", "not-an-int", "", "   ", "999999999", "inf", "1e9"],
+)
+def test_invalid_chunk_size_config_stays_usable(monkeypatch, configured):
+    """A misconfigured chunk size must not raise, count zero, or restore the quadratic encode cost."""
+    monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", configured)
+    try:
+        reloaded = importlib.reload(litellm.constants)
+        chunk_size = reloaded.TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
+        assert 1 <= chunk_size <= reloaded.TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS
+
+        encoding = tiktoken.get_encoding("cl100k_base")
+        count_tokens = _get_tiktoken_count_function(
+            lambda text: len(encoding.encode(text, disallowed_special=())),
+            chunk_size=chunk_size,
+        )
+        assert count_tokens("The quick brown fox jumps over the lazy dog. " * 40) > 0
+    finally:
+        monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS")
+        importlib.reload(litellm.constants)
+
+
+def test_valid_chunk_size_config_is_honoured(monkeypatch):
+    monkeypatch.setenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS", "2048")
+    try:
+        assert importlib.reload(litellm.constants).TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS == 2048
+    finally:
+        monkeypatch.delenv("TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS")
+        importlib.reload(litellm.constants)
+
+
 def test_token_counter_with_prefix():
     messages = [
         {"role": "user", "content": "Who won the world cup in 2022?"},