feat(rust_bridge): count prompt, input, query and documents bodies in Rust

The Rust token counter now mirrors _count_input_tokens key precedence
(messages, prompt, input, query/documents) so every LLM route that goes
through budget reservation gets the GIL-free count, not only /v1/messages
and /v1/chat/completions. Objects are serialised like json.dumps before
tokenizing; floats and unknown shapes still decline to Python. The body
model is optional so route-selected models can be matched by the caller

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-rust/crates/core/src/token_counter/mod.rs

@@ -2,6 +2,7 @@
 //! for the shapes it can count exactly. Everything else is declined so the host
 //! keeps its own counter as the reference.
 
+mod python_json;
 mod tools;
 pub mod types;
 
@@ -13,7 +14,9 @@ use crate::constants::{
     TOOL_CHOICE_NONE_TOKENS, TOOL_DEFINITIONS_TOKENS, TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT,
 };
 use tools::format_function_definitions;
-use types::{ContentBlock, ContentItem, CountableRequest, Message, MessageContent, ToolChoice};
+use types::{
+    ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice,
+};
 
 #[derive(Debug, ThisError, PartialEq, Eq)]
 pub enum TokenCountError {
@@ -29,7 +32,7 @@ pub enum TokenCountError {
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize)]
 pub struct InputTokenCount {
-    pub model: String,
+    pub model: Option<String>,
     pub input_tokens: usize,
 }
 
@@ -56,14 +59,37 @@ impl TokenCounter {
             .map_err(|error| TokenCountError::Encode(error.to_string()))
     }
 
+    /// Mirrors the host's key precedence: `messages`, then `prompt`, then
+    /// `input`, then `query` plus `documents`.
     pub fn count_request(
         &self,
         request: &CountableRequest,
     ) -> Result<InputTokenCount, TokenCountError> {
-        let messages = request
-            .messages
-            .as_deref()
-            .ok_or_else(|| TokenCountError::Unsupported("request has no messages".to_string()))?;
+        let input_tokens = if let Some(messages) = &request.messages {
+            self.count_messages(request, messages)?
+        } else if let Some(prompt) = &request.prompt {
+            self.count_text_value(prompt)?
+        } else if let Some(input) = &request.input {
+            self.count_text_value(input)?
+        } else if request.query.is_some() || request.documents.is_some() {
+            self.count_optional_text_value(request.query.as_ref())?
+                + self.count_optional_text_value(request.documents.as_ref())?
+        } else {
+            return Err(TokenCountError::Unsupported(
+                "request has no countable input".to_string(),
+            ));
+        };
+        Ok(InputTokenCount {
+            model: request.model.clone(),
+            input_tokens,
+        })
+    }
+
+    fn count_messages(
+        &self,
+        request: &CountableRequest,
+        messages: &[Message],
+    ) -> Result<usize, TokenCountError> {
         let message_tokens = messages
             .iter()
             .map(|message| self.count_message(message))
@@ -76,10 +102,35 @@ impl TokenCounter {
             request.tool_choice.as_ref(),
             includes_system_message,
         )?;
-        Ok(InputTokenCount {
-            model: request.model.clone(),
-            input_tokens: message_tokens + extra_tokens,
-        })
+        Ok(message_tokens + extra_tokens)
+    }
+
+    fn count_optional_text_value(
+        &self,
+        value: Option<&TextValue>,
+    ) -> Result<usize, TokenCountError> {
+        value.map_or(Ok(0), |value| self.count_text_value(value))
+    }
+
+    /// `str()` for scalars, `json.dumps()` for objects, lists flattened, nulls
+    /// skipped. Floats are declined because Python's `repr` and Rust's float
+    /// formatting disagree on exponents.
+    fn count_text_value(&self, value: &TextValue) -> Result<usize, TokenCountError> {
+        match value {
+            TextValue::Null => Ok(0),
+            TextValue::Bool(true) => self.count_text("True"),
+            TextValue::Bool(false) => self.count_text("False"),
+            TextValue::Integer(number) => self.count_text(&number.to_string()),
+            TextValue::Float(_) => Err(TokenCountError::Unsupported(
+                "float text values are counted by the python path".to_string(),
+            )),
+            TextValue::Text(text) => self.count_text(text),
+            TextValue::List(items) => items
+                .iter()
+                .map(|item| self.count_text_value(item))
+                .sum::<Result<usize, _>>(),
+            TextValue::Object(_) => self.count_text(&python_json::dumps(value)?),
+        }
     }
 
     fn count_message(&self, message: &Message) -> Result<usize, TokenCountError> {

litellm-rust/crates/core/src/token_counter/python_json.rs

@@ -0,0 +1,79 @@
+//! `json.dumps(value)` with Python's default arguments: `", "` and `": "`
+//! separators, `ensure_ascii=True`, and keys in insertion order.
+
+use std::fmt::Write;
+
+use super::TokenCountError;
+use super::types::TextValue;
+
+pub(super) fn dumps(value: &TextValue) -> Result<String, TokenCountError> {
+    let mut out = String::new();
+    write_value(&mut out, value)?;
+    Ok(out)
+}
+
+fn write_value(out: &mut String, value: &TextValue) -> Result<(), TokenCountError> {
+    match value {
+        TextValue::Null => out.push_str("null"),
+        TextValue::Bool(true) => out.push_str("true"),
+        TextValue::Bool(false) => out.push_str("false"),
+        TextValue::Integer(number) => write_number(out, number),
+        TextValue::Float(_) => {
+            return Err(TokenCountError::Unsupported(
+                "float repr is formatted by the python path".to_string(),
+            ));
+        }
+        TextValue::Text(text) => write_string(out, text),
+        TextValue::List(items) => {
+            out.push('[');
+            for (index, item) in items.iter().enumerate() {
+                if index > 0 {
+                    out.push_str(", ");
+                }
+                write_value(out, item)?;
+            }
+            out.push(']');
+        }
+        TextValue::Object(entries) => {
+            out.push('{');
+            for (index, (key, item)) in entries.iter().enumerate() {
+                if index > 0 {
+                    out.push_str(", ");
+                }
+                write_string(out, key);
+                out.push_str(": ");
+                write_value(out, item)?;
+            }
+            out.push('}');
+        }
+    }
+    Ok(())
+}
+
+fn write_number(out: &mut String, number: &i64) {
+    // Writing an integer into a String cannot fail.
+    let _ = write!(out, "{number}");
+}
+
+fn write_string(out: &mut String, text: &str) {
+    out.push('"');
+    for character in text.chars() {
+        match character {
+            '"' => out.push_str("\\\""),
+            '\\' => out.push_str("\\\\"),
+            '\n' => out.push_str("\\n"),
+            '\r' => out.push_str("\\r"),
+            '\t' => out.push_str("\\t"),
+            '\u{08}' => out.push_str("\\b"),
+            '\u{0c}' => out.push_str("\\f"),
+            ' '..='~' => out.push(character),
+            _ => {
+                let mut units = [0u16; 2];
+                for unit in character.encode_utf16(&mut units) {
+                    let _ = write!(out, "\\u{unit:04x}");
+                }
+            }
+        }
+    }
+    out.push('"');
+}

litellm-rust/crates/core/src/token_counter/tests.rs

@@ -45,23 +45,69 @@ const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5",
     "type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}],
   "tool_choice":"none"}"#;
 
+const COMPLETIONS_PROMPT: &str =
+    r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#;
+
+const COMPLETIONS_PROMPT_LIST: &str =
+    r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#;
+
+const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[
+  {"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]},
+  {"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#;
+
+const EMBEDDINGS_TOKEN_IDS: &str =
+    r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#;
+
+const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour",
+  "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#;
+
+/// Expected counts are pinned from
+/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`.
 #[rstest]
 #[case::text_only(SIMPLE, 14)]
 #[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)]
 #[case::openai_tools_named_choice(TOOLS_OPENAI, 123)]
 #[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)]
+#[case::completions_prompt(COMPLETIONS_PROMPT, 7)]
+#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)]
+#[case::responses_input_items(RESPONSES_INPUT, 62)]
+#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)]
+#[case::rerank_query_and_documents(RERANK, 41)]
 fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) {
     let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
     let count = counter().count_request(&request).expect("fixture counts");
     assert_eq!(
         count,
         InputTokenCount {
-            model: "claude-sonnet-4-5".to_string(),
+            model: Some("claude-sonnet-4-5".to_string()),
             input_tokens: expected,
         }
     );
 }
 
+#[rstest]
+#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)]
+#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)]
+#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)]
+#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)]
+fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) {
+    let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
+    let count = counter().count_request(&request).expect("fixture counts");
+    assert_eq!(count.input_tokens, expected);
+}
+
+#[test]
+fn objects_dump_like_python_json_dumps() {
+    let body = r#"{"model":"m","input":{"text":"caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~","n":-3,"ok":true,"no":false,"none":null,"list":[1,"a",{"z":[]}],"empty":{}}}"#;
+    let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
+    let dumped = python_json::dumps(request.input.as_ref().expect("input is present"))
+        .expect("fixture dumps");
+    assert_eq!(
+        dumped,
+        r#"{"text": "caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~", "n": -3, "ok": true, "no": false, "none": null, "list": [1, "a", {"z": []}], "empty": {}}"#
+    );
+}
+
 #[test]
 fn tool_definitions_render_like_python() {
     let request = CountableRequest::parse(TOOLS_OPENAI.as_bytes()).expect("fixture parses");
@@ -85,7 +131,6 @@ fn union_types_and_anthropic_schema_render_like_python() {
 
 #[rstest]
 #[case::not_json(b"not json" as &[u8])]
-#[case::missing_model(br#"{"messages":[]}"#)]
 #[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)]
 #[case::message_with_tool_calls(
     br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#
@@ -107,7 +152,9 @@ fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) {
 }
 
 #[rstest]
-#[case::no_messages(br#"{"model":"m"}"# as &[u8])]
+#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])]
+#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)]
+#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)]
 #[case::image_block(
     br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"#
 )]

litellm-rust/crates/core/src/token_counter/types.rs

@@ -1,17 +1,30 @@
 use indexmap::IndexMap;
-use serde::Deserialize;
+use serde::{Deserialize, Deserializer};
 
 use super::TokenCountError;
 
-/// The parts of a request body `litellm.token_counter` reads when a host counts
-/// input tokens for budget checks. Anything outside this shape is declined so
-/// the host can fall back to its own counter instead of silently miscounting.
+/// The parts of a request body the host's budget counter reads. Chat and
+/// Anthropic Messages bodies carry `messages`; completions carry `prompt`;
+/// Responses and embeddings carry `input`; rerank carries `query` and
+/// `documents`. The host checks key presence, not nullness, so an explicit
+/// `null` is kept distinct from an absent key. Anything outside this shape is
+/// declined so the host can fall back to its own counter instead of silently
+/// miscounting.
 #[derive(Clone, Debug, Deserialize, PartialEq)]
 pub struct CountableRequest {
-    pub model: String,
+    pub model: Option<String>,
+    #[serde(default, deserialize_with = "present_messages")]
     pub messages: Option<Vec<Message>>,
     pub tools: Option<Vec<ToolDefinition>>,
     pub tool_choice: Option<ToolChoice>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub prompt: Option<TextValue>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub input: Option<TextValue>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub query: Option<TextValue>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub documents: Option<TextValue>,
 }
 
 impl CountableRequest {
@@ -21,6 +34,32 @@ impl CountableRequest {
     }
 }
 
+fn present_messages<'de, D: Deserializer<'de>>(
+    deserializer: D,
+) -> Result<Option<Vec<Message>>, D::Error> {
+    Option::<Vec<Message>>::deserialize(deserializer)
+        .map(|messages| Some(messages.unwrap_or_default()))
+}
+
+fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<TextValue>, D::Error> {
+    TextValue::deserialize(deserializer).map(Some)
+}
+
+/// Free-form JSON the host counts as text: strings and integers via `str()`,
+/// objects via `json.dumps()`, lists flattened. Objects keep document order so
+/// the dumped text matches Python byte for byte.
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(untagged)]
+pub enum TextValue {
+    Null,
+    Bool(bool),
+    Integer(i64),
+    Float(f64),
+    Text(String),
+    List(Vec<TextValue>),
+    Object(IndexMap<String, TextValue>),
+}
+
 /// Python counts every string-valued key of a message, so any key beyond these
 /// makes the shape unsupported rather than silently uncounted.
 #[derive(Clone, Debug, Deserialize, PartialEq)]

litellm/rust_bridge/token_counter.py

@@ -28,7 +28,7 @@ class RustTokenCounterFactory(Protocol):
 
 @dataclass(frozen=True, slots=True)
 class InputTokenCount:
-    model: str
+    model: str | None
     input_tokens: int
 
 
@@ -49,7 +49,7 @@ TOKEN_COUNTER: Final = NativeBinding("TokenCounter", validate=_as_factory)
 
 
 def uses_anthropic_tokenizer(model: str) -> bool:
-    if litellm.disable_hf_tokenizer_download is True:
+    if litellm.disable_token_counter is True or litellm.disable_hf_tokenizer_download is True:
         return False
     return model in litellm.anthropic_models and "claude-3" not in model
 

tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py

@@ -199,13 +199,26 @@ def rust_counter(monkeypatch: pytest.MonkeyPatch):
 
 
 @pytest.mark.asyncio
-async def test_rust_count_replaces_python_tokenizing_for_anthropic_models(rust_counter: None) -> None:
+@pytest.mark.parametrize(
+    ("route", "request_body"),
+    (
+        ("/v1/messages", RUST_COUNTED_BODY),
+        ("/v1/chat/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "messages": ANTHROPIC_MESSAGES}),
+        ("/v1/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "prompt": "hi"}),
+        ("/v1/responses", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": "hi"}),
+        ("/v1/embeddings", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": ["hi"]}),
+        ("/v1/rerank", {"model": ANTHROPIC_TOKENIZER_MODEL, "query": "hi", "documents": ["a"]}),
+    ),
+)
+async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
+    rust_counter: None, route: str, request_body: dict
+) -> None:
     litellm.rust(True)
     rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
-    raw_body: Final = json.dumps(RUST_COUNTED_BODY).encode()
+    raw_body: Final = json.dumps(request_body).encode()
 
     counts: Final = await count_request_input_tokens(
-        request_body=RUST_COUNTED_BODY, route="/v1/messages", llm_router=None, raw_body=raw_body
+        request_body=request_body, route=route, llm_router=None, raw_body=raw_body
     )
 
     assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS}

tests/test_litellm/rust_bridge/test_token_counter.py

@@ -13,6 +13,7 @@ from typing import Final
 import pytest
 
 import litellm
+from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens
 from litellm.rust_bridge import bindings, configuration
 from litellm.rust_bridge import token_counter as bridge
 
@@ -140,8 +141,9 @@ def test_uses_anthropic_tokenizer_mirrors_python_tokenizer_selection(model: str,
     assert bridge.uses_anthropic_tokenizer(model) is expected
 
 
-def test_uses_anthropic_tokenizer_respects_hf_download_opt_out(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
+@pytest.mark.parametrize("flag", ("disable_hf_tokenizer_download", "disable_token_counter"))
+def test_uses_anthropic_tokenizer_respects_python_opt_outs(monkeypatch: pytest.MonkeyPatch, flag: str) -> None:
+    monkeypatch.setattr(litellm, flag, True)
 
     assert bridge.uses_anthropic_tokenizer(MODEL) is False
 
@@ -182,12 +184,27 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
         "model": MODEL,
         "messages": [{"role": "user", "content": "x " * 20_000}],
     },
+    {"model": MODEL, "prompt": "Write a haiku about ships.", "max_tokens": 20},
+    {"model": MODEL, "prompt": ["first prompt", "second prompt"]},
+    {
+        "model": MODEL,
+        "instructions": "be terse",
+        "input": [
+            {"role": "user", "content": [{"type": "input_text", "text": "Summarise caf\u00e9 menus \u2014 \"ok\"?\n"}]},
+            {"role": "assistant", "content": "Sure."},
+        ],
+    },
+    {"model": MODEL, "input": "a single embedding string"},
+    {"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"},
+    {"model": MODEL, "query": "best harbour", "documents": ["doc one", {"text": "doc two", "title": "T", "n": 3}]},
+    {"model": MODEL, "messages": None, "prompt": "messages key wins even when null"},
+    {"prompt": "model comes from the route"},
 )
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("request_body", PARITY_REQUESTS)
-async def test_native_count_matches_python_token_counter(
+async def test_native_count_matches_python_budget_counter(
     monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
 ) -> None:
     native: Final = pytest.importorskip("litellm.rust_bridge._native")
@@ -195,30 +212,31 @@ async def test_native_count_matches_python_token_counter(
     litellm.rust(True)
 
     rust_count: Final = await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode())
-    python_count: Final = litellm.token_counter(
-        model=MODEL,
-        messages=request_body["messages"],
-        tools=request_body.get("tools"),
-        tool_choice=request_body.get("tool_choice"),
-    )
+    python_count: Final = _count_input_tokens(request_body=request_body, model=MODEL)
 
     assert rust_count is not None
-    assert rust_count.model == MODEL
+    assert rust_count.model == request_body.get("model")
     assert rust_count.input_tokens == python_count
 
 
+DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
+    {
+        "model": MODEL,
+        "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}],
+    },
+    {"model": MODEL, "prompt": 1.5},
+    {"model": MODEL, "documents": [{"score": 0.5}]},
+    {"model": MODEL, "file": "audio.mp3"},
+)
+
+
 @pytest.mark.asyncio
-async def test_native_declines_image_content(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.parametrize("request_body", DECLINED_REQUESTS)
+async def test_native_declines_shapes_python_prices_differently(
+    monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
+) -> None:
     native: Final = pytest.importorskip("litellm.rust_bridge._native")
     monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
     litellm.rust(True)
-    body: Final = json.dumps(
-        {
-            "model": MODEL,
-            "messages": [
-                {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}
-            ],
-        }
-    ).encode()
 
-    assert await bridge.count_anthropic_input_tokens(body) is None
+    assert await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) is None