fix(rust): separate public and request token counter policies

litellm/rust_bridge/README.md

@@ -6,7 +6,7 @@ Every SDK API has one `NativeComponent` in the immutable `COMPONENTS` catalog. A
 
 `RustImplementationState` records whether Rust is unimplemented, experimental, or ready. `RolloutPolicy` independently selects unsupported, Python-only, Rust opt-in, Rust opt-out, or Rust-required execution. Optional Rust execution can fall back to Python. Rust-required execution cannot
 
-OCR completed delivery is ready and default-on. Messages, chat completions, token counting, and Responses WebSocket transport are experimental and opt-in. Other completed APIs remain Python-only. Bedrock transcription requires Rust because it has no Python implementation; Python-backed transcription providers remain on Python
+OCR completed delivery is ready and default-on. Messages, chat completions, raw-request input token counting, and Responses WebSocket transport are experimental and opt-in. The public `litellm.token_counter()` and other completed APIs remain Python-only. Bedrock transcription requires Rust because it has no Python implementation; Python-backed transcription providers remain on Python
 
 ```python
 execution = COMPONENT.resolve(
@@ -28,7 +28,11 @@ Provider failures, host callback failures, cancellation, conversion failures, an
 
 `invoke` and `ainvoke` return the native result or execute the supplied fallback directly. There is no public admission, prepare, accepts, or can-handle API
 
-Token counting follows the same component policy. Its one native counting entrypoint validates the tokenizer configuration and request body, obtains and caches the required tokenizer resource, then counts. Unsupported inputs decline, known resource loading failures report native unavailability, and unexpected counting failures propagate
+Token counting has two catalog entries. `UtilityName.TOKEN_COUNTER` declares the public `litellm.token_counter()` as Python-only, with no native exports. The public function continues to execute Python directly regardless of `litellm.rust(bool)` or `LITELLM_RUST`
+
+`UtilityName.REQUEST_INPUT_TOKEN_COUNTER` owns the experimental raw-request optimization. Budget reservation keeps its direct `litellm.rust_bridge.token_counter.count_input_tokens` import. That adapter uses `REQUEST_COMPONENT`; `COMPONENT` describes the public API
+
+The request optimization follows its component policy. Its one native counting entrypoint validates the tokenizer configuration and request body, obtains and caches the required tokenizer resource, then counts. Unsupported inputs decline, known resource loading failures report native unavailability, and unexpected counting failures propagate
 
 ## Package layout
 

litellm/rust_bridge/catalog.py

@@ -136,6 +136,11 @@ COMPONENTS: Final = MappingProxyType(
         ),
         UtilityName.TOKEN_COUNTER: _component(
             UtilityName.TOKEN_COUNTER,
+            _unimplemented(),
+            (),
+        ),
+        UtilityName.REQUEST_INPUT_TOKEN_COUNTER: _component(
+            UtilityName.REQUEST_INPUT_TOKEN_COUNTER,
             _experimental_completed,
             ("count_input_tokens",),
         ),

litellm/rust_bridge/configuration.py

@@ -25,6 +25,7 @@ class RouteName(str, Enum):
 
 class UtilityName(str, Enum):
     TOKEN_COUNTER = "token_counter"
+    REQUEST_INPUT_TOKEN_COUNTER = "request_input_token_counter"
 
 
 class RustImplementationState(str, Enum):

litellm/rust_bridge/token_counter/__init__.py

@@ -1,9 +1,12 @@
 from typing import Final
 
+from litellm.rust_bridge.token_counter.definition import COMPONENT, REQUEST_COMPONENT
 from litellm.rust_bridge.token_counter.types import InputTokenCount, RustTokenizer
 from litellm.rust_bridge.token_counter.value import TOKEN_COUNTER, count_input_tokens, rust_tokenizer
 
 __all__: Final = (
+    "COMPONENT",
+    "REQUEST_COMPONENT",
     "TOKEN_COUNTER",
     "InputTokenCount",
     "RustTokenizer",

litellm/rust_bridge/token_counter/definition.py

@@ -4,3 +4,4 @@ from litellm.rust_bridge.catalog import COMPONENTS
 from litellm.rust_bridge.configuration import UtilityName
 
 COMPONENT: Final = COMPONENTS[UtilityName.TOKEN_COUNTER]
+REQUEST_COMPONENT: Final = COMPONENTS[UtilityName.REQUEST_INPUT_TOKEN_COUNTER]

litellm/rust_bridge/token_counter/value.py

@@ -8,7 +8,7 @@ import litellm
 from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file
 from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting
 from litellm.rust_bridge.runtime import BridgeErrorContext, ainvoke
-from litellm.rust_bridge.token_counter.definition import COMPONENT
+from litellm.rust_bridge.token_counter.definition import REQUEST_COMPONENT
 from litellm.rust_bridge.token_counter.types import InputTokenCount, RustTokenCounter, RustTokenizer
 from litellm.utils import claude_json_str, huggingface_tokenizer_kind
 
@@ -19,11 +19,11 @@ def _as_counter(value: object) -> RustTokenCounter | None:
     return cast(RustTokenCounter, value) if callable(value) else None
 
 
-TOKEN_COUNTER: Final = COMPONENT.bind("count_input_tokens", validate=_as_counter)
+TOKEN_COUNTER: Final = REQUEST_COMPONENT.bind("count_input_tokens", validate=_as_counter)
 
 
 def rust_tokenizer(model: str) -> RustTokenizer | None:
-    execution: Final = COMPONENT.resolve()
+    execution: Final = REQUEST_COMPONENT.resolve()
     if execution.select(TOKEN_COUNTER) is None:
         return None
     kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model)
@@ -49,7 +49,7 @@ def _tokenizer_resource(tokenizer: str) -> str:
 
 
 async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:
-    execution: Final = COMPONENT.resolve()
+    execution: Final = REQUEST_COMPONENT.resolve()
     counter: Final = execution.select(TOKEN_COUNTER)
 
     async def python_fallback() -> None:
@@ -71,5 +71,5 @@ async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputToke
         else None,
         python_fallback=python_fallback,
         adapt=_INPUT_TOKEN_COUNT.validate_python,
-        context=BridgeErrorContext(route=COMPONENT.name.value, provider="", model=""),
+        context=BridgeErrorContext(route=REQUEST_COMPONENT.name.value, provider="", model=""),
     )

tests/test_litellm/rust_bridge/test_configuration_env.py

@@ -1,10 +1,21 @@
 from __future__ import annotations
 
+import json
+from collections.abc import Callable
+from types import ModuleType
+from typing import Final
+
 import pytest
 
+import litellm
+from litellm.litellm_core_utils import token_counter as python_counter
+from litellm.rust_bridge import bindings, configuration
+from litellm.rust_bridge import token_counter as bridge
 from litellm.rust_bridge.configuration import (
+    ExecutionDecision,
     _parse_env_bool,  # pyright: ignore[reportPrivateUsage]  # directly test env parsing contract
 )
+from litellm.rust_bridge.token_counter import COMPONENT
 
 
 @pytest.mark.parametrize(("value", "expected"), (("1", True), ("0", False), (" 1 ", True), (" 0 ", False)))
@@ -20,3 +31,78 @@ def test_parse_env_bool_preserves_unset_value() -> None:
 def test_parse_env_bool_rejects_unknown_value(value: str) -> None:
     with pytest.raises(ValueError, match="must be '1' or '0'"):
         _parse_env_bool(value)
+
+
+@pytest.mark.parametrize("environment", (None, "0", "1"))
+@pytest.mark.parametrize("override", (None, False, True))
+def test_public_token_counter_stays_python_only(
+    monkeypatch: pytest.MonkeyPatch, environment: str | None, override: bool | None
+) -> None:
+    def unexpected_native_load() -> ModuleType:
+        raise AssertionError("public token counting must not load Rust")
+
+    calls: Final[list[str]] = []
+
+    def python_count(text: str) -> int:
+        calls.append(text)
+        return len(text)
+
+    configuration.reset_rust_configuration()
+    if environment is None:
+        monkeypatch.delenv("LITELLM_RUST", raising=False)
+    else:
+        monkeypatch.setenv("LITELLM_RUST", environment)
+    if override is not None:
+        litellm.rust(override)
+    monkeypatch.setattr(bindings, "get_native_bridge", unexpected_native_load)
+    monkeypatch.setattr(python_counter, "_get_count_function", lambda model, custom_tokenizer: python_count)
+    try:
+        assert COMPONENT.resolve().decision is ExecutionDecision.PYTHON
+        assert litellm.token_counter(model="gpt-4o", text="hello") == 5
+        assert calls == ["hello"]
+    finally:
+        configuration.reset_rust_configuration()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("environment", (None, "0", "1"))
+@pytest.mark.parametrize("override", (None, False, True))
+async def test_budget_direct_import_follows_request_rollout(
+    monkeypatch: pytest.MonkeyPatch, environment: str | None, override: bool | None
+) -> None:
+    from litellm.proxy.spend_tracking.budget_reservation import count_request_input_tokens
+
+    configuration.reset_rust_configuration()
+    if environment is None:
+        monkeypatch.delenv("LITELLM_RUST", raising=False)
+    else:
+        monkeypatch.setenv("LITELLM_RUST", environment)
+    if override is not None:
+        litellm.rust(override)
+    model: Final = "gpt-4o"
+    messages: Final = [{"role": "user", "content": "hello"}]
+    body: Final = json.dumps({"model": model, "messages": messages}).encode()
+    native_calls: Final[list[bytes]] = []
+
+    async def counter(
+        body: bytes,
+        kind: str | None,
+        encoding: str,
+        disabled: bool,
+        legacy_accounting: bool,
+        resource_loader: Callable[[str], str],
+    ) -> object:
+        native_calls.append(body)
+        return {"model": model, "input_tokens": 42}
+
+    bridge.TOKEN_COUNTER.override(counter)
+    enabled: Final = override if override is not None else environment == "1"
+    try:
+        budget: Final = await count_request_input_tokens(
+            request_body=json.loads(body), route="/v1/messages", llm_router=None, raw_body=body
+        )
+        assert budget == {model: 42 if enabled else litellm.token_counter(model=model, messages=messages)}
+        assert native_calls == ([body] if enabled else [])
+    finally:
+        bridge.TOKEN_COUNTER.reset()
+        configuration.reset_rust_configuration()