Merge remote-tracking branch 'origin/main' into litellm_any_sweep_20261017

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -3,6 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
 """
 
 import asyncio
+import json
 from collections.abc import Mapping, Sequence
 from dataclasses import replace as dataclasses_replace
 from datetime import datetime, timedelta, timezone
@@ -227,7 +228,7 @@ class CheckBatchCost:
             "so it will no longer be polled"
         )
 
-    async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
+    async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
         """
         Atomically flip batch_processed from false to true, returning whether this pod won
         the row. Every pod and uvicorn worker schedules its own poller against the shared
@@ -241,10 +242,16 @@ class CheckBatchCost:
         """
         if not self._has_batch_processed_column:
             return True
+        from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json
+
         try:
             claimed: Final = await _managed_object_table(self.prisma_client).update_many(
                 where={"id": job.id, "batch_processed": False},
-                data={"batch_processed": True},
+                data={
+                    "batch_processed": True,
+                    "status": response.status if response.status != "completed" else "complete",
+                    "file_object": postgres_storable_json(response),
+                },
             )
         except Exception as db_err:
             verbose_proxy_logger.error(
@@ -253,7 +260,9 @@ class CheckBatchCost:
             return False
         return claimed > 0
 
-    async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
+    async def _release_job_claim(
+        self, job: "ManagedObjectRow", original_status: str | None, original_file_object: object
+    ) -> None:
         """Give a claimed row back once billing it failed, so a later poll cycle retries it.
 
         Safe to match on batch_processed=True: while this poller is active the retrieve
@@ -265,7 +274,11 @@ class CheckBatchCost:
         try:
             await _managed_object_table(self.prisma_client).update_many(
                 where={"id": job.id, "batch_processed": True},
-                data={"batch_processed": False},
+                data={
+                    "batch_processed": False,
+                    "status": original_status,
+                    "file_object": json.dumps(original_file_object),
+                },
             )
         except Exception as db_err:
             verbose_proxy_logger.error(
@@ -564,8 +577,6 @@ class CheckBatchCost:
 
     @staticmethod
     def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
-        import json
-
         from litellm.types.utils import LiteLLMBatch
 
         file_object: Final = job.file_object
@@ -718,8 +729,7 @@ class CheckBatchCost:
                         )
                     except Exception as _e:
                         verbose_proxy_logger.warning(
-                            f"CheckBatchCost: failed to create managed file ID for "
-                            f"{_file_attr}={_raw_file_id!r}: {_e}"
+                            f"CheckBatchCost: failed to create managed file ID for {_file_attr}={_raw_file_id!r}: {_e}"
                         )
 
         # Pass the deployment's router-registered pricing (litellm_params custom
@@ -788,7 +798,9 @@ class CheckBatchCost:
             custom_llm_provider=str(llm_provider) if llm_provider else None,
         )
 
-        if not await self._claim_job_for_costing(job):
+        original_status: Final = job.status
+        original_file_object: Final = job.file_object
+        if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})):
             verbose_proxy_logger.info(
                 f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
                 "in this window, so its cost is already being tracked there"
@@ -807,7 +819,7 @@ class CheckBatchCost:
                 batch_completion_cost=batch_result.completion_cost,
             )
         except Exception:
-            await self._release_job_claim(job)
+            await self._release_job_claim(job, original_status, original_file_object)
             raise
 
         # Record batch duration (completed_at - created_at)
@@ -1100,10 +1112,7 @@ class CheckBatchCost:
                 continue
 
             ## RETRIEVE THE BATCH JOB OUTPUT FILE
-            if (
-                response.status in PROVIDER_TERMINAL_BATCH_STATUSES
-                and response.output_file_id is not None
-            ):
+            if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None:
                 try:
                     tracked = await self._track_completed_batch_cost(
                         job=job,
@@ -1134,24 +1143,19 @@ class CheckBatchCost:
                 # Track this job for the final metrics summary
                 processed_models.append(tracked)
 
-                # mark the job as complete
-                try:
-                    update_data: dict = {
-                        "status": response.status if response.status != "completed" else "complete",
-                        "file_object": response.model_dump_json(),
-                    }
-                    if self._has_batch_processed_column:
-                        update_data["batch_processed"] = True
-                    await _managed_object_table(self.prisma_client).update(
-                        where={"id": job.id},
-                        data=update_data,
-                    )
-                except Exception as db_err:
-                    verbose_proxy_logger.error(
-                        f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
-                    )
-                    if not self._has_batch_processed_column:
-                        # nothing claimed the row, so it stays in the poll page and its lines wait for that retry
+                if not self._has_batch_processed_column:
+                    try:
+                        await _managed_object_table(self.prisma_client).update(
+                            where={"id": job.id},
+                            data={
+                                "status": response.status if response.status != "completed" else "complete",
+                                "file_object": response.model_dump_json(),
+                            },
+                        )
+                    except Exception as db_err:
+                        verbose_proxy_logger.error(
+                            f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
+                        )
                         continue
                 self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
 

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -416,7 +416,7 @@ def _storable_json_value(value: JsonValue) -> JsonValue:
     return value
 
 
-def _postgres_storable_json(file_object: BaseModel) -> str:
+def postgres_storable_json(file_object: BaseModel) -> str:
     dumped: Final = file_object.model_dump_json()
     if "\\u0000" not in dumped:
         return dumped
@@ -650,7 +650,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         )
         # FIX: Update status and file_object on every operation to keep state in sync
         update_columns: Final = {
-            "file_object": _postgres_storable_json(stored_object),
+            "file_object": postgres_storable_json(stored_object),
             "status": file_object.status,
             "updated_by": user_api_key_dict.user_id,
         }
@@ -666,7 +666,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             data={
                 "create": {
                     "unified_object_id": unified_object_id,
-                    "file_object": _postgres_storable_json(stored_object),
+                    "file_object": postgres_storable_json(stored_object),
                     "model_object_id": model_object_id,
                     "file_purpose": file_purpose,
                     "created_by": resolve_resource_owner_id(user_api_key_dict),
@@ -702,7 +702,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     "unified_object_id": response_id,
                     "model_object_id": decrypted_response_id(response_id),
                     "file_purpose": "response",
-                    "file_object": _postgres_storable_json(_without_model_output(response)),
+                    "file_object": postgres_storable_json(_without_model_output(response)),
                     "status": PENDING_RESPONSE_STATUSES[0],
                     "created_by": resolve_resource_owner_id(user_api_key_dict),
                     "team_id": user_api_key_dict.team_id,

litellm/batches/batch_utils.py

@@ -794,11 +794,17 @@ def _batch_response_was_successful(
     message batch results lines report ``result.type == "succeeded"``; Bedrock
     batch output lines report ``modelOutput`` (and no ``error``).
     """
-    if custom_llm_provider == "anthropic":
-        return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded"
     if custom_llm_provider == "bedrock":
         return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None
-    _response: Final[dict[str, object]] = cast(  # cast-ok: batch response bodies come from provider JSON
-        dict[str, object], batch_job_output_file.get("response", None) or {}
+    envelope_key: Final = "result" if custom_llm_provider == "anthropic" else "response"
+    envelope: Final = batch_job_output_file.get(envelope_key)
+    if not isinstance(envelope, Mapping):
+        if envelope is not None:
+            verbose_logger.warning("batch output %s must be an object, treating the line as unsuccessful", envelope_key)
+        return False
+    response: Final = cast(Mapping[object, object], envelope)  # cast-ok: mapping shape validated above
+    return (
+        response.get("type") == "succeeded"
+        if custom_llm_provider == "anthropic"
+        else response.get("status_code") == 200
     )
-    return _response.get("status_code", None) == 200

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -6166,3 +6166,25 @@ def resolve_structured_messages(
         if structured:
             return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured]
     return None
+
+
+def is_flat_responses_tool(tool: Mapping[str, object]) -> bool:
+    tool_type: Final = tool.get("type")
+    return tool_type in ("function", "custom") and tool_type not in tool and "input_schema" not in tool
+
+
+def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
+    from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
+    from litellm.responses.litellm_completion_transformation.transformation import (
+        LiteLLMCompletionResponsesConfig,
+    )
+
+    if is_flat_responses_tool(tool):
+        chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
+            [dict(tool)]  # pyright: ignore[reportArgumentType]  # a JSON dict in a flat Responses tool shape
+        )
+        return tuple(chat_tools)
+    anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(
+        [dict(tool)]  # pyright: ignore[reportArgumentType]  # a JSON dict in an Anthropic or chat tool shape
+    )
+    return tuple(anthropic_tools)

litellm/llms/baseten/chat.py

@@ -54,6 +54,7 @@ class BasetenConfig(OpenAIGPTConfig):
         return [
             "max_tokens",
             "max_completion_tokens",
+            "reasoning_effort",
             "response_format",
             "seed",
             "stop",

litellm/llms/bedrock/count_tokens/handler.py

@@ -48,7 +48,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
             verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model)
 
             # Get AWS region using existing LiteLLM function
-            aws_region_name: Final = self._get_aws_region_name(
+            aws_region_name: Final = self.get_aws_region_name(
                 optional_params=litellm_params,
                 model=resolved_model,
                 model_id=None,

litellm/llms/bedrock/count_tokens/kimi_k3.py

@@ -0,0 +1,502 @@
+"""
+Exact Bedrock token counts for Kimi K3, which has no CountTokens support.
+
+Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the
+official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256.
+"""
+
+import asyncio
+import base64
+import copy
+import hashlib
+import json
+import threading
+import time
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from concurrent.futures import Future
+from dataclasses import dataclass
+from enum import Enum
+from itertools import chain, groupby
+from pathlib import Path
+from typing import Annotated, Final, Literal, TypeAlias
+
+import regex
+import tiktoken
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator
+from typing_extensions import Self
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.prompt_templates.factory import (
+    chat_completion_tools_for,
+    is_flat_responses_tool,
+)
+from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
+
+KIMI_K3_REPO: Final = "moonshotai/Kimi-K3"
+KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06"
+KIMI_K3_RANKS_FILE: Final = "tiktoken.model"
+KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103"
+_PATTERN: Final = "|".join(
+    (
+        r"[\p{Han}]+",
+        r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
+        r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
+        r"\p{N}{1,3}",
+        r" ?[^\s\p{L}\p{N}]+[\r\n]*",
+        r"\s*[\r\n]+",
+        r"\s+(?!\S)",
+        r"\s+",
+    )
+)
+# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE
+# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece.
+_CHUNK: Final = regex.compile(
+    r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])"
+    r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*"
+)
+
+
+@dataclass(frozen=True, slots=True)
+class KimiK3TokenizerSource:
+    ranks_path: Callable[[], str]
+    sha256: str
+
+
+def _hugging_face_ranks_path() -> str:
+    import huggingface_hub
+
+    return huggingface_hub.hf_hub_download(
+        repo_id=KIMI_K3_REPO,
+        filename=KIMI_K3_RANKS_FILE,
+        revision=KIMI_K3_REVISION,
+        local_files_only=litellm.disable_hf_tokenizer_download is True,
+    )
+
+
+HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256)
+
+
+def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str:
+    try:
+        data: Final = Path(source.ranks_path()).read_bytes()
+    except Exception as e:  # noqa: BLE001  # a download or read failure of any kind falls back to the estimate
+        return type(e).__name__
+    if hashlib.sha256(data).hexdigest() != source.sha256:
+        return "SHA-256 mismatch"
+    rows: Final = (line.split() for line in data.splitlines() if line)
+    return tiktoken.Encoding(
+        name="kimi_k3",
+        pat_str=_PATTERN,
+        mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows},
+        special_tokens={},
+    )
+
+
+def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None:
+    loaded: Final = _loaded(source)
+    if isinstance(loaded, tiktoken.Encoding):
+        return loaded
+    verbose_logger.warning(
+        "Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. "
+        "Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
+        loaded,
+        KIMI_K3_RANKS_FILE,
+        KIMI_K3_REPO,
+        KIMI_K3_REVISION,
+    )
+    return None
+
+
+def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]":
+    encoding: Final = Future[tiktoken.Encoding | None]()
+    threading.Thread(
+        target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True
+    ).start()
+    return encoding
+
+
+@dataclass(frozen=True, slots=True)
+class _EncodingLoad:
+    encoding: "Future[tiktoken.Encoding | None]"
+    started_at: float
+    is_retry: bool
+
+    def retry_due(self, now: float, retry_after_seconds: float) -> bool:
+        return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds
+
+
+@dataclass(slots=True)
+class KimiK3Encodings:
+    source: KimiK3TokenizerSource
+    retry_after_seconds: float
+    clock: Callable[[], float] = time.monotonic
+    current: _EncodingLoad | None = None
+
+    def _load(self) -> _EncodingLoad:
+        now: Final = self.clock()
+        current: Final = self.current
+        if current is not None and not current.retry_due(now, self.retry_after_seconds):
+            return current
+        load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None)
+        self.current = load
+        return load
+
+    async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None:
+        load: Final = self._load()
+        remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock()
+        if load.encoding.done() or remaining <= 0:
+            return load.encoding.result() if load.encoding.done() else None
+        await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining)
+        if load.encoding.done():
+            return load.encoding.result()
+        verbose_logger.warning(
+            "Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. "
+            "Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
+            wait_seconds,
+            KIMI_K3_RANKS_FILE,
+            KIMI_K3_REPO,
+            KIMI_K3_REVISION,
+        )
+        return None
+
+
+class _Converse(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+
+class _Text(_Converse):
+    text: str
+
+
+class _CachePoint(_Converse):
+    cache_point: JsonValue = Field(alias="cachePoint")
+
+
+class _ToolUse(_Converse):
+    tool_use_id: str = Field(alias="toolUseId")
+    name: str
+    arguments: Mapping[str, JsonValue] = Field(alias="input")
+
+
+class _ToolUseBlock(_Converse):
+    tool_use: _ToolUse = Field(alias="toolUse")
+
+
+class _ToolResult(_Converse):
+    tool_use_id: str = Field(alias="toolUseId")
+    content: tuple[_Text, ...]
+    status: str | None = None
+
+
+class _ToolResultBlock(_Converse):
+    tool_result: _ToolResult = Field(alias="toolResult")
+
+
+class _ReasoningText(_Converse):
+    text: str
+    signature: str | None = None
+
+
+class _Reasoning(_Converse):
+    reasoning_text: _ReasoningText = Field(alias="reasoningText")
+
+
+class _ReasoningBlock(_Converse):
+    reasoning_content: _Reasoning = Field(alias="reasoningContent")
+
+
+class _UserMessage(_Converse):
+    role: Literal["user"]
+    content: tuple[_Text | _ToolResultBlock | _CachePoint, ...]
+
+
+class _AssistantMessage(_Converse):
+    role: Literal["assistant"]
+    content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...]
+
+    def calls(self) -> tuple[_ToolUse, ...]:
+        return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock))
+
+
+class _InputSchema(_Converse):
+    definition: JsonValue = Field(alias="json")
+
+
+class _ToolSpec(_Converse):
+    name: str
+    description: str | None = None
+    input_schema: _InputSchema = Field(alias="inputSchema")
+
+    def declared(self) -> Mapping[str, JsonValue]:
+        described: Final[Mapping[str, JsonValue]] = (
+            {} if self.description is None else {"description": self.description}
+        )
+        return {"name": self.name, **described, "parameters": self.input_schema.definition}
+
+
+class _ToolSpecBlock(_Converse):
+    tool_spec: _ToolSpec = Field(alias="toolSpec")
+
+
+class _ToolConfig(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    tools: tuple[_ToolSpecBlock | _CachePoint, ...]
+
+
+_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")]
+
+
+def _result_ids(message: _Message) -> frozenset[str]:
+    if isinstance(message, _AssistantMessage):
+        return frozenset()
+    return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock))
+
+
+def _call_ids(message: _Message | None) -> frozenset[str]:
+    if not isinstance(message, _AssistantMessage):
+        return frozenset()
+    return frozenset(call.tool_use_id for call in message.calls())
+
+
+class _ConverseRequest(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    messages: tuple[_Message, ...]
+    system: tuple[_Text | _CachePoint, ...] = ()
+    tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig")
+
+    @model_validator(mode="after")
+    def _results_answer_the_previous_calls(self) -> Self:
+        pairs: Final = zip((None, *self.messages), self.messages, strict=False)
+        if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs):
+            raise ValueError("a toolResult does not answer a toolUse of the previous assistant message")
+        return self
+
+
+class _Control(Enum):
+    OPEN = "<|open|>"
+    CLOSE = "<|close|>"
+    SEP = "<|sep|>"
+    END_OF_MSG = "<|end_of_msg|>"
+
+
+_Segment: TypeAlias = str | _Control
+_Segments: TypeAlias = tuple[_Segment, ...]
+
+
+def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments:
+    escaped: Final = (
+        (f" {key}", '="', value.replace("&", "&amp;").replace('"', "&quot;"), '"') for key, value in attrs
+    )
+    return (_Control.OPEN, tag, *chain.from_iterable(escaped), _Control.SEP)
+
+
+def _close(tag: str) -> _Segments:
+    return (_Control.CLOSE, tag, _Control.SEP)
+
+
+def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Segments:
+    return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG)
+
+
+def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments:
+    tools: Final = json.dumps(
+        [{"type": "function", "function": spec.declared()} for spec in specs],
+        ensure_ascii=False,
+        separators=(",", ":"),
+        sort_keys=True,
+    )
+    return _message(
+        (("role", "system"), ("type", "tool-declare")),
+        (f"# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n{tools}\n```",),
+    )
+
+
+def _text(blocks: Iterable[_Converse]) -> str:
+    return "".join(block.text for block in blocks if isinstance(block, _Text))
+
+
+def _json_type(value: JsonValue) -> str:
+    match value:
+        case bool():
+            return "boolean"
+        case None:
+            return "null"
+        case int() | float():
+            return "number"
+        case str():
+            return "string"
+        case dict():
+            return "object"
+        case _:
+            return "array"
+
+
+def _argument(key: str, value: JsonValue) -> _Segments:
+    text: Final = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
+    return (*_open("argument", (("key", key), ("type", _json_type(value)))), text, *_close("argument"))
+
+
+def _call(index: int, call: _ToolUse) -> _Segments:
+    arguments: Final = chain.from_iterable(_argument(key, value) for key, value in call.arguments.items())
+    return (*_open("call", (("tool", call.name), ("index", str(index)))), *arguments, *_close("call"))
+
+
+def _tool_calls(calls: Sequence[_ToolUse]) -> _Segments:
+    if not calls:
+        return ()
+    rendered: Final = chain.from_iterable(_call(index, call) for index, call in enumerate(calls, start=1))
+    return (*_open("tools"), *rendered, *_close("tools"))
+
+
+def _assistant_message(message: _AssistantMessage) -> _Segments:
+    reasoning: Final = "".join(
+        block.reasoning_content.reasoning_text.text for block in message.content if isinstance(block, _ReasoningBlock)
+    )
+    return _message(
+        (("role", "assistant"),),
+        (
+            *_open("think"),
+            reasoning,
+            *_close("think"),
+            *_open("response"),
+            _text(message.content),
+            *_close("response"),
+            *_tool_calls(message.calls()),
+        ),
+    )
+
+
+def _tool_messages(results: Sequence[_ToolResult], previous: _AssistantMessage) -> _Segments:
+    names: Final = {call.tool_use_id: call.name for call in previous.calls()}
+    return tuple(
+        chain.from_iterable(
+            _message(
+                (("role", "tool"), ("tool", names[result.tool_use_id]), ("index", str(index))),
+                (_text(result.content),),
+            )
+            for index, result in enumerate(results, start=1)
+        )
+    )
+
+
+def _user_run(blocks: Sequence[_Text | _ToolResultBlock], previous: _Message | None) -> _Segments:
+    results: Final = tuple(block.tool_result for block in blocks if isinstance(block, _ToolResultBlock))
+    if results and isinstance(previous, _AssistantMessage):
+        return _tool_messages(results, previous)
+    return _message((("role", "user"),), (_text(blocks),))
+
+
+def _user_messages(message: _UserMessage, previous: _Message | None) -> _Segments:
+    blocks: Final = tuple(block for block in message.content if not isinstance(block, _CachePoint))
+    runs: Final = groupby(blocks, key=lambda block: isinstance(block, _ToolResultBlock))
+    return tuple(chain.from_iterable(_user_run(tuple(run), previous) for _, run in runs))
+
+
+def _turn(previous: _Message | None, message: _Message) -> _Segments:
+    if isinstance(message, _AssistantMessage):
+        return _assistant_message(message)
+    return _user_messages(message, previous)
+
+
+def _rendered(request: _ConverseRequest) -> _Segments:
+    specs: Final = (
+        ()
+        if request.tool_config is None
+        else tuple(block.tool_spec for block in request.tool_config.tools if isinstance(block, _ToolSpecBlock))
+    )
+    system: Final = _text(request.system)
+    turns: Final = chain.from_iterable(
+        _turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False)
+    )
+    return (
+        *(_tool_declare(specs) if specs else ()),
+        *(_message((("role", "system"),), (system,)) if system else ()),
+        *turns,
+        *_open("message", (("role", "assistant"),)),
+        *_open("think"),
+    )
+
+
+def _text_tokens(encoding: tiktoken.Encoding, text: str) -> int:
+    spans: Final = (match.span() for match in _CHUNK.finditer(text))
+    return sum(len(encoding.encode_ordinary(text[start:end])) for start, end in spans)
+
+
+class _TextPart(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    type: Literal["text"]
+
+
+class _TextOnlyMessage(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    content: str | tuple[_TextPart, ...] | None = None
+
+
+def _text_only(message: AllMessageValues) -> bool:
+    try:
+        _TextOnlyMessage.model_validate(message)
+    except ValidationError:
+        return False
+    return True
+
+
+def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, object], ...]:
+    return chat_completion_tools_for(tool) if is_flat_responses_tool(tool) else (tool,)
+
+
+def _converse_request(
+    model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None
+) -> _ConverseRequest | None:
+    if not all(_text_only(message) for message in messages):
+        return None
+    try:
+        chat_tools: Final = tuple(
+            chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool)
+        )  # comprehension-ok: flattens each tool's translation
+        return _ConverseRequest.model_validate(
+            litellm.AmazonConverseConfig().transform_request(
+                model=model,
+                messages=copy.deepcopy(list(messages)),
+                optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {},
+                litellm_params={},
+                headers={},
+            )
+        )
+    except Exception:  # noqa: BLE001  # a body litellm cannot send to Bedrock, or cannot render, falls back to the estimate
+        return None
+
+
+def kimi_k3_tokens(
+    encoding: tiktoken.Encoding,
+    model: str,
+    messages: Sequence[AllMessageValues],
+    tools: Sequence[ChatCompletionToolParam] | None,
+) -> int | None:
+    request: Final = _converse_request(model, messages, tools)
+    if request is None:
+        return None
+    return sum(
+        1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request)
+    )
+
+
+def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Mapping[str, object]) -> int | None:
+    body: Final = {**fields, "model": model, "max_tokens": 1}
+    try:
+        translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
+            copy.deepcopy(body),  # pyright: ignore[reportArgumentType]  # raw /v1/messages fields
+            custom_llm_provider="bedrock",
+        )
+    except Exception:  # noqa: BLE001  # a request the adapter cannot translate falls back to the estimate
+        return None
+    match translated:
+        case {"tools": tools}:
+            return kimi_k3_tokens(encoding, model, translated["messages"], tools)
+        case _:
+            return kimi_k3_tokens(encoding, model, translated["messages"], None)
+
+
+def kimi_k3_template_tokens(message_count: int, has_tools: bool) -> int:
+    return 14 + 13 * max(message_count - 1, 0) + (40 if has_tools else 0)

litellm/llms/bedrock/count_tokens/mantle_handler.py

@@ -52,7 +52,7 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
     ) -> dict[str, JsonValue]:
         try:
             request: Final = _COUNT_REQUEST.validate_python(request_data)
-            aws_region_name: Final = self._get_aws_region_name(
+            aws_region_name: Final = self.get_aws_region_name(
                 optional_params=litellm_params, model=resolved_model, model_id=None
             )
             endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name)

litellm/llms/openai/responses/count_tokens/token_counter.py

@@ -124,6 +124,7 @@ class OpenAITokenCounter(BaseTokenCounter):
                 error=True,
                 error_message=request_error,
                 status_code=400,
+                rejected_before_provider_call=True,
             )
 
         deployment = deployment or {}

litellm/llms/vertex_ai/count_tokens/context.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal
 from pydantic import JsonValue, TypeAdapter, ValidationError
 from typing_extensions import assert_never
 
+from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for
 from litellm.types.llms.anthropic import AnthropicMessagesTextParam
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.llms.vertex_ai import ContentType
@@ -75,27 +76,6 @@ def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool:
     return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool
 
 
-def _is_responses_function_tool(tool: Mapping[str, JsonValue]) -> bool:
-    return tool.get("type") == "function" and "function" not in tool and "input_schema" not in tool
-
-
-def _chat_tool(tool: Mapping[str, JsonValue]) -> tuple[Mapping[str, object], ...]:
-    from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
-    from litellm.responses.litellm_completion_transformation.transformation import (
-        LiteLLMCompletionResponsesConfig,
-    )
-
-    if _is_responses_function_tool(tool):
-        chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
-            [dict(tool)]  # pyright: ignore[reportArgumentType]  # validated JSON dict in the Responses function tool shape
-        )
-        return tuple(chat_tools)
-    anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(
-        [dict(tool)]  # pyright: ignore[reportArgumentType]  # validated JSON dict in an Anthropic or chat tool shape
-    )
-    return tuple(anthropic_tools)
-
-
 def _map_function_tools(
     tools: object | None,
 ) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError:
@@ -112,7 +92,7 @@ def _map_function_tools(
     from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
 
     chat_tools: Final = tuple(
-        chat_tool for tool in function_tools for chat_tool in _chat_tool(tool)
+        chat_tool for tool in function_tools for chat_tool in chat_completion_tools_for(tool)
     )  # comprehension-ok: flattens each tool's translation
     try:
         mapped_tools: Final = VertexGeminiConfig().map_function(
@@ -163,6 +143,7 @@ def count_context_error_response(
         error=True,
         error_message=error.error_message,
         status_code=error.status_code,
+        rejected_before_provider_call=True,
     )
 
 

litellm/model_prices_and_context_window_backup.json

@@ -3028,7 +3028,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -3313,7 +3313,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -21144,6 +21144,7 @@
         "input_cost_per_token": 1.4e-07,
         "input_dbu_cost_per_token": 2e-06,
         "litellm_provider": "databricks",
+        "token_count_framing": 97,
         "max_input_tokens": 1000000,
         "max_output_tokens": 393216,
         "max_tokens": 393216,
@@ -21166,6 +21167,11 @@
         "supports_tool_choice": true,
         "supports_vision": false
     },
+    "databricks/system.ai.deepseek-v4-flash-0731": {
+        "litellm_provider": "databricks",
+        "mode": "chat",
+        "token_count_framing": 97
+    },
     "databricks/databricks-deepseek-v4-pro-0813": {
         "cache_creation_input_token_cost": 1.31999e-06,
         "cache_read_input_token_cost": 1.3202e-07,
@@ -22134,6 +22140,7 @@
         "input_cost_per_token": 2.99999e-06,
         "input_dbu_cost_per_token": 4.2857e-05,
         "litellm_provider": "databricks",
+        "token_count_framing": 100,
         "max_input_tokens": 1000000,
         "max_output_tokens": 1048576,
         "max_tokens": 1048576,
@@ -22157,6 +22164,11 @@
         "supports_tool_choice": true,
         "supports_vision": true
     },
+    "databricks/system.ai.kimi-k3": {
+        "litellm_provider": "databricks",
+        "mode": "chat",
+        "token_count_framing": 100
+    },
     "databricks/databricks-llama-4-maverick": {
         "cache_creation_input_token_cost": 5.0001e-07,
         "cache_read_input_token_cost": 5.0001e-07,
@@ -24695,7 +24707,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -33901,7 +33913,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -48928,7 +48940,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -71540,6 +71552,8 @@
         "input_cost_per_token_flex": 1.65e-06,
         "input_cost_per_token_priority": 5.775e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -71570,6 +71584,8 @@
         "input_cost_per_token_flex": 1.5e-06,
         "input_cost_per_token_priority": 5.25e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -71600,6 +71616,8 @@
         "input_cost_per_token_flex": 1.65e-06,
         "input_cost_per_token_priority": 5.775e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -73100,6 +73118,10 @@
         "cache_read_input_token_cost": 1e-07,
         "input_cost_per_token": 2e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -73173,6 +73195,10 @@
         "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
     },
     "au.anthropic.claude-sonnet-5-5": {
+        "token_count_model": {
+            "model": "au.anthropic.claude-opus-5-5",
+            "region": "ap-southeast-4"
+        },
         "bedrock_converse_supports_strict_tools": false,
         "cache_creation_input_token_cost": 2.75e-06,
         "cache_creation_input_token_cost_above_1hr": 4.4e-06,
@@ -73373,6 +73399,10 @@
         "cache_read_input_token_cost": 1e-07,
         "input_cost_per_token": 2e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "global.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -73582,6 +73612,10 @@
         "cache_read_input_token_cost": 1.1e-07,
         "input_cost_per_token": 2.2e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "us.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -74382,6 +74416,8 @@
         "input_cost_per_token_flex": 1.65e-06,
         "input_cost_per_token_priority": 5.775e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -74555,6 +74591,10 @@
         "cache_read_input_token_cost": 1e-08,
         "input_cost_per_token": 1e-07,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -74633,6 +74673,10 @@
         "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
     },
     "au.anthropic.claude-haiku-5-5": {
+        "token_count_model": {
+            "model": "au.anthropic.claude-opus-5-5",
+            "region": "ap-southeast-4"
+        },
         "input_cost_per_token_above_100k_tokens": 5.5e-07,
         "output_cost_per_token_above_100k_tokens": 2.75e-06,
         "cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
@@ -74863,6 +74907,10 @@
         "cache_read_input_token_cost": 1e-08,
         "input_cost_per_token": 1e-07,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "global.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -75010,6 +75058,10 @@
         "cache_read_input_token_cost": 1.1e-08,
         "input_cost_per_token": 1.1e-07,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "us.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -48,6 +48,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
     openai_error_type,
     with_litellm_call_id,
 )
+from litellm.proxy.common_utils.token_count_settings import count_source_fields
 from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli
 from litellm.router import Router
 
@@ -420,7 +421,7 @@ async def count_tokens(
         log_llm_api_exception(e, litellm_call_id)
         return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
     else:
-        return {"input_tokens": token_response.total_tokens}
+        return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)}
 
 
 @router.post(

litellm/proxy/common_utils/token_count_settings.py

@@ -0,0 +1,164 @@
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, replace
+from types import MappingProxyType
+from typing import Final, Literal
+
+from pydantic import BaseModel, ConfigDict, ValidationError
+
+import litellm
+from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
+from litellm.litellm_core_utils.core_helpers import as_str_mapping
+from litellm.proxy._types import ProxyException
+from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer
+from litellm.types.utils import TokenCountResponse
+
+CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404})
+TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0
+TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0
+TOKENIZER_LOAD_MAX_RUNNING: Final = 4
+
+
+class _UpstreamErrorDetail(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="allow")
+    message: str
+
+
+class _UpstreamErrorEnvelope(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="allow")
+    type: Literal["error"]
+    error: _UpstreamErrorDetail
+
+
+def _joined(message: str, hint: str) -> str:
+    return f"{message.rstrip().rstrip('.')}. {hint}"
+
+
+def _with_hint(message: str, hint: str) -> str:
+    try:
+        envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message)
+    except ValidationError:
+        return _joined(AnthropicExceptionMapping.extract_error_message(message), hint)
+    detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)})
+    return envelope.model_copy(update={"error": detail}).model_dump_json()
+
+
+@dataclass(frozen=True, slots=True)
+class TokenCountModel:
+    model: str
+    region: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class TokenCountSettings:
+    framing: int | None
+    count_model: TokenCountModel | None
+    tokenizer: TokenCountTokenizer | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ProviderCountFailure:
+    status_code: int
+    message: str
+    rejected_before_provider_call: bool = False
+
+    @property
+    def cannot_count_model(self) -> bool:
+        return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES
+
+    def with_hint(self, hint: str) -> "ProviderCountFailure":
+        return replace(self, message=_with_hint(self.message, hint))
+
+    def to_proxy_exception(self) -> ProxyException:
+        return ProxyException(
+            message=self.message,
+            type="token_counting_error",
+            param="model",
+            code=self.status_code,
+        )
+
+
+def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]:
+    try:
+        return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
+    except Exception:  # noqa: BLE001  # get_model_info raises bare Exception for unmapped models
+        return MappingProxyType({})
+
+
+def _framing(value: object) -> int | None:
+    match value:
+        case bool():
+            return None
+        case int() if value >= 0:
+            return value
+        case float() if value.is_integer() and value >= 0:
+            return int(value)
+        case _:
+            return None
+
+
+def _non_blank(value: object) -> str | None:
+    return value if isinstance(value, str) and value.strip() else None
+
+
+def _count_model(value: Mapping[str, object]) -> TokenCountModel | None:
+    model: Final = _non_blank(value.get("model"))
+    return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region")))
+
+
+def _geography(region: str) -> str:
+    return region.rsplit("-", 2)[0]
+
+
+def _within_deployment_geography(
+    count_model: TokenCountModel | None, deployment_region: Callable[[], str]
+) -> TokenCountModel | None:
+    if count_model is None or count_model.region is None or count_model.model.startswith("global."):
+        return count_model
+    if _geography(deployment_region()) == _geography(count_model.region):
+        return count_model
+    return replace(count_model, region=None)
+
+
+def _resolved_count_model(
+    configured: object, default: object, deployment_region: Callable[[], str]
+) -> TokenCountModel | None:
+    configured_mapping: Final = as_str_mapping(configured)
+    if configured_mapping is not None:
+        return _count_model(configured_mapping)
+    default_mapping: Final = as_str_mapping(default)
+    return (
+        None
+        if default_mapping is None
+        else _within_deployment_geography(_count_model(default_mapping), deployment_region)
+    )
+
+
+def _tokenizer(value: object) -> TokenCountTokenizer | None:
+    match value:
+        case "kimi_k3":
+            return "kimi_k3"
+        case _:
+            return None
+
+
+def resolve_token_count_settings(
+    configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str]
+) -> TokenCountSettings:
+    framing: Final = _framing(configured.get("token_count_framing"))
+    return TokenCountSettings(
+        framing=framing if framing is not None else _framing(default.get("token_count_framing")),
+        count_model=_resolved_count_model(
+            configured.get("token_count_model"), default.get("token_count_model"), deployment_region
+        ),
+        tokenizer=_tokenizer(default.get("token_count_tokenizer")),
+    )
+
+
+def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields:
+    match response.count_source:
+        case None:
+            return {}
+        case "token_count_model":
+            return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used}
+        case source:
+            return {"litellm_count_source": source}

litellm/proxy/google_endpoints/endpoints.py

@@ -21,6 +21,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (  # noqa: F401  # leg
     _read_request_body,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     read_request_body,
 )
+from litellm.proxy.common_utils.token_count_settings import count_source_fields
 from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
 
 
@@ -309,10 +310,12 @@ async def google_count_tokens(request: Request, model_name: str):
             ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500),
         )
     original_response: Final[dict] = token_response.original_response or {}
-    return TokenCountDetailsResponse(
-        totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0,
-        promptTokensDetails=original_response.get("promptTokensDetails", []),
-    )
+    counted: Final[TokenCountDetailsResponse] = {
+        "totalTokens": original_response.get("totalTokens") or token_response.total_tokens or 0,
+        "promptTokensDetails": original_response.get("promptTokensDetails", []),
+        **count_source_fields(token_response),
+    }
+    return counted
 
 
 # ============================================================

litellm/proxy/proxy_server.py

@@ -55,6 +55,7 @@ from typing import (
 )
 
 import anyio
+import anyio.lowlevel
 import websockets
 import websockets.exceptions
 from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config
@@ -99,6 +100,13 @@ from litellm.litellm_core_utils.litellm_logging import (
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
 from litellm.litellm_core_utils.token_counter import offload_token_count
+from litellm.llms.bedrock.count_tokens.kimi_k3 import (
+    HUGGING_FACE_SOURCE,
+    KimiK3Encodings,
+    kimi_k3_anthropic_tokens,
+    kimi_k3_template_tokens,
+    kimi_k3_tokens,
+)
 from litellm.proxy._types import (
     UI_TEAM_ID,
     AwsErrorTypeProxyException,
@@ -149,6 +157,15 @@ from litellm.proxy.common_utils.realtime_utils import (  # noqa: F401, RUF100  #
     _realtime_request_body,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     realtime_request_body,
 )
+from litellm.proxy.common_utils.token_count_settings import (
+    TOKENIZER_LOAD_MAX_RUNNING,
+    TOKENIZER_LOAD_REUSE_SECONDS,
+    TOKENIZER_LOAD_WAIT_SECONDS,
+    ProviderCountFailure,
+    TokenCountSettings,
+    cost_map_entry,
+    resolve_token_count_settings,
+)
 from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
 from litellm.router_utils.access_windows import access_windows_config_error
 from litellm.router_utils.add_retry_fallback_headers import (
@@ -304,6 +321,7 @@ from litellm.constants import (
     CLI_SSO_SESSION_TTL_SECONDS,
     DAYS_IN_A_MONTH,
     DEFAULT_HEALTH_CHECK_INTERVAL,
+    DEFAULT_MAX_LRU_CACHE_SIZE,
     DEFAULT_MODEL_CREATED_AT_TIME,
     GLOBAL_PROXY_SPEND_CACHE_KEY,
     LITELLM_PROXY_ADMIN_NAME,
@@ -337,6 +355,7 @@ from litellm.litellm_core_utils.core_helpers import (
     drop_params_flag,
     get_litellm_metadata_from_kwargs,
     get_parent_otel_span_from_kwargs,
+    iter_items,
 )
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -993,7 +1012,13 @@ from litellm.types.secret_managers.main import (
     KeyManagementSettings,
     KeyManagementSystem,
 )
-from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload
+from litellm.types.utils import (
+    CredentialItem,
+    CustomHuggingfaceTokenizer,
+    RawRequestTypedDict,
+    SelectTokenizerResponse,
+    StandardLoggingPayload,
+)
 from litellm.types.utils import ModelInfo as ModelMapInfo
 from litellm.utils import add_custom_logger_callback_to_specific_event
 
@@ -14673,25 +14698,36 @@ async def _try_provider_token_count(
     provider_counter: "BaseTokenCounter",
     custom_llm_provider: str | None,
     model_to_use: str,
+    region: str | None,
     messages: list | None,
     contents: list | None,
-    deployment: dict[str, object] | None,
+    deployment: Mapping[str, object] | None,
     request_model: str,
     tools: list | None = None,
     system: str | None = None,
     messages_format: TokenCountMessagesFormat = "openai",
     output_format: JsonValue = None,
     context_management: CountContextManagement | None = None,
-) -> Optional["TokenCountResponse"]:
-    """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting."""
+) -> Union["TokenCountResponse", ProviderCountFailure, None]:
+    """Attempt provider-specific token counting. Returns None to fall through to local counting."""
     if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider):
         return None
+    deployment_in_region: Final = (
+        None
+        if deployment is None
+        else dict(deployment)
+        if region is None
+        else {
+            **deployment,
+            "litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), "aws_region_name": region},
+        }
+    )
     try:
         result: Final = await provider_counter.count_tokens(
             model_to_use=model_to_use or "",
             messages=messages,
             contents=contents,
-            deployment=deployment,
+            deployment=deployment_in_region,
             request_model=request_model,
             tools=tools,
             system=system,
@@ -14700,24 +14736,325 @@ async def _try_provider_token_count(
             context_management=context_management,
         )
     except httpx.HTTPStatusError as e:
-        error_message: Final = getattr(e, "message", None) or str(e)
-        status_code: Final = getattr(e, "status_code", None) or e.response.status_code
-        raise ProxyException(
-            message=error_message,
-            type="token_counting_error",
-            param="model",
-            code=status_code,
+        return ProviderCountFailure(
+            status_code=getattr(e, "status_code", None) or e.response.status_code,
+            message=getattr(e, "message", None) or str(e),
         )
     if result is not None and result.error is True:
-        raise ProxyException(
+        return ProviderCountFailure(
+            status_code=result.status_code or 500,
             message=result.error_message or "Token counting failed",
-            type="token_counting_error",
-            param="model",
-            code=result.status_code or 500,
+            rejected_before_provider_call=result.rejected_before_provider_call,
         )
     return result
 
 
+async def _provider_token_count(
+    count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]],
+    model_to_use: str,
+    settings: TokenCountSettings,
+    local_estimate_allowed: bool,
+) -> Union["TokenCountResponse", ProviderCountFailure, None]:
+    """Count with the provider, using the token_count_model in its region when one is set.
+
+    Returns None when the count should be estimated locally, or the provider's failure when the framed local
+    estimate stands in for it and that failure is the answer if the estimate cannot read the request."""
+    count_model: Final = settings.count_model
+    counted: Final = await (
+        count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region)
+    )
+    if counted is None:
+        return None
+    if isinstance(counted, TokenCountResponse):
+        if count_model is None:
+            return counted.model_copy(update={"count_source": "provider"})
+        return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model})
+    failure: Final = (
+        counted
+        if count_model is None or count_model.region is None
+        else counted.with_hint(
+            f"This model counts tokens with {count_model.model} in {count_model.region}. "
+            "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info"
+        )
+    )
+    if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed:
+        return failure
+    raise failure.to_proxy_exception()
+
+
+@dataclass(frozen=True, slots=True)
+class _LocalTokenEstimate:
+    total_tokens: int
+    tokenizer_type: str
+
+
+def _is_fork_guard_error(error: BaseException) -> bool:
+    from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking
+
+    return isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking))
+
+
+def _is_tokenizer_panic(error: BaseException) -> bool:
+    return type(error).__module__ == "pyo3_runtime" and type(error).__name__ == "PanicException"
+
+
+def _answers_provider_error(error: BaseException) -> bool:
+    return (isinstance(error, Exception) or _is_tokenizer_panic(error)) and not _is_fork_guard_error(error)
+
+
+def _provider_error_after_failed_estimate(
+    provider_failure: ProviderCountFailure, error: BaseException
+) -> ProxyException:
+    verbose_proxy_logger.warning(
+        "Local token estimate failed after the provider's count error %s, answering that error: %s",
+        provider_failure.status_code,
+        type(error).__name__,
+    )
+    return provider_failure.to_proxy_exception()
+
+
+async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int:
+    try:
+        return await estimate
+    except BaseException as e:
+        if not _answers_provider_error(e):
+            raise
+        raise _provider_error_after_failed_estimate(provider_failure, e) from e
+
+
+@dataclass(frozen=True, slots=True)
+class _CustomTokenizerKey:
+    identifier: str
+    revision: str
+    auth_token: str | None
+
+
+class _TokenizerLoadsBusy(Exception):
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class _TokenizerLoadOutcome:
+    tokenizer: SelectTokenizerResponse | BaseException
+    finished_at: float
+
+
+@dataclass(frozen=True, slots=True)
+class _SharedTokenizerLoad:
+    outcome: asyncio.Task[_TokenizerLoadOutcome]
+    started_at: float
+
+    def reusable(self, now: float) -> bool:
+        return not self.outcome.done() or now - self.outcome.result().finished_at < TOKENIZER_LOAD_REUSE_SECONDS
+
+
+async def _loaded_tokenizer(
+    model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter
+) -> SelectTokenizerResponse | BaseException:
+    try:
+        return await asyncify(litellm.utils.select_tokenizer, limiter=threads)(
+            model=model, custom_tokenizer=custom_tokenizer
+        )
+    except asyncio.CancelledError:
+        raise
+    except BaseException as e:  # noqa: BLE001  # every count sharing this load decides what it answers, pyo3 panics included
+        return e
+
+
+async def _tokenizer_load_outcome(
+    model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter
+) -> _TokenizerLoadOutcome:
+    tokenizer: Final = await _loaded_tokenizer(model, custom_tokenizer, threads)
+    return _TokenizerLoadOutcome(tokenizer=tokenizer, finished_at=asyncio.get_running_loop().time())
+
+
+@dataclass(slots=True)
+class _TokenizerLoads:
+    threads: anyio.CapacityLimiter
+    running: int = 0
+
+    def start(
+        self, model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, loop: asyncio.AbstractEventLoop
+    ) -> _SharedTokenizerLoad | None:
+        if self.running >= self.threads.total_tokens:
+            return None
+        self.running += 1
+        outcome: Final = loop.create_task(_tokenizer_load_outcome(model, custom_tokenizer, self.threads))
+        outcome.add_done_callback(self._finished)
+        return _SharedTokenizerLoad(outcome=outcome, started_at=loop.time())
+
+    def _finished(self, _outcome: asyncio.Task[_TokenizerLoadOutcome]) -> None:
+        self.running -= 1
+
+
+_TOKENIZER_LOADS: Final = anyio.lowlevel.RunVar[_TokenizerLoads]("litellm_proxy_tokenizer_loads")
+
+
+def _tokenizer_loads_for_this_loop() -> _TokenizerLoads:
+    existing: Final = _TOKENIZER_LOADS.get(None)
+    if existing is not None:
+        return existing
+    created: Final = _TokenizerLoads(threads=anyio.CapacityLimiter(TOKENIZER_LOAD_MAX_RUNNING))
+    _TOKENIZER_LOADS.set(created)
+    return created
+
+
+@dataclass(slots=True)
+class _TokenizerLoadSlot:
+    model: str
+    custom_tokenizer: CustomHuggingfaceTokenizer | None
+    current: _SharedTokenizerLoad | None = None
+
+    def load(self, loads: _TokenizerLoads, loop: asyncio.AbstractEventLoop) -> _SharedTokenizerLoad | None:
+        current: Final = self.current
+        if current is not None and current.reusable(loop.time()):
+            return current
+        self.current = loads.start(self.model, self.custom_tokenizer, loop)
+        return self.current
+
+
+@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
+def _tokenizer_load_slot(
+    model: str, custom_tokenizer: _CustomTokenizerKey | None, loop: asyncio.AbstractEventLoop
+) -> _TokenizerLoadSlot:
+    return _TokenizerLoadSlot(
+        model=model,
+        custom_tokenizer=None
+        if custom_tokenizer is None
+        else CustomHuggingfaceTokenizer(
+            identifier=custom_tokenizer.identifier,
+            revision=custom_tokenizer.revision,
+            auth_token=custom_tokenizer.auth_token,
+        ),
+    )
+
+
+async def _tokenizer_after_provider_error(
+    model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None
+) -> SelectTokenizerResponse | BaseException:
+    loop: Final = asyncio.get_running_loop()
+    try:
+        load: Final = _tokenizer_load_slot(
+            model,
+            None
+            if custom_tokenizer is None
+            else _CustomTokenizerKey(
+                identifier=custom_tokenizer["identifier"],
+                revision=custom_tokenizer["revision"],
+                auth_token=custom_tokenizer["auth_token"],
+            ),
+            loop,
+        ).load(_tokenizer_loads_for_this_loop(), loop)
+        if load is None:
+            return _TokenizerLoadsBusy()
+        outcome: Final = await asyncio.wait_for(
+            asyncio.shield(load.outcome),
+            timeout=max(load.started_at + TOKENIZER_LOAD_WAIT_SECONDS - loop.time(), 0),
+        )
+    except Exception as e:  # noqa: BLE001  # a tokenizer that cannot be read, loaded in time, or loaded at all answers the provider's error
+        return e
+    return outcome.tokenizer
+
+
+async def _tokenizer_or_provider_error(
+    model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure
+) -> SelectTokenizerResponse:
+    tokenizer: Final = await _tokenizer_after_provider_error(model, custom_tokenizer)
+    if not isinstance(tokenizer, BaseException):
+        return tokenizer
+    if not _answers_provider_error(tokenizer):
+        raise tokenizer.with_traceback(None)
+    raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer
+
+
+async def _local_token_estimate(
+    model: str,
+    prompt: str | None,
+    messages: Sequence[AllMessageValues] | None,
+    tools: list[ChatCompletionToolParam] | None,
+    custom_tokenizer: CustomHuggingfaceTokenizer | None,
+    provider_failure: ProviderCountFailure | None,
+) -> _LocalTokenEstimate:
+    tokenizer: Final = (
+        await asyncify(litellm.utils.select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer)
+        if provider_failure is None
+        else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure)
+    )
+    count: Final = offload_token_count(litellm.token_counter)(
+        model=model,
+        text=prompt,
+        messages=messages,
+        tools=tools,
+        custom_tokenizer=tokenizer,
+    )
+    total_tokens: Final = await (
+        count if provider_failure is None else _estimate_or_provider_error(count, provider_failure)
+    )
+    return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=str(tokenizer["type"]))
+
+
+_KIMI_K3_ENCODINGS: Final = KimiK3Encodings(
+    source=HUGGING_FACE_SOURCE, retry_after_seconds=TOKENIZER_LOAD_REUSE_SECONDS
+)
+
+
+async def _kimi_k3_count(
+    request: TokenCountRequest,
+    messages_format: TokenCountMessagesFormat,
+    model: str,
+    messages: Sequence[AllMessageValues],
+    tools: Sequence[ChatCompletionToolParam] | None,
+) -> int | None:
+    encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS)
+    if encoding is None:
+        return None
+    match messages_format:
+        case "anthropic":
+            return await offload_token_count(kimi_k3_anthropic_tokens)(
+                encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True)
+            )
+        case "openai":
+            return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools)
+
+
+class _TokenCountDeployment(TypedDict, total=False):
+    litellm_params: ReadOnly[Mapping[str, object]]
+    model_info: ReadOnly[Mapping[str, object]]
+
+
+_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment)
+
+
+def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None:
+    group: Final = () if router is None else router.get_all_deployments(model_name=model_group)
+    return group[0] if group else None
+
+
+def _non_empty_str(value: object) -> str | None:
+    return value if isinstance(value, str) and value else None
+
+
+def _deployment_aws_region(deployment: _TokenCountDeployment) -> str:
+    from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+
+    params: Final = deployment.get("litellm_params", {})
+    return BaseAWSLLM().get_aws_region_name(
+        {"aws_region_name": params.get("aws_region_name")}, model=_non_empty_str(params.get("model"))
+    )
+
+
+def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]:
+    model_info: Final = deployment.get("model_info", {})
+    params: Final = deployment.get("litellm_params", {})
+    base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model"))
+    if base_model is not None:
+        return base_model, None
+    model: Final = _non_empty_str(params.get("model"))
+    if model is None:
+        return request_model, None
+    return model, _non_empty_str(params.get("custom_llm_provider"))
+
+
 _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset(
     {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value}
 )
@@ -14728,6 +15065,14 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool:
     return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True)
 
 
+_ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"})
+
+
+def _estimable_contents(contents: Sequence[Mapping[str, object]]) -> bool:
+    parts: Final = chain.from_iterable(iter_items(content.get("parts") or ()) for content in contents)
+    return all(set(as_str_mapping(part) or ()) <= _ESTIMABLE_GOOGLE_PART_KEYS for part in parts)
+
+
 def _router_default_litellm_params(router: Router) -> Mapping[str, object]:
     defaults: Final[Mapping[str, object] | None] = getattr(router, "default_litellm_params", None)
     if defaults is None:
@@ -14831,6 +15176,14 @@ async def count_request_tokens(
     model_to_use: str = (
         litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model
     )
+    deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python(
+        deployment or _model_group_deployment(llm_router, request.model) or {}
+    )
+    settings: Final = await asyncify(resolve_token_count_settings)(
+        configured=deployment_view.get("model_info", {}),
+        default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)),
+        deployment_region=partial(_deployment_aws_region, deployment_view),
+    )
 
     # Try provider-specific token counting first - only for non-direct requests (from provider endpoints)
     provider_counter: BaseTokenCounter | None = None
@@ -14850,24 +15203,33 @@ async def count_request_tokens(
             code=400,
         )
     prompt_only: Final = bool(prompt) and conversation_empty
-    if provider_counter is not None:
-        prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
-        result: Final = await _try_provider_token_count(
-            provider_counter=provider_counter,
-            custom_llm_provider=custom_llm_provider,
+    contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents))
+    prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
+    ladder: Final = (
+        None
+        if provider_counter is None
+        else await _provider_token_count(
+            count_with=partial(
+                _try_provider_token_count,
+                provider_counter,
+                custom_llm_provider,
+                messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages,
+                contents=None if prompt_as_message else contents,
+                deployment=deployment,
+                request_model=request.model,
+                tools=tools,
+                system=system,
+                messages_format=messages_format,
+                output_format=output_format,
+                context_management=_count_context_management(context_management, request.model, request_litellm_params),
+            ),
             model_to_use=model_to_use,
-            messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages,
-            contents=None if prompt_as_message else contents,
-            deployment=deployment,
-            request_model=request.model,
-            tools=tools,
-            system=system,
-            messages_format=messages_format,
-            output_format=output_format,
-            context_management=_count_context_management(context_management, request.model, request_litellm_params),
+            settings=settings,
+            local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable,
         )
-        if result is not None:
-            return result
+    )
+    if isinstance(ladder, TokenCountResponse):
+        return ladder
 
     # Check if token counter is disabled before fallback
     if litellm.disable_token_counter is True:
@@ -14885,11 +15247,6 @@ async def count_request_tokens(
             CustomHuggingfaceTokenizer | None,
             model_info.get("custom_tokenizer", None),
         )
-    _tokenizer_used: Final = await asyncify(litellm.utils.select_tokenizer)(
-        model=model_to_use, custom_tokenizer=custom_tokenizer
-    )
-
-    tokenizer_used: Final = str(_tokenizer_used["type"])
     system_message: Final = _system_message(system)
     prompt_turn: Final = prompt_only and not (system_message is None and not tools)
     typed_messages: Final = cast(  # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes
@@ -14902,18 +15259,38 @@ async def count_request_tokens(
     counted_tools: Final = cast(  # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats
         list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None
     )
-    total_tokens: Final = await offload_token_count(litellm.token_counter)(
+    framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0
+    kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None
+    exact: Final = (
+        await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools)
+        if kimi_k3_messages is not None and custom_tokenizer is None
+        else None
+    )
+    if exact is not None:
+        return TokenCountResponse(
+            total_tokens=exact + framing,
+            request_model=request.model,
+            model_used=model_to_use,
+            tokenizer_type="huggingface_tokenizer",
+            count_source="local_estimate",
+        )
+    estimate: Final = await _local_token_estimate(
         model=model_to_use,
-        text=prompt if counted_messages is None else None,
+        prompt=prompt if counted_messages is None else None,
         messages=counted_messages,
         tools=counted_tools,
-        custom_tokenizer=_tokenizer_used,
+        custom_tokenizer=custom_tokenizer,
+        provider_failure=ladder,
+    )
+    template: Final = (
+        0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools))
     )
     return TokenCountResponse(
-        total_tokens=total_tokens,
+        total_tokens=estimate.total_tokens + template + framing,
         request_model=request.model,
         model_used=model_to_use,
-        tokenizer_type=tokenizer_used,
+        tokenizer_type=estimate.tokenizer_type,
+        count_source="local_estimate",
     )
 
 

litellm/proxy/response_api_endpoints/endpoints.py

@@ -67,6 +67,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
     error_status_code,
     openai_error_type_for_status,
 )
+from litellm.proxy.common_utils.token_count_settings import count_source_fields
 from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
 from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id
 from litellm.proxy.litellm_pre_call_utils import (
@@ -86,6 +87,7 @@ from litellm.types.llms.openai import (
     ResponsesAPIResponse,
 )
 from litellm.types.responses.main import DeleteResponseResult
+from litellm.types.token_count import TokenCountSourceFields
 from litellm.types.utils import TokenCountResponse
 
 if TYPE_CHECKING:
@@ -1373,7 +1375,7 @@ class _ResponsesApiErrorBody(TypedDict):
     error: ReadOnly[_ResponsesApiErrorDetail]
 
 
-class _ResponsesInputTokensResult(TypedDict):
+class _ResponsesInputTokensResult(TokenCountSourceFields):
     object: ReadOnly[str]
     input_tokens: ReadOnly[int]
 
@@ -1610,6 +1612,7 @@ async def responses_input_tokens(
     result: Final[_ResponsesInputTokensResult] = {
         "object": "response.input_tokens",
         "input_tokens": token_response.total_tokens,
+        **count_source_fields(token_response),
     }
     return result
 

litellm/types/llms/vertex_ai.py

@@ -9,6 +9,7 @@ from typing_extensions import (
 )
 
 from litellm.types.llms.openai import ChatCompletionFileObject, EmbeddingInput
+from litellm.types.token_count import TokenCountSourceFields
 
 # Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit
 # opt-in for combined embeddings — a provider-specific extension of the
@@ -320,12 +321,13 @@ class UsageMetadata(TypedDict, total=False):
     candidatesTokensDetails: list[PromptTokensDetails]  # Alternative key name used in some responses
 
 
-class TokenCountDetailsResponse(TypedDict):
+class TokenCountDetailsResponse(TokenCountSourceFields):
     """
     Response structure for token count details with modality breakdown.
 
     Example:
-        {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]}
+        {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}],
+         'litellm_count_source': 'provider'}
     """
 
     totalTokens: int

litellm/types/token_count.py

@@ -0,0 +1,16 @@
+from typing import Literal, TypeAlias
+
+from typing_extensions import NotRequired, ReadOnly, TypedDict
+
+TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"]
+TokenCountTokenizer: TypeAlias = Literal["kimi_k3"]
+
+
+class TokenCountModelInfo(TypedDict):
+    model: ReadOnly[str]
+    region: ReadOnly[NotRequired[str]]
+
+
+class TokenCountSourceFields(TypedDict, total=False):
+    litellm_count_source: ReadOnly[TokenCountSource]
+    litellm_count_model: ReadOnly[str]

litellm/types/utils.py

@@ -58,6 +58,7 @@ from litellm.types.llms.base import (
     LiteLLMPydanticObjectBase,
 )
 from litellm.types.mcp import MCPServerCostInfo
+from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer
 from litellm.types.workload_identity import (
     ANTHROPIC_WIF_KWARGS_KEYS,
     OPENAI_WIF_KWARGS_KEYS,
@@ -230,6 +231,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
     bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
     bedrock_converse_supports_strict_tools: bool | None
     supports_regex_lookaround: ReadOnly[bool | None]
+    token_count_framing: ReadOnly[int | None]
+    token_count_model: ReadOnly[TokenCountModelInfo | None]
+    token_count_tokenizer: ReadOnly[TokenCountTokenizer | None]
 
 
 class SearchContextCostPerQuery(TypedDict, total=False):
@@ -4051,11 +4055,14 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset(
     {"supports_regex_lookaround", "supports_assistant_prefill", "supports_bedrock_native_structured_output"}
 )
 
+DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"})
+
 SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
     frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
     - frozenset(CustomPricingLiteLLMParams.model_fields)
     - DEPLOYMENT_SCOPED_PRICING_FIELDS
     - DEPLOYMENT_SCOPED_CAPABILITY_FIELDS
+    - DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS
 )
 
 
@@ -4492,6 +4499,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase):
     request_model: str
     model_used: str
     tokenizer_type: str
+    count_source: TokenCountSource | None = None
     original_response: dict | None = None
     """
     Original Response from upstream API call - if an API call was made for token counting
@@ -4502,6 +4510,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase):
     HTTP status code from the token counting API (e.g., 200 for success, 429 for rate limit, 400 for bad request)
     """
     status_code: int | None = None
+    rejected_before_provider_call: bool = Field(default=False, exclude=True)
 
 
 class CustomHuggingfaceTokenizer(TypedDict):

litellm/utils.py

@@ -6694,6 +6694,9 @@ def get_model_info_helper(
                 uses_embed_content=_model_info.get("uses_embed_content", None),
                 supports_image_size=_model_info.get("supports_image_size", None),
                 supported_audio_formats=_model_info.get("supported_audio_formats", None),
+                token_count_framing=_model_info.get("token_count_framing", None),
+                token_count_model=_model_info.get("token_count_model", None),
+                token_count_tokenizer=_model_info.get("token_count_tokenizer", None),
                 vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None),
             )
             for cost_key, cost_value in _model_info.items():

model_prices_and_context_window.json

@@ -3028,7 +3028,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -3313,7 +3313,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -21144,6 +21144,7 @@
         "input_cost_per_token": 1.4e-07,
         "input_dbu_cost_per_token": 2e-06,
         "litellm_provider": "databricks",
+        "token_count_framing": 97,
         "max_input_tokens": 1000000,
         "max_output_tokens": 393216,
         "max_tokens": 393216,
@@ -21166,6 +21167,11 @@
         "supports_tool_choice": true,
         "supports_vision": false
     },
+    "databricks/system.ai.deepseek-v4-flash-0731": {
+        "litellm_provider": "databricks",
+        "mode": "chat",
+        "token_count_framing": 97
+    },
     "databricks/databricks-deepseek-v4-pro-0813": {
         "cache_creation_input_token_cost": 1.31999e-06,
         "cache_read_input_token_cost": 1.3202e-07,
@@ -22134,6 +22140,7 @@
         "input_cost_per_token": 2.99999e-06,
         "input_dbu_cost_per_token": 4.2857e-05,
         "litellm_provider": "databricks",
+        "token_count_framing": 100,
         "max_input_tokens": 1000000,
         "max_output_tokens": 1048576,
         "max_tokens": 1048576,
@@ -22157,6 +22164,11 @@
         "supports_tool_choice": true,
         "supports_vision": true
     },
+    "databricks/system.ai.kimi-k3": {
+        "litellm_provider": "databricks",
+        "mode": "chat",
+        "token_count_framing": 100
+    },
     "databricks/databricks-llama-4-maverick": {
         "cache_creation_input_token_cost": 5.0001e-07,
         "cache_read_input_token_cost": 5.0001e-07,
@@ -24695,7 +24707,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -33901,7 +33913,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -48928,7 +48940,7 @@
         "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
         "cache_read_input_token_cost_above_200k_tokens": 6e-07,
         "litellm_provider": "bedrock_converse",
-        "max_input_tokens": 1000000,
+        "max_input_tokens": 200000,
         "max_output_tokens": 64000,
         "max_tokens": 64000,
         "mode": "chat",
@@ -71540,6 +71552,8 @@
         "input_cost_per_token_flex": 1.65e-06,
         "input_cost_per_token_priority": 5.775e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -71570,6 +71584,8 @@
         "input_cost_per_token_flex": 1.5e-06,
         "input_cost_per_token_priority": 5.25e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -71600,6 +71616,8 @@
         "input_cost_per_token_flex": 1.65e-06,
         "input_cost_per_token_priority": 5.775e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -73100,6 +73118,10 @@
         "cache_read_input_token_cost": 1e-07,
         "input_cost_per_token": 2e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -73173,6 +73195,10 @@
         "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
     },
     "au.anthropic.claude-sonnet-5-5": {
+        "token_count_model": {
+            "model": "au.anthropic.claude-opus-5-5",
+            "region": "ap-southeast-4"
+        },
         "bedrock_converse_supports_strict_tools": false,
         "cache_creation_input_token_cost": 2.75e-06,
         "cache_creation_input_token_cost_above_1hr": 4.4e-06,
@@ -73373,6 +73399,10 @@
         "cache_read_input_token_cost": 1e-07,
         "input_cost_per_token": 2e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "global.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -73582,6 +73612,10 @@
         "cache_read_input_token_cost": 1.1e-07,
         "input_cost_per_token": 2.2e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "us.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -74382,6 +74416,8 @@
         "input_cost_per_token_flex": 1.65e-06,
         "input_cost_per_token_priority": 5.775e-06,
         "litellm_provider": "bedrock_converse",
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3",
         "max_input_tokens": 1000000,
         "max_output_tokens": 131072,
         "max_tokens": 131072,
@@ -74555,6 +74591,10 @@
         "cache_read_input_token_cost": 1e-08,
         "input_cost_per_token": 1e-07,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -74633,6 +74673,10 @@
         "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json"
     },
     "au.anthropic.claude-haiku-5-5": {
+        "token_count_model": {
+            "model": "au.anthropic.claude-opus-5-5",
+            "region": "ap-southeast-4"
+        },
         "input_cost_per_token_above_100k_tokens": 5.5e-07,
         "output_cost_per_token_above_100k_tokens": 2.75e-06,
         "cache_creation_input_token_cost_above_100k_tokens": 6.875e-07,
@@ -74863,6 +74907,10 @@
         "cache_read_input_token_cost": 1e-08,
         "input_cost_per_token": 1e-07,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "global.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,
@@ -75010,6 +75058,10 @@
         "cache_read_input_token_cost": 1.1e-08,
         "input_cost_per_token": 1.1e-07,
         "litellm_provider": "bedrock_converse",
+        "token_count_model": {
+            "model": "us.anthropic.claude-opus-5-5",
+            "region": "us-east-1"
+        },
         "supports_tool_search": true,
         "max_input_tokens": 1000000,
         "max_output_tokens": 128000,

model_prices_and_context_window.schema.json

@@ -1210,6 +1210,36 @@
             "additionalProperties": false
           }
         },
+        "token_count_framing": {
+          "type": "integer",
+          "minimum": 0,
+          "description": "Tokens added once per request to a local count estimate, for the provider's chat framing."
+        },
+        "token_count_model": {
+          "type": "object",
+          "required": [
+            "model"
+          ],
+          "additionalProperties": false,
+          "properties": {
+            "model": {
+              "type": "string",
+              "pattern": "\\S"
+            },
+            "region": {
+              "type": "string",
+              "pattern": "\\S"
+            }
+          },
+          "description": "Model that counts tokens for this one, and the AWS region its count API runs in."
+        },
+        "token_count_tokenizer": {
+          "type": "string",
+          "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.",
+          "enum": [
+            "kimi_k3"
+          ]
+        },
         "tpm": {
           "type": "integer",
           "minimum": 0,

tests/e2e/llm_translation/test_count_tokens_source_e2e.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+
+import base64
+import io
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Final
+
+import pytest
+from anthropic import Anthropic, BadRequestError, Omit, omit
+from anthropic.types import DocumentBlockParam, MessageParam, TextBlockParam, ToolParam
+from e2e_config import unique_marker
+from e2e_http import unwrap
+from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
+from lifecycle import ResourceManager
+from models import (
+    ChatMessage,
+    CountTokensResponse,
+    LiteLLMParamsBody,
+    ModelInfoBody,
+    ModelNewBody,
+    ResponsesInputTokensResponse,
+    TokenCounterBody,
+    TokenCountModelBody,
+)
+from proxy_client import ProxyClient
+from reportlab.lib.pagesizes import LETTER
+from reportlab.pdfgen.canvas import Canvas
+from sdk_clients import NO_PROXY_CACHE, SdkClients
+
+pytestmark = pytest.mark.e2e
+
+BEDROCK_SONNET: Final = "bedrock/us.anthropic.claude-sonnet-5-5"
+BEDROCK_HAIKU: Final = "bedrock/us.anthropic.claude-haiku-5-5"
+COST_MAP_COUNT_MODEL: Final = "us.anthropic.claude-opus-5-5"
+CONFIGURED_COUNT_MODEL: Final = "us.anthropic.claude-sonnet-5"
+ANTHROPIC_SONNET: Final = "anthropic/claude-sonnet-5-5"
+BEDROCK_KIMI: Final = "bedrock/us.moonshotai.kimi-k3"
+BEDROCK_NOVA: Final = "bedrock/us.amazon.nova-pro-v1:0"
+OPENAI_GPT: Final = "openai/gpt-5.6"
+
+COUNT_TOLERANCE: Final = 0.05
+MAX_TOKENS: Final = 16
+PDF_LINES_PER_PAGE: Final = 45
+
+LARGE_SYSTEM: Final = " ".join(
+    f"Rule {n}: answer in plain English, cite clause {n} when it applies, and never invent a figure."
+    for n in range(1, 301)
+)
+LARGE_TOOL: Final[ToolParam] = {
+    "name": "search_orders",
+    "description": " ".join(
+        f"Filter {n} narrows the order search by the attribute recorded in field_{n}." for n in range(1, 151)
+    ),
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            f"field_{n}": {"type": "string", "description": f"Order attribute number {n}, matched exactly."}
+            for n in range(1, 61)
+        },
+        "required": ["field_1"],
+    },
+}
+
+
+@dataclass(frozen=True, slots=True)
+class Served:
+    model: str
+    key: str
+
+
+@dataclass(frozen=True, slots=True)
+class PromptShape:
+    name: str
+    system: str | None = None
+    tool: ToolParam | None = None
+    pdf_pages: int = 0
+
+
+SHAPES: Final = (
+    PromptShape(name="baseline"),
+    PromptShape(name="large_system", system=LARGE_SYSTEM),
+    PromptShape(name="large_tool", tool=LARGE_TOOL),
+    PromptShape(name="system_and_tool", system=LARGE_SYSTEM, tool=LARGE_TOOL),
+    PromptShape(name="pdf_1_page", pdf_pages=1),
+    PromptShape(name="pdf_12_pages", pdf_pages=12),
+)
+SHAPE_IDS: Final = tuple(shape.name for shape in SHAPES)
+TEXT_SHAPES: Final = SHAPES[:2]
+TEXT_SHAPE_IDS: Final = SHAPE_IDS[:2]
+
+
+def _bedrock(model: str) -> LiteLLMParamsBody:
+    return LiteLLMParamsBody(
+        model=model,
+        aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
+        aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
+        aws_region_name="os.environ/AWS_REGION",
+    )
+
+
+def _serve(proxy: ProxyClient, backend: LiteLLMParamsBody, model_info: ModelInfoBody | None = None) -> Iterator[Served]:
+    resources: Final = ResourceManager(client=proxy)
+    try:
+        model: Final = f"e2e-count-source-{unique_marker()}"
+        model_id: Final = proxy.register_model(
+            ModelNewBody(model_name=model, litellm_params=backend, model_info=model_info or ModelInfoBody())
+        )
+        resources.defer(lambda: proxy.delete_model(model_id))
+        yield Served(model=model, key=resources.key())
+    finally:
+        resources.teardown()
+
+
+@pytest.fixture(scope="module")
+def bedrock_sonnet(proxy: ProxyClient) -> Iterator[Served]:
+    yield from _serve(proxy, _bedrock(BEDROCK_SONNET))
+
+
+@pytest.fixture(scope="module")
+def bedrock_haiku(proxy: ProxyClient) -> Iterator[Served]:
+    yield from _serve(proxy, _bedrock(BEDROCK_HAIKU))
+
+
+@pytest.fixture(scope="module")
+def bedrock_sonnet_with_configured_count_model(proxy: ProxyClient) -> Iterator[Served]:
+    yield from _serve(
+        proxy,
+        _bedrock(BEDROCK_SONNET),
+        ModelInfoBody(token_count_model=TokenCountModelBody(model=CONFIGURED_COUNT_MODEL)),
+    )
+
+
+@pytest.fixture(scope="module")
+def anthropic_sonnet(proxy: ProxyClient) -> Iterator[Served]:
+    yield from _serve(proxy, LiteLLMParamsBody(model=ANTHROPIC_SONNET, api_key="os.environ/ANTHROPIC_API_KEY"))
+
+
+@pytest.fixture(scope="module")
+def bedrock_kimi(proxy: ProxyClient) -> Iterator[Served]:
+    yield from _serve(proxy, _bedrock(BEDROCK_KIMI))
+
+
+@pytest.fixture(scope="module")
+def bedrock_nova(proxy: ProxyClient) -> Iterator[Served]:
+    yield from _serve(proxy, _bedrock(BEDROCK_NOVA))
+
+
+@pytest.fixture(scope="module")
+def openai_gpt(proxy: ProxyClient) -> Iterator[Served]:
+    yield from _serve(proxy, LiteLLMParamsBody(model=OPENAI_GPT, api_key="os.environ/OPENAI_API_KEY"))
+
+
+def _pdf(pages: int) -> str:
+    buffer: Final = io.BytesIO()
+    canvas: Final = Canvas(buffer, pagesize=LETTER, invariant=True)
+    for page in range(1, pages + 1):
+        for line in range(1, PDF_LINES_PER_PAGE + 1):
+            canvas.drawString(
+                72, 740 - 14 * line, f"Page {page}, line {line}: revenue, costs and headcount for region {line}."
+            )
+        canvas.showPage()
+    canvas.save()
+    return base64.b64encode(buffer.getvalue()).decode()
+
+
+def _user_turn(shape: PromptShape) -> MessageParam:
+    question: Final[TextBlockParam] = {"type": "text", "text": f"Reply with one word. {unique_marker()}"}
+    if shape.pdf_pages == 0:
+        return {"role": "user", "content": [question]}
+    document: Final[DocumentBlockParam] = {
+        "type": "document",
+        "source": {"type": "base64", "media_type": "application/pdf", "data": _pdf(shape.pdf_pages)},
+    }
+    return {"role": "user", "content": [document, question]}
+
+
+def _count_and_bill(client: Anthropic, model: str, shape: PromptShape) -> tuple[CountTokensResponse, int]:
+    messages: Final = [_user_turn(shape)]
+    system: Final[str | Omit] = omit if shape.system is None else shape.system
+    tools: Final[list[ToolParam] | Omit] = omit if shape.tool is None else [shape.tool]
+    counted: Final = CountTokensResponse.model_validate_json(
+        client.messages.with_raw_response.count_tokens(model=model, messages=messages, system=system, tools=tools).text
+    )
+    usage: Final = client.messages.create(
+        model=model,
+        max_tokens=MAX_TOKENS,
+        messages=messages,
+        system=system,
+        tools=tools,
+        extra_body=NO_PROXY_CACHE,
+    ).usage
+    billed: Final = usage.input_tokens + (usage.cache_creation_input_tokens or 0) + (usage.cache_read_input_tokens or 0)
+    return counted, billed
+
+
+def _within_tolerance(counted: int, billed: int) -> bool:
+    return abs(counted - billed) <= billed * COUNT_TOLERANCE
+
+
+def _assert_counted_by_the_cost_map_count_model(served: Served, sdk: SdkClients, shape: PromptShape) -> None:
+    counted, billed = _count_and_bill(sdk.anthropic(served.key), served.model, shape)
+
+    assert counted.litellm_count_source == "token_count_model", counted
+    assert counted.litellm_count_model == COST_MAP_COUNT_MODEL, counted
+    assert _within_tolerance(counted.input_tokens, billed), (counted, billed)
+
+
+class TestBedrockClaudeCountsWithCostMapCountModel:
+    @pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS)
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.BEDROCK,),
+            models=(BEDROCK_SONNET, COST_MAP_COUNT_MODEL),
+            capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_sonnet_count_names_the_count_model_and_matches_billed_input(
+        self, bedrock_sonnet: Served, sdk: SdkClients, shape: PromptShape
+    ) -> None:
+        _assert_counted_by_the_cost_map_count_model(bedrock_sonnet, sdk, shape)
+
+    @pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS)
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.BEDROCK,),
+            models=(BEDROCK_HAIKU, COST_MAP_COUNT_MODEL),
+            capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_haiku_count_names_the_count_model_and_matches_billed_input(
+        self, bedrock_haiku: Served, sdk: SdkClients, shape: PromptShape
+    ) -> None:
+        _assert_counted_by_the_cost_map_count_model(bedrock_haiku, sdk, shape)
+
+
+class TestAnthropicClaudeCountsWithProvider:
+    @pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS)
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.ANTHROPIC,),
+            models=(ANTHROPIC_SONNET,),
+            capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_count_comes_from_the_provider_and_matches_billed_input(
+        self, anthropic_sonnet: Served, sdk: SdkClients, shape: PromptShape
+    ) -> None:
+        counted, billed = _count_and_bill(sdk.anthropic(anthropic_sonnet.key), anthropic_sonnet.model, shape)
+
+        assert counted.litellm_count_source == "provider", counted
+        assert counted.litellm_count_model is None, counted
+        assert _within_tolerance(counted.input_tokens, billed), (counted, billed)
+
+
+class TestBedrockKimiCountsWithFramedLocalEstimate:
+    @pytest.mark.parametrize("shape", TEXT_SHAPES, ids=TEXT_SHAPE_IDS)
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.BEDROCK,),
+            models=(BEDROCK_KIMI,),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_text_estimate_matches_billed_input(self, bedrock_kimi: Served, sdk: SdkClients, shape: PromptShape) -> None:
+        counted, billed = _count_and_bill(sdk.anthropic(bedrock_kimi.key), bedrock_kimi.model, shape)
+
+        assert counted.litellm_count_source == "local_estimate", counted
+        assert _within_tolerance(counted.input_tokens, billed), (counted, billed)
+
+
+class TestConfiguredCountModelOverridesCostMap:
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.BEDROCK,),
+            models=(BEDROCK_SONNET, CONFIGURED_COUNT_MODEL),
+        )
+    )
+    def test_count_names_the_model_info_count_model(
+        self, bedrock_sonnet_with_configured_count_model: Served, sdk: SdkClients
+    ) -> None:
+        served: Final = bedrock_sonnet_with_configured_count_model
+        raw: Final = sdk.anthropic(served.key).messages.with_raw_response.count_tokens(
+            model=served.model, messages=[_user_turn(SHAPES[0])]
+        )
+        counted: Final = CountTokensResponse.model_validate_json(raw.text)
+
+        assert counted.litellm_count_source == "token_count_model", counted
+        assert counted.litellm_count_model == CONFIGURED_COUNT_MODEL, counted
+        assert counted.input_tokens > 0, counted
+
+
+class TestProviderCountErrorWithoutCountModel:
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.BEDROCK,),
+            models=(BEDROCK_NOVA,),
+        )
+    )
+    def test_count_keeps_the_provider_400_instead_of_estimating(self, bedrock_nova: Served, sdk: SdkClients) -> None:
+        with pytest.raises(BadRequestError) as rejected:
+            _ = sdk.anthropic(bedrock_nova.key).messages.count_tokens(
+                model=bedrock_nova.model, messages=[_user_turn(SHAPES[0])]
+            )
+
+        assert rejected.value.status_code == 400, rejected.value
+
+
+class TestCountSourceOnOtherCountEndpoints:
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.OPENAI,),
+            models=(OPENAI_GPT,),
+        )
+    )
+    def test_responses_input_tokens_says_the_provider_counted(self, openai_gpt: Served, sdk: SdkClients) -> None:
+        raw: Final = sdk.openai(openai_gpt.key).responses.input_tokens.with_raw_response.count(
+            model=openai_gpt.model, input=f"Reply with one word. {unique_marker()}"
+        )
+        counted: Final = ResponsesInputTokensResponse.model_validate_json(raw.text)
+
+        assert counted.litellm_count_source == "provider", counted
+        assert counted.input_tokens > 0, counted
+
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.COUNT_TOKENS,
+            providers=(Provider.BEDROCK,),
+            models=(BEDROCK_KIMI,),
+        )
+    )
+    def test_token_counter_says_kimi_was_estimated_locally(self, proxy: ProxyClient, bedrock_kimi: Served) -> None:
+        counted: Final = unwrap(
+            proxy.token_counter(
+                bedrock_kimi.key,
+                TokenCounterBody(
+                    model=bedrock_kimi.model,
+                    messages=[ChatMessage(role="user", content=f"Reply with one word. {unique_marker()}")],
+                ),
+            )
+        )
+
+        assert counted.count_source == "local_estimate", counted
+        assert counted.total_tokens > 0, counted

tests/e2e/models.py

@@ -677,13 +677,39 @@ class AnthropicMessagesResponse(BaseModel):
     usage: Usage | None = None
 
 
-class CountTokensResponse(BaseModel):
+TokenCountSource = Literal["provider", "token_count_model", "local_estimate"]
+
+
+class LitellmCountSource(BaseModel, frozen=True):
+    litellm_count_source: TokenCountSource | None = None
+    litellm_count_model: str | None = None
+
+
+class CountTokensResponse(LitellmCountSource, frozen=True):
     """`/v1/messages/count_tokens` answer. `input_tokens` is required so a 200
     whose body lacks it fails validation instead of passing vacuously."""
 
     input_tokens: int
 
 
+class ResponsesInputTokensResponse(LitellmCountSource, frozen=True):
+    object: str
+    input_tokens: int
+
+
+class TokenCounterBody(BaseModel, frozen=True):
+    model: str
+    messages: list[ChatMessage]
+
+
+class TokenCounterResponse(BaseModel, frozen=True):
+    total_tokens: int
+    request_model: str
+    model_used: str
+    tokenizer_type: str
+    count_source: TokenCountSource | None = None
+
+
 class AnthropicErrorBody(BaseModel):
     type: str
     message: str
@@ -1392,6 +1418,11 @@ class LiteLLMParamsBody(BaseModel):
 ModelMode = Literal["chat", "batch", "realtime", "image_generation"]
 
 
+class TokenCountModelBody(BaseModel, frozen=True):
+    model: str
+    region: str | None = None
+
+
 class ModelInfoBody(BaseModel):
     # id is left unset so the proxy assigns a unique model_id per deployment.
     # Pinning it to the model_name made re-registrations of a fixed-name model
@@ -1404,6 +1435,7 @@ class ModelInfoBody(BaseModel):
     team_id: str | None = None
     allowed_fails: int | None = None
     allowed_fails_policy: dict[str, int] | None = None
+    token_count_model: TokenCountModelBody | None = None
 
 
 class ModelNewBody(BaseModel):

tests/e2e/proxy_client.py

@@ -100,6 +100,8 @@ from models import (
     TeamNewBody,
     TeamNewResponse,
     TeamUpdateBody,
+    TokenCounterBody,
+    TokenCounterResponse,
     ToolsetCreateBody,
     ToolsetRow,
     ToolsetUpdateBody,
@@ -1049,6 +1051,15 @@ class ProxyClient:
             response_type=CountTokensResponse,
         )
 
+    @step("Count tokens with /utils/token_counter for {body.model}")
+    def token_counter(self, key: str, body: TokenCounterBody) -> Result[TokenCounterResponse]:
+        return self.transport.post(
+            "/utils/token_counter",
+            headers=self.transport.bearer(key),
+            json=body,
+            response_type=TokenCounterResponse,
+        )
+
     @step("Send a /v1/messages request to {body.model}")
     def messages(
         self, key: str, body: AnthropicMessagesBody, *, session_id: str | None = None

tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py

@@ -104,7 +104,7 @@ def test_anthropic_count_tokens_forwards_body_and_refuses_a_key_without_model_ac
 
             client: Final = _client(gateway, gateway.key)
             result: Final = _sdk_count(client, model, spelling, _COUNT_BETA)
-            assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS}
+            assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS, "litellm_count_source": "provider"}
             requests: Final = wire.drain()
             assert len(requests) == 1, requests
 
@@ -236,5 +236,5 @@ def test_anthropic_count_tokens_merges_the_client_beta_with_the_token_counting_b
                 num_retries=0,
             )
             result: Final = _sdk_count(_client(gateway, gateway.key), model, spelling, _CLIENT_BETA)
-            assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS}
+            assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS, "litellm_count_source": "provider"}
             assert len(wire.drain()) == 1

tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py

@@ -245,7 +245,10 @@ def test_enabled_gateway_count_tokens_alias_matches_native_count_tokens_route(en
             "/claude_code_gateway/v1/messages/count_tokens?beta=true", json=body, headers=headers
         )
         assert gateway_response.status_code == 200, gateway_response.text
-        assert _JSON_OBJECT.validate_json(gateway_response.content) == {"input_tokens": 41}, gateway_response.text
+        assert _JSON_OBJECT.validate_json(gateway_response.content) == {
+            "input_tokens": 41,
+            "litellm_count_source": "provider",
+        }, gateway_response.text
         native_response: Final = enabled_gateway.client.post(
             "/v1/messages/count_tokens?beta=true", json=body, headers=headers
         )

tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py

@@ -274,7 +274,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run(
 ) -> None:
     with _listening(deployment) as peer:
         response: Final = _count(gateway, count_request(deployment.model, case))
-        assert _payload(response) == {"input_tokens": _COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
         assert _bodies(peer, deployment.provider) == (expected_count_body(_MODEL, case),)
 
 
@@ -329,7 +329,11 @@ def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gate
             "/v1/responses/input_tokens",
             {"model": deployment.model, "input": USER_TEXT, "instructions": INSTRUCTION},
         )
-        assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text
+        assert _payload(response) == {
+            "object": "response.input_tokens",
+            "input_tokens": _COUNT,
+            "litellm_count_source": "provider",
+        }, response.text
         assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
 
 
@@ -342,7 +346,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite
             "/v1/responses/input_tokens",
             {"model": deployment.model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION},
         )
-        assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text
+        assert _payload(response) == {
+            "object": "response.input_tokens",
+            "input_tokens": _COUNT,
+            "litellm_count_source": "provider",
+        }, response.text
         assert _bodies(peer, deployment.provider) == (
             {"model": _MODEL, "messages": [USER], "system": [LIFTED, {"type": "text", "text": REMINDER}]},
         )
@@ -373,7 +381,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system(
 ) -> None:
     with _listening(deployment) as peer:
         response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "tools": TOOLS})
-        assert _payload(response) == {"input_tokens": _COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
         assert _bodies(peer, deployment.provider) == ({**_STRING_BODY, "tools": TOOLS},)
 
 
@@ -386,7 +394,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place(
             "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP],
         }
         response: Final = _count(gateway, body)
-        assert _payload(response) == {"input_tokens": _COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
         assert _bodies(peer, deployment.provider) == (
             {"model": _MODEL, "messages": [USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], "system": [LIFTED]},
         )
@@ -418,7 +426,7 @@ def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Dep
         answers: Final = tuple(
             _payload(_count(gateway, count_request(deployment.model, STRING_CASE))) for _ in range(2)
         )
-        assert answers == ({"input_tokens": _COUNT},) * 2
+        assert answers == ({"input_tokens": _COUNT, "litellm_count_source": "provider"},) * 2
         assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2
 
 
@@ -432,7 +440,7 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_
         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
+        assert _payload(follow_up) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, follow_up.text
         assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
 
 
@@ -447,7 +455,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value(
             content=f'{{"model": "{deployment.model}", "messages": {first}, "messages": {last}}}',
             headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
         )
-        assert _payload(response) == {"input_tokens": _COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text
         assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
 
 

tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py

@@ -378,7 +378,7 @@ def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot(
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",)
         assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
 
@@ -394,7 +394,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mant
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = _count(gateway, _bare(model))
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
 
@@ -429,7 +429,11 @@ def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: Owne
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
-        assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
+        assert _payload(response) == {
+            "totalTokens": _MANTLE_COUNT,
+            "promptTokensDetails": [],
+            "litellm_count_source": "provider",
+        }, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},)
 
@@ -445,7 +449,11 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock
     ):
         model: Final = _deployment(scenario, runtime.url, _SONNET)
         response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
-        assert _payload(response) == {"totalTokens": _RUNTIME_COUNT, "promptTokensDetails": []}, response.text
+        assert _payload(response) == {
+            "totalTokens": _RUNTIME_COUNT,
+            "promptTokensDetails": [],
+            "litellm_count_source": "provider",
+        }, response.text
         assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",)
         assert mantle.drain() == ()
 
@@ -460,7 +468,7 @@ def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy,
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
 
@@ -520,7 +528,11 @@ def test_google_count_tokens_forwards_system_and_tools_to_mantle(
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
-        assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
+        assert _payload(response) == {
+            "totalTokens": _MANTLE_COUNT,
+            "promptTokensDetails": [],
+            "litellm_count_source": "provider",
+        }, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (
             _mantle_body(
@@ -565,7 +577,11 @@ def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle(
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
-        assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
+        assert _payload(response) == {
+            "totalTokens": _MANTLE_COUNT,
+            "promptTokensDetails": [],
+            "litellm_count_source": "provider",
+        }, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (
             _mantle_body(
@@ -595,7 +611,11 @@ def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy
         response: Final = gateway.request(
             "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"}
         )
-        assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {
+            "object": "response.input_tokens",
+            "input_tokens": _MANTLE_COUNT,
+            "litellm_count_source": "provider",
+        }, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
 
@@ -640,7 +660,10 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx
         gateway.scenario() as scenario,
     ):
         model: Final = _deployment(scenario, runtime.url, _SONNET)
-        assert _payload(_count(gateway, _full(model))) == {"input_tokens": _RUNTIME_COUNT}
+        assert _payload(_count(gateway, _full(model))) == {
+            "input_tokens": _RUNTIME_COUNT,
+            "litellm_count_source": "provider",
+        }
         detailed: Final = gateway.request(
             "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"}
         )
@@ -728,7 +751,11 @@ def test_responses_input_tokens_with_instructions_counts_through_mantle(
             "/v1/responses/input_tokens",
             {"model": model, "input": "Count this message", "instructions": "Be terse"},
         )
-        assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {
+            "object": "response.input_tokens",
+            "input_tokens": _MANTLE_COUNT,
+            "litellm_count_source": "provider",
+        }, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_mantle_body(system=[{"type": "text", "text": "Be terse"}]),)
 
@@ -745,7 +772,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle(
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = _count(gateway, count_request(model, case))
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (expected_count_body(_OPUS_BASE, case),)
 
@@ -858,7 +885,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite
             "/v1/responses/input_tokens",
             {"model": model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION},
         )
-        assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {
+            "object": "response.input_tokens",
+            "input_tokens": _MANTLE_COUNT,
+            "litellm_count_source": "provider",
+        }, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED, {"type": "text", "text": REMINDER}]),)
 
@@ -874,7 +905,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS})
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED], tools=_TOOLS),)
 
@@ -894,7 +925,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant
             "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP],
         }
         response: Final = _count(gateway, body)
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (
             _mantle_body(messages=[USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], system=[LIFTED]),
@@ -947,7 +978,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti
     ):
         model: Final = _deployment(scenario, runtime.url)
         answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2))
-        assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2
+        assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2
         assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2
         assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2
 
@@ -968,7 +999,9 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_
         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
+        assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, (
+            follow_up.text
+        )
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
 
@@ -990,7 +1023,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_m
             content=f'{{"model": "{model}", "messages": {first}, "messages": {last}}}',
             headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
         )
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,)
 
@@ -1015,7 +1048,9 @@ def test_disabled_token_counter_counts_a_leading_system_through_mantle(
         )
         with wire_server(_mantle(_strict), port=mantle_port) as strict:
             counted: Final = _count(owned.gateway, count_request(_OWNED_OPUS, STRING_CASE))
-            assert _payload(counted) == {"input_tokens": _MANTLE_COUNT}, counted.text
+            assert _payload(counted) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, (
+                counted.text
+            )
             assert _mantle_bodies(strict) == (_MANTLE_LIFTED,)
         assert len(_runtime_count_targets(runtime)) == 1
 
@@ -1220,7 +1255,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle(
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = _count(gateway, {**_bare(model), "system": system})
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_mantle_body(system=system),)
 
@@ -1241,7 +1276,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value
             content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}',
             headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
         )
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         (sent,) = mantle.drain()
         assert _mantle_requests((sent,)) == (_MANTLE_FULL,)
@@ -1259,7 +1294,7 @@ def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_b
     ):
         model: Final = _deployment(scenario, runtime.url)
         response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""})
-        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text
         assert len(_runtime_count_targets(runtime)) == 1
         assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
 
@@ -1275,7 +1310,7 @@ def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time(
     ):
         model: Final = _deployment(scenario, runtime.url)
         answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2))
-        assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2
+        assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2
         assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2
         assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2
 
@@ -1304,9 +1339,15 @@ def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_lo
             assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text
             assert _mantle_bodies(refusing) == (_MANTLE_FULL,)
         with wire_server(_mantle(), port=mantle_port) as counting:
-            assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {"input_tokens": _MANTLE_COUNT}
+            assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {
+                "input_tokens": _MANTLE_COUNT,
+                "litellm_count_source": "provider",
+            }
             assert _mantle_bodies(counting) == (_MANTLE_FULL,)
-            assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {"input_tokens": _RUNTIME_COUNT}
+            assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {
+                "input_tokens": _RUNTIME_COUNT,
+                "litellm_count_source": "provider",
+            }
             rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA))
             assert rejected.status_code == 400, rejected.text
             assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text

tests/integration/providers/test_count_tokens_model_wire.py

@@ -0,0 +1,719 @@
+import json
+import re
+from collections.abc import Callable, Iterator, Mapping
+from itertools import chain
+from pathlib import Path
+from typing import Final
+from uuid import uuid4
+
+import anthropic
+import httpx
+import pytest
+from anthropic.types import MessageCountTokensToolParam, MessageTokensCount
+from integration._support.bedrock_runtime_peer import target_of
+from integration._support.client import Gateway, Scenario, gateway_from_environment
+from integration._support.ports import allocate_port
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+_SONNET: Final = "us.anthropic.claude-sonnet-5-5"
+_SONNET_BASE: Final = "anthropic.claude-sonnet-5-5"
+_OPUS: Final = "us.anthropic.claude-opus-5-5"
+_OPUS_BASE: Final = "anthropic.claude-opus-5-5"
+_NOVA: Final = "amazon.nova-lite-v1:0"
+_KIMI: Final = "us.moonshotai.kimi-k3"
+_KIMI_BASE: Final = "moonshotai.kimi-k3"
+_PROFILE_ID: Final = "lit9563count"
+_PROFILE_ARN: Final = f"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/{_PROFILE_ID}"
+_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json"
+_REGION: Final = "us-west-2"
+_COUNT_REGION: Final = "us-east-1"
+_ACCESS_KEY: Final = "AKIAINTEGRATIONCOUNT"
+_SECRET_KEY: Final = "integration-count-secret"
+_PROVIDER_COUNT: Final = 1345
+_COUNT_MODEL_COUNT: Final = 2718
+_FRAMING: Final = 50
+_UNSUPPORTED: Final = "The provided model doesn't support counting tokens."
+_COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$")
+_SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})")
+_SCRIPTED_COUNT_MODEL_STATUS: Final = re.compile(r"count_model=(\d{3})")
+_CREDENTIAL: Final = re.compile(r"Credential=([^/]+)/\d+/([^/]+)/")
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]])
+
+_SYSTEM: Final = "You are a terse assistant that answers in one sentence"
+_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = {
+    "type": "object",
+    "properties": {"city": {"type": "string", "description": "City to look up"}},
+    "required": ["city"],
+}
+_SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [
+    {
+        "name": "get_weather",
+        "description": "Look up the current weather for a city",
+        "input_schema": {
+            "type": "object",
+            "properties": {"city": {"type": "string", "description": "City to look up"}},
+            "required": ["city"],
+        },
+    }
+]
+_COUNT_MODEL_SOURCE: Final = {"litellm_count_source": "token_count_model", "litellm_count_model": _OPUS}
+_OPENAI_PRIMARY: Final = "gpt-count-primary"
+_OPENAI_COUNT_MODEL: Final = "gpt-count-model"
+_OPENAI_COUNT_TARGET: Final = "/v1/responses/input_tokens"
+_COUNTED_FIELDS: Final = ("input", "instructions", "tools", "text")
+_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
+_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = {
+    "name": "get_weather",
+    "description": "Look up the current weather for a city",
+    "input_schema": _WEATHER_SCHEMA,
+}
+_TOOL_TURNS: Final[list[JsonValue]] = [
+    {"role": "user", "content": "What is the weather in Paris?"},
+    {
+        "role": "assistant",
+        "content": [{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}}],
+    },
+    {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "Sunny"}]},
+]
+_LOCAL_SOURCE: Final = {"litellm_count_source": "local_estimate"}
+
+
+def _prompt(status: int, count_model_status: int | None = None) -> str:
+    scripted: Final = "" if count_model_status is None else f" count_model={count_model_status}"
+    return f"Count this message status={status}{scripted}"
+
+
+def _converse_body(prompt: str) -> dict[str, JsonValue]:
+    return {
+        "input": {
+            "converse": {
+                "messages": [{"role": "user", "content": [{"text": prompt}]}],
+                "system": [{"text": _SYSTEM}],
+                "toolConfig": {
+                    "tools": [
+                        {
+                            "toolSpec": {
+                                "name": "get_weather",
+                                "description": "Look up the current weather for a city",
+                                "inputSchema": {"json": _WEATHER_SCHEMA},
+                            }
+                        }
+                    ]
+                },
+            }
+        }
+    }
+
+
+def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply:
+    return Reply(status=status, body=json.dumps(payload).encode())
+
+
+def _counted_model(request: Request) -> str:
+    counted: Final = _COUNT_TARGET.match(target_of(request))
+    return target_of(request) if counted is None else counted.group(1)
+
+
+def _scripted(pattern: re.Pattern[str], request: Request, default: int) -> int:
+    scripted: Final = pattern.search(request.body.decode(errors="replace"))
+    return default if scripted is None else int(scripted.group(1))
+
+
+def _runtime(request: Request) -> Reply:
+    model: Final = _counted_model(request)
+    if model == _OPUS_BASE:
+        count_model_status: Final = _scripted(_SCRIPTED_COUNT_MODEL_STATUS, request, 200)
+        if count_model_status == 200:
+            return _json_reply(200, {"inputTokens": _COUNT_MODEL_COUNT})
+        return _json_reply(count_model_status, {"message": f"scripted runtime {count_model_status} for {model}"})
+    status: Final = _scripted(_SCRIPTED_STATUS, request, 500)
+    if status == 200:
+        return _json_reply(200, {"inputTokens": _PROVIDER_COUNT})
+    return _json_reply(status, {"message": _UNSUPPORTED if status == 400 else f"scripted runtime {status}"})
+
+
+def _mantle(request: Request) -> Reply:
+    model: Final = _JSON_OBJECT.validate_json(request.body).get("model")
+    return _json_reply(
+        404, {"type": "error", "error": {"type": "not_found_error", "message": f"scripted mantle 404 for {model}"}}
+    )
+
+
+def _signer(request: Request) -> tuple[str, str]:
+    credential: Final = _CREDENTIAL.search(request.headers.get("authorization", ""))
+    return ("", "") if credential is None else (credential.group(1), credential.group(2))
+
+
+def _runtime_counts(wire: Wire) -> tuple[tuple[str, tuple[str, str], dict[str, JsonValue]], ...]:
+    return tuple(
+        (_counted_model(request), _signer(request), _JSON_OBJECT.validate_json(request.body))
+        for request in wire.drain()
+    )
+
+
+def _runtime_models(wire: Wire) -> tuple[str, ...]:
+    return tuple(_counted_model(request) for request in wire.drain())
+
+
+def _mantle_models(wire: Wire) -> tuple[JsonValue, ...]:
+    return tuple(_JSON_OBJECT.validate_json(request.body).get("model") for request in wire.drain())
+
+
+def _url(port: int) -> str:
+    return f"http://127.0.0.1:{port}"
+
+
+def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]:
+    return {
+        "model": f"bedrock/{model}",
+        "api_base": api_base,
+        "aws_access_key_id": _ACCESS_KEY,
+        "aws_secret_access_key": _SECRET_KEY,
+        "aws_region_name": _REGION,
+    }
+
+
+@pytest.fixture(scope="module")
+def runtime_port() -> int:
+    return allocate_port()
+
+
+@pytest.fixture(scope="module")
+def mantle_port() -> int:
+    return allocate_port()
+
+
+@pytest.fixture(scope="module")
+def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -> Iterator[OwnedProxy]:
+    with (
+        gateway_from_environment() as gateway,
+        owned_proxy_process(
+            gateway,
+            tmp_path_factory.mktemp("count-model"),
+            {
+                "BEDROCK_MANTLE_API_BASE": _url(mantle_port),
+                "LITELLM_LOCAL_MODEL_COST_MAP": "True",
+                "HF_HUB_OFFLINE": "1",
+            },
+            workers=1,
+            remove_environment=("AWS_BEARER_TOKEN_BEDROCK",),
+        ) as owned,
+    ):
+        yield owned
+
+
+def _deployment(
+    scenario: Scenario, runtime_port: int, model: str = _SONNET, model_info: Mapping[str, JsonValue] | None = None
+) -> str:
+    return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port)))
+
+
+def _count(gateway: Gateway, model: str, status: int, count_model_status: int | None = None) -> MessageTokensCount:
+    with anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) as client:
+        return client.messages.count_tokens(
+            model=model,
+            messages=[{"role": "user", "content": _prompt(status, count_model_status)}],
+            system=_SYSTEM,
+            tools=_SDK_TOOLS,
+        )
+
+
+def _messages_estimate(gateway: Gateway, model: str) -> tuple[int, JsonValue]:
+    counted: Final = _count(gateway, model, 400)
+    return counted.input_tokens, (counted.model_extra or {}).get("litellm_count_source")
+
+
+def _utils_estimate(gateway: Gateway, model: str) -> tuple[int, JsonValue]:
+    response: Final = gateway.request(
+        "POST", "/utils/token_counter", {"model": model, "messages": [{"role": "user", "content": _prompt(400)}]}
+    )
+    assert response.status_code == 200, response.text
+    body: Final = _JSON_OBJECT.validate_json(response.content)
+    total: Final = body["total_tokens"]
+    assert isinstance(total, int), body
+    return total, body.get("count_source")
+
+
+def _shipped_framing(model: str) -> JsonValue:
+    return _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())[model].get("token_count_framing")
+
+
+def test_shipped_token_count_model_counts_the_request_in_its_region_without_trying_the_model(
+    counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        counted: Final = _count(gateway, _deployment(scenario, runtime_port), 404)
+        assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted
+        assert _runtime_counts(runtime) == ((_OPUS_BASE, (_ACCESS_KEY, _COUNT_REGION), _converse_body(_prompt(404))),)
+        assert _mantle_models(mantle) == ()
+
+
+@pytest.mark.parametrize(
+    ("count_model_status", "error", "message"),
+    [
+        (401, anthropic.AuthenticationError, f"scripted runtime 401 for {_OPUS_BASE}"),
+        (429, anthropic.RateLimitError, f"scripted runtime 429 for {_OPUS_BASE}"),
+        (404, anthropic.NotFoundError, f"scripted mantle 404 for {_OPUS_BASE}"),
+    ],
+)
+def test_token_count_model_failure_answers_its_error_and_names_the_region_to_enable(
+    counting_proxy: OwnedProxy,
+    runtime_port: int,
+    mantle_port: int,
+    count_model_status: int,
+    error: type[anthropic.APIStatusError],
+    message: str,
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        with pytest.raises(error) as refused:
+            _count(gateway, _deployment(scenario, runtime_port), 200, count_model_status=count_model_status)
+        assert message in str(refused.value) and f"{_OPUS} in {_COUNT_REGION}" in str(refused.value), refused.value
+        assert _runtime_models(runtime) == (_OPUS_BASE,)
+        assert _mantle_models(mantle) == ((_OPUS_BASE,) if count_model_status == 404 else ())
+
+
+def test_framing_is_added_to_the_local_estimate_when_the_runtime_cannot_count(
+    counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        framed: Final = _count(
+            gateway, _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": _FRAMING}), 400
+        )
+        unframed: Final = _count(gateway, _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": 0}), 400)
+        assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed)
+        assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed)
+        assert _runtime_models(runtime) == (_NOVA, _NOVA)
+        assert _mantle_models(mantle) == ()
+
+
+def test_framing_on_one_deployment_leaves_a_sibling_of_the_same_model_raising_the_400(
+    counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        framed: Final = _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": _FRAMING})
+        sibling: Final = _deployment(scenario, runtime_port, _NOVA)
+        assert _count(gateway, framed, 400).model_extra == _LOCAL_SOURCE
+        with pytest.raises(anthropic.BadRequestError, match=re.escape(_UNSUPPORTED)):
+            _count(gateway, sibling, 400)
+        assert _runtime_models(runtime) == (_NOVA, _NOVA)
+        assert _mantle_models(mantle) == ()
+
+
+def test_empty_configured_token_count_model_counts_the_model_itself_in_the_deployment_region(
+    counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        counted: Final = _count(gateway, _deployment(scenario, runtime_port, model_info={"token_count_model": {}}), 200)
+        assert (counted.input_tokens, counted.model_extra) == (_PROVIDER_COUNT, {"litellm_count_source": "provider"})
+        assert _runtime_counts(runtime) == ((_SONNET_BASE, (_ACCESS_KEY, _REGION), _converse_body(_prompt(200))),)
+        assert _mantle_models(mantle) == ()
+
+
+@pytest.mark.parametrize("base_model_in", ["model_info", "litellm_params"])
+def test_inference_profile_arn_counts_with_the_token_count_model_of_its_base_model(
+    counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int, base_model_in: str
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    base_model: Final[dict[str, JsonValue]] = {"base_model": f"bedrock/{_SONNET}"}
+    with (
+        wire_server(_runtime, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = scenario.model(
+            model_info=base_model if base_model_in == "model_info" else None,
+            api_key=None,
+            **_litellm_params(_PROFILE_ARN, _url(runtime_port)),
+            **(base_model if base_model_in == "litellm_params" else {}),
+        )
+        counted: Final = _count(gateway, model, 404)
+        assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted
+        assert _runtime_models(runtime) == (_OPUS_BASE,)
+        assert _mantle_models(mantle) == ()
+
+
+@pytest.mark.parametrize(
+    ("estimate", "runtime_calls"),
+    [(_messages_estimate, (_KIMI_BASE, _KIMI_BASE)), (_utils_estimate, ())],
+    ids=["messages_count_tokens", "utils_token_counter"],
+)
+def test_shipped_framing_is_added_to_the_estimate_and_a_configured_zero_turns_it_off(
+    counting_proxy: OwnedProxy,
+    runtime_port: int,
+    mantle_port: int,
+    estimate: Callable[[Gateway, str], tuple[int, JsonValue]],
+    runtime_calls: tuple[str, ...],
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        shipped: Final = estimate(gateway, _deployment(scenario, runtime_port, _KIMI))
+        unframed: Final = estimate(gateway, _deployment(scenario, runtime_port, _KIMI, {"token_count_framing": 0}))
+        assert (shipped[1], unframed[1]) == ("local_estimate", "local_estimate"), (shipped, unframed)
+        framing: Final = _shipped_framing(_KIMI)
+        assert isinstance(framing, int) and framing > 0, framing
+        assert shipped[0] - unframed[0] == framing, (shipped, unframed)
+        assert _runtime_models(runtime) == runtime_calls
+        assert _mantle_models(mantle) == ()
+
+
+def _openai_count(request: Request) -> Reply:
+    model: Final = _JSON_OBJECT.validate_json(request.body).get("model")
+    if request.target == _OPENAI_COUNT_TARGET and model == _OPENAI_COUNT_MODEL:
+        return _json_reply(200, {"object": "response.input_tokens", "input_tokens": _COUNT_MODEL_COUNT})
+    return _json_reply(
+        404,
+        {"error": {"message": f"scripted openai 404 for {model}", "type": "invalid_request_error", "code": None}},
+    )
+
+
+def _counted_fields(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+    return {key: body[key] for key in _COUNTED_FIELDS if key in body}
+
+
+def _openai_counts(wire: Wire) -> tuple[tuple[str, JsonValue, dict[str, JsonValue]], ...]:
+    bodies: Final = tuple((request.target, _JSON_OBJECT.validate_json(request.body)) for request in wire.drain())
+    return tuple((target, body.get("model"), _counted_fields(body)) for target, body in bodies)
+
+
+def _openai_deployment(scenario: Scenario, wire: Wire, model_info: Mapping[str, JsonValue]) -> str:
+    return scenario.model(model=f"openai/{_OPENAI_PRIMARY}", api_base=f"{wire.url}/v1", model_info=model_info)
+
+
+def test_token_count_model_counts_the_anthropic_body_converted_like_the_provider_count(
+    counting_proxy: OwnedProxy,
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_openai_count) as openai, gateway.scenario() as scenario:
+        model: Final = _openai_deployment(scenario, openai, {"token_count_model": {"model": _OPENAI_COUNT_MODEL}})
+        response: Final = gateway.request(
+            "POST",
+            "/v1/messages/count_tokens",
+            {
+                "model": model,
+                "system": _SYSTEM,
+                "tools": [_ANTHROPIC_TOOL],
+                "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}},
+                "messages": _TOOL_TURNS,
+            },
+        )
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content) == {
+            "input_tokens": _COUNT_MODEL_COUNT,
+            "litellm_count_source": "token_count_model",
+            "litellm_count_model": _OPENAI_COUNT_MODEL,
+        }
+        counts: Final = _openai_counts(openai)
+        assert tuple((target, counted) for target, counted, _ in counts) == (
+            (_OPENAI_COUNT_TARGET, _OPENAI_COUNT_MODEL),
+        ), counts
+        primary: Final = counts[0][2]
+        assert primary["instructions"] == _SYSTEM, primary
+        text: Final = primary.get("text")
+        assert isinstance(text, dict) and isinstance(text.get("format"), dict), primary
+        input_items: Final = primary["input"]
+        assert isinstance(input_items, list), primary
+        item_types: Final = {item.get("type") for item in input_items if isinstance(item, dict)}
+        assert {"function_call", "function_call_output"} <= item_types, input_items
+
+
+@pytest.mark.parametrize(
+    ("body", "named"),
+    [
+        (
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": "Summarize the notes."},
+                            {
+                                "type": "search_result",
+                                "source": "https://example.com/notes",
+                                "title": "Notes",
+                                "content": [],
+                            },
+                        ],
+                    }
+                ]
+            },
+            "search_result",
+        ),
+        (
+            {"messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]},
+            "messages.1.role: Field required",
+        ),
+    ],
+    ids=["unconvertible_block", "missing_role"],
+)
+def test_framing_keeps_the_400_for_a_body_the_openai_count_rejects_before_calling_openai(
+    counting_proxy: OwnedProxy, body: Mapping[str, JsonValue], named: str
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_openai_count) as openai, gateway.scenario() as scenario:
+        model: Final = _openai_deployment(scenario, openai, {"token_count_framing": _FRAMING})
+        rejected: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **body})
+        assert rejected.status_code == 400, rejected.text
+        error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"]
+        assert isinstance(error, dict) and named in str(error.get("message")), error
+        assert openai.drain() == ()
+        estimated: Final = gateway.request(
+            "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": "hello"}]}
+        )
+        assert estimated.status_code == 200, estimated.text
+        assert _JSON_OBJECT.validate_json(estimated.content).get("litellm_count_source") == "local_estimate"
+        assert tuple(counted for _, counted, _ in _openai_counts(openai)) == (_OPENAI_PRIMARY,)
+
+
+def _runtime_cannot_count(_request: Request) -> Reply:
+    return _json_reply(400, {"message": _UNSUPPORTED})
+
+
+_REPLAYED_WEB_SEARCH: Final[list[JsonValue]] = [
+    {"role": "user", "content": "What is the newest LiteLLM release?"},
+    {
+        "role": "assistant",
+        "content": [
+            {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "LiteLLM"}},
+            {
+                "type": "web_search_tool_result",
+                "tool_use_id": "srvtoolu_01",
+                "content": [
+                    {
+                        "type": "web_search_result",
+                        "url": "https://docs.litellm.ai/release_notes",
+                        "title": "Release notes",
+                        "encrypted_content": "",
+                    }
+                ],
+            },
+            {"type": "text", "text": "The newest release is in the release notes."},
+        ],
+    },
+    {"role": "user", "content": "Thanks."},
+]
+_BOOLEAN_SUBSCHEMA_TOOL: Final[dict[str, JsonValue]] = {
+    "name": "lookup",
+    "description": "Look up a value",
+    "input_schema": {"type": "object", "properties": {"query": True}},
+}
+_UNREADABLE_BODIES: Final[dict[str, dict[str, JsonValue]]] = {
+    "web_search_replay": {"messages": _REPLAYED_WEB_SEARCH},
+    "boolean_subschema_tool": {
+        "messages": [{"role": "user", "content": "Look up the value."}],
+        "tools": [_BOOLEAN_SUBSCHEMA_TOOL],
+    },
+}
+_COUNT_ENDPOINTS: Final[dict[str, tuple[str, str]]] = {
+    "messages_count_tokens": ("/v1/messages/count_tokens", "litellm_count_source"),
+    "utils_token_counter": ("/utils/token_counter?call_endpoint=true", "count_source"),
+}
+
+
+def _count_answer(response: httpx.Response, source_key: str) -> tuple[int, str]:
+    if response.status_code not in (200, 400):
+        return response.status_code, response.text
+    body: Final = _JSON_OBJECT.validate_json(response.content)
+    error: Final = body.get("error")
+    if response.status_code == 400 and isinstance(error, dict):
+        return 400, str(error.get("message"))
+    return response.status_code, str(body.get(source_key))
+
+
+@pytest.mark.parametrize("endpoint", sorted(_COUNT_ENDPOINTS))
+@pytest.mark.parametrize("unreadable", sorted(_UNREADABLE_BODIES))
+def test_shipped_framing_never_answers_a_500_for_a_body_the_local_count_cannot_read(
+    counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int, unreadable: str, endpoint: str
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    path, source_key = _COUNT_ENDPOINTS[endpoint]
+    with (
+        wire_server(_runtime_cannot_count, port=runtime_port) as runtime,
+        wire_server(_mantle, port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        response: Final = gateway.request(
+            "POST", path, {"model": _deployment(scenario, runtime_port, _KIMI), **_UNREADABLE_BODIES[unreadable]}
+        )
+        status, answer = _count_answer(response, source_key)
+        assert (status, answer) == (200, "local_estimate") or (status == 400 and _UNSUPPORTED in answer), (
+            response.status_code,
+            response.text,
+        )
+        assert _runtime_models(runtime) == (_KIMI_BASE,)
+        assert _mantle_models(mantle) == ()
+
+
+_GEMINI_MODEL: Final = "gemini-4-argon"
+_GEMINI_COUNT_MODEL: Final = "gemini-3.8-flash"
+_VERTEX_PROJECT: Final = "scripted-project"
+_COMPACT_UNDER_THE_MINIMUM_TRIGGER: Final[dict[str, JsonValue]] = {
+    "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
+}
+_CLEAR_OLDER_TOOL_RESULTS: Final[dict[str, JsonValue]] = {
+    "edits": [
+        {
+            "type": "clear_tool_uses_20250919",
+            "trigger": {"type": "input_tokens", "value": 500},
+            "keep": {"type": "tool_uses", "value": 1},
+        }
+    ]
+}
+_READ_FILE_TOOL: Final[dict[str, JsonValue]] = {
+    "name": "read_file",
+    "description": "Read a file",
+    "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
+}
+_FILE_TEXT: Final = "line of a long file that the agent read earlier. " * 180
+_FILES: Final = 3
+
+
+def _google_cannot_count(_request: Request) -> Reply:
+    return _json_reply(400, {"error": {"code": 400, "message": "scripted count 400", "status": "INVALID_ARGUMENT"}})
+
+
+def _google_counts_only_the_count_model(request: Request) -> Reply:
+    if request.target.endswith(f"/{_GEMINI_COUNT_MODEL}:countTokens"):
+        return _json_reply(200, {"totalTokens": _COUNT_MODEL_COUNT})
+    return _google_cannot_count(request)
+
+
+def _google_deployment(
+    gateway: Gateway, scenario: Scenario, wire: Wire, provider: str, model_info: dict[str, JsonValue]
+) -> str:
+    if provider == "gemini":
+        return scenario.model(
+            model=f"gemini/{_GEMINI_MODEL}", api_base=wire.url, api_key="scripted-gemini-key", model_info=model_info
+        )
+    return scenario.model(
+        model=f"vertex_ai/{_GEMINI_MODEL}",
+        api_base=wire.url,
+        api_key=None,
+        vertex_project=_VERTEX_PROJECT,
+        vertex_location="us-east5",
+        vertex_credentials=service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")),
+        model_info=model_info,
+    )
+
+
+def _google_count_targets(wire: Wire) -> tuple[bool, ...]:
+    return tuple(request.target.endswith(":countTokens") for request in wire.drain())
+
+
+def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]:
+    return (
+        {"role": "user", "content": f"Read file {index}."},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}}
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "tool_result",
+                    "tool_use_id": f"toolu_{index}",
+                    "content": f"<file {marker}-{index}>{_FILE_TEXT}",
+                }
+            ],
+        },
+    )
+
+
+def _agent_messages(marker: str) -> list[JsonValue]:
+    rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES)))
+    return [*rounds, {"role": "user", "content": "Summarize what you read."}]
+
+
+def _counted_files(request: Request, marker: str) -> tuple[bool, tuple[bool, ...]]:
+    sent: Final = json.dumps(_JSON_OBJECT.validate_json(request.body))
+    return (
+        request.target.endswith(f"/{_GEMINI_COUNT_MODEL}:countTokens"),
+        tuple(f"<file {marker}-{index}>" in sent for index in range(_FILES)),
+    )
+
+
+@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"])
+def test_framing_keeps_the_400_for_a_context_edit_the_gemini_count_rejects_before_calling_google(
+    counting_proxy: OwnedProxy, provider: str
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    hello: Final[list[JsonValue]] = [{"role": "user", "content": "hello"}]
+    with wire_server(_google_cannot_count) as google, gateway.scenario() as scenario:
+        model: Final = _google_deployment(gateway, scenario, google, provider, {"token_count_framing": _FRAMING})
+        rejected: Final = gateway.request(
+            "POST",
+            "/v1/messages/count_tokens",
+            {"model": model, "messages": hello, "context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER},
+        )
+        assert rejected.status_code == 400, rejected.text
+        assert _google_count_targets(google) == ()
+        estimated: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, "messages": hello})
+        assert estimated.status_code == 200, estimated.text
+        assert _JSON_OBJECT.validate_json(estimated.content).get("litellm_count_source") == "local_estimate"
+        assert _google_count_targets(google) == (True,)
+
+
+@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"])
+def test_token_count_model_counts_the_body_left_after_the_context_management_edits(
+    counting_proxy: OwnedProxy, provider: str
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    marker: Final = uuid4().hex
+    with wire_server(_google_counts_only_the_count_model) as google, gateway.scenario() as scenario:
+        model: Final = _google_deployment(
+            gateway, scenario, google, provider, {"token_count_model": {"model": _GEMINI_COUNT_MODEL}}
+        )
+        counted: Final = gateway.request(
+            "POST",
+            "/v1/messages/count_tokens",
+            {
+                "model": model,
+                "tools": [_READ_FILE_TOOL],
+                "messages": _agent_messages(marker),
+                "context_management": _CLEAR_OLDER_TOOL_RESULTS,
+            },
+        )
+        assert counted.status_code == 200, counted.text
+        assert _JSON_OBJECT.validate_json(counted.content) == {
+            "input_tokens": _COUNT_MODEL_COUNT,
+            "litellm_count_source": "token_count_model",
+            "litellm_count_model": _GEMINI_COUNT_MODEL,
+        }
+        assert tuple(_counted_files(request, marker) for request in google.drain()) == ((True, (False, False, True)),)

tests/integration/providers/test_malformed_tools_wire.py

@@ -46,6 +46,7 @@ _CLAUDE_BACKEND: Final = "claude-haiku-4-5"
 _GEMINI_BACKEND: Final = "gemini-2.5-flash"
 _PROVIDER_KEY: Final = "scripted-malformed-tools-key"
 _COUNT: Final = 37
+_PROVIDER_SOURCE: Final[Mapping[str, JsonValue]] = MappingProxyType({"litellm_count_source": "provider"})
 _CONFIG_STORE: Final = "vs_integration_config_store"
 _OTHER_STORE: Final = "vs_malformed_tools_other_store"
 _CLAUDE_CODE_AGENT: Final = "claude-cli/2.1.0 (external, cli)"
@@ -821,7 +822,7 @@ def test_count_tokens_forwards_well_formed_tools(rig: _Rig) -> None:
         "POST", "/v1/messages/count_tokens", _count_body(rig, marker, tools=[MESSAGES.tool])
     )
     assert response.status_code == 200, response.text
-    assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text
+    assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT, **_PROVIDER_SOURCE}, response.text
     assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool]
 
 
@@ -853,6 +854,7 @@ def test_responses_input_tokens_counts_well_formed_tools(rig: _Rig) -> None:
     assert _JSON_OBJECT.validate_json(response.content) == {
         "object": "response.input_tokens",
         "input_tokens": _COUNT,
+        **_PROVIDER_SOURCE,
     }, response.text
     sent: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens")
     assert "get_weather" in json.dumps(sent["tools"]), sent
@@ -1396,7 +1398,7 @@ def test_count_tokens_drops_the_bad_entry_a_deployment_list_names(rig: _Rig, fie
         {**_count_body(rig, marker, tools=[MESSAGES.tool, 5]), "model": field_groups.anthropic},
     )
     assert response.status_code == 200, response.text
-    assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text
+    assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT, **_PROVIDER_SOURCE}, response.text
     assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool]
 
 
@@ -1656,7 +1658,7 @@ _CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType(
             named_as=MESSAGES,
             send=_count_tokens,
             target="/v1/messages/count_tokens",
-            answered=_counted({"input_tokens": _COUNT}),
+            answered=_counted({"input_tokens": _COUNT, **_PROVIDER_SOURCE}),
             kept_by_a_tools_list=("get_weather",),
             bad_entry=_MESSAGES_BAD_ENTRY,
         ),
@@ -1666,7 +1668,7 @@ _CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType(
             named_as=MESSAGES,
             send=_input_tokens,
             target="/v1/messages/count_tokens",
-            answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT}),
+            answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT, **_PROVIDER_SOURCE}),
             kept_by_a_tools_list=("get_weather",),
             bad_entry=_RESPONSES_BAD_ENTRY,
         ),

tests/integration/providers/test_responses_bridge_message_cache_control_wire.py

@@ -1118,6 +1118,7 @@ def test_input_tokens_counts_marked_items_without_their_marks(gateway: Gateway)
     assert _JSON_OBJECT.validate_json(response.text) == {
         "object": "response.input_tokens",
         "input_tokens": _COUNTED_TOKENS,
+        "litellm_count_source": "provider",
     }
     assert counted.target.endswith("/count_tokens"), counted.target
     assert anthropic_labels(counted) == []
@@ -1199,6 +1200,7 @@ def test_input_tokens_reads_each_input_shape_without_item_marks(gateway: Gateway
     assert _JSON_OBJECT.validate_json(response.text) == {
         "object": "response.input_tokens",
         "input_tokens": _COUNTED_TOKENS,
+        "litellm_count_source": "provider",
     }
     assert [anthropic_labels(request) for request in counted] == [[]]
 

tests/integration/providers/test_vertex_partner_count_tokens_wire.py

@@ -396,7 +396,7 @@ def test_messages_count_tokens_forwards_system_and_tools(gateway: Gateway, syste
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url)
         response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system, "tools": _TOOLS},)
 
 
@@ -404,7 +404,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_bare_body(gateway:
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url)
         response: Final = _count(gateway, _bare(model))
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         assert _counted_bodies(wire) == (_PEER_BARE,)
 
 
@@ -443,7 +443,11 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway
         response: Final = gateway.request(
             "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"}
         )
-        assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {
+            "object": "response.input_tokens",
+            "input_tokens": _PEER_COUNT,
+            "litellm_count_source": "provider",
+        }, response.text
         assert _counted_bodies(wire) == (_PEER_BARE,)
 
 
@@ -498,7 +502,7 @@ def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> Non
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}")
         response: Final = _count(gateway, _bare(model))
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         (sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET)
         assert sent["messages"] == _MESSAGES, sent
 
@@ -616,7 +620,7 @@ def test_messages_count_tokens_forwards_an_empty_tools_list(gateway: Gateway) ->
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url)
         response: Final = _count(gateway, {**_bare(model), "tools": []})
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": []},)
 
 
@@ -624,7 +628,7 @@ def test_messages_count_tokens_forwards_an_empty_system_string(gateway: Gateway)
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url)
         response: Final = _count(gateway, {**_bare(model), "system": ""})
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         assert _counted_bodies(wire) == ({**_PEER_BARE, "system": ""},)
 
 
@@ -632,7 +636,7 @@ def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url)
         response: Final = _count(gateway, {**_bare(model), "system": None, "tools": None})
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         assert _counted_bodies(wire) == (_PEER_BARE,)
 
 
@@ -650,7 +654,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway)
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url)
         response: Final = _count(gateway, {**_bare(model), "system": system})
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system},)
 
 
@@ -663,7 +667,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value
             content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}',
             headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
         )
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         (sent,) = _count_requests(wire.drain())
         assert _count_bodies((sent,)) == (_PEER_FULL,)
         assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body
@@ -696,7 +700,7 @@ def test_count_tokens_location_override_targets_the_count_region(gateway: Gatewa
             gateway, scenario, wire.url, vertex_location="global", vertex_count_tokens_location="europe-west1"
         )
         response: Final = _count(gateway, _bare(model))
-        assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text
+        assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text
         target: Final = _COUNT_TARGET.replace(f"/locations/{_LOCATION}/", "/locations/europe-west1/")
         assert _count_bodies(wire.drain(), target) == (_PEER_BARE,)
 
@@ -705,7 +709,7 @@ def test_messages_count_tokens_repeated_request_reaches_the_peer_each_time(gatew
     with wire_server(_peer()) as wire, gateway.scenario() as scenario:
         model: Final = _deployment(gateway, scenario, wire.url)
         answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2))
-        assert answers == ({"input_tokens": _PEER_COUNT},) * 2
+        assert answers == ({"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"},) * 2
         assert _counted_bodies(wire) == (_PEER_BARE,) * 2
 
 
@@ -722,7 +726,7 @@ def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_l
         )
         with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
             counted: Final = _count(owned.gateway, _full(_OWNED_MODEL))
-            assert _payload(counted) == {"input_tokens": _PEER_COUNT}, counted.text
+            assert _payload(counted) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, counted.text
             rejected: Final = _count(
                 owned.gateway, {**_full(_OWNED_MODEL), "messages": [{"role": "user", "content": _REJECT_TEXT}]}
             )

tests/integration/spend/test_cache_and_quota.py

@@ -453,9 +453,11 @@ def test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completi
             [{"role": "user", "content": prompt}]
         ]
         assert count.status_code == 200, count.text
-        assert count.json() == {"totalTokens": 12, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}]}, (
-            count.text
-        )
+        assert count.json() == {
+            "totalTokens": 12,
+            "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}],
+            "litellm_count_source": "provider",
+        }, count.text
         provider_calls: Final = wire.drain()
         assert [(call.method, call.target) for call in provider_calls] == [
             ("POST", "/v1beta/models/gemini-3.8-flash:countTokens")

tests/unit/batches/test_batch_utils.py

@@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable.
 import asyncio, json, time
 import logging
 from types import MappingProxyType
+from typing import Final, Literal
 
 import httpx
 import pytest
@@ -3347,3 +3348,30 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
         assert mock_batch._hidden_params["batch_successful_requests"] == 8
         assert mock_batch._hidden_params["batch_failed_requests"] == 0
         assert mock_batch.usage == expected_usage
+
+
+@pytest.mark.parametrize("provider", ("openai", "anthropic"))
+@pytest.mark.parametrize("envelope", (1, "invalid", [1], None, {}))
+@pytest.mark.parametrize("include_valid", (True, False))
+def test_invalid_batch_envelopes_preserve_recoverable_cost(
+    provider: Literal["openai", "anthropic"], envelope: object, include_valid: bool
+) -> None:
+    invalid: Final[dict[str, object]] = {"response" if provider == "openai" else "result": envelope}
+    valid: Final[dict[str, object]] = (
+        _success_row(model="batch-accounting-control", usage=_usage(100, 20))
+        if provider == "openai"
+        else {
+            "result": {
+                "type": "succeeded",
+                "message": {"model": "batch-accounting-control", "usage": {"input_tokens": 100, "output_tokens": 20}},
+            }
+        }
+    )
+    entries: Final = (valid, invalid, valid) if include_valid else (invalid,)
+    result: Final = bu._aggregate_batch_cost_usage_models(
+        entries, provider, model_info={"input_cost_per_token_batches": 0.001, "output_cost_per_token_batches": 0.002}
+    )
+    assert result.cost == pytest.approx(0.28 if include_valid else 0)
+    assert result.usage.prompt_tokens == (200 if include_valid else 0)
+    assert result.usage.completion_tokens == (40 if include_valid else 0)
+    assert (result.successful_requests, result.failed_requests) == (2 if include_valid else 0, 1)

tests/unit/enterprise/proxy/test_managed_files_access_check.py

@@ -9,6 +9,8 @@ with deployment credentials, bypassing the managed files access-control hooks.
 """
 
 import base64
+import json
+from typing import Final
 import pytest
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -367,13 +369,14 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti
         # managed_files_obj.afile_content should NOT have been called
         mock_managed_files_hook.afile_content.assert_not_called()
 
-        # Verify the DB update writes batch_processed, status, and file_object
-        mock_prisma.db.litellm_managedobjecttable.update.assert_called_once()
-        update_call_kwargs = (
-            mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs
+        claim_data: Final = next(
+            call.kwargs["data"]
+            for call in mock_prisma.db.litellm_managedobjecttable.update_many.call_args_list
+            if call.kwargs["where"].get("id") == mock_job.id
         )
-        assert update_call_kwargs["data"]["batch_processed"] is True
-        assert update_call_kwargs["data"]["status"] == "complete"
-        assert (
-            "file_object" in update_call_kwargs["data"]
-        ), "file_object must be written to DB so list_batches reads updated status"
+        assert claim_data["batch_processed"] is True
+        assert claim_data["status"] == "complete"
+        persisted_batch: Final = json.loads(claim_data["file_object"])
+        assert persisted_batch["status"] == "completed"
+        assert persisted_batch["usage"]["prompt_tokens"] == 10
+        assert persisted_batch["usage"]["completion_tokens"] == 5

tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py

@@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
             "request_model": "azure_ai/claude-sonnet-4-5",
             "model_used": "claude-sonnet-4-5",
             "tokenizer_type": "azure_ai_anthropic_api",
+            "count_source": None,
             "original_response": payload,
             "error": False,
             "error_message": None,
@@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v
             "request_model": "azure_ai/claude-sonnet-4-5",
             "model_used": "claude-sonnet-4-5",
             "tokenizer_type": "azure_ai_anthropic_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": str(rejection.value),
@@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam
             "request_model": "azure_ai/claude-sonnet-4-5",
             "model_used": "claude-sonnet-4-5",
             "tokenizer_type": "azure_ai_anthropic_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": str(failed_lookup.value),
@@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_
             "request_model": "azure_ai/claude-sonnet-4-5",
             "model_used": "claude-sonnet-4-5",
             "tokenizer_type": "azure_ai_anthropic_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": upstream_text,
@@ -180,6 +184,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_
             "request_model": "azure_ai/claude-sonnet-4-5",
             "model_used": "claude-sonnet-4-5",
             "tokenizer_type": "azure_ai_anthropic_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": f"CountTokens processing error: {undecodable.value}",

tests/unit/llms/baseten/chat/test_baseten_completions.py

@@ -18,10 +18,7 @@ class TestBasetenRouting:
         )
 
         # Model API (non-8-character)
-        assert (
-            config.get_api_base_for_model("openai/gpt-oss-120b")
-            == "https://inference.baseten.co/v1"
-        )
+        assert config.get_api_base_for_model("openai/gpt-oss-120b") == "https://inference.baseten.co/v1"
 
 
 class TestBasetenModelAPI:
@@ -47,12 +44,37 @@ class TestBasetenModelAPI:
         assert result["top_p"] == 0.9
 
         # Test provider info
-        api_base, api_key = config._get_openai_compatible_provider_info(
-            None, "test-key"
-        )
+        api_base, api_key = config._get_openai_compatible_provider_info(None, "test-key")
         assert api_base == "https://inference.baseten.co/v1"
         assert api_key == "test-key"
 
 
+class TestBasetenReasoningEffort:
+    def test_reasoning_effort_passthrough_with_drop_params(self):
+        config = BasetenConfig()
+
+        result = config.map_openai_params(
+            non_default_params={"reasoning_effort": "high", "max_tokens": 100},
+            optional_params={},
+            model="deepseek-ai/DeepSeek-V4-Flash-0731",
+            drop_params=True,
+        )
+
+        assert result["reasoning_effort"] == "high"
+        assert result["max_tokens"] == 100
+
+    def test_reasoning_effort_none_passthrough(self):
+        config = BasetenConfig()
+
+        result = config.map_openai_params(
+            non_default_params={"reasoning_effort": "none"},
+            optional_params={},
+            model="moonshotai/Kimi-K3",
+            drop_params=True,
+        )
+
+        assert result["reasoning_effort"] == "none"
+
+
 if __name__ == "__main__":
     pytest.main([__file__])

tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py

@@ -0,0 +1,524 @@
+import asyncio
+import base64
+import hashlib
+import json
+import logging
+import re
+import threading
+from collections.abc import Iterable
+from itertools import product
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+import respx
+import tiktoken
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.llms.bedrock.count_tokens.kimi_k3 import (
+    KimiK3Encodings,
+    KimiK3TokenizerSource,
+    kimi_k3_anthropic_tokens,
+    kimi_k3_tokens,
+    load_kimi_k3_encoding,
+)
+from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig
+
+MODEL: Final = "us.moonshotai.kimi-k3"
+_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>")
+_WEATHER: Final = {
+    "type": "function",
+    "function": {
+        "name": "get_weather",
+        "description": "Get the weather in °C",
+        "parameters": {
+            "type": "object",
+            "properties": {"city": {"type": "string"}, "days": {"type": "integer"}},
+            "required": ["city"],
+        },
+    },
+}
+
+
+def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]:
+    return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}}
+
+
+def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource:
+    path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens)))
+    return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest())
+
+
+def _single_bytes() -> tuple[bytes, ...]:
+    return tuple(bytes((value,)) for value in range(256))
+
+
+@pytest.fixture(scope="module")
+def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource:
+    return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes())
+
+
+@pytest.fixture(scope="module")
+def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding:
+    loaded: Final = load_kimi_k3_encoding(byte_source)
+    assert loaded is not None
+    return loaded
+
+
+@pytest.fixture(scope="module")
+def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding:
+    pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2))
+    source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs))
+    loaded: Final = load_kimi_k3_encoding(source)
+    assert loaded is not None
+    return loaded
+
+
+def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int:
+    texts: Final = _CONTROL.split(rendered)
+    return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts)
+
+
+_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>'
+
+
+# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched
+# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16
+@pytest.mark.parametrize(
+    ("messages", "tools", "rendered"),
+    [
+        pytest.param(
+            [{"role": "user", "content": "Hello there"}],
+            None,
+            '<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
+            id="user",
+        ),
+        pytest.param(
+            [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}],
+                },
+                {"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]},
+            ],
+            None,
+            '<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>'
+            '<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
+            id="text_parts_join",
+        ),
+        pytest.param(
+            [
+                {"role": "user", "content": "What is 2+2?"},
+                {"role": "assistant", "content": "4"},
+                {"role": "user", "content": "And 3+3?"},
+            ],
+            None,
+            '<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>'
+            '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
+            "<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>"
+            '<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
+            id="assistant_turn",
+        ),
+        pytest.param(
+            [
+                {"role": "user", "content": "Weather in Paris and Oslo?"},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})],
+                },
+                {"role": "tool", "tool_call_id": "c2", "content": "Snow"},
+                {"role": "tool", "tool_call_id": "c1", "content": "Sun"},
+            ],
+            [_WEATHER],
+            '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
+            "Here are the available tools, described in JSONSchema.\n\n```json\n"
+            '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
+            '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
+            '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
+            '<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>'
+            '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
+            "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
+            '<|open|>call tool="get_weather" index="1"<|sep|>'
+            '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>'
+            '<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>'
+            '<|open|>call tool="get_weather" index="2"<|sep|>'
+            '<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>'
+            "<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
+            '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>'
+            '<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>'
+            + _GENERATION_PROMPT,
+            id="tool_results_out_of_call_order",
+        ),
+        pytest.param(
+            [
+                {"role": "user", "content": "Weather?"},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        _call(
+                            "c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}}
+                        )
+                    ],
+                },
+                {"role": "tool", "tool_call_id": "c1", "content": "Rain"},
+            ],
+            [_WEATHER],
+            '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
+            "Here are the available tools, described in JSONSchema.\n\n```json\n"
+            '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
+            '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
+            '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
+            '<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>'
+            '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
+            "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
+            '<|open|>call tool="get_weather" index="1"<|sep|>'
+            '<|open|>argument key="a&amp;b" type="string"<|sep|>x"y<|close|>argument<|sep|>'
+            '<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>'
+            '<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>'
+            '<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>'
+            '<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>'
+            "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
+            '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>'
+            + _GENERATION_PROMPT,
+            id="argument_types_and_escaping",
+        ),
+    ],
+)
+@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"])
+def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request(
+    request: pytest.FixtureRequest,
+    encoding_fixture: str,
+    messages: list[dict[str, object]],
+    tools: list[dict[str, object]] | None,
+    rendered: str,
+):
+    encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture)
+
+    assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered)
+
+
+def _parallel_calls(count: int) -> list[dict[str, object]]:
+    ids: Final = tuple(f"c{number}" for number in range(1, count + 1))
+    return [
+        {"role": "user", "content": "Weather?"},
+        {"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]},
+        *({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids),
+    ]
+
+
+def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding):
+    nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER])
+    ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER])
+    assert nine is not None and ten is not None
+
+    assert ten - nine == _template_count(
+        byte_encoding,
+        '<|open|>call tool="get_weather" index="10"<|sep|>'
+        '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>'
+        '<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>',
+    )
+
+
+@pytest.mark.parametrize(
+    "text",
+    [
+        pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"),
+        pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"),
+        pytest.param("--- " * 600, id="symbols"),
+        pytest.param("def f(x):\n    return x  # it's 42\n\n" * 80, id="code"),
+        pytest.param("word   next\t\ttab\n \n" * 120, id="whitespace"),
+        pytest.param("naïve café, “quotes” & <tags> 中文汉字，日本語のテキスト 😀🎉\r\n" * 40, id="unicode"),
+        pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"),
+        pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"),
+    ],
+)
+def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str):
+    def counted(content: str) -> int | None:
+        return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None)
+
+    short: Final = counted("x")
+    assert short is not None
+
+    assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text))
+
+
+@pytest.mark.parametrize(
+    ("messages", "tools"),
+    [
+        pytest.param(
+            [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "What is this?"},
+                        {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
+                    ],
+                }
+            ],
+            None,
+            id="image",
+        ),
+        pytest.param(
+            [{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}],
+            [_WEATHER],
+            id="tool_result_without_its_call",
+        ),
+    ],
+)
+def test_request_the_template_cannot_render_is_not_counted(
+    pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None
+):
+    assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None
+
+
+@pytest.mark.parametrize(
+    "tool",
+    [
+        pytest.param(
+            {
+                "type": "function",
+                "name": "get_weather",
+                "description": "Get the weather in °C",
+                "parameters": _WEATHER["function"]["parameters"],
+            },
+            id="responses_api",
+        ),
+        pytest.param(
+            {
+                "name": "get_weather",
+                "description": "Get the weather in °C",
+                "input_schema": _WEATHER["function"]["parameters"],
+            },
+            id="anthropic_messages",
+        ),
+    ],
+)
+def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool(
+    pair_encoding: tiktoken.Encoding, tool: dict[str, object]
+):
+    messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
+
+    counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool])
+
+    assert counted is not None
+    assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
+
+
+@pytest.mark.parametrize(
+    "tool",
+    [
+        pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"),
+        pytest.param({"type": "web_search"}, id="responses_web_search"),
+        pytest.param(
+            {
+                "type": "mcp",
+                "server_label": "docs",
+                "server_url": "https://mcp.example/mcp",
+                "require_approval": "never",
+            },
+            id="mcp",
+        ),
+    ],
+)
+def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone(
+    pair_encoding: tiktoken.Encoding, tool: dict[str, object]
+):
+    messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
+
+    counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool])
+
+    assert counted is not None
+    assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
+
+
+def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding):
+    messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
+    custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"}
+    sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
+        [dict(custom_tool)]
+    )
+
+    without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
+
+    counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool])
+
+    assert counted is not None and without_it is not None
+    assert counted > without_it
+    assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference])
+
+
+def test_counting_a_request_with_an_image_url_never_downloads_the_image(
+    pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setattr(litellm, "user_url_validation", False)
+    image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200))
+    messages: Final = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What is this?"},
+                {"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}},
+            ],
+        }
+    ]
+
+    assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None
+    assert not image.called
+
+
+def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding):
+    anthropic: Final = kimi_k3_anthropic_tokens(
+        pair_encoding,
+        MODEL,
+        {
+            "system": "Be terse.",
+            "messages": [
+                {"role": "user", "content": "Weather in Paris?"},
+                {
+                    "role": "assistant",
+                    "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}],
+                },
+                {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]},
+            ],
+            "tools": [
+                {
+                    "name": "get_weather",
+                    "description": "Get the weather in °C",
+                    "input_schema": _WEATHER["function"]["parameters"],
+                }
+            ],
+        },
+    )
+    openai: Final = kimi_k3_tokens(
+        pair_encoding,
+        MODEL,
+        [
+            {"role": "system", "content": "Be terse."},
+            {"role": "user", "content": "Weather in Paris?"},
+            {"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]},
+            {"role": "tool", "tool_call_id": "c1", "content": "Sun"},
+        ],
+        [_WEATHER],
+    )
+
+    assert openai is not None
+    assert anthropic == openai
+
+
+@pytest.mark.parametrize("failure", ["checksum", "missing_file"])
+def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded(
+    byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str
+):
+    source: Final = (
+        KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest())
+        if failure == "checksum"
+        else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256)
+    )
+    verbose_logger.addHandler(caplog.handler)
+    try:
+        with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
+            loaded: Final = load_kimi_k3_encoding(source)
+    finally:
+        verbose_logger.removeHandler(caplog.handler)
+
+    assert loaded is None
+    assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records)
+
+
+@pytest.mark.asyncio
+async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done(
+    byte_source: KimiK3TokenizerSource,
+):
+    released: Final = threading.Event()
+
+    def slow_ranks_path() -> str:
+        released.wait(10)
+        return byte_source.ranks_path()
+
+    encodings: Final = KimiK3Encodings(
+        source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256),
+        retry_after_seconds=60,
+        clock=lambda: 0.0,
+    )
+
+    assert await encodings.encoding(wait_seconds=0) is None
+    released.set()
+    loaded: Final = await encodings.encoding(wait_seconds=5)
+    assert loaded is not None
+    assert await encodings.encoding(wait_seconds=0) is loaded
+
+
+class _Clock:
+    def __init__(self) -> None:
+        self.now = 0.0
+
+    def __call__(self) -> float:
+        return self.now
+
+
+async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None:
+    for _ in range(500):
+        if (loaded := await encodings.encoding(wait_seconds=0)) is not None:
+            return loaded
+        await asyncio.sleep(0.01)
+    return None
+
+
+@pytest.mark.asyncio
+async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window(
+    byte_source: KimiK3TokenizerSource, tmp_path: Path
+):
+    clock: Final = _Clock()
+    reads: Final[list[float]] = []
+    released: Final = threading.Event()
+
+    def ranks_path() -> str:
+        reads.append(clock.now)
+        if len(reads) == 1:
+            return str(tmp_path / "missing")
+        released.wait(10)
+        return byte_source.ranks_path()
+
+    encodings: Final = KimiK3Encodings(
+        source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256),
+        retry_after_seconds=60,
+        clock=clock,
+    )
+
+    assert await encodings.encoding(wait_seconds=5) is None
+    clock.now = 59.0
+    assert await encodings.encoding(wait_seconds=5) is None
+    clock.now = 60.0
+    assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None
+    released.set()
+    assert await _loaded_eventually(encodings) is not None
+    assert reads == [0.0, 60.0]
+
+
+@pytest.mark.asyncio
+async def test_tokenizer_load_still_running_after_the_wait_is_logged(
+    byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture
+):
+    released: Final = threading.Event()
+
+    def hung_ranks_path() -> str:
+        released.wait(10)
+        return byte_source.ranks_path()
+
+    encodings: Final = KimiK3Encodings(
+        source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60
+    )
+    verbose_logger.addHandler(caplog.handler)
+    try:
+        with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
+            loaded: Final = await encodings.encoding(wait_seconds=0.05)
+    finally:
+        verbose_logger.removeHandler(caplog.handler)
+        released.set()
+
+    assert loaded is None
+    assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records)

tests/unit/llms/openai/responses/count_tokens/test_token_counter.py

@@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
             "request_model": "openai/gpt-4o",
             "model_used": "gpt-4o",
             "tokenizer_type": "openai_api",
+            "count_source": None,
             "original_response": payload,
             "error": False,
             "error_message": None,
@@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v
             "request_model": "openai/gpt-4o",
             "model_used": "gpt-4o",
             "tokenizer_type": "openai_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": str(rejection.value),
@@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam
             "request_model": "openai/gpt-4o",
             "model_used": "gpt-4o",
             "tokenizer_type": "openai_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": str(failed_lookup.value),
@@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_
             "request_model": "openai/gpt-4o",
             "model_used": "gpt-4o",
             "tokenizer_type": "openai_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": upstream_text,
@@ -382,6 +386,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_
             "request_model": "openai/gpt-4o",
             "model_used": "gpt-4o",
             "tokenizer_type": "openai_api",
+            "count_source": None,
             "original_response": None,
             "error": True,
             "error_message": f"CountTokens processing error: {undecodable.value}",

tests/unit/proxy/common_utils/test_check_batch_cost.py

@@ -17,6 +17,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
 from fastapi import HTTPException
 
+from litellm.types.utils import LiteLLMBatch
+
 if TYPE_CHECKING:
     from litellm.batches.batch_utils import BatchCostUsageResult
 
@@ -668,10 +670,16 @@ class TestCheckBatchCost:
         assert check_batch_cost_instance._has_batch_processed_column is True
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
 
-        mock_response = MagicMock()
-        mock_response.status = "completed"
-        mock_response.output_file_id = "file-output-123"
-        mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
+        mock_response: Final = LiteLLMBatch(
+            id="batch-1",
+            status="completed",
+            output_file_id="file-output-123",
+            created_at=1,
+            completion_window="24h",
+            endpoint="/v1/chat/completions",
+            input_file_id="file-input-1",
+            object="batch",
+        )
 
         mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
         mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
@@ -730,10 +738,10 @@ class TestCheckBatchCost:
 
             await check_batch_cost_instance.check_batch_cost()
 
-        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, (
             "Expected update() to be called exactly once for the completed job"
         )
-        update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["batch_processed"] is True, (
             "update() must include batch_processed=True when column is present"
         )
@@ -863,7 +871,7 @@ class TestCheckBatchCost:
 
         mock_update_database.assert_awaited_once()
         assert mock_update_database.call_args.kwargs["response_cost"] == 0.01
-        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, (
             "the job must still be marked processed once cost tracking succeeds"
         )
 
@@ -1285,11 +1293,16 @@ class TestCheckBatchCost:
         assert check_batch_cost_instance._has_batch_processed_column is True
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
 
-        mock_response = MagicMock()
-        mock_response.status = terminal_status
-        mock_response.output_file_id = "file-output-123"
-        mock_response.error_file_id = None
-        mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}'
+        mock_response: Final = LiteLLMBatch(
+            id="batch-1",
+            status=terminal_status,
+            output_file_id="file-output-123",
+            created_at=1,
+            completion_window="24h",
+            endpoint="/v1/chat/completions",
+            input_file_id="file-input-1",
+            object="batch",
+        )
 
         mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
         mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
@@ -1352,8 +1365,8 @@ class TestCheckBatchCost:
             f"{terminal_status} batch with an output file must fetch results and be billed"
         )
         mock_logging_obj.async_success_handler.assert_awaited_once()
-        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
-        update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["batch_processed"] is True
         assert update_data["status"] == terminal_status, (
             f"billed {terminal_status} batch must keep its real terminal status in the DB"
@@ -1847,13 +1860,16 @@ class TestUnmanagedVertexRouting:
         router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash"
         router.get_model_ids.return_value = ["deploy-1"]
 
-        mock_response = MagicMock()
-        mock_response.status = "completed"
-        mock_response.output_file_id = "gs://bucket/out/predictions.jsonl"
-        mock_response.error_file_id = None
-        mock_response.completed_at = None
-        mock_response.created_at = None
-        mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}'
+        mock_response: Final = LiteLLMBatch(
+            id="8823717160934178816",
+            status="completed",
+            output_file_id="gs://bucket/out/predictions.jsonl",
+            created_at=1,
+            completion_window="24h",
+            endpoint="/v1/chat/completions",
+            input_file_id="file-input-1",
+            object="batch",
+        )
         router.aretrieve_batch = AsyncMock(return_value=mock_response)
         router.get_deployment_credentials_with_provider = MagicMock(
             return_value={"vertex_project": "p", "vertex_location": "us-central1"}
@@ -1921,8 +1937,8 @@ class TestUnmanagedVertexRouting:
         mock_logging_obj.async_success_handler.assert_awaited_once()
         assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01
 
-        assert prisma.db.litellm_managedobjecttable.update.call_count == 1
-        update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3
+        update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["batch_processed"] is True
         assert update_data["status"] == "complete"
 
@@ -2065,13 +2081,16 @@ class TestUnmanagedBedrockRouting:
         router.resolve_model_name_from_model_id.return_value = "claude-sonnet-4"
         router.get_model_ids.return_value = ["deploy-1"]
 
-        mock_response = MagicMock()
-        mock_response.status = "completed"
-        mock_response.output_file_id = "s3://bucket/out/predictions.jsonl"
-        mock_response.error_file_id = None
-        mock_response.completed_at = None
-        mock_response.created_at = None
-        mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}'
+        mock_response: Final = LiteLLMBatch(
+            id=self._ARN,
+            status="completed",
+            output_file_id="s3://bucket/out/predictions.jsonl",
+            created_at=1,
+            completion_window="24h",
+            endpoint="/v1/chat/completions",
+            input_file_id="file-input-1",
+            object="batch",
+        )
         router.aretrieve_batch = AsyncMock(return_value=mock_response)
         router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"})
 
@@ -2135,8 +2154,8 @@ class TestUnmanagedBedrockRouting:
         mock_logging_obj.async_success_handler.assert_awaited_once()
         assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02
 
-        assert prisma.db.litellm_managedobjecttable.update.call_count == 1
-        update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3
+        update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["batch_processed"] is True
         assert update_data["status"] == "complete"
 
@@ -3003,13 +3022,20 @@ class TestMultiPodBatchCostClaim:
 
     @staticmethod
     def _router():
-        response = MagicMock()
-        response.status = "completed"
-        response.output_file_id = _CLAIM_OUTPUT_FILE_ID
-        response.error_file_id = None
-        response.created_at = 1
-        response.completed_at = 2
-        response.model_dump_json.return_value = '{"id":"batch-456","status":"completed"}'
+        from litellm.types.utils import LiteLLMBatch
+
+        response: Final = LiteLLMBatch(
+            id="batch-456",
+            status="completed",
+            output_file_id=_CLAIM_OUTPUT_FILE_ID,
+            error_file_id=None,
+            created_at=1,
+            completed_at=2,
+            completion_window="24h",
+            endpoint="/v1/chat/completions",
+            input_file_id="file-input-1",
+            object="batch",
+        )
 
         deployment = MagicMock()
         deployment.litellm_params.custom_llm_provider = "openai"
@@ -3089,8 +3115,9 @@ class TestMultiPodBatchCostClaim:
             ("openai/batch-test-model", 2, 0, "batch-test-model"),
         ),
     )
+    @pytest.mark.parametrize("failure_envelope", (None, 1, "invalid", [1]))
     async def test_batch_accounting_preserves_models_and_mixed_results(
-        self, deployment_model: str, successes: int, failures: int, expected_model: str
+        self, deployment_model: str, successes: int, failures: int, expected_model: str, failure_envelope: object
     ) -> None:
         import base64
         from datetime import datetime
@@ -3125,7 +3152,7 @@ class TestMultiPodBatchCostClaim:
             },
             "error": None,
         }
-        failure: Final = {"response": None, "error": {"code": "batch_expired", "message": "expired"}}
+        failure: Final = {"response": failure_envelope, "error": {"code": "batch_expired", "message": "expired"}}
         rows: Final = (success,) * successes + (failure,) * failures
         content: Final = b"\n".join(json.dumps(row).encode() for row in rows)
 
@@ -3217,6 +3244,30 @@ class TestMultiPodBatchCostClaim:
         with patch("litellm.proxy.proxy_server.scheduler", scheduler):
             await guard._check_file_deletion_allowed(file_id)
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("status", ("completed", "failed", "expired", "cancelled"))
+    @pytest.mark.parametrize("metadata_text", ("before\x00after", r"\u0000"))
+    async def test_cost_claim_persists_terminal_metadata_without_a_later_write(
+        self, status: str, metadata_text: str
+    ) -> None:
+        row: Final = _FakeManagedObjectRow()
+        journal: Final = []
+        prisma: Final = self._prisma(row, journal)
+        router: Final = self._router()
+        router.aretrieve_batch.return_value = router.aretrieve_batch.return_value.model_copy(
+            update={"status": status, "metadata": {"annotation": metadata_text}}
+        )
+        prisma.db.litellm_managedobjecttable.update.side_effect = RuntimeError("final write unavailable")
+        worker: Final = self._instance(prisma, router)
+        with self._billing_patches(journal) as logging_obj:
+            await worker.check_batch_cost()
+            await worker.check_batch_cost()
+        assert row.batch_processed is True
+        assert row.status == ("complete" if status == "completed" else status)
+        assert json.loads(row.file_object)["status"] == status
+        assert json.loads(row.file_object)["metadata"] == {"annotation": metadata_text.replace("\x00", "")}
+        assert logging_obj.async_success_handler.await_count == 1
+
     @pytest.mark.asyncio
     async def test_winning_pod_claims_the_row_between_fetching_and_billing(self):
         """The claim flips batch_processed false -> true after the results are in hand and
@@ -3228,13 +3279,9 @@ class TestMultiPodBatchCostClaim:
         with self._billing_patches(journal) as logging_obj:
             await self._instance(prisma, self._router()).check_batch_cost()
 
-        assert journal == ["fetch", "claim", "bill", "finalize"]
-        assert self._claim_calls(prisma) == [
-            {
-                "where": {"id": "job-claim-1", "batch_processed": False},
-                "data": {"batch_processed": True},
-            }
-        ]
+        assert journal == ["fetch", "claim", "bill"]
+        assert row.status == "complete"
+        assert json.loads(row.file_object)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5}
         logging_obj.async_success_handler.assert_awaited_once()
         assert row.batch_processed is True
 
@@ -3254,18 +3301,17 @@ class TestMultiPodBatchCostClaim:
 
         assert journal == ["fetch"]
         logging_obj.async_success_handler.assert_not_awaited()
-        assert self._claim_calls(prisma) == [
-            {
-                "where": {"id": "job-claim-1", "batch_processed": False},
-                "data": {"batch_processed": True},
-            }
-        ]
+        assert row.status == "in_progress"
+        assert json.loads(row.file_object)["status"] == "in_progress"
 
     @pytest.mark.asyncio
-    async def test_a_failed_spend_log_write_releases_the_claim(self):
+    @pytest.mark.parametrize("stored_as_json_string", (True, False))
+    async def test_a_failed_spend_log_write_releases_the_claim(self, stored_as_json_string: bool) -> None:
         """A transient failure while billing a claimed batch must hand the row back, or its
         spend is silently lost instead of being retried on the next cycle."""
         row = _FakeManagedObjectRow()
+        original_file_object: Final = row.file_object if stored_as_json_string else json.loads(row.file_object)
+        row.file_object = original_file_object
         journal = []
         prisma = self._prisma(row, journal)
 
@@ -3274,10 +3320,16 @@ class TestMultiPodBatchCostClaim:
 
         assert journal == ["fetch", "claim", "bill", "release"]
         assert row.batch_processed is False
-        assert self._claim_calls(prisma)[-1] == {
-            "where": {"id": "job-claim-1", "batch_processed": True},
-            "data": {"batch_processed": False},
-        }
+        assert row.status == "in_progress"
+        assert json.loads(row.file_object) == original_file_object
+
+        with self._billing_patches(journal) as logging_obj:
+            await self._instance(prisma, self._router()).check_batch_cost()
+            await self._instance(prisma, self._router()).check_batch_cost()
+
+        assert row.batch_processed is True
+        assert row.status == "complete"
+        logging_obj.async_success_handler.assert_awaited_once()
 
     @pytest.mark.asyncio
     async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self):
@@ -3309,7 +3361,7 @@ class TestMultiPodBatchCostClaim:
         with self._billing_patches(survivor_journal) as survivor_logging:
             await self._instance(survivor_prisma, self._router()).check_batch_cost()
 
-        assert survivor_journal == ["fetch", "claim", "bill", "finalize"]
+        assert survivor_journal == ["fetch", "claim", "bill"]
         survivor_logging.async_success_handler.assert_awaited_once()
         assert row.batch_processed is True
 
@@ -3340,7 +3392,7 @@ class TestMultiPodBatchCostClaim:
             finish_fetch.set()
             await asyncio.wait_for(costing, timeout=5)
 
-        assert journal == ["fetch", "claim", "bill", "finalize"]
+        assert journal == ["fetch", "claim", "bill"]
         assert row.batch_processed is True
         await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID)
 
@@ -3465,7 +3517,6 @@ class TestBatchLineItemsFromThePoller:
             f"read {_CLAIM_OUTPUT_FILE_ID}",
             "claim",
             "bill",
-            "finalize",
             f"read {_CLAIM_OUTPUT_FILE_ID}",
             "read file-input-1",
             "lines",
@@ -3539,7 +3590,7 @@ class TestBatchLineItemsFromThePoller:
             await instance.check_batch_cost()
             await self._stored(instance)
 
-        assert journal[-4:] == ["finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"]
+        assert journal[-4:] == ["bill", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"]
         ((line,),) = [payload["lines"] for payload in logger.payloads]
         assert line["request"] == json.loads(self._INPUT)
         assert line["response"] == json.loads(self._OUTPUT)
@@ -3558,7 +3609,7 @@ class TestBatchLineItemsFromThePoller:
             await instance.check_batch_cost()
             await self._stored(instance)
 
-        assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "finalize", "lines"]
+        assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "lines"]
         assert row.batch_processed is True
         assert row.status == "complete"
 
@@ -3574,7 +3625,7 @@ class TestBatchLineItemsFromThePoller:
             await instance.check_batch_cost()
             await self._stored(instance)
 
-        assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"]
+        assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"]
 
     @pytest.mark.asyncio
     async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is_retired(self, monkeypatch):
@@ -3751,7 +3802,6 @@ class TestBatchLineItemsFromThePoller:
                 f"read {_CLAIM_OUTPUT_FILE_ID}",
                 "claim",
                 "bill",
-                "finalize",
                 f"read {_CLAIM_OUTPUT_FILE_ID}",
                 "read file-input-1",
                 "lines started",
@@ -3834,7 +3884,6 @@ class TestBatchLineItemsFromThePoller:
             f"read {_CLAIM_OUTPUT_FILE_ID}",
             "claim",
             "bill",
-            "finalize",
             "store batch-earlier",
             "store batch-456",
         ]
@@ -3863,7 +3912,7 @@ class TestBatchLineItemsFromThePoller:
             await instance.check_batch_cost()
             await self._stored(instance)
 
-        assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"]
+        assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"]
         assert logger.payloads == []
         assert (row.batch_processed, row.status) == (True, "complete")
         assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text
@@ -3887,7 +3936,7 @@ class TestBatchLineItemsFromThePoller:
             await instance.check_batch_cost()
             await self._stored(instance)
 
-        assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"]
+        assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"]
         assert logger.payloads == []
         assert (row.batch_processed, row.status) == (True, "complete")
         assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text

tests/unit/proxy/common_utils/test_token_count_settings.py

@@ -0,0 +1,209 @@
+import json
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+from litellm.proxy.common_utils.token_count_settings import (
+    ProviderCountFailure,
+    TokenCountModel,
+    TokenCountSettings,
+    count_source_fields,
+    resolve_token_count_settings,
+)
+from litellm.types.utils import TokenCountResponse
+
+_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1")
+_DEFAULT: Final = {
+    "token_count_framing": 97,
+    "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"},
+}
+_COST_MAP: Final = Path(__file__).parents[4] / "model_prices_and_context_window.json"
+
+
+def _in_us_west_2() -> str:
+    return "us-west-2"
+
+
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [
+        ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
+        (
+            {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}},
+            TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")),
+        ),
+        (
+            {"token_count_model": {"model": "custom-model"}},
+            TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
+        ),
+        (
+            {"token_count_model": {"model": "custom-model", "region": " "}},
+            TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
+        ),
+        ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)),
+        (
+            {"token_count_framing": None, "token_count_model": None},
+            TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
+        ),
+        (
+            {"token_count_framing": True, "token_count_model": "custom-model"},
+            TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
+        ),
+        ({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)),
+        ({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
+        ({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
+        ({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
+        ({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)),
+        ({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
+        (
+            {"token_count_model": {"model": "  ", "region": "us-east-1"}},
+            TokenCountSettings(framing=97, count_model=None),
+        ),
+    ],
+)
+def test_config_overrides_the_cost_map_default_key_by_key(
+    configured: dict[str, object], expected: TokenCountSettings
+) -> None:
+    assert (
+        resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2)
+        == expected
+    )
+
+
+@pytest.mark.parametrize(
+    ("deployment_region", "configured", "expected_region"),
+    [
+        ("us-west-2", {}, "us-east-1"),
+        ("eu-central-1", {}, None),
+        ("ca-central-1", {}, None),
+        ("us-gov-west-1", {}, None),
+        (
+            "eu-central-1",
+            {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}},
+            "us-east-1",
+        ),
+    ],
+)
+def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography(
+    deployment_region: str, configured: dict[str, object], expected_region: str | None
+) -> None:
+    settings: Final = resolve_token_count_settings(
+        configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region
+    )
+    assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region)
+
+
+def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None:
+    settings: Final = resolve_token_count_settings(
+        configured={},
+        default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}},
+        deployment_region=lambda: "eu-west-1",
+    )
+    assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1")
+
+
+def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None:
+    settings: Final = resolve_token_count_settings(
+        configured={},
+        default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}},
+        deployment_region=lambda: "ap-southeast-2",
+    )
+    assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4")
+
+
+def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None:
+    assert resolve_token_count_settings(
+        configured={}, default={}, deployment_region=_in_us_west_2
+    ) == TokenCountSettings(framing=None, count_model=None)
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)],
+)
+def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None:
+    settings: Final = resolve_token_count_settings(
+        configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2
+    )
+    assert settings.tokenizer == expected
+
+
+def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse:
+    return TokenCountResponse(
+        total_tokens=7,
+        request_model="sonnet",
+        model_used=model_used,
+        tokenizer_type="bedrock_api",
+        count_source=source,
+    )
+
+
+@pytest.mark.parametrize(
+    ("source", "expected"),
+    [
+        (None, {}),
+        ("provider", {"litellm_count_source": "provider"}),
+        ("local_estimate", {"litellm_count_source": "local_estimate"}),
+        (
+            "token_count_model",
+            {"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"},
+        ),
+    ],
+)
+def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted(
+    source: str | None, expected: dict[str, str]
+) -> None:
+    assert count_source_fields(_counted(source)) == expected
+
+
+def test_every_shipped_token_count_model_is_another_mapped_model_of_the_same_provider_in_a_named_region() -> None:
+    cost_map: Final = json.loads(_COST_MAP.read_text())
+    count_models: Final = {
+        name: entry["token_count_model"]
+        for name, entry in cost_map.items()
+        if isinstance(entry, dict) and "token_count_model" in entry
+    }
+    assert count_models
+    assert {
+        name: count_model
+        for name, count_model in count_models.items()
+        if cost_map.get(count_model["model"], {}).get("litellm_provider") != cost_map[name]["litellm_provider"]
+        or count_model["model"] == name
+        or not count_model.get("region")
+    } == {}
+
+
+_UPSTREAM_404: Final = {
+    "type": "error",
+    "request_id": "req_upstream",
+    "error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"},
+}
+
+
+@pytest.mark.parametrize(
+    ("message", "expected"),
+    [
+        (
+            json.dumps(_UPSTREAM_404),
+            {
+                **_UPSTREAM_404,
+                "error": {
+                    "type": "not_found_error",
+                    "message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1",
+                },
+            },
+        ),
+        (json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"),
+        ("scripted 404", "scripted 404. Enable eu-west-1"),
+        (
+            "Credential should be scoped to a valid region. ",
+            "Credential should be scoped to a valid region. Enable eu-west-1",
+        ),
+    ],
+)
+def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None:
+    hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1")
+
+    assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected
+    assert hinted.status_code == 404

tests/unit/proxy/google_endpoints/test_google_api_endpoints.py

@@ -4,10 +4,11 @@ Test to verify the Google GenAI proxy API endpoints
 """
 
 from typing import Final
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
+from litellm.types.utils import TokenCountResponse
 
 
 def _build_test_client():
@@ -88,9 +89,7 @@ def test_google_stream_generate_content_endpoint():
         assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True
         assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True
         assert init_kwargs["data"]["model"] == "test-model"
-        assert init_kwargs["data"]["contents"] == [
-            {"role": "user", "parts": [{"text": "Hello"}]}
-        ]
+        assert init_kwargs["data"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
 
 
 def test_google_generate_content_data_flows_through_processor():
@@ -123,9 +122,7 @@ def test_google_generate_content_data_flows_through_processor():
         data = mock_init.call_args.kwargs["data"]
         assert data["model"] == "test-model"
         assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
-        assert data["systemInstruction"] == {
-            "parts": [{"text": "Your name is Doodle."}]
-        }
+        assert data["systemInstruction"] == {"parts": [{"text": "Your name is Doodle."}]}
         # generationConfig arrives intact here; the rename to `config` is
         # done downstream in route_request (see test_route_llm_request).
         assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"]
@@ -160,12 +157,17 @@ def test_google_count_tokens_unchanged():
     except ImportError as e:
         pytest.skip(f"Skipping test due to missing dependency: {e}")
 
-    fake_response = MagicMock()
-    fake_response.original_response = {
-        "totalTokens": 7,
-        "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
-    }
-    fake_response.total_tokens = 7
+    fake_response = TokenCountResponse(
+        total_tokens=7,
+        request_model="test-model",
+        model_used="test-model",
+        tokenizer_type="gemini_api",
+        original_response={
+            "totalTokens": 7,
+            "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
+        },
+        count_source="provider",
+    )
 
     with patch(
         "litellm.proxy.proxy_server.token_counter",
@@ -180,6 +182,7 @@ def test_google_count_tokens_unchanged():
         assert response.status_code == 200
         body = response.json()
         assert body["totalTokens"] == 7
+        assert body["litellm_count_source"] == "provider"
 
 
 def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope():
@@ -212,9 +215,14 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
     except ImportError as e:
         pytest.skip(f"Skipping test due to missing dependency: {e}")
 
-    fake_response: Final = MagicMock()
-    fake_response.original_response = {"input_tokens": 2167}
-    fake_response.total_tokens = 2167
+    fake_response: Final = TokenCountResponse(
+        total_tokens=2167,
+        request_model="claude-opus-4-8",
+        model_used="us.anthropic.claude-opus-5-5",
+        tokenizer_type="bedrock_api",
+        original_response={"input_tokens": 2167},
+        count_source="token_count_model",
+    )
 
     with patch(
         "litellm.proxy.proxy_server.token_counter",
@@ -230,3 +238,5 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
         body: Final = response.json()
         assert body["totalTokens"] == 2167
         assert body["promptTokensDetails"] == []
+        assert body["litellm_count_source"] == "token_count_model"
+        assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5"

tests/unit/proxy/management_endpoints/test_model_management_endpoints.py

@@ -4184,6 +4184,39 @@ class TestModelInfoCostMapEchoFilter:
         assert "litellm_provider" not in info
         assert "supported_openai_params" not in info
 
+    @pytest.mark.parametrize(
+        ("model", "field", "edited"),
+        [
+            ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}),
+            ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0),
+        ],
+    )
+    def test_token_count_setting_is_stored_only_when_the_operator_changed_it(
+        self, model: str, field: str, edited: object
+    ) -> None:
+        import litellm
+        from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
+
+        entry: Final = litellm.get_model_info(model)
+        assert entry[field] not in (None, edited), entry
+        db_model: Final = Deployment(
+            model_name="count-settings",
+            litellm_params=LiteLLM_Params(model=model),
+            model_info=ModelInfo(id="dep-echo-count"),
+        )
+
+        echoed: Final = update_db_model(
+            db_model=db_model,
+            updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count"})),
+        )
+        changed: Final = update_db_model(
+            db_model=db_model,
+            updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count", field: edited})),
+        )
+
+        assert field not in json.loads(echoed["model_info"])
+        assert json.loads(changed["model_info"])[field] == edited
+
     def test_metadata_without_a_cost_map_key_is_persisted(self):
         import litellm
 

tests/unit/proxy/proxy_server/test_routes_utils.py

@@ -54,6 +54,7 @@ def test_token_counter_happy_path(client, auth_as, patched_token_counter):
         "request_model": "gpt-4",
         "model_used": "gpt-4",
         "tokenizer_type": "openai_tokenizer",
+        "count_source": "local_estimate",
         "original_response": None,
         "error": False,
         "error_message": None,

tests/unit/test_router_model_cost_isolation.py

@@ -584,6 +584,40 @@ def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None:
         _restore_model_cost_entries(model_keys)
 
 
+def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None:
+    """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on
+    the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them."""
+    from litellm.proxy.common_utils.token_count_settings import cost_map_entry
+
+    backend_model: Final = "bedrock/us.anthropic.claude-sonnet-5-5"
+    deploy_id: Final = "sonnet-deploy-keep-token-count"
+    configured: Final = {
+        "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"},
+        "token_count_framing": 7,
+    }
+    builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured}
+    assert all(builtin[field] != value for field, value in configured.items()), builtin
+    model_keys: Final = {
+        key: copy.deepcopy(litellm.model_cost.get(key))
+        for key in (deploy_id, backend_model, backend_model.removeprefix("bedrock/"))
+    }
+    try:
+        Router(
+            model_list=[
+                {
+                    "model_name": "sonnet-keep-token-count",
+                    "litellm_params": {"model": backend_model},
+                    "model_info": {"id": deploy_id, **configured},
+                }
+            ],
+        )
+
+        assert {field: litellm.model_cost[deploy_id][field] for field in configured} == configured
+        assert {field: cost_map_entry(backend_model).get(field) for field in configured} == builtin
+    finally:
+        _restore_model_cost_entries(model_keys)
+
+
 def test_should_store_full_pricing_under_deployment_model_id():
     """
     Per-deployment pricing (including zero) should be stored and
@@ -3884,9 +3918,7 @@ def test_router_capacity_widens_to_a_larger_base_model_window(
     )
     _invalidate_model_cost_lowercase_map()
     router: Final = _capacity_test_router(route, base, base_location, explicit_limit)
-    expected: Final = (
-        explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192)
-    )
+    expected: Final = explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192)
     assert _effective_max_input_tokens(router) == expected
     assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0)
     _invalidate_model_cost_lowercase_map()

tests/unit/test_utils.py

@@ -1138,6 +1138,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
                 "supports_box_questions": {"type": "boolean"},
                 "supports_vision": {"type": "boolean"},
                 "supports_web_search": {"type": "boolean"},
+                "token_count_framing": {"type": "integer"},
+                "token_count_model": {
+                    "type": "object",
+                    "properties": {"model": {"type": "string"}, "region": {"type": "string"}},
+                    "required": ["model"],
+                },
+                "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]},
                 "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"},
                 "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"},
                 "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"},

ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx

@@ -10,6 +10,7 @@ import { chooseSelectOption, renderWithProviders as render } from "../../../../.
 vi.mock("@/components/networking", () => ({
   apiClient: { get: vi.fn() },
   createAgentCall: vi.fn(),
+  getProxyBaseUrl: vi.fn(() => "http://gw.test:4000/"),
   getAgentCreateMetadata: vi.fn(),
   getAgentsList: vi.fn(),
   keyCreateForAgentCall: vi.fn(),

ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx

@@ -8,6 +8,7 @@ import type { AgentCreateInfo } from "@/components/networking";
 
 vi.mock("@/components/networking", () => ({
   createAgentCall: vi.fn(),
+  getProxyBaseUrl: vi.fn(() => "http://gw.test:4000/"),
   getAgentCreateMetadata: vi.fn(),
   getAgentsList: vi.fn(),
   keyCreateForAgentCall: vi.fn(),

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -1914,7 +1914,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
 
       {apiKey && (
         <Dialog open={isModalVisible} onOpenChange={(open) => !open && handleCancel()}>
-          <DialogContent className="max-h-[calc(100dvh-2rem)] overflow-y-auto">
+          <DialogContent className="max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-2xl">
             <div className="grid grid-cols-1 gap-2 w-full">
               <DialogTitle className="text-lg font-medium text-foreground">Save your Key</DialogTitle>
               {apiKey != null ? (

ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx

@@ -1,13 +1,43 @@
+import { cloneElement, type MouseEventHandler, type ReactElement, type ReactNode } from "react";
 import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { render, screen, act } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
-import CreatedKeyDisplay from "./CreatedKeyDisplay";
+import CreatedKeyDisplay, { buildAgentInstructions } from "./CreatedKeyDisplay";
 
 import { toast } from "@/lib/toast";
 
+const mockCopy = vi.hoisted(() => vi.fn<(text: string) => void>());
+
+vi.mock("@/components/networking", () => ({
+  getProxyBaseUrl: vi.fn(() => "http://gw.test:4000/"),
+}));
+
+vi.mock("react-copy-to-clipboard", () => ({
+  CopyToClipboard: ({
+    text,
+    onCopy,
+    children,
+  }: {
+    text: string;
+    onCopy?: (text: string, result: boolean) => void;
+    children: ReactNode;
+  }) => {
+    const child = children as ReactElement<{ onClick?: MouseEventHandler<HTMLButtonElement> }>;
+    return cloneElement(child, {
+      onClick: (event) => {
+        child.props.onClick?.(event);
+        mockCopy(text);
+        onCopy?.(text, true);
+      },
+    });
+  },
+}));
+
 describe("CreatedKeyDisplay", () => {
   beforeEach(() => {
     vi.useFakeTimers({ shouldAdvanceTime: true });
+    mockCopy.mockReset();
+    vi.mocked(toast.success).mockClear();
   });
 
   afterEach(() => {
@@ -69,4 +99,87 @@ describe("CreatedKeyDisplay", () => {
 
     expect(screen.getByRole("button", { name: /copy virtual key/i })).toBeInTheDocument();
   });
+
+  it("shows the gateway URL without its trailing slash and copies it", async () => {
+    const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
+    render(<CreatedKeyDisplay apiKey="sk-test-123" />);
+
+    expect(screen.getByText("http://gw.test:4000")).toBeInTheDocument();
+    await user.click(screen.getByRole("button", { name: "Copy URL" }));
+
+    expect(mockCopy).toHaveBeenCalledWith("http://gw.test:4000");
+    expect(toast.success).toHaveBeenCalledWith("Gateway URL copied to clipboard");
+  });
+
+  it("shows one shared instruction card with both agent logos and expands on demand", async () => {
+    const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
+    render(<CreatedKeyDisplay apiKey="sk-test-123" />);
+
+    expect(screen.getByRole("img", { name: "Claude Code" })).toBeInTheDocument();
+    expect(screen.getByRole("img", { name: "Codex" })).toBeInTheDocument();
+    expect(screen.getByText("Agent instructions")).toBeInTheDocument();
+    expect(screen.queryByText("Set up your agent")).not.toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: "Claude Code" })).not.toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: "Codex" })).not.toBeInTheDocument();
+
+    const showButton = screen.getByRole("button", { name: "Show" });
+    expect(showButton).toHaveAttribute("aria-expanded", "false");
+    expect(screen.queryByText(/POST http:\/\/gw\.test:4000\/v1\/responses/)).not.toBeInTheDocument();
+
+    await user.click(showButton);
+
+    expect(screen.getByRole("button", { name: "Hide" })).toHaveAttribute("aria-expanded", "true");
+    expect(screen.getByText(/POST http:\/\/gw\.test:4000\/v1\/responses/)).toBeInTheDocument();
+    expect(screen.getByText(/ANTHROPIC_BASE_URL=http:\/\/gw\.test:4000/)).toBeInTheDocument();
+    expect(screen.getByText(/OPENAI_BASE_URL=http:\/\/gw\.test:4000\/v1/)).toBeInTheDocument();
+  });
+
+  it("copies complete shared instructions while the card is collapsed", async () => {
+    const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
+    render(<CreatedKeyDisplay apiKey="sk-test-123" />);
+
+    expect(screen.getByRole("button", { name: "Show" })).toHaveAttribute("aria-expanded", "false");
+    expect(screen.queryByText(/POST http:\/\/gw\.test:4000\/v1\/responses/)).not.toBeInTheDocument();
+    await user.click(screen.getByRole("button", { name: "Copy instructions" }));
+
+    const copiedInstructions = mockCopy.mock.calls.at(-1)?.[0] ?? "";
+    expect(copiedInstructions).toContain("Gateway URL: http://gw.test:4000");
+    expect(copiedInstructions).toContain("Virtual key: sk-test-123");
+    expect(copiedInstructions).toContain("POST http://gw.test:4000/v1/chat/completions");
+    expect(copiedInstructions).toContain("POST http://gw.test:4000/v1/responses");
+    expect(copiedInstructions).toContain("POST http://gw.test:4000/v1/messages");
+    expect(copiedInstructions).toContain("GET http://gw.test:4000/v1/models");
+    expect(copiedInstructions).toContain(
+      "To run Claude Code through this gateway, set ANTHROPIC_BASE_URL=http://gw.test:4000 and ANTHROPIC_AUTH_TOKEN=sk-test-123.",
+    );
+    expect(copiedInstructions).toContain(
+      "To run Codex through this gateway, set OPENAI_BASE_URL=http://gw.test:4000/v1 and OPENAI_API_KEY=sk-test-123.",
+    );
+    expect(screen.getByRole("button", { name: "Copied!" })).toBeInTheDocument();
+    expect(toast.success).toHaveBeenCalledWith("Instructions copied");
+  });
+});
+
+describe("buildAgentInstructions", () => {
+  it("builds shared instructions with both agent settings and a normalized URL", () => {
+    expect(buildAgentInstructions("http://gw.test:4000/", "sk-test-123")).toBe(
+      [
+        "Use this LiteLLM AI gateway for model inference.",
+        "",
+        "Gateway URL: http://gw.test:4000",
+        "Virtual key: sk-test-123",
+        "",
+        "LiteLLM is a gateway that exposes OpenAI and Anthropic compatible APIs. Send the virtual key as a Bearer token (Authorization: Bearer sk-test-123) to:",
+        "- POST http://gw.test:4000/v1/chat/completions (OpenAI Chat Completions)",
+        "- POST http://gw.test:4000/v1/responses (OpenAI Responses)",
+        "- POST http://gw.test:4000/v1/messages (Anthropic Messages)",
+        "- GET http://gw.test:4000/v1/models (models this key can use)",
+        "",
+        "To run Claude Code through this gateway, set ANTHROPIC_BASE_URL=http://gw.test:4000 and ANTHROPIC_AUTH_TOKEN=sk-test-123.",
+        "To run Codex through this gateway, set OPENAI_BASE_URL=http://gw.test:4000/v1 and OPENAI_API_KEY=sk-test-123.",
+        "",
+        "Docs for everything else LiteLLM supports: https://docs.litellm.ai/docs/",
+      ].join("\n"),
+    );
+  });
 });

ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx

@@ -1,5 +1,9 @@
 import React, { useState } from "react";
 import { CopyToClipboard } from "react-copy-to-clipboard";
+import { ChevronDown, ChevronUp } from "lucide-react";
+import anthropicLogo from "../../../public/assets/logos/anthropic.svg";
+import openaiLogo from "../../../public/assets/logos/openai_small.svg";
+import { getProxyBaseUrl } from "@/components/networking";
 import { Button } from "@/components/ui/button";
 import { toast } from "@/lib/toast";
 
@@ -7,12 +11,41 @@ interface CreatedKeyDisplayProps {
   apiKey: string;
 }
 
+export const buildAgentInstructions = (baseUrl: string, apiKey: string): string => {
+  const normalizedBaseUrl = baseUrl.replace(/\/+$/, "");
+  const endpointIntroduction =
+    "LiteLLM is a gateway that exposes OpenAI and Anthropic compatible APIs. " +
+    `Send the virtual key as a Bearer token (Authorization: Bearer ${apiKey}) to:`;
+
+  return [
+    "Use this LiteLLM AI gateway for model inference.",
+    "",
+    `Gateway URL: ${normalizedBaseUrl}`,
+    `Virtual key: ${apiKey}`,
+    "",
+    endpointIntroduction,
+    `- POST ${normalizedBaseUrl}/v1/chat/completions (OpenAI Chat Completions)`,
+    `- POST ${normalizedBaseUrl}/v1/responses (OpenAI Responses)`,
+    `- POST ${normalizedBaseUrl}/v1/messages (Anthropic Messages)`,
+    `- GET ${normalizedBaseUrl}/v1/models (models this key can use)`,
+    "",
+    `To run Claude Code through this gateway, set ANTHROPIC_BASE_URL=${normalizedBaseUrl} and ANTHROPIC_AUTH_TOKEN=${apiKey}.`,
+    `To run Codex through this gateway, set OPENAI_BASE_URL=${normalizedBaseUrl}/v1 and OPENAI_API_KEY=${apiKey}.`,
+    "",
+    "Docs for everything else LiteLLM supports: https://docs.litellm.ai/docs/",
+  ].join("\n");
+};
+
 /**
  * Shared component for displaying a newly-created virtual key.
  * Used on the Virtual Keys page and in the Add Agent wizard.
  */
 const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
   const [copied, setCopied] = useState(false);
+  const [instructionsCopied, setInstructionsCopied] = useState(false);
+  const [showInstructions, setShowInstructions] = useState(false);
+  const baseUrl = getProxyBaseUrl().replace(/\/+$/, "");
+  const instructions = buildAgentInstructions(baseUrl, apiKey);
 
   const handleCopy = () => {
     setCopied(true);
@@ -20,6 +53,12 @@ const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
     setTimeout(() => setCopied(false), 2000);
   };
 
+  const handleCopyInstructions = () => {
+    setInstructionsCopied(true);
+    toast.success("Instructions copied");
+    setTimeout(() => setInstructionsCopied(false), 2000);
+  };
+
   return (
     <div>
       <p className="mb-2">
@@ -36,6 +75,47 @@ const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
       <CopyToClipboard text={apiKey} onCopy={handleCopy}>
         <Button className="mt-3">{copied ? "Copied!" : "Copy Virtual Key"}</Button>
       </CopyToClipboard>
+
+      <p className="text-sm text-muted-foreground mt-4 mb-1">Gateway URL:</p>
+      <div className="bg-muted rounded-md p-2.5 mb-2.5 flex items-center justify-between gap-2">
+        <pre className="m-0 whitespace-normal break-words text-foreground">{baseUrl}</pre>
+        <CopyToClipboard text={baseUrl} onCopy={() => toast.success("Gateway URL copied to clipboard")}>
+          <Button type="button" variant="outline" className="shrink-0">
+            Copy URL
+          </Button>
+        </CopyToClipboard>
+      </div>
+
+      <div className="mt-4 rounded-md border border-border p-3">
+        <div className="flex flex-wrap items-center justify-between gap-2">
+          <div className="flex items-center gap-2">
+            <p className="text-sm font-medium">Agent instructions</p>
+            <img src={anthropicLogo.src} alt="Claude Code" title="Claude Code" className="size-5 object-contain" />
+            <img src={openaiLogo.src} alt="Codex" title="Codex" className="size-5 object-contain" />
+          </div>
+          <div className="flex items-center gap-2">
+            <CopyToClipboard text={instructions} onCopy={handleCopyInstructions}>
+              <Button type="button">{instructionsCopied ? "Copied!" : "Copy instructions"}</Button>
+            </CopyToClipboard>
+            <Button
+              type="button"
+              variant="outline"
+              aria-expanded={showInstructions}
+              onClick={() => setShowInstructions((expanded) => !expanded)}
+            >
+              {showInstructions ? "Hide" : "Show"}
+              {showInstructions ? (
+                <ChevronUp aria-hidden="true" className="size-4" />
+              ) : (
+                <ChevronDown aria-hidden="true" className="size-4" />
+              )}
+            </Button>
+          </div>
+        </div>
+        {showInstructions && (
+          <pre className="mt-3 whitespace-pre-wrap break-words text-sm text-foreground">{instructions}</pre>
+        )}
+      </div>
     </div>
   );
 };

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -46847,9 +46847,17 @@ export interface components {
          * @description Response structure for token count details with modality breakdown.
          *
          *     Example:
-         *         {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]}
+         *         {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}],
+         *          'litellm_count_source': 'provider'}
          */
         TokenCountDetailsResponse: {
+            /** Litellm Count Model */
+            litellm_count_model?: string;
+            /**
+             * Litellm Count Source
+             * @enum {string}
+             */
+            litellm_count_source?: "provider" | "token_count_model" | "local_estimate";
             /** Prompttokensdetails */
             promptTokensDetails: components["schemas"]["PromptTokensDetails"][];
             /** Totaltokens */
@@ -46878,6 +46886,8 @@ export interface components {
         };
         /** TokenCountResponse */
         TokenCountResponse: {
+            /** Count Source */
+            count_source?: ("provider" | "token_count_model" | "local_estimate") | null;
             /**
              * Error
              * @default false