Merge remote-tracking branch 'origin/stable/1.104.x' into litellm_backport_mcp_sse_1_104

ci_cd/generate_model_prices_schema.py

@@ -19,6 +19,7 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
 NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
 BOOLEAN: JsonSchema = {"type": "boolean"}
 STRING: JsonSchema = {"type": "string"}
+NON_BLANK_STRING: Final[JsonSchema] = {"type": "string", "pattern": r"\S"}
 TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
 WEEKDAY_PATTERN: Final = (
     r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
@@ -103,6 +104,16 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
         "type": "object",
         "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
     },
+    "token_count_model": {
+        "type": "object",
+        "required": ["model"],
+        "additionalProperties": False,
+        "properties": {
+            "model": NON_BLANK_STRING,
+            "region": NON_BLANK_STRING,
+        },
+        "description": "Model that counts tokens for this one, and the AWS region its count API runs in.",
+    },
 }
 
 ARRAY_KEYS: dict[str, JsonSchema] = {
@@ -189,6 +200,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = {
         **NONNEG_INTEGER,
         "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
     },
+    "token_count_framing": {
+        **NONNEG_INTEGER,
+        "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
+    },
     "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
     "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
 }
@@ -244,6 +259,11 @@ def cost_schema(key: str) -> JsonSchema:
 
 def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
     return {
+        "token_count_tokenizer": {
+            "type": "string",
+            "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.",
+            "enum": ["kimi_k3"],
+        },
         "litellm_provider": {
             "type": "string",
             "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",

litellm/integrations/gitlab/gitlab_prompt_manager.py

@@ -215,7 +215,7 @@ class GitLabTemplateManager:
                 elif value.replace(".", "").isdigit():
                     try:
                         result[key] = float(value)
-                    except Exception:
+                    except ValueError:
                         result[key] = value
                 else:
                     result[key] = value.strip("\"'")

litellm/litellm_core_utils/core_helpers.py

@@ -3,9 +3,9 @@
 import copy
 import logging
 import re
-from collections.abc import Collection, Iterable, Mapping
+from collections.abc import Collection, Iterable, Iterator, Mapping
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable
 
 import httpx
 from pydantic import TypeAdapter, ValidationError
@@ -51,6 +51,28 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non
 _DROP_PARAMS_BOOL: Final = TypeAdapter(bool)
 
 
+def as_str_mapping(value: object) -> Mapping[str, object] | None:
+    return value if isinstance(value, Mapping) else None
+
+
+@runtime_checkable
+class _SupportsIter(Protocol):
+    def __iter__(self) -> Iterator[object]: ...
+
+
+@runtime_checkable
+class _SupportsGetItem(Protocol):
+    def __getitem__(self, index: int, /) -> object: ...
+
+
+def iter_items(value: object) -> Iterable[object]:
+    if isinstance(value, _SupportsIter):
+        return value
+    if isinstance(value, _SupportsGetItem):
+        return iter(value)
+    raise TypeError(f"'{type(value).__name__}' object is not iterable")
+
+
 def normalize_drop_params(value: object) -> bool | None:
     if value is None or isinstance(value, bool):
         return value

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -5580,3 +5580,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.experimental_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/anthropic/count_tokens/token_counter.py

@@ -5,10 +5,13 @@ Anthropic Token Counter implementation using the CountTokens API.
 import os
 from typing import Any, Final
 
+from pydantic import JsonValue
+
+import litellm
 from litellm._logging import verbose_logger
 from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
 from litellm.llms.base_llm.base_utils import BaseTokenCounter
-from litellm.types.utils import LlmProviders, TokenCountResponse
+from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
 
 # Global handler instance - reuse across all token counting requests
 anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler()
@@ -30,8 +33,10 @@ class AnthropicTokenCounter(BaseTokenCounter):
         contents: list[dict[str, Any]] | None,
         deployment: dict[str, Any] | None = None,
         request_model: str = "",
-        tools: list[dict[str, Any]] | None = None,
-        system: Any | None = None,
+        tools: list[dict[str, JsonValue]] | None = None,
+        system: JsonValue = None,
+        messages_format: TokenCountMessagesFormat = "anthropic",
+        output_format: JsonValue = None,
     ) -> TokenCountResponse | None:
         """
         Count tokens using Anthropic's CountTokens API.
@@ -54,14 +59,22 @@ class AnthropicTokenCounter(BaseTokenCounter):
         deployment = deployment or {}
         litellm_params: Final = deployment.get("litellm_params", {})
 
-        # Get Anthropic API key from deployment config or environment
-        api_key = litellm_params.get("api_key")
+        api_key: Final = (
+            litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key or os.getenv("ANTHROPIC_API_KEY")
+        )
         if not api_key:
-            api_key = os.getenv("ANTHROPIC_API_KEY")
-
-        if not api_key:
-            verbose_logger.warning("No Anthropic API key found for token counting")
-            return None
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="anthropic_api",
+                error=True,
+                error_message=(
+                    f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be "
+                    "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment."
+                ),
+                status_code=401,
+            )
 
         try:
             result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request(

litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py

@@ -460,7 +460,7 @@ class LiteLLMAnthropicMessagesAdapter:
         *,
         custom_llm_provider: str | None = None,
         preserve_midturn_system: bool = False,
-    ) -> list:
+    ) -> list[AllMessageValues]:
         new_messages: Final[list[AllMessageValues]] = []
         replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)
         leading_count: Final = next(

litellm/llms/anthropic/experimental_pass_through/convertible_blocks.py

@@ -0,0 +1,106 @@
+from collections.abc import Iterator, Mapping, Sequence
+from itertools import chain
+from types import MappingProxyType
+from typing import Final
+
+from pydantic import JsonValue, TypeAdapter, ValidationError
+from typing_extensions import assert_never
+
+from litellm.llms.anthropic.experimental_pass_through.messages.utils import normalize_native_messages
+from litellm.types.utils import TokenCountMessagesFormat
+
+_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType(
+    {
+        "user": frozenset({"text", "image", "document", "tool_result"}),
+        "assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}),
+    }
+)
+_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"})
+_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"})
+
+
+_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...])
+
+
+def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]:
+    try:
+        return _JSON_MESSAGES.validate_python(messages)
+    except ValidationError:
+        return ()
+
+
+def _source_label(block_type: str, source: JsonValue) -> str | None:
+    if not isinstance(source, dict):
+        return f"{block_type} without a source"
+    match source.get("type"):
+        case "base64" if source.get("data"):
+            return None
+        case "url" if source.get("url"):
+            return None
+        case "base64" | "url" as source_type:
+            return f"{block_type} with an empty {source_type} source"
+        case str() as source_type:
+            return f"{block_type} with a {source_type} source"
+        case _:
+            return f"{block_type} with an untyped source"
+
+
+def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None:
+    if not isinstance(block, dict):
+        return None
+    block_type: Final = block.get("type")
+    if not isinstance(block_type, str):
+        return "a block without a type"
+    if block_type not in convertible:
+        return block_type
+    return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None
+
+
+def _tool_result_labels(block: JsonValue) -> Iterator[str]:
+    if not isinstance(block, dict) or block.get("type") != "tool_result":
+        return iter(())
+    inner: Final = block.get("content")
+    if not isinstance(inner, list):
+        return iter(())
+    return (
+        f"{label} inside a tool_result"
+        for item in inner
+        if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None
+    )
+
+
+def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]:
+    label: Final = _block_label(block, convertible)
+    return chain(() if label is None else (label,), _tool_result_labels(block))
+
+
+def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]:
+    role: Final = message.get("role")
+    convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None
+    content: Final = message.get("content")
+    if convertible is None or not isinstance(content, list):
+        return iter(())
+    return chain.from_iterable(_content_labels(block, convertible) for block in content)
+
+
+def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
+    normalized: Final = _json_messages(normalize_native_messages(messages))
+    return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized))))
+
+
+def unconvertible_blocks_error(
+    messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat
+) -> str | None:
+    match messages_format:
+        case "openai":
+            return None
+        case "anthropic":
+            labels: Final = unconvertible_anthropic_blocks(messages or ())
+            return (
+                "count_tokens can't convert these content blocks for this model, so the count would leave them out: "
+                + ", ".join(labels)
+                if labels
+                else None
+            )
+        case _:
+            assert_never(messages_format)

litellm/llms/anthropic/experimental_pass_through/messages/utils.py

@@ -113,6 +113,19 @@ def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMe
     return ChatCompletionSystemMessage(role="system", content=system)
 
 
+def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]:
+    """Apply the block rewrites every /v1/messages request gets before it reaches a provider."""
+    from litellm.llms.anthropic.common_utils import (
+        flatten_unencrypted_web_search_results_in_anthropic_messages,
+        sanitize_tool_use_ids_in_anthropic_messages,
+        strip_empty_content_blocks_from_anthropic_messages,
+    )
+
+    return flatten_unencrypted_web_search_results_in_anthropic_messages(
+        sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages))
+    )
+
+
 @lru_cache(maxsize=1)
 def _anthropic_messages_optional_param_keys() -> frozenset[str]:
     """

litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py

@@ -500,6 +500,19 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
             result["summary"] = "detailed"
         return result
 
+    @staticmethod
+    def system_blocks_to_instructions(system: object) -> str:
+        if not isinstance(system, list):
+            return ""
+        return "\n".join(
+            text
+            for block in system
+            if isinstance(block, dict)
+            and block.get("type") == "text"
+            and isinstance(text := block.get("text"), str)
+            and text
+        )
+
     def translate_request(
         self,
         anthropic_request: AnthropicMessagesRequest,
@@ -549,9 +562,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
             if isinstance(system, str):
                 responses_kwargs["instructions"] = system
             elif isinstance(system, list):
-                responses_kwargs["instructions"] = "\n".join(
-                    filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text"))
-                )
+                responses_kwargs["instructions"] = self.system_blocks_to_instructions(system)
 
         # max_tokens -> max_output_tokens
         max_tokens: Final = anthropic_request.get("max_tokens")

litellm/llms/azure_ai/anthropic/count_tokens/handler.py

@@ -28,7 +28,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
         self,
         model: str,
         messages: list[dict[str, Any]],
-        api_key: str,
+        api_key: str | None,
         api_base: str,
         litellm_params: dict[str, Any] | None = None,
         timeout: float | httpx.Timeout | None = None,
@@ -78,6 +78,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
                 api_key=api_key,
                 litellm_params=litellm_params,
             )
+            if not {"x-api-key", "Authorization"} & headers.keys():
+                raise AnthropicError(
+                    status_code=401,
+                    message=(
+                        f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be "
+                        "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or "
+                        "AZURE_API_KEY in the proxy environment."
+                    ),
+                )
 
             # Use LiteLLM's async httpx client
             async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI)

litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py

@@ -5,17 +5,35 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API.
 import os
 from typing import Any, Final
 
+from pydantic import JsonValue
+
+import litellm
 from litellm._logging import verbose_logger
 from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
     AzureAIAnthropicCountTokensHandler,
 )
 from litellm.llms.base_llm.base_utils import BaseTokenCounter
-from litellm.types.utils import LlmProviders, TokenCountResponse
+from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
 
 # Global handler instance - reuse across all token counting requests
 azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler()
 
 
+def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse:
+    return TokenCountResponse(
+        total_tokens=0,
+        request_model=request_model,
+        model_used=model_to_use,
+        tokenizer_type="azure_ai_anthropic_api",
+        error=True,
+        error_message=(
+            f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. "
+            "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment."
+        ),
+        status_code=500,
+    )
+
+
 class AzureAIAnthropicTokenCounter(BaseTokenCounter):
     """Token counter implementation for Azure AI Anthropic provider using the CountTokens API."""
 
@@ -34,6 +52,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
         request_model: str = "",
         tools: list[dict[str, object]] | None = None,
         system: object | None = None,
+        messages_format: TokenCountMessagesFormat = "anthropic",
+        output_format: JsonValue = None,
     ) -> TokenCountResponse | None:
         """
         Count tokens using Azure AI Anthropic's CountTokens API.
@@ -61,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
         if not api_key:
             api_key = os.getenv("AZURE_AI_API_KEY")
 
-        # Get API base from deployment config or environment
-        api_base = litellm_params.get("api_base")
-        if not api_base:
-            api_base = os.getenv("AZURE_AI_API_BASE")
-
-        if not api_key:
-            verbose_logger.warning("No Azure AI API key found for token counting")
-            return None
+        api_base: Final = (
+            litellm_params.get("api_base")
+            or os.getenv("AZURE_AI_API_BASE")
+            or os.getenv("AZURE_API_BASE")
+            or litellm.api_base
+        )
 
         if not api_base:
-            verbose_logger.warning("No Azure AI API base found for token counting")
-            return None
+            return _missing_api_base_response(request_model, model_to_use)
 
         try:
             result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request(

litellm/llms/azure_ai/anthropic/count_tokens/transformation.py

@@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
 
     def get_required_headers(
         self,
-        api_key: str,
+        api_key: str | None,
         litellm_params: dict[str, object] | None = None,
     ) -> dict[str, str]:
         """
@@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
             "Content-Type": "application/json",
             "anthropic-version": "2023-06-01",
             "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
-            "x-api-key": api_key,  # Azure AI Anthropic requires this header
+            **({"x-api-key": api_key} if api_key else {}),
         }
 
         # Also set up Azure auth headers for flexibility
         litellm_params = litellm_params or {}
-        if "api_key" not in litellm_params:
+        if api_key and "api_key" not in litellm_params:
             litellm_params["api_key"] = api_key
 
         litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params)
@@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
         # Get Azure auth headers (api-key or Authorization)
         azure_headers = BaseAzureLLM._base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj)
 
-        # Merge Azure auth headers
         headers.update(azure_headers)
+        if "api-key" in azure_headers:
+            headers.setdefault("x-api-key", azure_headers["api-key"])
 
         return headers
 

litellm/llms/base_llm/base_utils.py

@@ -8,11 +8,11 @@ from abc import ABC, abstractmethod
 from typing import Any, Final
 
 from openai.lib import _parsing, _pydantic
-from pydantic import BaseModel
+from pydantic import BaseModel, JsonValue
 
 from litellm._logging import verbose_logger
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
-from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse
+from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse
 
 
 class BaseTokenCounter(ABC):
@@ -26,6 +26,8 @@ class BaseTokenCounter(ABC):
         request_model: str = "",
         tools: list[dict[str, Any]] | None = None,
         system: Any | None = None,
+        messages_format: TokenCountMessagesFormat = "openai",
+        output_format: JsonValue = None,
     ) -> TokenCountResponse | None:
         pass
 

litellm/llms/bedrock/count_tokens/bedrock_token_counter.py

@@ -5,11 +5,13 @@ Bedrock Token Counter implementation using the CountTokens API.
 from collections.abc import Mapping, Sequence
 from typing import Any, Final
 
+from pydantic import JsonValue
+
 from litellm._logging import verbose_logger
 from litellm.llms.base_llm.base_utils import BaseTokenCounter
 from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
 from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
-from litellm.types.utils import LlmProviders, TokenCountResponse
+from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
 
 
 class BedrockTokenCounter(BaseTokenCounter):
@@ -33,6 +35,8 @@ class BedrockTokenCounter(BaseTokenCounter):
         request_model: str = "",
         tools: Sequence[Mapping[str, object]] | None = None,
         system: object | None = None,
+        messages_format: TokenCountMessagesFormat = "anthropic",
+        output_format: JsonValue = None,
     ) -> TokenCountResponse | None:
         """
         Count tokens using AWS Bedrock's CountTokens API.

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.experimental_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/gemini/common_utils.py

@@ -6,6 +6,7 @@ from collections.abc import Mapping, Sequence
 from typing import Any, Final
 
 import httpx
+from pydantic import JsonValue
 
 import litellm
 from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@@ -13,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import TokenCountResponse
+from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
 
 GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = {
     "1:1": 1 / 1,
@@ -491,21 +492,54 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
         request_model: str = "",
         tools: list[dict[str, object]] | None = None,
         system: object | None = None,
+        messages_format: TokenCountMessagesFormat = "openai",
+        output_format: JsonValue = None,
     ) -> TokenCountResponse | None:
         import copy
 
         from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
+        from litellm.llms.vertex_ai.count_tokens.context import (
+            VertexGeminiCountContextError,
+            count_context_error_response,
+            gemini_count_contents,
+            inference_request_view,
+            provider_error_response,
+            unconvertible_messages_error,
+            vertex_gemini_count_context,
+        )
+
+        request_view: Final = inference_request_view(messages or [], system, messages_format)
+        count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools)
+        if isinstance(count_context, VertexGeminiCountContextError):
+            return count_context_error_response(count_context, request_model, model_to_use)
+        blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format)
+        if blocks_error is not None:
+            return count_context_error_response(blocks_error, request_model, model_to_use)
 
         deployment = deployment or {}
         count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {}))
+        resolved_contents: Final = (
+            contents
+            if contents is not None
+            else await gemini_count_contents(
+                request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request
+            )
+        )
+        if isinstance(resolved_contents, VertexGeminiCountContextError):
+            return count_context_error_response(resolved_contents, request_model, model_to_use)
         count_tokens_params: Final = {
             "model": model_to_use,
-            "contents": contents,
+            "contents": resolved_contents,
+            "system_instruction": count_context.system_instruction,
+            "tools": count_context.tools,
         }
         count_tokens_params_request.update(count_tokens_params)
-        result: Final = await GoogleAIStudioTokenCounter().acount_tokens(
-            **count_tokens_params_request,
-        )
+        try:
+            result: Final = await GoogleAIStudioTokenCounter().acount_tokens(
+                **count_tokens_params_request,
+            )
+        except litellm.APIError as error:
+            return provider_error_response(error, request_model, model_to_use)
 
         if result is not None:
             return TokenCountResponse(

litellm/llms/gemini/count_tokens/handler.py

@@ -1,9 +1,10 @@
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
 
 import litellm
-from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
 from litellm.types.utils import LlmProviders
 
 if TYPE_CHECKING:
@@ -28,12 +29,12 @@ class GoogleAIStudioTokenCounter:
         """
         import copy
 
-        from google.genai.types import FunctionResponse
-
         # Handle None or empty contents
         if not contents:
             return contents
 
+        from google.genai.types import FunctionResponse
+
         cleaned_contents: Final = copy.deepcopy(contents)
 
         for content in cleaned_contents:
@@ -54,6 +55,28 @@ class GoogleAIStudioTokenCounter:
         base_url: Final = api_base or "https://generativelanguage.googleapis.com"
         return f"{base_url}/v1beta/models/{model}:countTokens"
 
+    def __init__(self, client: AsyncHTTPHandler | None = None) -> None:
+        super().__init__()
+        self._client: Final = client
+
+    def _build_count_tokens_request_body(
+        self,
+        model: str,
+        contents: object,
+        system_instruction: Mapping[str, object] | None,
+        tools: Sequence[Mapping[str, object]] | None,
+    ) -> Mapping[str, object]:
+        if system_instruction is None and not tools:
+            return {"contents": contents}
+
+        generate_content_request: Final = {
+            "model": model if model.startswith("models/") else f"models/{model}",
+            "contents": contents,
+            **({"systemInstruction": dict(system_instruction)} if system_instruction is not None else {}),
+            **({"tools": [dict(tool) for tool in tools]} if tools else {}),
+        }
+        return {"generateContentRequest": generate_content_request}
+
     async def validate_environment(
         self,
         api_base: str | None = None,
@@ -84,6 +107,8 @@ class GoogleAIStudioTokenCounter:
         api_key: str | None = None,
         api_base: str | None = None,
         timeout: float | httpx.Timeout | None = None,
+        system_instruction: Mapping[str, object] | None = None,
+        tools: Sequence[Mapping[str, object]] | None = None,
         **kwargs: object,
     ) -> dict[str, Any]:
         """
@@ -130,14 +155,19 @@ class GoogleAIStudioTokenCounter:
 
         # Prepare request body - clean up contents to remove unsupported fields
         cleaned_contents: Final = self._clean_contents_for_gemini_api(contents)
-        request_body: Final = {"contents": cleaned_contents}
+        request_body: Final = self._build_count_tokens_request_body(
+            model=model,
+            contents=cleaned_contents,
+            system_instruction=system_instruction,
+            tools=tools,
+        )
 
-        async_httpx_client: Final = get_async_httpx_client(
-            llm_provider=LlmProviders.GEMINI,
+        async_httpx_client: Final = (
+            self._client if self._client is not None else get_async_httpx_client(llm_provider=LlmProviders.GEMINI)
         )
 
         try:
-            response: Final = await async_httpx_client.post(url=url, headers=headers, json=request_body)
+            response: Final = await async_httpx_client.post(url=url, headers=headers, json=dict(request_body))
 
             # Check for HTTP errors
             response.raise_for_status()

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

@@ -5,6 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint.
 """
 
 import json
+from collections.abc import Mapping, Sequence
 from typing import Any, Final
 
 import httpx
@@ -30,9 +31,10 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
         api_key: str,
         api_base: str | None = None,
         timeout: float | httpx.Timeout | None = None,
-        tools: list[dict[str, Any]] | None = None,
+        tools: Sequence[Mapping[str, object]] | None = None,
         instructions: str | None = None,
-    ) -> dict[str, Any]:
+        text: Mapping[str, object] | None = None,
+    ) -> dict[str, object]:
         """
         Handle a token counting request to OpenAI's Responses API.
 
@@ -52,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
                 input=input,
                 tools=tools,
                 instructions=instructions,
+                text=text,
             )
 
             endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base)

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

@@ -2,24 +2,47 @@
 OpenAI Token Counter implementation using the Responses API /input_tokens endpoint.
 """
 
-import os
+from collections.abc import Mapping, Sequence
 from typing import Any, Final
 
+import openai
+from pydantic import JsonValue
+from typing_extensions import assert_never
+
 from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.llms.anthropic.experimental_pass_through.convertible_blocks import unconvertible_blocks_error
+from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import (
+    LiteLLMAnthropicToResponsesAPIAdapter,
+)
 from litellm.llms.base_llm.base_utils import BaseTokenCounter
-from litellm.llms.openai.common_utils import OpenAIError
+from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials
 from litellm.llms.openai.responses.count_tokens.handler import (
     OpenAICountTokensHandler,
 )
 from litellm.llms.openai.responses.count_tokens.transformation import (
     OpenAICountTokensConfig,
+    ResponsesCountRequestParts,
 )
-from litellm.types.utils import LlmProviders, TokenCountResponse
+from litellm.llms.openai.workload_identity import (
+    OpenAIWorkloadIdentityConfig,
+    get_workload_identity_bearer_token,
+    resolve_openai_workload_identity_config,
+)
+from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
 
 # Global handler instance - reuse across all token counting requests
 openai_count_tokens_handler: Final = OpenAICountTokensHandler()
 
 
+def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None:
+    api_base: Final = litellm_params.get("api_base")
+    return resolve_openai_workload_identity_config(
+        api_key=None,
+        api_base=api_base if isinstance(api_base, str) else None,
+    )
+
+
 class OpenAITokenCounter(BaseTokenCounter):
     """Token counter implementation for OpenAI provider using the Responses API."""
 
@@ -29,6 +52,37 @@ class OpenAITokenCounter(BaseTokenCounter):
     ) -> bool:
         return custom_llm_provider == LlmProviders.OPENAI.value
 
+    @staticmethod
+    def _system_instructions(system: object) -> str | None:
+        if system is None or isinstance(system, str):
+            return system
+        return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None
+
+    @staticmethod
+    def _responses_request_parts(
+        model_to_use: str,
+        messages: Sequence[Mapping[str, object]],
+        tools: Sequence[Mapping[str, object]] | None,
+        system: object,
+        messages_format: TokenCountMessagesFormat,
+        output_format: JsonValue,
+    ) -> ResponsesCountRequestParts:
+        match messages_format:
+            case "anthropic":
+                return OpenAICountTokensConfig.anthropic_messages_to_responses_input(
+                    model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format
+                )
+            case "openai":
+                input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
+                    [dict(message) for message in messages]
+                )
+                fallback_instructions: Final = OpenAITokenCounter._system_instructions(system)
+                return OpenAICountTokensConfig.count_request_parts(
+                    input_items, instructions if instructions is not None else fallback_instructions, tools
+                )
+            case _:
+                assert_never(messages_format)
+
     async def count_tokens(
         self,
         model_to_use: str,
@@ -38,6 +92,8 @@ class OpenAITokenCounter(BaseTokenCounter):
         request_model: str = "",
         tools: list[dict[str, object]] | None = None,
         system: object = None,
+        messages_format: TokenCountMessagesFormat = "openai",
+        output_format: JsonValue = None,
     ) -> TokenCountResponse | None:
         """
         Count tokens using OpenAI's Responses API /input_tokens endpoint.
@@ -45,39 +101,76 @@ class OpenAITokenCounter(BaseTokenCounter):
         if not messages:
             return None
 
+        malformed_error: Final = (
+            OpenAICountTokensConfig.malformed_anthropic_messages_error(messages)
+            if messages_format == "anthropic"
+            else None
+        )
+        request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format)
+        if request_error is not None:
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="openai_api",
+                error=True,
+                error_message=request_error,
+                status_code=400,
+                rejected_before_provider_call=True,
+            )
+
         deployment = deployment or {}
         litellm_params: Final = deployment.get("litellm_params", {})
 
-        # Get OpenAI API key from deployment config or environment
-        api_key = litellm_params.get("api_key")
-        if not api_key:
-            api_key = os.getenv("OPENAI_API_KEY")
+        credentials: Final = get_openai_credentials(
+            api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key")
+        )
+        credential: Final = credentials.api_key or _workload_identity_config(litellm_params)
 
-        if not api_key:
-            verbose_logger.warning("No OpenAI API key found for token counting")
-            return None
+        if credential is None:
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="openai_api",
+                error=True,
+                error_message=(
+                    f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be "
+                    "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy "
+                    "environment."
+                ),
+                status_code=401,
+            )
 
-        api_base: Final = litellm_params.get("api_base")
+        api_base: Final = credentials.api_base
 
-        # Convert chat messages to Responses API input format
-        input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
-
-        # Use system param if instructions not extracted from messages
-        if instructions is None and system is not None:
-            instructions = system if isinstance(system, str) else str(system)
+        input_items, instructions, counted_tools, text = self._responses_request_parts(
+            model_to_use=model_to_use,
+            messages=messages,
+            tools=tools,
+            system=system,
+            messages_format=messages_format,
+            output_format=output_format,
+        )
 
         # If no input items were produced (e.g., system-only messages), fall back to local counting
         if not input_items:
             return None
 
         try:
+            api_key: Final = (
+                credential
+                if isinstance(credential, str)
+                else await asyncify(get_workload_identity_bearer_token)(credential)
+            )
             result: Final = await openai_count_tokens_handler.handle_count_tokens_request(
                 model=model_to_use,
-                input=input_items if input_items is not None else [],
+                input=input_items,
                 api_key=api_key,
                 api_base=api_base,
-                tools=tools,
+                tools=counted_tools,
                 instructions=instructions,
+                text=text,
             )
 
             if result is not None:
@@ -99,6 +192,17 @@ class OpenAITokenCounter(BaseTokenCounter):
                 error_message=e.message,
                 status_code=e.status_code,
             )
+        except openai.AuthenticationError as e:
+            verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message)
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="openai_api",
+                error=True,
+                error_message=e.message,
+                status_code=401,
+            )
         except Exception as e:
             verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e)
             return TokenCountResponse(

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

@@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp
 """
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Final, Literal
+from typing import Any, Final, Literal, TypeAlias
 
+from pydantic import JsonValue, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
+from litellm.llms.anthropic.experimental_pass_through.messages.utils import normalize_native_messages
+from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import (
+    LiteLLMAnthropicToResponsesAPIAdapter,
+)
+from litellm.responses.utils import ResponsesAPIRequestUtils
+from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema
+
 
 class ResponsesInputTextPart(TypedDict):
     type: ReadOnly[Literal["input_text"]]
@@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response
 
 ResponsesContentRole = Literal["user", "assistant"]
 
+ResponsesCountRequestParts: TypeAlias = tuple[
+    tuple[Mapping[str, object], ...],
+    str | None,
+    tuple[Mapping[str, object], ...] | None,
+    Mapping[str, object] | None,
+]
+_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts)
+_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object])
+
+
+class _AnthropicCountMessage(TypedDict):
+    role: ReadOnly[str]
+
+
+_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...])
+_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema)
+
+
+def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None:
+    try:
+        return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format)
+    except ValidationError:
+        return None
+
 
 def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None:
     url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url
@@ -118,8 +150,9 @@ class OpenAICountTokensConfig:
         self,
         model: str,
         input: str | Sequence[object],
-        tools: list[dict[str, Any]] | None = None,
+        tools: Sequence[Mapping[str, object]] | None = None,
         instructions: str | None = None,
+        text: Mapping[str, object] | None = None,
     ) -> dict[str, object]:
         """
         Transform request to OpenAI Responses API token counting format.
@@ -135,7 +168,10 @@ class OpenAICountTokensConfig:
             request["instructions"] = instructions
 
         if tools is not None:
-            request["tools"] = self._transform_tools_for_responses_api(tools)
+            request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools])
+
+        if text is not None:
+            request["text"] = text
 
         return request
 
@@ -183,7 +219,7 @@ class OpenAICountTokensConfig:
     @staticmethod
     def messages_to_responses_input(
         messages: list[dict[str, Any]],
-    ) -> tuple:
+    ) -> tuple[list[dict[str, object]], str | None]:
         """
         Convert standard chat messages format to OpenAI Responses API input format.
 
@@ -246,3 +282,47 @@ class OpenAICountTokensConfig:
 
         instructions: Final = "\n".join(instructions_parts) if instructions_parts else None
         return input_items, instructions
+
+    @staticmethod
+    def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None:
+        try:
+            _ANTHROPIC_COUNT_MESSAGES.validate_python(messages)
+        except ValidationError as e:
+            return "; ".join(
+                f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors()
+            )
+        return None
+
+    @staticmethod
+    def anthropic_messages_to_responses_input(
+        model: str,
+        messages: Sequence[Mapping[str, object]],
+        system: object,
+        tools: Sequence[Mapping[str, object]] | None,
+        output_format: JsonValue = None,
+    ) -> ResponsesCountRequestParts:
+        """Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills."""
+        anthropic_request: Final = AnthropicMessagesRequest(
+            model=model,
+            messages=normalize_native_messages(messages),
+            system=system if isinstance(system, (str, list)) else None,
+            tools=[dict(tool) for tool in tools] if tools else None,
+            output_format=_anthropic_output_schema(output_format),
+        )
+        translated: Final = _TRANSLATED_REQUEST.validate_python(
+            LiteLLMAnthropicToResponsesAPIAdapter().translate_request(
+                anthropic_request, include_encrypted_reasoning=False
+            )
+        )
+        return OpenAICountTokensConfig.count_request_parts(
+            ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(translated.get("input")),  # pyright: ignore[reportPrivateUsage]  # the input restore inference runs
+            translated.get("instructions"),
+            translated.get("tools"),
+            translated.get("text"),
+        )
+
+    @staticmethod
+    def count_request_parts(
+        input_items: object, instructions: object, tools: object, text: object = None
+    ) -> ResponsesCountRequestParts:
+        return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text))

litellm/llms/vertex_ai/common_utils.py

@@ -7,7 +7,7 @@ from types import MappingProxyType
 from typing import Any, Final, Literal, cast, get_type_hints
 
 import httpx
-from pydantic import TypeAdapter, ValidationError
+from pydantic import JsonValue, TypeAdapter, ValidationError
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 import litellm
@@ -22,7 +22,7 @@ from litellm.types.llms.vertex_ai import (
     PartType,
     Schema,
 )
-from litellm.types.utils import TokenCountResponse
+from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
 from litellm.utils import supports_response_schema, supports_system_messages
 
 VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS: Final = frozenset(
@@ -1237,9 +1237,14 @@ class VertexAITokenCounter(BaseTokenCounter):
         request_model: str = "",
         tools: list[dict[str, object]] | None = None,
         system: object | None = None,
+        messages_format: TokenCountMessagesFormat = "openai",
+        output_format: JsonValue = None,
     ) -> TokenCountResponse | None:
         import copy
 
+        from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
+            VertexAIError as PartnerVertexAIError,
+        )
         from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
             VertexAIPartnerModels,
         )
@@ -1269,14 +1274,32 @@ class VertexAITokenCounter(BaseTokenCounter):
                 "vertex_ai_credentials"
             )
 
-            result = await partner_models_handler.count_tokens(
-                model=model_to_use,
-                messages=messages or [],
-                litellm_params=partner_litellm_params,
-                vertex_project=vertex_project,
-                vertex_location=vertex_location,
-                vertex_credentials=vertex_credentials,
-            )
+            try:
+                result = await partner_models_handler.count_tokens(
+                    model=model_to_use,
+                    messages=messages or [],
+                    litellm_params=partner_litellm_params,
+                    vertex_project=vertex_project,
+                    vertex_location=vertex_location,
+                    vertex_credentials=vertex_credentials,
+                    system=system,
+                    tools=tools,
+                )
+            except (PartnerVertexAIError, httpx.HTTPStatusError) as e:
+                status_code: Final = e.response.status_code
+                error_message: Final = e.message if isinstance(e, PartnerVertexAIError) else e.response.text
+                verbose_logger.warning(
+                    "Vertex AI partner CountTokens API error: status=%s, message=%s", status_code, error_message
+                )
+                return TokenCountResponse(
+                    total_tokens=0,
+                    request_model=request_model,
+                    model_used=model_to_use,
+                    tokenizer_type="vertex_ai_partner_models",
+                    error=True,
+                    error_message=error_message,
+                    status_code=status_code,
+                )
 
             if result is not None:
                 return TokenCountResponse(
@@ -1287,34 +1310,55 @@ class VertexAITokenCounter(BaseTokenCounter):
                     original_response=result,
                 )
         else:
-            from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter
-            from litellm.llms.vertex_ai.gemini.transformation import (
-                _gemini_convert_messages_with_history,  # pyright: ignore[reportPrivateUsage]  # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini
+            from litellm.llms.vertex_ai.count_tokens.context import (
+                VertexGeminiCountContextError,
+                count_context_error_response,
+                gemini_count_contents,
+                inference_request_view,
+                provider_error_response,
+                unconvertible_messages_error,
+                vertex_gemini_count_context,
             )
+            from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter
+
+            request_view: Final = inference_request_view(messages or [], system, messages_format)
+            count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools)
+            if isinstance(count_context, VertexGeminiCountContextError):
+                return count_context_error_response(count_context, request_model, model_to_use)
+            blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format)
+            if blocks_error is not None:
+                return count_context_error_response(blocks_error, request_model, model_to_use)
 
             resolved_contents: Final = (
                 contents
                 if contents is not None
-                else _gemini_convert_messages_with_history(
-                    messages=messages or []  # mutable-ok: fallback for None messages; helper signature requires list
+                else await gemini_count_contents(
+                    request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request
                 )
             )
+            if isinstance(resolved_contents, VertexGeminiCountContextError):
+                return count_context_error_response(resolved_contents, request_model, model_to_use)
 
             count_tokens_params: Final = {
                 "model": model_to_use,
                 "contents": resolved_contents,
+                "system_instruction": count_context.system_instruction,
+                "tools": count_context.tools,
             }
             count_tokens_params_request.update(count_tokens_params)
-            result = await VertexAITokenCounter().acount_tokens(
-                **count_tokens_params_request,
-            )
+            try:
+                result = await VertexAITokenCounter().acount_tokens(
+                    **count_tokens_params_request,
+                )
+            except litellm.APIError as error:
+                return provider_error_response(error, request_model, model_to_use)
 
-            if result is not None and "totalTokens" in result:
+            if "totalTokens" in result:
                 return TokenCountResponse(
                     total_tokens=result["totalTokens"],
                     request_model=request_model,
                     model_used=model_to_use,
-                    tokenizer_type=result.get("tokenizer_used", ""),
+                    tokenizer_type=result.get("tokenizer_used", "gemini"),
                     original_response=result,
                 )
 

litellm/llms/vertex_ai/count_tokens/context.py

@@ -0,0 +1,242 @@
+import copy
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from types import MappingProxyType
+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
+from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
+
+if TYPE_CHECKING:
+    from litellm.exceptions import APIError
+    from litellm.llms.anthropic.experimental_pass_through.context_management.result import PolyfillResult
+
+
+@dataclass(frozen=True, slots=True)
+class VertexGeminiCountContext:
+    system_instruction: Mapping[str, object] | None
+    tools: tuple[Mapping[str, object], ...] | None
+
+
+@dataclass(frozen=True, slots=True)
+class VertexGeminiCountContextError:
+    error_message: str
+
+
+_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...])
+_TOOLS: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...])
+_VERTEX_TOOL_KEYS: Final = MappingProxyType(
+    {
+        "function_declarations": "functionDeclarations",
+        "url_context": "urlContext",
+        "code_execution": "codeExecution",
+    }
+)
+
+
+def _system_parts(
+    system: object | None,
+) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError:
+    if system is None:
+        return ()
+    if isinstance(system, str):
+        return ({"text": system},) if system else ()
+    try:
+        blocks: Final = _SYSTEM_BLOCKS.validate_python(system)
+    except ValidationError:
+        return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.")
+    return tuple({"text": block["text"]} for block in blocks)
+
+
+_INVALID_TOOLS: Final = VertexGeminiCountContextError(
+    error_message="tools must be a list of objects, and each function tool needs a JSON object schema."
+)
+
+
+def _has_object_schema(tool: Mapping[str, JsonValue]) -> bool:
+    function: Final = tool.get("function")
+    schema: Final = (
+        function.get("parameters") if isinstance(function, dict) else tool.get("input_schema", tool.get("parameters"))
+    )
+    return schema is None or isinstance(schema, dict)
+
+
+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 _map_function_tools(
+    tools: object | None,
+) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError:
+    try:
+        tool_values: Final = _TOOLS.validate_python(tools or ())
+    except ValidationError:
+        return _INVALID_TOOLS
+    function_tools: Final = tuple(tool for tool in tool_values if _is_function_tool(tool))
+    if not function_tools:
+        return None
+    if not all(_has_object_schema(tool) for tool in function_tools):
+        return _INVALID_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_completion_tools_for(tool)
+    )  # comprehension-ok: flattens each tool's translation
+    try:
+        mapped_tools: Final = VertexGeminiConfig()._map_function(  # pyright: ignore[reportPrivateUsage]  # the tool mapping inference uses
+            value=[dict(tool) for tool in chat_tools],
+            optional_params={},
+        )
+    except (AttributeError, TypeError, ValueError):
+        return _INVALID_TOOLS
+    return tuple(
+        MappingProxyType({_VERTEX_TOOL_KEYS.get(key, key): value for key, value in tool.items()})
+        for tool in mapped_tools
+    )
+
+
+def vertex_gemini_count_context(
+    system: object | None,
+    tools: object | None,
+) -> VertexGeminiCountContext | VertexGeminiCountContextError:
+    system_parts: Final = _system_parts(system)
+    if isinstance(system_parts, VertexGeminiCountContextError):
+        return system_parts
+    function_tools: Final = _map_function_tools(tools)
+    if isinstance(function_tools, VertexGeminiCountContextError):
+        return function_tools
+    return VertexGeminiCountContext(
+        system_instruction=MappingProxyType({"parts": system_parts}) if system_parts else None,
+        tools=function_tools,
+    )
+
+
+def unconvertible_messages_error(
+    messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat
+) -> VertexGeminiCountContextError | None:
+    from litellm.llms.anthropic.experimental_pass_through.convertible_blocks import unconvertible_blocks_error
+
+    error_message: Final = unconvertible_blocks_error(messages, messages_format)
+    return None if error_message is None else VertexGeminiCountContextError(error_message=error_message)
+
+
+def count_context_error_response(
+    error: VertexGeminiCountContextError, request_model: str, model_used: str
+) -> TokenCountResponse:
+    return TokenCountResponse(
+        total_tokens=0,
+        request_model=request_model,
+        model_used=model_used,
+        tokenizer_type="gemini",
+        error=True,
+        error_message=error.error_message,
+        status_code=400,
+        rejected_before_provider_call=True,
+    )
+
+
+def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse:
+    return TokenCountResponse(
+        total_tokens=0,
+        request_model=request_model,
+        model_used=model_used,
+        tokenizer_type="gemini",
+        error=True,
+        error_message=error.message,
+        status_code=error.status_code,
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class InferenceRequestView:
+    messages: Sequence[Mapping[str, object]]
+    system: object | None
+
+
+def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None":
+    from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
+        apply_client_compaction_block_history,
+    )
+
+    chat_messages: Final = [dict(message) for message in messages]
+    if system is None or isinstance(system, str):
+        return apply_client_compaction_block_history(messages=chat_messages, system=system)
+    try:
+        blocks: Final = _SYSTEM_BLOCKS.validate_python(system)
+    except ValidationError:
+        return None
+    return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks])
+
+
+def inference_request_view(
+    messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat
+) -> InferenceRequestView:
+    """Apply the history rewrites /v1/messages runs before it converts a body for Gemini."""
+    from litellm.llms.anthropic.experimental_pass_through.messages.utils import normalize_native_messages
+
+    match messages_format:
+        case "openai":
+            return InferenceRequestView(messages=messages, system=system)
+        case "anthropic":
+            normalized: Final = normalize_native_messages(messages)
+            compacted: Final = _client_compaction(normalized, system)
+            if compacted is None:
+                return InferenceRequestView(messages=normalized, system=system)
+            return InferenceRequestView(messages=compacted.messages, system=compacted.system)
+        case _:
+            assert_never(messages_format)
+
+
+def _chat_messages(
+    messages: Sequence[Mapping[str, object]],
+    messages_format: TokenCountMessagesFormat,
+    model: str,
+    custom_llm_provider: Literal["vertex_ai", "gemini"],
+) -> Sequence[AllMessageValues]:
+    from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
+
+    match messages_format:
+        case "anthropic":
+            return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
+                messages=list(messages),  # pyright: ignore[reportArgumentType]  # raw Anthropic body, read the way /v1/messages inference reads it
+                model=model,
+                custom_llm_provider=custom_llm_provider,
+            )
+        case "openai":
+            return copy.deepcopy(list(messages))  # pyright: ignore[reportReturnType]  # chat messages from the request, copied because the AI Studio conversion rewrites them in place
+        case _:
+            assert_never(messages_format)
+
+
+async def gemini_count_contents(
+    messages: Sequence[Mapping[str, object]],
+    messages_format: TokenCountMessagesFormat,
+    model: str,
+    custom_llm_provider: Literal["vertex_ai", "gemini"],
+    litellm_params: Mapping[str, object],
+) -> Sequence[ContentType] | VertexGeminiCountContextError:
+    from litellm.exceptions import BadRequestError
+    from litellm.litellm_core_utils.asyncify import asyncify
+    from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
+    from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
+    from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media
+    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
+
+    config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig()
+    try:
+        chat_messages: Final = await async_inline_remote_media(
+            list(_chat_messages(messages, messages_format, model, custom_llm_provider)),
+            should_inline=should_inline_remote_media(custom_llm_provider),
+        )
+        return await asyncify(config._transform_messages)(  # pyright: ignore[reportPrivateUsage]  # the message transform inference runs
+            messages=chat_messages, model=model, litellm_params=dict(litellm_params)
+        )
+    except BadRequestError as error:
+        return VertexGeminiCountContextError(error_message=error.message)

litellm/llms/vertex_ai/count_tokens/handler.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping, Sequence
 from typing import Any, Final
 
 from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
@@ -5,6 +6,19 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 
 
 class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase):
+    def _build_count_tokens_request_body(
+        self,
+        model: str,
+        contents: object,
+        system_instruction: Mapping[str, object] | None,
+        tools: Sequence[Mapping[str, object]] | None,
+    ) -> Mapping[str, object]:
+        return {
+            "contents": contents,
+            **({"systemInstruction": dict(system_instruction)} if system_instruction else {}),
+            **({"tools": [dict(tool) for tool in tools]} if tools else {}),
+        }
+
     async def validate_environment(
         self,
         api_base: str | None = None,
@@ -16,29 +30,29 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase):
         """
         Returns a Tuple of headers and url for the Vertex AI countTokens endpoint.
         """
-        litellm_params = litellm_params or {}
-        vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=litellm_params)
-        vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params)
-        vertex_location: Final = self.get_vertex_ai_location(litellm_params=litellm_params)
-        _auth_header, vertex_project = await self._ensure_access_token_async(
+        resolved_litellm_params: Final = litellm_params or {}
+        vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=resolved_litellm_params)
+        vertex_project: Final = self.get_vertex_ai_project(litellm_params=resolved_litellm_params)
+        vertex_location: Final = self.get_vertex_ai_location(litellm_params=resolved_litellm_params)
+        _auth_header, resolved_vertex_project = await self._ensure_access_token_async(
             credentials=vertex_credentials,
             project_id=vertex_project,
             custom_llm_provider="vertex_ai",
         )
 
-        auth_header, api_base = self._get_token_and_url(
+        auth_header, url = self._get_token_and_url(
             model=model,
             gemini_api_key=None,
             auth_header=_auth_header,
-            vertex_project=vertex_project,
+            vertex_project=resolved_vertex_project,
             vertex_location=vertex_location,
             vertex_credentials=vertex_credentials,
             stream=False,
             custom_llm_provider="vertex_ai",
-            api_base=None,
+            api_base=api_base,
             mode="count_tokens",
         )
         headers = {
             "Authorization": f"Bearer {auth_header}",
         }
-        return headers, api_base
+        return headers, url

litellm/llms/vertex_ai/gemini/transformation.py

@@ -7,7 +7,7 @@ Why separate file? Make it easy to see how transformation works
 import json
 import os
 import re
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
 from urllib.parse import quote
 
@@ -1327,6 +1327,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool:
     )
 
 
+def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]:
+    return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
+
+
 async def async_transform_request_body(
     gemini_api_key: str | None,
     messages: list[AllMessageValues],
@@ -1369,7 +1373,7 @@ async def async_transform_request_body(
     )
 
     inlined_messages: Final = await async_inline_remote_media(
-        messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
+        messages, should_inline=should_inline_remote_media(custom_llm_provider)
     )
 
     if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):

litellm/llms/vertex_ai/vertex_ai_partner_models/main.py

@@ -2,6 +2,7 @@
 ## API Handler for calling Vertex AI Partner Models
 from collections.abc import Callable
 from enum import Enum
+from types import MappingProxyType
 from typing import Final
 
 import httpx
@@ -270,6 +271,8 @@ class VertexAIPartnerModels(VertexBase):
         vertex_project=None,
         vertex_location=None,
         vertex_credentials=None,
+        system: object | None = None,
+        tools: list[dict[str, object]] | None = None,
     ):
         """
         Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.)
@@ -303,6 +306,9 @@ class VertexAIPartnerModels(VertexBase):
             request_data: Final = {
                 "model": model,
                 "messages": messages,
+                **MappingProxyType(
+                    {key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
+                ),
             }
 
             # Prepare litellm_params with credentials

litellm/model_prices_and_context_window_backup.json

@@ -19337,7 +19337,13 @@
         "supports_prompt_caching": true,
         "supports_reasoning": true,
         "supports_tool_choice": true,
-        "supports_vision": false
+        "supports_vision": false,
+        "token_count_framing": 97
+    },
+    "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,
@@ -20328,7 +20334,13 @@
         "supports_prompt_caching": true,
         "supports_reasoning": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 100
+    },
+    "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,
@@ -75642,7 +75654,9 @@
         "supports_response_schema": true,
         "supports_system_messages": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3"
     },
     "global.moonshotai.kimi-k3": {
         "cache_creation_input_token_cost": 3.75e-06,
@@ -75662,7 +75676,9 @@
         "supports_response_schema": true,
         "supports_system_messages": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3"
     },
     "us.moonshotai.kimi-k3": {
         "cache_creation_input_token_cost": 4.125e-06,
@@ -75682,7 +75698,9 @@
         "supports_response_schema": true,
         "supports_system_messages": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3"
     },
     "openrouter/xiaomi/mimo-v2.6-flash": {
         "cache_read_input_token_cost": 2.8e-09,

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -2,10 +2,12 @@
 Unified /v1/messages endpoint - (Anthropic Spec)
 """
 
+from collections.abc import Mapping
 from typing import Final
 
 from fastapi import APIRouter, Depends, HTTPException, Request, Response
 from fastapi.responses import JSONResponse
+from pydantic import JsonValue
 
 import litellm
 from litellm.anthropic_interface.exceptions import (
@@ -27,6 +29,7 @@ from litellm.proxy.common_request_processing import (
     create_response,
     log_llm_api_exception,
     proxy_exception_from_http_exception,
+    proxy_exception_from_route_error,
     resolve_litellm_call_id,
 )
 from litellm.proxy.common_utils.error_body_call_id import error_body_call_id
@@ -38,7 +41,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
     openai_error_type,
     with_litellm_call_id,
 )
-from litellm.types.utils import TokenCountResponse
+from litellm.proxy.common_utils.token_count_settings import count_source_fields
 
 router: Final = APIRouter()
 
@@ -286,6 +289,14 @@ async def anthropic_response(
         )
 
 
+def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue:
+    output_format: Final = data.get("output_format")
+    output_config: Final = data.get("output_config")
+    if isinstance(output_format, dict) or not isinstance(output_config, dict):
+        return output_format
+    return output_config.get("format")
+
+
 @router.post(
     "/v1/messages/count_tokens",
     tags=["[beta] Anthropic Messages Token Counting"],
@@ -315,7 +326,7 @@ async def count_tokens(
     
     Returns: {"input_tokens": <number>}
     """
-    from litellm.proxy.proxy_server import token_counter as internal_token_counter
+    from litellm.proxy.proxy_server import count_request_tokens
 
     litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
     try:
@@ -327,10 +338,10 @@ async def count_tokens(
         messages: Final = data.get("messages", [])
 
         if not model_name:
-            raise HTTPException(status_code=400, detail={"error": "model parameter is required"})
+            raise HTTPException(status_code=400, detail="model parameter is required")
 
         if not messages:
-            raise HTTPException(status_code=400, detail={"error": "messages parameter is required"})
+            raise HTTPException(status_code=400, detail="messages parameter is required")
 
         # Create TokenCountRequest for the internal endpoint
         from litellm.proxy._types import TokenCountRequest
@@ -342,35 +353,21 @@ async def count_tokens(
             system=data.get("system"),
         )
 
-        # Call the internal token counter function with direct request flag set to False
-        token_response: Final = await internal_token_counter(
+        token_response: Final = await count_request_tokens(
             request=token_request,
             call_endpoint=True,
+            messages_format="anthropic",
+            output_format=_requested_output_format(data),
         )
-        _token_response_dict: dict = {}
-        if isinstance(token_response, TokenCountResponse):
-            _token_response_dict = token_response.model_dump()
-        elif isinstance(token_response, dict):
-            _token_response_dict = token_response
-
-        # Convert the internal response to Anthropic API format
-        return {"input_tokens": _token_response_dict.get("total_tokens", 0)}
-
-    except HTTPException:
-        raise
+    except HTTPException as e:
+        return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request)
     except ProxyException as e:
-        status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500
-        detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error(
-            status_code=status_code,
-            raw_message=e.message,
-        )
-        raise HTTPException(
-            status_code=status_code,
-            detail=detail,
-        )
+        return _anthropic_error_json_response(e, request)
     except Exception as e:
         log_llm_api_exception(e, litellm_call_id)
-        raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})
+        return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
+    else:
+        return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)}
 
 
 @router.post(

litellm/proxy/auth/auth_checks.py

@@ -4352,6 +4352,24 @@ def _check_model_access_helper(
     return True
 
 
+def _owning_team_names_for_team_deployment(
+    model: str,
+    llm_router: Router | None,
+    team_id: str | None,
+    team_model_aliases: Mapping[str, str] | None,
+) -> tuple[str, ...]:
+    if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"):
+        return ()
+    deployment: Final = llm_router.get_deployment_by_model_group_name(model_group_name=model)
+    public_name: Final = (
+        deployment.model_info.team_public_model_name
+        if deployment is not None and deployment.model_info.team_id == team_id
+        else None
+    )
+    legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model)
+    return (*((public_name,) if public_name else ()), *legacy_aliases)
+
+
 def _can_object_call_model(
     model: str | list[str],
     llm_router: Router | None,
@@ -4423,6 +4441,13 @@ def _can_object_call_model(
             *((global_or_router_alias_target,) if global_or_router_alias_target else ()),
         )
     )
+    owning_team_names: Final = (
+        ()
+        if key_alias_applied
+        else _owning_team_names_for_team_deployment(
+            model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases
+        )
+    )
 
     ## check model access for alias + underlying model - allow if either is in allowed models
     for m in potential_models:
@@ -4434,6 +4459,11 @@ def _can_object_call_model(
             team_id=team_id,
         ):
             return True
+    if any(
+        _check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id)
+        for name in owning_team_names
+    ):
+        return True
 
     internal_message: Final = (
         f"{object_type} not allowed to access model. This {object_type} can only access models={models}. "

litellm/proxy/common_request_processing.py

@@ -609,6 +609,27 @@ def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, s
     )
 
 
+def proxy_exception_from_route_error(exc: Exception) -> ProxyException:
+    if isinstance(exc, ProxyException):
+        return exc
+    if isinstance(exc, HTTPException):
+        return proxy_exception_from_http_exception(exc, {})
+    error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR)
+    message: Final = attribute_of(exc, "message", str(exc))
+    carried_code: Final = attribute_of(exc, "code")
+    provider_fields: Final = attribute_of(exc, "provider_specific_fields")
+    return ProxyException(
+        message=redact_internal_details_from_client_message(
+            strip_bug_report_notice(message) if isinstance(message, str) else str(exc)
+        ),
+        type=openai_error_type(exc, error_status),
+        param=openai_error_param(exc),
+        code=error_status,
+        openai_code=carried_code if isinstance(carried_code, str) else None,
+        provider_specific_fields=provider_fields if isinstance(provider_fields, dict) else None,
+    )
+
+
 def _collect_response_file_search_vector_store_ids(data: Mapping[str, object]) -> set[str]:
     vector_store_ids: Final[set[str]] = set()
     tools: Final = data.get("tools")

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

@@ -1,19 +1,139 @@
+from collections.abc import Mapping
+from itertools import chain
+from types import MappingProxyType
 from typing import Final
 
-from fastapi import APIRouter, Depends, Request, Response
-from fastapi.responses import ORJSONResponse
+from fastapi import APIRouter, Depends, HTTPException, Request, Response
+from fastapi.responses import JSONResponse, ORJSONResponse
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator
 
+from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types
 from litellm.proxy._types import *
+from litellm.proxy._types import ProxyException, TokenCountRequest
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
-from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_request_processing import (
+    ProxyBaseLLMRequestProcessing,
+    log_llm_api_exception,
+    proxy_exception_from_http_exception,
+    resolve_litellm_call_id,
+)
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
+from litellm.proxy.common_utils.token_count_settings import count_source_fields
 from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
 
+
+class _GoogleCountTokensTextPart(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    text: str
+
+
+class _GoogleCountTokensSystemInstruction(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    parts: tuple[_GoogleCountTokensTextPart, ...]
+
+    @field_validator("parts", mode="before")
+    @classmethod
+    def normalize_parts(cls, parts: object) -> object:
+        return (parts,) if isinstance(parts, Mapping) else parts
+
+
+class _GoogleCountTokensFunctionDeclaration(BaseModel):
+    model_config = ConfigDict(frozen=True, populate_by_name=True)
+    name: str
+    description: str | None = None
+    parameters: JsonValue | None = None
+    parameters_json_schema: JsonValue | None = Field(default=None, alias="parametersJsonSchema")
+
+
+class _GoogleCountTokensTool(BaseModel):
+    model_config = ConfigDict(frozen=True, populate_by_name=True)
+    function_declarations: tuple[_GoogleCountTokensFunctionDeclaration, ...] = Field(
+        default=(), alias="functionDeclarations"
+    )
+
+
+class _GoogleCountTokensRequest(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True)
+    contents: tuple[Mapping[str, JsonValue], ...] = ()
+    system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction")
+    tools: tuple[_GoogleCountTokensTool, ...] | None = ()
+
+
 router: Final = APIRouter(
     tags=["google genai endpoints"],
 )
 
 
+def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object:
+    if declaration.parameters_json_schema is not None:
+        return declaration.parameters_json_schema
+    if declaration.parameters is not None:
+        return normalize_json_schema_types(declaration.parameters)
+    return {"type": "object", "properties": {}}
+
+
+_GOOGLE_RPC_STATUS: Final = MappingProxyType(
+    {
+        400: "INVALID_ARGUMENT",
+        401: "UNAUTHENTICATED",
+        403: "PERMISSION_DENIED",
+        404: "NOT_FOUND",
+        409: "ABORTED",
+        429: "RESOURCE_EXHAUSTED",
+        499: "CANCELLED",
+        501: "UNIMPLEMENTED",
+        503: "UNAVAILABLE",
+        504: "DEADLINE_EXCEEDED",
+    }
+)
+
+
+def _google_error_response(request: Request, error: ProxyException) -> JSONResponse:
+    from litellm.proxy.proxy_server import (
+        _close_dangling_otel_server_span,  # pyright: ignore[reportPrivateUsage]  # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does
+    )
+
+    status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500
+    _close_dangling_otel_server_span(request, status_code, exc=error)
+    status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION")
+    return JSONResponse(
+        status_code=status_code,
+        content={"error": {"code": status_code, "message": error.message, "status": status}},
+        headers=error.headers,
+    )
+
+
+def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest:
+    from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
+
+    validated_data: Final = TypeAdapter(_GoogleCountTokensRequest).validate_python(data)
+    contents: Final = [dict(content) for content in validated_data.contents]
+    messages: Final = (
+        GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", [])
+    )
+    system: Final = (
+        [{"type": "text", "text": part.text} for part in validated_data.system_instruction.parts]
+        if validated_data.system_instruction is not None
+        else None
+    )
+    declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools or ())
+    tools: Final = [
+        {
+            "name": declaration.name,
+            **({"description": declaration.description} if declaration.description is not None else {}),
+            "input_schema": _input_schema(declaration),
+        }
+        for declaration in declarations
+    ]
+    return TokenCountRequest(
+        model=model_name,
+        contents=contents,
+        messages=messages,
+        system=system,
+        tools=tools or None,
+    )
+
+
 @router.post(
     "/v1beta/models/{model_name:path}:generateContent",
     dependencies=[Depends(user_api_key_auth)],
@@ -165,50 +285,39 @@ async def google_count_tokens(request: Request, model_name: str):
     }
     ```
     """
-    from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
     from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
     from litellm.proxy.proxy_server import token_counter as internal_token_counter
 
     data: Final = await _read_request_body(request=request)
-    contents: Final = data.get("contents", [])
-    # Create TokenCountRequest for the internal endpoint
-    from litellm.proxy._types import TokenCountRequest
+    try:
+        token_request: Final = _google_count_tokens_request(model_name, data)
+    except ValidationError:
+        return _google_error_response(
+            request,
+            ProxyException(
+                message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400
+            ),
+        )
 
-    # Translate contents to openai format messages using the adapter
-    messages = GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", [])
-
-    token_request: Final = TokenCountRequest(
-        model=model_name,
-        contents=contents,
-        messages=messages,  # compatibility when use openai-like endpoint
-    )
-
-    # Call the internal token counter function with direct request flag set to False
-    token_response: Final = await internal_token_counter(
-        request=token_request,
-        call_endpoint=True,
-    )
-    if token_response is not None:
-        # cast the response to the well known format
-        original_response: Final[dict] = token_response.original_response or {}
-        if original_response:
-            return TokenCountDetailsResponse(
-                totalTokens=original_response.get("totalTokens", 0),
-                promptTokensDetails=original_response.get("promptTokensDetails", []),
-            )
-        else:
-            return TokenCountDetailsResponse(
-                totalTokens=token_response.total_tokens or 0,
-                promptTokensDetails=[],
-            )
-
-    #########################################################
-    # Return the response in the well known format
-    #########################################################
-    return TokenCountDetailsResponse(
-        totalTokens=0,
-        promptTokensDetails=[],
-    )
+    try:
+        token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True)
+    except ProxyException as error:
+        return _google_error_response(request, error)
+    except HTTPException as error:
+        return _google_error_response(request, proxy_exception_from_http_exception(error, {}))
+    except Exception as error:
+        log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id")))
+        return _google_error_response(
+            request,
+            ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500),
+        )
+    original_response: Final[dict] = token_response.original_response or {}
+    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

@@ -26,6 +26,7 @@ from collections.abc import (
     MutableMapping,
     Sequence,
 )
+from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from itertools import chain
 from types import MappingProxyType, UnionType
@@ -47,6 +48,7 @@ from typing import (
 )
 
 import anyio
+import anyio.lowlevel
 import websockets
 import websockets.exceptions
 from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError
@@ -87,6 +89,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,
     CallbackDelete,
@@ -133,6 +142,15 @@ from litellm.proxy.common_utils.callback_utils import (
     strip_callback_config,
 )
 from litellm.proxy.common_utils.realtime_utils import _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 (
@@ -162,6 +180,7 @@ from litellm.types.utils import (
     ModelResponseStream,
     StreamingChoices,
     TextCompletionResponse,
+    TokenCountMessagesFormat,
     TokenCountResponse,
     echoed_cost_map_pricing_fields,
     is_server_derived_pricing_key,
@@ -277,6 +296,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,
@@ -307,6 +327,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import (
 from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
 from litellm.litellm_core_utils.core_helpers import (
     _get_parent_otel_span_from_kwargs,
+    as_str_mapping,
     drop_params_flag,
     get_litellm_metadata_from_kwargs,
 )
@@ -879,7 +900,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
 
@@ -13613,52 +13640,362 @@ 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, Any] | None,
+    deployment: Mapping[str, object] | None,
     request_model: str,
     tools: list | None = None,
     system: str | None = None,
-) -> Optional["TokenCountResponse"]:
-    """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting."""
+    messages_format: TokenCountMessagesFormat = "openai",
+    output_format: JsonValue = None,
+) -> 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,
+            messages_format=messages_format,
+            output_format=output_format,
         )
     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:
-        if litellm.disable_token_counter is True:
-            raise ProxyException(
-                message=result.error_message or "Token counting failed",
-                type="token_counting_error",
-                param="model",
-                code=result.status_code or 500,
-            )
-        verbose_proxy_logger.warning(
-            "Provider token counting failed (%s): %s. Falling back to local tokenizer.",
-            result.status_code,
-            result.error_message,
+        return ProviderCountFailure(
+            status_code=result.status_code or 500,
+            message=result.error_message or "Token counting failed",
+            rejected_before_provider_call=result.rejected_before_provider_call,
         )
-        return None
     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_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 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}
+)
+
+
 def _system_message(system: object) -> ChatCompletionSystemMessage | None:
     if not isinstance(system, (str, list)) or not system:
         return None
@@ -13681,8 +14018,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
     Returns:
         TokenCountResponse
     """
-    global llm_router
+    return await count_request_tokens(request=request, call_endpoint=call_endpoint)
 
+
+async def count_request_tokens(
+    request: TokenCountRequest,
+    call_endpoint: bool = False,
+    messages_format: TokenCountMessagesFormat = "openai",
+    output_format: JsonValue = None,
+) -> TokenCountResponse:
     prompt: Final = request.prompt
     messages: Final = request.messages
     contents: Final = request.contents
@@ -13718,8 +14062,16 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
             litellm_model_name = litellm_model_name.split("/", 1)[1]
 
     model_to_use: str = (
-        litellm_model_name or request.model
-    )  # use litellm model name, if it's not avalable then fallback to request.model
+        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
@@ -13730,20 +14082,33 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
         if _model is not None:
             model_to_use = _model
 
-    if provider_counter is not None:
-        result: Final = await _try_provider_token_count(
-            provider_counter=provider_counter,
-            custom_llm_provider=custom_llm_provider,
+    conversation_empty: Final = not messages and not contents
+    prompt_only: Final = bool(prompt) and conversation_empty
+    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,
+            ),
             model_to_use=model_to_use,
-            messages=messages,
-            contents=contents,
-            deployment=deployment,
-            request_model=request.model,
-            tools=tools,
-            system=system,
+            settings=settings,
+            local_estimate_allowed=litellm.disable_token_counter is not True,
         )
-        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:
@@ -13761,14 +14126,11 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
             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
-        Sequence[AllMessageValues] | None, messages
+        Sequence[AllMessageValues] | None,
+        [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages),
     )
     counted_messages: Final = (
         typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
@@ -13776,18 +14138,38 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
     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,
+        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

@@ -35,6 +35,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     _safe_set_request_parsed_body,
 )
+from litellm.proxy.common_utils.token_count_settings import count_source_fields
 from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
 from litellm.types.llms.openai import (
     REASONING_EFFORT,
@@ -42,6 +43,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:
@@ -1067,7 +1069,7 @@ class _ResponsesApiErrorBody(TypedDict):
     error: ReadOnly[_ResponsesApiErrorDetail]
 
 
-class _ResponsesInputTokensResult(TypedDict):
+class _ResponsesInputTokensResult(TokenCountSourceFields):
     object: ReadOnly[str]
     input_tokens: ReadOnly[int]
 
@@ -1198,6 +1200,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

@@ -7,6 +7,7 @@ from typing_extensions import (
 )
 
 from litellm.types.llms.openai import 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
@@ -308,12 +309,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

@@ -9,6 +9,7 @@ from typing import (
     Any,
     Final,
     Literal,
+    TypeAlias,
     get_args,
 )
 
@@ -54,6 +55,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_core_utils.core_helpers import map_finish_reason, process_response_headers
 from . import litellm_params as _litellm_params
@@ -210,6 +212,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
     vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None]
     bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
     bedrock_converse_supports_strict_tools: 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):
@@ -3825,10 +3830,13 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
 
 DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"})
 
+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_TOKEN_COUNT_FIELDS
 )
 
 
@@ -4217,11 +4225,15 @@ class LiteLLMLoggingBaseClass:
         pass
 
 
+TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"]
+
+
 class TokenCountResponse(LiteLLMPydanticObjectBase):
     total_tokens: int
     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
@@ -4232,6 +4244,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

@@ -6327,6 +6327,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

@@ -19337,7 +19337,13 @@
         "supports_prompt_caching": true,
         "supports_reasoning": true,
         "supports_tool_choice": true,
-        "supports_vision": false
+        "supports_vision": false,
+        "token_count_framing": 97
+    },
+    "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,
@@ -20328,7 +20334,13 @@
         "supports_prompt_caching": true,
         "supports_reasoning": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 100
+    },
+    "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,
@@ -75642,7 +75654,9 @@
         "supports_response_schema": true,
         "supports_system_messages": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3"
     },
     "global.moonshotai.kimi-k3": {
         "cache_creation_input_token_cost": 3.75e-06,
@@ -75662,7 +75676,9 @@
         "supports_response_schema": true,
         "supports_system_messages": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3"
     },
     "us.moonshotai.kimi-k3": {
         "cache_creation_input_token_cost": 4.125e-06,
@@ -75682,7 +75698,9 @@
         "supports_response_schema": true,
         "supports_system_messages": true,
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "token_count_framing": 64,
+        "token_count_tokenizer": "kimi_k3"
     },
     "openrouter/xiaomi/mimo-v2.6-flash": {
         "cache_read_input_token_cost": 2.8e-09,

model_prices_and_context_window.schema.json

@@ -1083,6 +1083,18 @@
             "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_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,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm"
-version = "1.104.3"
+version = "1.104.4"
 description = "Library to easily interface with LLM API providers"
 readme = "README.md"
 requires-python = ">=3.10, <3.15"
@@ -355,7 +355,7 @@ litellm-enterprise = { workspace = true }
 members = ["enterprise", "litellm-proxy-extras"]
 
 [tool.commitizen]
-version = "1.104.3"
+version = "1.104.4"
 version_files = [
     "pyproject.toml:^version",
 ]

tests/integration/_support/process.py

@@ -46,12 +46,16 @@ def signal_group(group: int, action: int) -> None:
         pass
 
 
+def graceful_stop_seconds() -> float:
+    return 30.0
+
+
 def stop_root_process(process: subprocess.Popen[bytes]) -> bool:
     if process.poll() is not None:
         return True
     process.terminate()
     try:
-        process.wait(timeout=30)
+        process.wait(timeout=graceful_stop_seconds())
     except subprocess.TimeoutExpired:
         return False
     return True

tests/integration/_support/provider.py

@@ -0,0 +1,49 @@
+"""The fake provider shared by every integration test that takes the `provider` fixture.
+
+Deployments in `proxy_config.yaml` point at `PROVIDER_URL`, so one server answers for all of them. A test
+queues the replies it expects with `expect` and reads what the proxy sent with `received`. Tests run one at a
+time against it; the `provider` fixture checks nothing is left over between tests.
+The proxy's model-info refresh is answered here.
+"""
+
+from __future__ import annotations
+
+import os
+from collections import deque
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Final
+
+from tests.integration._support.wire import Reply, Request, Wire, wire_server
+
+PROVIDER_PORT: Final = int(os.environ.get("INTEGRATION_PROVIDER_PORT", "8191"))
+PROVIDER_URL: Final = f"http://127.0.0.1:{PROVIDER_PORT}"
+MODEL_DISCOVERY: Final = ("GET", "/v1/models")
+_UNQUEUED: Final = Reply(status=500, body=b'{"error": "the shared fake provider has no reply queued for this request"}')
+
+
+@dataclass(slots=True)
+class SharedProvider:
+    wire: Wire
+    replies: deque[Reply]
+    last_test: str | None = field(default=None)
+
+    def expect(self, *replies: Reply) -> None:
+        self.replies.extend(replies)
+
+    def received(self) -> tuple[Request, ...]:
+        return tuple(request for request in self.wire.drain() if (request.method, request.target) != MODEL_DISCOVERY)
+
+
+@contextmanager
+def shared_provider() -> Iterator[SharedProvider]:
+    replies: Final[deque[Reply]] = deque()
+
+    def respond(request: Request) -> Reply:
+        if (request.method, request.target) == MODEL_DISCOVERY:
+            return Reply(body=b'{"object":"list","data":[]}')
+        return replies.popleft() if replies else _UNQUEUED
+
+    with wire_server(respond, port=PROVIDER_PORT) as wire:
+        yield SharedProvider(wire, replies)

tests/integration/_support/vertex.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import json
+from typing import Final
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+
+def service_account_json(project: str, token_url: str) -> str:
+    private_key: Final = (
+        rsa.generate_private_key(public_exponent=65537, key_size=2048)
+        .private_bytes(
+            serialization.Encoding.PEM,
+            serialization.PrivateFormat.PKCS8,
+            serialization.NoEncryption(),
+        )
+        .decode()
+    )
+    return json.dumps(
+        {
+            "type": "service_account",
+            "project_id": project,
+            "private_key_id": "scripted",
+            "private_key": private_key,
+            "client_email": f"scripted@{project}.iam.gserviceaccount.com",
+            "client_id": "0",
+            "auth_uri": f"{token_url}/_oauth/authorize",
+            "token_uri": f"{token_url}/_oauth/token",
+        }
+    )

tests/integration/conftest.py

@@ -15,6 +15,7 @@ from redis import Redis
 from tests.integration._support.client import Gateway, eventually, gateway_from_environment
 from tests.integration._support.generation import LIFECYCLE_SETTINGS
 from tests.integration._support.manifest import OWNED_DIRECTORIES
+from tests.integration._support.provider import SharedProvider, shared_provider
 from tests.integration._support.routing import RoutingPlugin
 
 COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
@@ -126,6 +127,35 @@ def gateway() -> Iterator[Gateway]:
         yield value
 
 
+@pytest.fixture(scope="session")
+def shared_provider_server() -> Iterator[SharedProvider]:
+    if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1":
+        pytest.fail("the shared fake provider needs tests to run one at a time; this group runs under pytest-xdist")
+    with shared_provider() as server:
+        yield server
+        late: Final = server.received()
+        assert late == (), (
+            f"the shared fake provider got {[item.target for item in late]} after {server.last_test} finished"
+        )
+
+
+@pytest.fixture
+def provider(shared_provider_server: SharedProvider, request: pytest.FixtureRequest) -> Iterator[SharedProvider]:
+    stray: Final = shared_provider_server.received()
+    shared_provider_server.replies.clear()
+    assert stray == (), (
+        f"the shared fake provider got {[item.target for item in stray]} "
+        f"after {shared_provider_server.last_test} finished"
+    )
+    yield shared_provider_server
+    shared_provider_server.last_test = request.node.nodeid
+    unused: Final = len(shared_provider_server.replies)
+    unread: Final = shared_provider_server.received()
+    shared_provider_server.replies.clear()
+    assert unused == 0, f"{unused} queued provider replies were never requested"
+    assert unread == (), f"the test never read the provider requests {[item.target for item in unread]}"
+
+
 @pytest.fixture
 def peer(gateway: Gateway) -> Iterator[Gateway]:
     url: Final = os.environ["INTEGRATION_PEER_URL"]

tests/integration/providers/_cache_control_marks_support.py

@@ -0,0 +1,479 @@
+import json
+import re
+import uuid
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import yaml
+from integration._support.client import Gateway, eventually, object_value
+from integration._support.database import read_rows
+from integration._support.wire import Reply, Request
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
+
+# https://platform.claude.com/docs/en/build-with-claude/prompt-caching (read 2026-09-28): at most 4 blocks with cache_control
+ANTHROPIC_CACHE_CONTROL_CAP: Final = 4
+# https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html (read 2026-09-28): at most 4 cache checkpoints
+BEDROCK_CACHE_CHECKPOINT_CAP: Final = 4
+
+ANTHROPIC_MODEL: Final = "claude-opus-5-5"
+BEDROCK_MODEL: Final = "anthropic.claude-opus-5-5"
+UNMAPPED_BEDROCK_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/scripted"
+PROVIDER_KEY: Final = "synthetic-provider-key"
+SYSTEM: Final = "Use the provided tool results to answer the user."
+ASK: Final = "Look up the weather in London, Paris, and Tokyo."
+CITIES: Final = ("London", "Paris", "Tokyo")
+EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"}
+POINTS: Final[list[JsonValue]] = [
+    {"location": "message", "role": "system"},
+    {"location": "message", "index": -1},
+]
+TOOL: Final[dict[str, JsonValue]] = {
+    "type": "function",
+    "function": {
+        "name": "lookup_weather",
+        "description": "Look up the weather in a city.",
+        "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+    },
+}
+SYSTEM_LABEL: Final = f"system:{SYSTEM}"
+ASK_LABEL: Final = f"user:text:{ASK}"
+_MARKER: Final = re.compile(rb"marker-([0-9a-f]{32})")
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+
+@dataclass(frozen=True, slots=True)
+class Mark:
+    label: str
+    ttl: str | None
+
+
+def new_marker() -> str:
+    return uuid.uuid4().hex
+
+
+def final_text(marker: str) -> str:
+    return f"Summarize the results in one word. marker-{marker}"
+
+
+def final_label(marker: str) -> str:
+    return f"user:text:{final_text(marker)}"
+
+
+def call_id(city: str) -> str:
+    return f"call_weather_{city.lower()}"
+
+
+def tool_use_label(city: str) -> str:
+    return f"assistant:tool_use:{call_id(city)}"
+
+
+def tool_call(city: str, **fields: JsonValue) -> dict[str, JsonValue]:
+    return {
+        "id": call_id(city),
+        "type": "function",
+        "function": {"name": "lookup_weather", "arguments": json.dumps({"city": city})},
+        **fields,
+    }
+
+
+def marked_calls(mark: JsonValue = EPHEMERAL, cities: Sequence[str] = CITIES) -> list[JsonValue]:
+    return [tool_call(city, cache_control=mark) for city in cities]
+
+
+def client_marked() -> list[str]:
+    return [ASK_LABEL, *(tool_use_label(city) for city in CITIES)]
+
+
+def ask(*, marked: bool) -> dict[str, JsonValue]:
+    block: Final[dict[str, JsonValue]] = {"type": "text", "text": ASK}
+    return {"role": "user", "content": [{**block, "cache_control": EPHEMERAL} if marked else block]}
+
+
+def tool_results(calls: Sequence[JsonValue]) -> list[JsonValue]:
+    return [
+        {"role": "tool", "tool_call_id": call["id"], "content": "sunny."}
+        for call in calls
+        if isinstance(call, dict) and str(call.get("id", "")).startswith("call_")
+    ]
+
+
+def conversation(
+    marker: str,
+    calls: Sequence[JsonValue],
+    *,
+    ask_marked: bool = True,
+    assistant: dict[str, JsonValue] | None = None,
+) -> list[JsonValue]:
+    turn: Final[dict[str, JsonValue]] = {
+        "role": "assistant",
+        "content": "",
+        "tool_calls": list(calls),
+        **(assistant or {}),
+    }
+    return [
+        {"role": "system", "content": SYSTEM},
+        ask(marked=ask_marked),
+        turn,
+        *tool_results(calls),
+        {"role": "user", "content": final_text(marker)},
+    ]
+
+
+def chat_body(model: str, messages: Sequence[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]:
+    return {"model": model, "messages": list(messages), "tools": [TOOL], "max_tokens": 64, **fields}
+
+
+def messages_body(model: str, marker: str, *, stream: bool) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "max_tokens": 64,
+        "stream": stream,
+        "system": [{"type": "text", "text": SYSTEM}],
+        "tools": [
+            {
+                "name": "lookup_weather",
+                "description": "Look up the weather in a city.",
+                "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+            }
+        ],
+        "messages": [
+            {"role": "user", "content": [{"type": "text", "text": ASK, "cache_control": EPHEMERAL}]},
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "tool_use",
+                        "id": call_id(city),
+                        "name": "lookup_weather",
+                        "input": {"city": city},
+                        "cache_control": EPHEMERAL,
+                    }
+                    for city in CITIES
+                ],
+            },
+            {
+                "role": "user",
+                "content": [
+                    *({"type": "tool_result", "tool_use_id": call_id(city), "content": "sunny."} for city in CITIES),
+                    {"type": "text", "text": final_text(marker)},
+                ],
+            },
+        ],
+    }
+
+
+def responses_body(model: str, marker: str) -> dict[str, JsonValue]:
+    items: Final[list[JsonValue]] = [
+        {"role": "user", "content": [{"type": "input_text", "text": ASK, "cache_control": EPHEMERAL}]},
+        *(
+            {
+                "type": "function_call",
+                "call_id": call_id(city),
+                "name": "lookup_weather",
+                "arguments": json.dumps({"city": city}),
+                "cache_control": EPHEMERAL,
+            }
+            for city in CITIES
+        ),
+        *({"type": "function_call_output", "call_id": call_id(city), "output": "sunny."} for city in CITIES),
+        {"role": "user", "content": final_text(marker)},
+    ]
+    return {
+        "model": model,
+        "instructions": SYSTEM,
+        "max_output_tokens": 64,
+        "tools": [{"type": "function", "name": "lookup_weather", "parameters": {"type": "object"}}],
+        "input": items,
+    }
+
+
+def marker_of(request: Request) -> str:
+    found: Final = _MARKER.findall(request.body)
+    return found[-1].decode() if found else "unmarked"
+
+
+class _AnthropicBlock(BaseModel):
+    model_config = ConfigDict(extra="allow", frozen=True)
+    type: str
+    text: str | None = None
+    id: str | None = None
+    tool_use_id: str | None = None
+    name: str | None = None
+    cache_control: JsonValue = None
+
+
+class _AnthropicMessage(BaseModel):
+    model_config = ConfigDict(extra="allow", frozen=True)
+    role: str
+    content: str | tuple[_AnthropicBlock, ...]
+
+
+class _AnthropicTool(BaseModel):
+    model_config = ConfigDict(extra="allow", frozen=True)
+    name: str = ""
+    cache_control: JsonValue = None
+
+
+class _AnthropicBody(BaseModel):
+    model_config = ConfigDict(extra="allow", frozen=True)
+    system: str | tuple[_AnthropicBlock, ...] = ()
+    messages: tuple[_AnthropicMessage, ...] = ()
+    tools: tuple[_AnthropicTool, ...] = ()
+    cache_control: JsonValue = None
+    stream: bool = False
+
+
+def _ttl(cache_control: JsonValue) -> str | None:
+    if not isinstance(cache_control, dict):
+        return None
+    ttl: Final = cache_control.get("ttl")
+    return ttl if isinstance(ttl, str) else None
+
+
+def _block_label(role: str, block: _AnthropicBlock) -> str:
+    detail: Final = block.text if block.type == "text" else block.id or block.tool_use_id or ""
+    return f"{role}:{block.type}:{detail}"
+
+
+def _message_blocks(body: _AnthropicBody) -> Iterator[tuple[str, _AnthropicBlock]]:
+    for message in body.messages:
+        if isinstance(message.content, tuple):
+            yield from ((message.role, block) for block in message.content)
+
+
+def _anthropic_body(request: Request) -> _AnthropicBody:
+    return _AnthropicBody.model_validate_json(request.body)
+
+
+def anthropic_marks(request: Request) -> tuple[Mark, ...]:
+    body: Final = _anthropic_body(request)
+    system: Final = body.system if isinstance(body.system, tuple) else ()
+    return (
+        *(Mark(f"tool:{tool.name}", _ttl(tool.cache_control)) for tool in body.tools if tool.cache_control is not None),
+        *(
+            Mark(f"system:{block.text}", _ttl(block.cache_control))
+            for block in system
+            if block.cache_control is not None
+        ),
+        *(
+            Mark(_block_label(role, block), _ttl(block.cache_control))
+            for role, block in _message_blocks(body)
+            if block.cache_control is not None
+        ),
+        *((Mark("request", _ttl(body.cache_control)),) if body.cache_control is not None else ()),
+    )
+
+
+def anthropic_labels(request: Request) -> list[str]:
+    return [mark.label for mark in anthropic_marks(request)]
+
+
+def _anthropic_error(message: str) -> Reply:
+    return Reply(
+        status=400,
+        body=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": message}}).encode(),
+    )
+
+
+def _ttl_out_of_order(marks: Sequence[Mark]) -> bool:
+    first_short: Final = next((index for index, mark in enumerate(marks) if mark.ttl != "1h"), len(marks))
+    return any(mark.ttl == "1h" for mark in marks[first_short:])
+
+
+def _usage() -> dict[str, JsonValue]:
+    return {"input_tokens": 12, "output_tokens": 1, "cache_creation_input_tokens": 12, "cache_read_input_tokens": 0}
+
+
+def _anthropic_message(identity: str) -> dict[str, JsonValue]:
+    return {
+        "id": identity,
+        "type": "message",
+        "role": "assistant",
+        "model": ANTHROPIC_MODEL,
+        "content": [{"type": "text", "text": "sunny"}],
+        "stop_reason": "end_turn",
+        "stop_sequence": None,
+        "usage": _usage(),
+    }
+
+
+def _event(name: str, data: dict[str, JsonValue]) -> bytes:
+    return f"event: {name}\ndata: {json.dumps(data)}\n\n".encode()
+
+
+def _anthropic_stream(identity: str) -> tuple[bytes, ...]:
+    return (
+        _event(
+            "message_start",
+            {"type": "message_start", "message": {**_anthropic_message(identity), "content": [], "stop_reason": None}},
+        ),
+        _event(
+            "content_block_start",
+            {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+        ),
+        _event(
+            "content_block_delta",
+            {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "sunny"}},
+        ),
+        _event("content_block_stop", {"type": "content_block_stop", "index": 0}),
+        _event(
+            "message_delta",
+            {
+                "type": "message_delta",
+                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                "usage": {"output_tokens": 1},
+            },
+        ),
+        _event("message_stop", {"type": "message_stop"}),
+    )
+
+
+def anthropic_peer(request: Request) -> Reply:
+    marks: Final = anthropic_marks(request)
+    if len(marks) > ANTHROPIC_CACHE_CONTROL_CAP:
+        return _anthropic_error(
+            f"A maximum of {ANTHROPIC_CACHE_CONTROL_CAP} blocks with cache_control may be provided. Found {len(marks)}."
+        )
+    if _ttl_out_of_order(marks):
+        return _anthropic_error("a ttl='1h' cache_control block must not come after a ttl='5m' cache_control block")
+    identity: Final = f"msg_{marker_of(request)}"
+    if _anthropic_body(request).stream:
+        return Reply(content_type="text/event-stream", chunks=_anthropic_stream(identity))
+    return Reply(body=json.dumps(_anthropic_message(identity)).encode())
+
+
+def _bedrock_label(role: str, block: JsonValue) -> str:
+    if not isinstance(block, dict):
+        return f"{role}:start"
+    if isinstance(block.get("text"), str):
+        return f"{role}:text:{block['text']}"
+    for kind in ("toolUse", "toolResult"):
+        inner = block.get(kind)
+        if isinstance(inner, dict):
+            return f"{role}:{kind}:{inner.get('toolUseId')}"
+    spec: Final = block.get("toolSpec")
+    return f"tool:{spec.get('name')}" if isinstance(spec, dict) else f"{role}:other"
+
+
+def _cache_points(role: str, blocks: JsonValue) -> Iterator[str]:
+    listed: Final = blocks if isinstance(blocks, list) else []
+    for previous, block in zip([None, *listed], listed):
+        if isinstance(block, dict) and "cachePoint" in block:
+            yield _bedrock_label(role, previous)
+
+
+def _bedrock_sections(body: dict[str, JsonValue]) -> Iterator[tuple[str, JsonValue]]:
+    tool_config: Final = body.get("toolConfig")
+    yield ("tool", tool_config.get("tools") if isinstance(tool_config, dict) else None)
+    yield ("system", body.get("system"))
+    messages: Final = body.get("messages")
+    for message in messages if isinstance(messages, list) else []:
+        if isinstance(message, dict):
+            yield (str(message.get("role")), message.get("content"))
+
+
+def bedrock_labels(request: Request) -> list[str]:
+    body: Final = _JSON_OBJECT.validate_json(request.body)
+    return [label for role, blocks in _bedrock_sections(body) for label in _cache_points(role, blocks)]
+
+
+def bedrock_peer(request: Request) -> Reply:
+    found: Final = len(bedrock_labels(request))
+    if found > BEDROCK_CACHE_CHECKPOINT_CAP:
+        return Reply(
+            status=400,
+            headers={"x-amzn-errortype": "ValidationException"},
+            body=json.dumps(
+                {
+                    "message": f"A maximum of {BEDROCK_CACHE_CHECKPOINT_CAP} cache checkpoints may be provided. Found {found}."
+                }
+            ).encode(),
+        )
+    return Reply(
+        body=json.dumps(
+            {
+                "output": {"message": {"role": "assistant", "content": [{"text": "sunny"}]}},
+                "stopReason": "end_turn",
+                "usage": {
+                    "inputTokens": 12,
+                    "outputTokens": 1,
+                    "totalTokens": 13,
+                    "cacheWriteInputTokens": 12,
+                    "cacheReadInputTokens": 0,
+                },
+                "metrics": {"latencyMs": 1},
+            }
+        ).encode()
+    )
+
+
+def _without_cache_points(value: JsonValue) -> JsonValue:
+    if isinstance(value, dict):
+        return {key: _without_cache_points(inner) for key, inner in value.items()}
+    if isinstance(value, list):
+        return [
+            _without_cache_points(inner) for inner in value if not (isinstance(inner, dict) and "cachePoint" in inner)
+        ]
+    return value
+
+
+def sent_without_cache_points(requests: Sequence[Request], marker: str) -> list[JsonValue]:
+    return [
+        _without_cache_points(_JSON_OBJECT.validate_json(request.body.replace(marker.encode(), b"marker")))
+        for request in requests
+    ]
+
+
+def gateway_injected(response_id: str) -> bool:
+    rows: Final = eventually(
+        lambda: read_rows(
+            """SELECT metadata->>'litellm_gateway_injected_cache' AS injected FROM "LiteLLM_SpendLogs" """
+            "WHERE request_id=%s",
+            (response_id,),
+        ),
+        lambda values: len(values) == 1,
+        seconds=70,
+    )
+    return rows[0]["injected"] is not None
+
+
+def post_chat(gateway: Gateway, body: dict[str, JsonValue], *, key: str | None = None) -> tuple[int, str, str]:
+    response: Final = gateway.request("POST", "/v1/chat/completions", body, key=key)
+    identity: Final = _JSON_OBJECT.validate_json(response.content).get("id") if response.status_code == 200 else None
+    return response.status_code, str(identity), response.text
+
+
+def anthropic_deployment(name: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]:
+    return {
+        "model_name": name,
+        "litellm_params": {
+            "model": f"anthropic/{ANTHROPIC_MODEL}",
+            "api_base": api_base,
+            "api_key": PROVIDER_KEY,
+            **fields,
+        },
+    }
+
+
+def owned_config(
+    directory: Path,
+    model_list: Sequence[JsonValue],
+    *,
+    litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}),
+    router_settings: Mapping[str, JsonValue] = MappingProxyType({}),
+) -> Path:
+    config: Final = _JSON_OBJECT.validate_python(
+        yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    )
+    merged: Final = {
+        **config,
+        "model_list": list(model_list),
+        "litellm_settings": {**object_value(config["litellm_settings"]), **litellm_settings},
+        "router_settings": {**object_value(config["router_settings"]), "num_retries": 0, **router_settings},
+    }
+    path: Final = directory / f"cache-control-marks-{uuid.uuid4().hex}.yaml"
+    path.write_text(yaml.safe_dump(merged))
+    return path

tests/integration/providers/test_count_tokens_error_shapes_wire.py

@@ -0,0 +1,215 @@
+import json
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Final
+
+import anthropic
+import openai
+import pytest
+from integration._support.client import Gateway, gateway_from_environment
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}]
+_PROVIDER_KEY_VARIABLES: Final = (
+    "ANTHROPIC_API_KEY",
+    "ANTHROPIC_AUTH_TOKEN",
+    "AZURE_AI_API_KEY",
+    "AZURE_API_KEY",
+    "AZURE_OPENAI_API_KEY",
+    "AZURE_AD_TOKEN",
+    "AZURE_CLIENT_ID",
+    "AZURE_CLIENT_SECRET",
+    "AZURE_TENANT_ID",
+    "OPENAI_API_KEY",
+)
+
+
+@pytest.fixture(scope="module")
+def provider_error_config(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    config: Final = tmp_path_factory.mktemp("count-tokens-provider-errors") / "config.yaml"
+    base: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
+    config.write_text(base.read_text() + "\nlitellm_settings:\n  disable_token_counter: true\n")
+    return config
+
+
+@pytest.fixture(scope="module")
+def keyless_proxy(tmp_path_factory: pytest.TempPathFactory, provider_error_config: Path) -> Iterator[OwnedProxy]:
+    with (
+        gateway_from_environment() as gateway,
+        owned_proxy_process(
+            gateway,
+            tmp_path_factory.mktemp("count-tokens-error-shapes"),
+            {},
+            config=provider_error_config,
+            remove_environment=_PROVIDER_KEY_VARIABLES,
+            workers=1,
+        ) as owned,
+    ):
+        yield owned
+
+
+@pytest.fixture(scope="module")
+def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
+    with (
+        gateway_from_environment() as gateway,
+        owned_proxy_process(
+            gateway,
+            tmp_path_factory.mktemp("count-tokens-azure-environment-key"),
+            {"AZURE_API_KEY": "synthetic-azure-env-key"},
+            remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"),
+            workers=1,
+        ) as owned,
+    ):
+        yield owned
+
+
+def _anthropic(gateway: Gateway) -> anthropic.Anthropic:
+    return anthropic.Anthropic(
+        api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client
+    )
+
+
+def _never_called(request: Request) -> Reply:
+    raise AssertionError(f"count reached the provider without a key: {request.target}")
+
+
+@pytest.mark.parametrize(
+    ("deployment", "missing"),
+    (
+        ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "),
+        ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "),
+    ),
+    ids=("anthropic", "azure_ai_anthropic"),
+)
+def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401(
+    keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str
+) -> None:
+    gateway: Final = keyless_proxy.gateway
+    with wire_server(_never_called) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0)
+        with pytest.raises(anthropic.AuthenticationError) as raised:
+            _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
+        assert raised.value.body == {
+            "type": "error",
+            "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]},
+        }
+        assert missing in raised.value.message, raised.value.message
+        assert wire.drain() == ()
+
+
+def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key(
+    azure_environment_key_proxy: OwnedProxy,
+) -> None:
+    gateway: Final = azure_environment_key_proxy.gateway
+
+    def count(request: Request) -> Reply:
+        assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers
+        return Reply(body=json.dumps({"input_tokens": 58}).encode())
+
+    with wire_server(count) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model="azure_ai/claude-sonnet-4-6",
+            api_base=wire.url,
+            api_key=None,
+            num_retries=0,
+        )
+        result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
+        assert result.input_tokens == 58
+
+
+def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None:
+    gateway: Final = keyless_proxy.gateway
+
+    def count(request: Request) -> Reply:
+        assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers
+        assert "x-api-key" not in request.headers, request.headers
+        return Reply(body=json.dumps({"input_tokens": 61}).encode())
+
+    with wire_server(count) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model="azure_ai/claude-sonnet-4-6",
+            api_base=wire.url,
+            api_key=None,
+            azure_ad_token="synthetic-entra-token",
+            num_retries=0,
+        )
+        result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
+        assert result.input_tokens == 61
+        assert len(wire.drain()) == 1
+
+
+def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None:
+    gateway: Final = keyless_proxy.gateway
+    upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}}
+    with (
+        wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = scenario.model(
+            model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0
+        )
+        missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model})
+        assert missing_messages.status_code == 400, missing_messages.text
+        assert _JSON_OBJECT.validate_json(missing_messages.content) == {
+            "type": "error",
+            "error": {"type": "invalid_request_error", "message": "messages parameter is required"},
+        }
+
+        rejected: Final = gateway.request(
+            "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES}
+        )
+        assert rejected.status_code == 400, rejected.text
+        body: Final = _JSON_OBJECT.validate_json(rejected.content)
+        assert set(body) <= {"type", "error", "request_id"}, body
+        assert body["error"] == upstream_error["error"], body
+        assert len(wire.drain()) == 1
+
+
+def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None:
+    gateway: Final = keyless_proxy.gateway
+    with wire_server(_never_called) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0
+        )
+        client: Final = openai.OpenAI(
+            api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client
+        )
+        with pytest.raises(openai.AuthenticationError) as raised:
+            client.responses.input_tokens.count(model=model, input="Count this synthetic request.")
+        assert f"No OpenAI API key is set for model {model}" in raised.value.message
+        assert wire.drain() == ()
+
+
+def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None:
+    gateway: Final = keyless_proxy.gateway
+    upstream_error: Final = {
+        "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"}
+    }
+    with (
+        wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = scenario.model(
+            model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0
+        )
+        rejected: Final = gateway.request(
+            "POST",
+            f"/v1beta/models/{model}:countTokens",
+            {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]},
+        )
+        assert rejected.status_code == 400, rejected.text
+        error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"]
+        assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error
+        assert "API key not valid" in str(error["message"]), error
+        assert len(wire.drain()) == 1
+
+        malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"})
+        assert malformed.status_code == 400, malformed.text
+        assert _JSON_OBJECT.validate_json(malformed.content) == {
+            "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"}
+        }

tests/integration/providers/test_openai_count_tokens_tools_wire.py

@@ -0,0 +1,393 @@
+import json
+from collections.abc import Iterator, Mapping
+from typing import Final
+
+import pytest
+from integration._support.client import Scenario, gateway_from_environment
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+_COUNT: Final = 4242
+_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {}}
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+
+@pytest.fixture(scope="module")
+def counting_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
+    with (
+        gateway_from_environment() as gateway,
+        owned_proxy_process(gateway, tmp_path_factory.mktemp("openai-count-tools"), {}, workers=1) as owned,
+    ):
+        yield owned
+
+
+def _deployment(scenario: Scenario, api_base: str) -> str:
+    return scenario.model(
+        model="openai/gpt-4o-mini",
+        api_base=f"{api_base}/v1",
+        api_key="sk-integration-openai",
+    )
+
+
+def _count_response(_request: Request) -> Reply:
+    return Reply(body=json.dumps({"input_tokens": _COUNT}).encode())
+
+
+_PDF_BASE64: Final = "JVBERi0xLjQK"
+_PNG_BASE64: Final = "aW1hZ2U="
+_LOOKUP_TOOL: Final[dict[str, JsonValue]] = {
+    "name": "lookup",
+    "description": "Find a record",
+    "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
+}
+_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = {
+    "id": "resp_count_parity",
+    "object": "response",
+    "created_at": 1,
+    "status": "completed",
+    "model": "gpt-4o-mini",
+    "output": [
+        {
+            "type": "message",
+            "id": "msg_count_parity",
+            "status": "completed",
+            "role": "assistant",
+            "content": [{"type": "output_text", "text": "OK", "annotations": []}],
+        }
+    ],
+    "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
+}
+_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal"
+_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
+_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = {
+    "structured-output": {
+        "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}},
+        "messages": [{"role": "user", "content": "Answer in JSON."}],
+    },
+    "affinity-wrapped-reasoning": {
+        "messages": [
+            {"role": "user", "content": "hi"},
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"},
+                    {"type": "text", "text": "OK"},
+                ],
+            },
+            {"role": "user", "content": "again"},
+        ],
+    },
+    "base64-pdf": {
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Summarize the attached contract."},
+                    {
+                        "type": "document",
+                        "source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64},
+                    },
+                ],
+            }
+        ],
+    },
+    "base64-image": {
+        "system": [{"type": "text", "text": "Describe images briefly."}],
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}},
+                    {"type": "text", "text": "What is this?"},
+                ],
+            }
+        ],
+    },
+    "replayed-web-search": {
+        "messages": [
+            {"role": "user", "content": "What is the newest LiteLLM release?"},
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "server_tool_use",
+                        "id": "srvtoolu_01",
+                        "name": "web_search",
+                        "input": {"query": "newest LiteLLM release"},
+                    },
+                    {
+                        "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": "",
+                                "snippet": "LiteLLM ships a weekly release candidate and a stable line.",
+                            }
+                        ],
+                    },
+                    {"type": "text", "text": "The newest release is in the release notes."},
+                ],
+            },
+            {"role": "user", "content": "Thanks."},
+        ],
+    },
+    "tool-turns": {
+        "system": "Use the lookup tool.",
+        "tools": [_LOOKUP_TOOL],
+        "messages": [
+            {"role": "user", "content": "Find record 14."},
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "Looking it up."},
+                    {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}},
+                ],
+            },
+            {
+                "role": "user",
+                "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}],
+            },
+        ],
+    },
+}
+
+
+def _responses_peer(request: Request) -> Reply:
+    if request.target == "/v1/responses/input_tokens":
+        return _count_response(request)
+    if request.target == "/v1/responses":
+        return Reply(body=json.dumps(_RESPONSES_REPLY).encode())
+    return Reply(status=404, body=json.dumps({"error": request.target}).encode())
+
+
+def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]:
+    matching: Final = tuple(request for request in wire_requests if request.target == target)
+    assert len(matching) == 1, [request.target for request in wire_requests]
+    return _JSON_OBJECT.validate_json(matching[0].body)
+
+
+def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None}
+
+
+def _count_body(wire: Wire) -> dict[str, JsonValue]:
+    requests: Final = wire.drain()
+    assert len(requests) == 1, requests
+    request: Final = requests[0]
+    assert request.method == "POST"
+    assert request.target == "/v1/responses/input_tokens", request.target
+    return _JSON_OBJECT.validate_json(request.body)
+
+
+def test_openai_count_tokens_uses_deployment_api_base(counting_proxy: OwnedProxy) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_count_response) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/messages/count_tokens",
+            {"model": model, "messages": [{"role": "user", "content": "hello"}]},
+        )
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT
+        assert "tools" not in _count_body(wire)
+
+
+def test_google_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_count_response) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            f"/v1beta/models/{model}:countTokens",
+            {
+                "contents": [{"role": "user", "parts": [{"text": "hello"}]}],
+                "tools": [
+                    {
+                        "functionDeclarations": [
+                            {
+                                "name": "lookup",
+                                "description": "Find a record",
+                                "parameters": _TOOL_SCHEMA,
+                            }
+                        ]
+                    }
+                ],
+            },
+        )
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT
+        assert _count_body(wire)["tools"] == [
+            {
+                "type": "function",
+                "name": "lookup",
+                "description": "Find a record",
+                "parameters": _TOOL_SCHEMA,
+                "strict": False,
+            }
+        ]
+
+
+def test_messages_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_count_response) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/messages/count_tokens",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": "hello"}],
+                "tools": [
+                    {
+                        "name": "lookup",
+                        "description": "Find a record",
+                        "input_schema": _TOOL_SCHEMA,
+                    }
+                ],
+            },
+        )
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT
+        assert _count_body(wire)["tools"] == [
+            {
+                "type": "function",
+                "name": "lookup",
+                "description": "Find a record",
+                "parameters": _TOOL_SCHEMA,
+                "strict": False,
+            }
+        ]
+
+
+def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedProxy) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_count_response) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            "/utils/token_counter",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": "hello"}],
+                "tools": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "lookup",
+                            "description": "Find a record",
+                            "parameters": _TOOL_SCHEMA,
+                        },
+                    }
+                ],
+            },
+            params={"call_endpoint": "true"},
+        )
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT
+        assert _count_body(wire)["tools"] == [
+            {
+                "type": "function",
+                "name": "lookup",
+                "description": "Find a record",
+                "parameters": _TOOL_SCHEMA,
+            }
+        ]
+
+
+@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES))
+def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends(
+    counting_proxy: OwnedProxy, case: str
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_responses_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        body: Final = {"model": model, **_ANTHROPIC_BODIES[case]}
+        generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16})
+        counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body)
+        assert generated.status_code == 200, generated.text
+        assert counted.status_code == 200, counted.text
+        assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT
+        requests: Final = wire.drain()
+        inference_body: Final = _body_for(requests, "/v1/responses")
+        count_body: Final = _body_for(requests, "/v1/responses/input_tokens")
+        assert _counted_fields(count_body) == _counted_fields(inference_body)
+        assert count_body["model"] == inference_body["model"]
+
+
+_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = {
+    "messages": [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Summarize the notes."},
+                {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}},
+                {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []},
+            ],
+        },
+        {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]},
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "tool_result",
+                    "tool_use_id": "toolu_01",
+                    "content": [
+                        {
+                            "type": "search_result",
+                            "source": "https://example.com/hits",
+                            "title": "Hits",
+                            "content": [{"type": "text", "text": "hit"}],
+                        }
+                    ],
+                }
+            ],
+        },
+    ],
+}
+
+
+def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_responses_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY})
+        assert response.status_code == 400, response.text
+        assert "document with a text source, search_result, search_result inside a tool_result" in response.text
+        assert wire.drain() == ()
+
+
+def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_count_response) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/messages/count_tokens",
+            {"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]},
+        )
+        assert response.status_code == 400, response.text
+        assert "messages.1.role: Field required" in response.text
+        assert wire.drain() == ()
+
+
+def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None:
+    gateway: Final = counting_proxy.gateway
+    with wire_server(_count_response) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            f"/v1beta/models/{model}:countTokens",
+            {
+                "systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]},
+                "contents": [{"role": "user", "parts": [{"text": "hello"}]}],
+            },
+        )
+        assert response.status_code == 200, response.text
+        assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English."

tests/integration/providers/test_vertex_gemini_count_tokens_wire.py

@@ -0,0 +1,630 @@
+import asyncio
+import json
+import threading
+from collections.abc import Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
+from types import MappingProxyType
+from typing import Final
+from uuid import uuid4
+
+import anthropic
+import httpx
+import pytest
+from integration._support.client import Gateway, Scenario
+from integration._support.process import graceful_stop_seconds
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+_MODEL: Final = "gemini-4-argon"
+_PROJECT: Final = "scripted-project"
+_LOCATION: Final = "us-east5"
+_COUNT: Final = 777
+_UNIT: Final = " red blue green yellow orange purple black white silver gold"
+_USER: Final = f"<filler>{_UNIT * 100}</filler>\nReturn exactly OK."
+_MESSAGES: Final[JsonValue] = [{"role": "user", "content": _USER}]
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_JSON_LIST: Final = TypeAdapter(list[JsonValue])
+_REPLY: Final[Mapping[str, JsonValue]] = {
+    "candidates": [
+        {
+            "content": {"role": "model", "parts": [{"text": "OK"}]},
+            "finishReason": "STOP",
+        }
+    ],
+    "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11},
+}
+
+
+class _GeminiPeer:
+    def __init__(self) -> None:
+        self.status: int = 200
+        self.requests: tuple[Request, ...] = ()
+        self._lock: Final = threading.Lock()
+
+    def set_status(self, status: int) -> None:
+        self.status = status
+
+    def __call__(self, request: Request) -> Reply:
+        with self._lock:
+            self.requests = (*self.requests, request)
+            status: Final = self.status
+        if request.target.endswith(":countTokens"):
+            return Reply(status=status, body=json.dumps({"totalTokens": _COUNT}).encode())
+        if request.target.endswith(":generateContent"):
+            return Reply(body=json.dumps(_REPLY).encode())
+        return Reply(status=404, body=json.dumps({"error": request.target}).encode())
+
+    def drain(self) -> tuple[Request, ...]:
+        with self._lock:
+            requests: Final = self.requests
+            self.requests = ()
+        return requests
+
+
+def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str:
+    return scenario.model(
+        model=f"vertex_ai/{_MODEL}",
+        api_base=api_base,
+        api_key=None,
+        vertex_project=_PROJECT,
+        vertex_location=_LOCATION,
+        vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
+    )
+
+
+def _messages(model: str, system: JsonValue | None, tools: JsonValue | None) -> Mapping[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": _MESSAGES,
+        **({"system": system} if system is not None else {}),
+        **({"tools": tools} if tools is not None else {}),
+    }
+
+
+def _json_body(request: Request) -> Mapping[str, JsonValue]:
+    return _JSON_OBJECT.validate_json(request.body)
+
+
+def _count_requests(wire: Wire) -> tuple[Request, ...]:
+    return tuple(request for request in wire.drain() if request.target.endswith(":countTokens"))
+
+
+def _count_tokens_tools(generated_tools: JsonValue) -> JsonValue:
+    return _JSON_LIST.validate_python(
+        [
+            {
+                {
+                    "function_declarations": "functionDeclarations",
+                    "url_context": "urlContext",
+                    "code_execution": "codeExecution",
+                }.get(key, key): value
+                for key, value in tool.items()
+            }
+            for tool in TypeAdapter(list[Mapping[str, JsonValue]]).validate_python(generated_tools)
+        ]
+    )
+
+
+def _count_tokens_request_from_generate_content(
+    generated_body: Mapping[str, JsonValue],
+) -> Mapping[str, JsonValue]:
+    system_instruction: Final = generated_body.get("systemInstruction", generated_body.get("system_instruction"))
+    generated_tools: Final = generated_body.get("tools")
+    return MappingProxyType(
+        {
+            "contents": generated_body["contents"],
+            **({"systemInstruction": system_instruction} if system_instruction is not None else {}),
+            **({"tools": _count_tokens_tools(generated_tools)} if generated_tools is not None else {}),
+        }
+    )
+
+
+def _tools(description: str) -> JsonValue:
+    return _JSON_LIST.validate_python(
+        [
+            {
+                "name": "noop",
+                "description": description,
+                "input_schema": {"type": "object", "properties": {}},
+            }
+        ]
+    )
+
+
+@pytest.mark.parametrize(
+    ("large_system", "large_tool"),
+    ((False, False), (True, False), (False, True), (True, True)),
+    ids=("baseline", "large-system", "large-tool", "large-system-and-tool"),
+)
+def test_messages_and_count_tokens_send_matching_gemini_context(
+    gateway: Gateway, large_system: bool, large_tool: bool
+) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        system: Final = (
+            f"<cache-buster>{uuid4().hex}</cache-buster>\n"
+            "This is a synthetic context-budget diagnostic. Return exactly OK." + (_UNIT * 100 if large_system else "")
+            if large_system
+            else None
+        )
+        tools: Final = _tools(_UNIT * 100) if large_tool else None
+        body: Final = _messages(model, system, tools)
+        generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16})
+        counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body)
+        assert generated.status_code == 200, generated.text
+        assert counted.status_code == 200, counted.text
+        assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text
+        requests: Final = peer.drain()
+        generated_body: Final = _json_body(
+            next(request for request in requests if request.target.endswith(":generateContent"))
+        )
+        count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens")))
+        assert count_body == _count_tokens_request_from_generate_content(generated_body)
+
+
+_LOOKUP_TOOL: Final[JsonValue] = {
+    "name": "lookup",
+    "description": "Look up a value",
+    "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
+}
+_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType(
+    {
+        "base64-pdf": {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "Summarize the document."},
+                        {
+                            "type": "document",
+                            "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
+                        },
+                    ],
+                }
+            ],
+        },
+        "base64-image": {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}},
+                        {"type": "text", "text": "What is this?"},
+                    ],
+                }
+            ],
+        },
+        "tool-turns": {
+            "tools": [_LOOKUP_TOOL],
+            "messages": [
+                {"role": "user", "content": "Look up the value."},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {"type": "text", "text": "Looking it up."},
+                        {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}},
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}],
+                },
+            ],
+        },
+    }
+)
+
+
+def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str:
+    return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters)
+
+
+def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue:
+    if "contents" in count_body:
+        return count_body["contents"]
+    return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"]
+
+
+@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
+@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES))
+def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = (
+            _deployment(gateway, scenario, wire.url)
+            if provider == "vertex_ai"
+            else _ai_studio_deployment(scenario, wire.url)
+        )
+        body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]}
+        generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16})
+        counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body)
+        assert generated.status_code == 200, generated.text
+        assert counted.status_code == 200, counted.text
+        assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text
+        requests: Final = peer.drain()
+        generated_body: Final = _json_body(
+            next(request for request in requests if request.target.endswith(":generateContent"))
+        )
+        count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens")))
+        assert _counted_contents(count_body) == generated_body["contents"]
+
+
+_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "Summarize the notes."},
+            {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}},
+            {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []},
+        ],
+    },
+    {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]},
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "tool_result",
+                "tool_use_id": "toolu_01",
+                "content": [
+                    {
+                        "type": "search_result",
+                        "source": "https://example.com/hits",
+                        "title": "Hits",
+                        "content": [{"type": "text", "text": "hit"}],
+                    }
+                ],
+            }
+        ],
+    },
+]
+
+
+@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
+def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = (
+            _deployment(gateway, scenario, wire.url)
+            if provider == "vertex_ai"
+            else _ai_studio_deployment(scenario, wire.url)
+        )
+        response: Final = gateway.request(
+            "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES}
+        )
+        assert response.status_code == 400, response.text
+        assert "document with a text source, search_result, search_result inside a tool_result" in response.text
+        assert peer.drain() == ()
+
+
+def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        system: Final[JsonValue] = [
+            {"type": "text", "text": "First instruction"},
+            {"type": "text", "text": "Second instruction"},
+        ]
+        response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, system, None))
+        assert response.status_code == 200, response.text
+        requests: Final = _count_requests(wire)
+        assert len(requests) == 1
+        assert _json_body(requests[0]) == {
+            "contents": [{"role": "user", "parts": [{"text": _USER}]}],
+            "systemInstruction": {"parts": [{"text": "First instruction"}, {"text": "Second instruction"}]},
+        }
+
+
+def test_count_tokens_without_system_or_tools_sends_contents_only(gateway: Gateway) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None))
+        assert response.status_code == 200, response.text
+        requests: Final = _count_requests(wire)
+        assert len(requests) == 1
+        assert _json_body(requests[0]) == {"contents": [{"role": "user", "parts": [{"text": _USER}]}]}
+
+
+_SDK_COUNT_ARGS: Final[Mapping[str, JsonValue]] = MappingProxyType(
+    {
+        "messages": [{"role": "user", "content": _USER}],
+        "system": "Follow the system instruction.",
+        "tools": [
+            {
+                "name": "noop",
+                "description": "Return no output",
+                "input_schema": {"type": "object", "properties": {}},
+            }
+        ],
+    }
+)
+
+
+def _sync_sdk_count_tokens(gateway: Gateway, model: str) -> int:
+    client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
+    return client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS).input_tokens  # pyright: ignore[reportArgumentType]  # SDK TypedDicts reject a JsonValue mapping
+
+
+async def _async_sdk_count_tokens(gateway: Gateway, model: str) -> int:
+    client: Final = anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
+    return (await client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS)).input_tokens  # pyright: ignore[reportArgumentType]  # SDK TypedDicts reject a JsonValue mapping
+
+
+@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync-sdk", "async-sdk"))
+def test_anthropic_sdk_count_tokens_forwards_system_and_tools(gateway: Gateway, asynchronous: bool) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        input_tokens: Final = (
+            asyncio.run(_async_sdk_count_tokens(gateway, model))
+            if asynchronous
+            else _sync_sdk_count_tokens(gateway, model)
+        )
+        assert input_tokens == _COUNT
+        requests: Final = _count_requests(wire)
+        assert len(requests) == 1
+        request_body: Final = _json_body(requests[0])
+        assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]}
+        assert request_body["tools"] == [
+            {
+                "functionDeclarations": [
+                    {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}}
+                ]
+            }
+        ]
+
+
+def test_google_count_tokens_forwards_system_and_tools_to_vertex(gateway: Gateway) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        body: Final[Mapping[str, JsonValue]] = {
+            "contents": [
+                {"role": "user", "parts": [{"text": _USER}]},
+                {"role": "model", "parts": [{"functionCall": {"name": "noop", "args": {"query": "value"}}}]},
+                {
+                    "role": "user",
+                    "parts": [{"functionResponse": {"name": "noop", "response": {"result": "ok"}}}],
+                },
+            ],
+            "systemInstruction": {"parts": [{"text": "Follow this instruction."}]},
+            "tools": [
+                {
+                    "functionDeclarations": [
+                        {
+                            "name": "noop",
+                            "description": "Return no output",
+                            "parameters": {"type": "OBJECT", "properties": {}},
+                        }
+                    ]
+                }
+            ],
+        }
+        response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT, response.text
+        requests: Final = _count_requests(wire)
+        assert len(requests) == 1
+        request_body: Final = _json_body(requests[0])
+        assert request_body["contents"] == body["contents"]
+        assert request_body["systemInstruction"] == body["systemInstruction"]
+        assert request_body["tools"] == [
+            {
+                "functionDeclarations": [
+                    {
+                        "name": "noop",
+                        "description": "Return no output",
+                        "parameters": {"type": "object"},
+                    }
+                ]
+            }
+        ]
+
+
+def test_responses_input_tokens_maps_flat_function_tools_to_vertex(gateway: Gateway) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        body: Final[Mapping[str, JsonValue]] = {
+            "model": model,
+            "instructions": "Follow the system instruction.",
+            "input": [
+                {
+                    "type": "function_call",
+                    "call_id": "call_lookup",
+                    "name": "lookup",
+                    "arguments": '{"query":"value"}',
+                },
+                {
+                    "type": "function_call_output",
+                    "call_id": "call_lookup",
+                    "output": "result",
+                },
+            ],
+            "tools": [
+                {
+                    "type": "function",
+                    "name": "lookup",
+                    "description": "Look up a value",
+                    "parameters": {"type": "object", "properties": {}},
+                }
+            ],
+        }
+        response: Final = gateway.request("POST", "/v1/responses/input_tokens", body)
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT, response.text
+        requests: Final = _count_requests(wire)
+        assert len(requests) == 1
+        request_body: Final = _json_body(requests[0])
+        assert request_body["tools"] == [
+            {
+                "functionDeclarations": [
+                    {
+                        "name": "lookup",
+                        "description": "Look up a value",
+                        "parameters": {"type": "object"},
+                    }
+                ]
+            }
+        ]
+
+
+def test_utils_token_counter_maps_openai_function_tools_to_vertex(gateway: Gateway) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        tool_schema: Final[JsonValue] = {"type": "object", "properties": {}}
+        body: Final[Mapping[str, JsonValue]] = {
+            "model": model,
+            "messages": _MESSAGES,
+            "system": "Follow the system instruction.",
+            "tools": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "noop",
+                        "description": "Return no output",
+                        "parameters": tool_schema,
+                    },
+                }
+            ],
+        }
+        response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"})
+        assert response.status_code == 200, response.text
+        assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT, response.text
+        requests: Final = _count_requests(wire)
+        assert len(requests) == 1
+        request_body: Final = _json_body(requests[0])
+        assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]}
+        assert request_body["tools"] == [
+            {
+                "functionDeclarations": [
+                    {
+                        "name": "noop",
+                        "description": "Return no output",
+                        "parameters": {"type": "object"},
+                    }
+                ]
+            }
+        ]
+
+
+def test_count_tokens_surfaces_peer_errors_and_rejects_non_text_system(gateway: Gateway) -> None:
+    peer: Final = _GeminiPeer()
+    peer.set_status(400)
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None))
+        assert response.status_code == 400, response.text
+        assert len(_count_requests(wire)) == 1
+
+    invalid_peer: Final = _GeminiPeer()
+    with wire_server(invalid_peer) as invalid_wire, gateway.scenario() as scenario:
+        invalid_model: Final = _deployment(gateway, scenario, invalid_wire.url)
+        invalid_response: Final = gateway.request(
+            "POST",
+            "/v1/messages/count_tokens",
+            _messages(
+                invalid_model,
+                [{"type": "image", "source": {"type": "base64", "data": "AAAA"}}],
+                None,
+            ),
+        )
+        assert invalid_response.status_code == 400, invalid_response.text
+        assert _count_requests(invalid_wire) == ()
+
+
+@pytest.mark.parametrize(
+    "tools",
+    (
+        [{"name": "noop", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}],
+        [{"name": "noop", "input_schema": "not-a-schema"}],
+    ),
+    ids=("malformed-property", "string-schema"),
+)
+def test_count_tokens_rejects_unconvertible_tools_without_request(gateway: Gateway, tools: JsonValue) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, tools))
+        assert response.status_code == 400, response.text
+        assert "tools" in response.text
+        assert _count_requests(wire) == ()
+
+
+def _count_wave(clients: Sequence[httpx.Client], body: Mapping[str, JsonValue]) -> tuple[tuple[int, JsonValue], ...]:
+    def count(client: httpx.Client) -> tuple[int, JsonValue]:
+        response: Final = client.post("/v1/messages/count_tokens", json=dict(body))
+        payload: Final = _JSON_OBJECT.validate_json(response.content)
+        return response.status_code, payload.get("input_tokens")
+
+    with ThreadPoolExecutor(max_workers=len(clients)) as executor:
+        return tuple(executor.map(count, clients))
+
+
+def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gateway) -> None:
+    peer: Final = _GeminiPeer()
+    with wire_server(peer) as wire, gateway.scenario() as scenario, ExitStack() as stack:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        clients: Final = tuple(
+            stack.enter_context(
+                httpx.Client(
+                    base_url=str(gateway.client.base_url),
+                    headers={"Authorization": f"Bearer {gateway.key}"},
+                    timeout=30,
+                    trust_env=False,
+                )
+            )
+            for _ in range(3)
+        )
+        body: Final = _messages(model, None, None)
+        assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT))
+        peer.set_status(503)
+        assert _count_wave(clients, body) == ((503, None), (503, None), (503, None))
+        peer.set_status(200)
+        assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT))
+        requests: Final = _count_requests(wire)
+        assert len(requests) == 9
+        assert (
+            tuple(_json_body(request) for request in requests)
+            == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9
+        )
+
+
+@pytest.mark.parametrize(
+    ("conversation", "call_endpoint"),
+    (
+        ({"messages": _MESSAGES}, True),
+        ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]}, True),
+        ({"messages": _MESSAGES}, False),
+    ),
+    ids=("provider-count-messages", "provider-count-contents", "local-count-messages"),
+)
+@pytest.mark.parametrize("drop_params", (False, True), ids=("no-drop", "drop"))
+def test_utils_token_counter_counts_the_conversation_regardless_of_prompt_and_drop_params(
+    gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool, drop_params: bool
+) -> None:
+    peer: Final = _GeminiPeer()
+    params: Final = {"call_endpoint": "true" if call_endpoint else "false"}
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params)
+        mixed: Final = gateway.request(
+            "POST", "/utils/token_counter", {"model": model, "prompt": "hi", **conversation}, params=params
+        )
+        mixed_requests: Final = _count_requests(wire)
+        alone: Final = gateway.request("POST", "/utils/token_counter", {"model": model, **conversation}, params=params)
+        alone_requests: Final = _count_requests(wire)
+    assert mixed.status_code == 200, mixed.text
+    assert alone.status_code == 200, alone.text
+    mixed_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(mixed.content)["total_tokens"])
+    alone_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(alone.content)["total_tokens"])
+    assert mixed_tokens == alone_tokens > 0
+    assert len(mixed_requests) == len(alone_requests) == int(call_endpoint)
+    assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [
+        _counted_contents(_json_body(request)) for request in alone_requests
+    ]

tests/integration/providers/test_vertex_partner_count_tokens_wire.py

@@ -0,0 +1,858 @@
+import json
+import re
+import signal
+import socket
+import threading
+import uuid
+from collections.abc import Callable, Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
+from dataclasses import dataclass
+from pathlib import Path
+from queue import SimpleQueue
+from types import MappingProxyType
+from typing import Final
+
+import httpx
+import psutil
+import pytest
+import yaml
+from integration._support.client import Gateway, Scenario, eventually
+from integration._support.process import 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
+from tokenizers import Regex, Tokenizer, models, pre_tokenizers
+
+_BACKEND: Final = "claude-sonnet-4-6"
+_PROJECT: Final = "scripted-project"
+_LOCATION: Final = "us-east5"
+_MODELS_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models"
+_COUNT_TARGET: Final = f"{_MODELS_PATH}/count-tokens:rawPredict"
+_MESSAGE_TARGET: Final = f"{_MODELS_PATH}/{_BACKEND}:rawPredict"
+_STREAM_TARGET: Final = f"{_MODELS_PATH}/{_BACKEND}:streamRawPredict?alt=sse"
+_PEER_COUNT: Final = 4242
+_REJECTION: Final = "scripted partner rejection"
+_REJECT_TEXT: Final = "The peer must reject this message"
+_REPLY_TEXT: Final = "scripted reply"
+_OWNED_MODEL: Final = "partner-claude"
+_OWNED_UNREACHABLE_MODEL: Final = "partner-claude-unreachable"
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_GLM_BACKEND: Final = "zai-org/glm-5.2-maas"
+_MISTRAL_BACKEND: Final = "mistral-medium-3"
+_LLAMA_BACKEND: Final = "meta/llama-4-maverick-17b-128e-instruct-maas"
+_LLAMA_COUNT_TARGET: Final = (
+    f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/meta/models/count-tokens:rawPredict"
+)
+_COUNTED_TEXT: Final = "Count this message"
+_TOKENIZER_REPO: Final = "scripted/char-tokenizer"
+_TOKENIZER_TARGET: Final = f"/{_TOKENIZER_REPO}/resolve/main/tokenizer.json"
+_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'})
+_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
+
+_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _COUNTED_TEXT}]
+_SYSTEM: Final = "You are a terse assistant that answers in one sentence"
+_SYSTEM_BLOCKS: Final[list[JsonValue]] = [
+    {"type": "text", "text": "You are a terse assistant"},
+    {"type": "text", "text": "Answer in one sentence"},
+]
+_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = {
+    "type": "object",
+    "properties": {"city": {"type": "string", "description": "City to look up"}},
+    "required": ["city"],
+}
+_TOOLS: Final[list[JsonValue]] = [
+    {"name": "get_weather", "description": "Look up the current weather for a city", "input_schema": _WEATHER_SCHEMA}
+]
+_OPENAI_TOOLS: Final[list[JsonValue]] = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Look up the current weather for a city",
+            "parameters": _WEATHER_SCHEMA,
+        },
+    }
+]
+_RESPONSES_TOOLS: Final[list[JsonValue]] = [
+    {
+        "type": "function",
+        "name": "get_weather",
+        "description": "Look up the current weather for a city",
+        "parameters": _WEATHER_SCHEMA,
+    }
+]
+_PEER_BARE: Final[dict[str, JsonValue]] = {"model": _BACKEND, "messages": _MESSAGES}
+_PEER_FULL: Final[dict[str, JsonValue]] = {**_PEER_BARE, "system": _SYSTEM, "tools": _TOOLS}
+_GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "Count this"}]}]}
+_GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}]
+
+_REPLY: Final[dict[str, JsonValue]] = {
+    "id": "msg_scripted",
+    "type": "message",
+    "role": "assistant",
+    "model": _BACKEND,
+    "content": [{"type": "text", "text": _REPLY_TEXT}],
+    "stop_reason": "end_turn",
+    "stop_sequence": None,
+    "usage": {"input_tokens": 5, "output_tokens": 3},
+}
+_EVENTS: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = (
+    (
+        "message_start",
+        {
+            "type": "message_start",
+            "message": {**_REPLY, "content": [], "stop_reason": None, "usage": {"input_tokens": 5, "output_tokens": 1}},
+        },
+    ),
+    ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
+    (
+        "content_block_delta",
+        {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}},
+    ),
+    ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+    (
+        "message_delta",
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": 3},
+        },
+    ),
+    ("message_stop", {"type": "message_stop"}),
+)
+_SSE: Final = tuple(f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in _EVENTS)
+
+
+def _counted(_request: Request) -> Reply:
+    return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode())
+
+
+def _rejected(status: int, error_type: str = "invalid_request_error") -> Reply:
+    return Reply(
+        status=status,
+        body=json.dumps({"type": "error", "error": {"type": error_type, "message": _REJECTION}}).encode(),
+    )
+
+
+def _rejecting(status: int, error_type: str = "invalid_request_error") -> Callable[[Request], Reply]:
+    def count(_request: Request) -> Reply:
+        return _rejected(status, error_type)
+
+    return count
+
+
+def _anthropic_message(message: JsonValue) -> bool:
+    return isinstance(message, dict) and message.get("role") in ("user", "assistant")
+
+
+def _anthropic_tool(tool: JsonValue) -> bool:
+    return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict)
+
+
+def _strict(request: Request) -> Reply:
+    body: Final = _JSON_OBJECT.validate_json(request.body)
+    messages: Final = body.get("messages")
+    tools: Final = body.get("tools", [])
+    accepted: Final = (
+        isinstance(messages, list)
+        and all(map(_anthropic_message, messages))
+        and isinstance(body.get("system", ""), (str, list))
+        and isinstance(tools, list)
+        and all(map(_anthropic_tool, tools))
+    )
+    return _counted(request) if accepted else _rejected(400)
+
+
+def _rejecting_marked_messages(request: Request) -> Reply:
+    return _rejected(400) if _REJECT_TEXT in request.body.decode() else _counted(request)
+
+
+def _peer(count: Callable[[Request], Reply] = _counted) -> Callable[[Request], Reply]:
+    def respond(request: Request) -> Reply:
+        if request.target.endswith("/count-tokens:rawPredict"):
+            return count(request)
+        if request.target.endswith(":streamRawPredict?alt=sse"):
+            return Reply(content_type="text/event-stream", chunks=_SSE)
+        if request.target.endswith(f"/{_BACKEND}:rawPredict"):
+            return Reply(body=json.dumps(_REPLY).encode())
+        return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode())
+
+    return respond
+
+
+def _count_requests(requests: Sequence[Request]) -> tuple[Request, ...]:
+    return tuple(request for request in requests if "count-tokens" in request.target)
+
+
+def _count_bodies(requests: Sequence[Request], target: str = _COUNT_TARGET) -> tuple[dict[str, JsonValue], ...]:
+    counts: Final = _count_requests(requests)
+    for request in counts:
+        assert (request.method, request.target) == ("POST", target), request.target
+        assert request.headers["authorization"] == "Bearer scripted-token", request.headers
+    return tuple(_JSON_OBJECT.validate_json(request.body) for request in counts)
+
+
+def _counted_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
+    return _count_bodies(wire.drain())
+
+
+def _bare(model: str) -> dict[str, JsonValue]:
+    return {"model": model, "messages": _MESSAGES}
+
+
+def _full(model: str) -> dict[str, JsonValue]:
+    return {**_bare(model), "system": _SYSTEM, "tools": _TOOLS}
+
+
+def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **overrides: JsonValue) -> str:
+    return scenario.model(
+        **{
+            "model": f"vertex_ai/{_BACKEND}",
+            "api_base": api_base,
+            "api_key": None,
+            "vertex_project": _PROJECT,
+            "vertex_location": _LOCATION,
+            "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
+            **overrides,
+        }
+    )
+
+
+def _count(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response:
+    return gateway.request("POST", "/v1/messages/count_tokens", body)
+
+
+def _payload(response: httpx.Response) -> dict[str, JsonValue]:
+    assert response.status_code == 200, response.text
+    return _JSON_OBJECT.validate_json(response.content)
+
+
+def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int:
+    response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"})
+    payload: Final = _payload(response)
+    total: Final = payload["total_tokens"]
+    assert payload["tokenizer_type"] != "vertex_ai_partner_models", response.text
+    assert isinstance(total, int) and 0 < total != _PEER_COUNT, response.text
+    return total
+
+
+@dataclass(frozen=True, slots=True)
+class _CountRoute:
+    path: Callable[[str], str]
+    body: Callable[[str], dict[str, JsonValue]]
+    rejection: Callable[[str], dict[str, JsonValue]]
+    count_field: str
+    params: Mapping[str, str] = MappingProxyType({})
+
+    def send(self, gateway: Gateway, model: str) -> httpx.Response:
+        return gateway.request("POST", self.path(model), self.body(model), params=self.params)
+
+
+_MESSAGES_ROUTE: Final = _CountRoute(
+    path=lambda _model: "/v1/messages/count_tokens",
+    body=_bare,
+    rejection=lambda message: {"type": "error", "error": {"type": "invalid_request_error", "message": message}},
+    count_field="input_tokens",
+)
+_RESPONSES_ROUTE: Final = _CountRoute(
+    path=lambda _model: "/v1/responses/input_tokens",
+    body=lambda model: {"model": model, "input": _COUNTED_TEXT},
+    rejection=lambda message: {
+        "error": {"message": message, "type": "invalid_request_error", "param": None, "code": None}
+    },
+    count_field="input_tokens",
+)
+_GOOGLE_ROUTE: Final = _CountRoute(
+    path=lambda model: f"/v1beta/models/{model}:countTokens",
+    body=lambda _model: {"contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]},
+    rejection=lambda message: {"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}},
+    count_field="totalTokens",
+)
+_UTILS_ROUTE: Final = _CountRoute(
+    path=lambda _model: "/utils/token_counter",
+    body=_bare,
+    rejection=lambda message: {
+        "error": {"message": message, "type": "token_counting_error", "param": "model", "code": "400"}
+    },
+    count_field="total_tokens",
+    params=MappingProxyType({"call_endpoint": "true"}),
+)
+_COUNT_ROUTES: Final = (_MESSAGES_ROUTE, _RESPONSES_ROUTE, _GOOGLE_ROUTE, _UTILS_ROUTE)
+
+
+def _error_message(payload: Mapping[str, JsonValue]) -> str:
+    error: Final = payload.get("error")
+    assert isinstance(error, dict), payload
+    message: Final = error.get("message")
+    assert isinstance(message, str), payload
+    return message
+
+
+def _route_count(gateway: Gateway, route: _CountRoute, model: str) -> tuple[int, JsonValue]:
+    response: Final = route.send(gateway, model)
+    return response.status_code, _JSON_OBJECT.validate_json(response.content).get(route.count_field)
+
+
+def _char_tokenizer() -> bytes:
+    tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]"))
+    tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated")
+    return tokenizer.to_str().encode()
+
+
+def _tokenizer_hub(tokenizer: bytes) -> Callable[[Request], Reply]:
+    def respond(request: Request) -> Reply:
+        if request.target == _TOKENIZER_TARGET:
+            return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS)
+        return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode())
+
+    return respond
+
+
+def _closed_port_url() -> str:
+    with socket.socket() as reserve:
+        reserve.bind(("127.0.0.1", 0))
+        return f"http://127.0.0.1:{reserve.getsockname()[1]}"
+
+
+def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]:
+    return tuple(
+        stack.enter_context(httpx.Client(base_url=base_url, timeout=30, trust_env=False)) for _ in range(count)
+    )
+
+
+def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]:
+    response: Final = client.post(
+        "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"}
+    )
+    return response.status_code, _JSON_OBJECT.validate_json(response.content).get("input_tokens")
+
+
+def _generated_then_counted(
+    client: httpx.Client, key: str, model: str, body: Mapping[str, JsonValue]
+) -> tuple[int, int, JsonValue]:
+    generated: Final = client.post(
+        "/v1/messages",
+        json={
+            "model": model,
+            "max_tokens": 16,
+            "messages": [{"role": "user", "content": f"Generate before counting {uuid.uuid4().hex}"}],
+        },
+        headers={"Authorization": f"Bearer {key}"},
+    )
+    return generated.status_code, *_counted_on(client, key, body)
+
+
+def _local_port(client: httpx.Client) -> int:
+    with client.stream("GET", "/health/liveliness") as response:
+        port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1])
+        response.read()
+    assert response.status_code == 200, response.text
+    return port
+
+
+def _counted_or_dropped(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue] | None:
+    try:
+        return _counted_on(client, key, body)
+    except httpx.TransportError:
+        return None
+
+
+def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]:
+    return frozenset(
+        connection.raddr.port
+        for connection in psutil.Process(pid).net_connections(kind="tcp")
+        if connection.raddr and connection.laddr.port == proxy_port
+    )
+
+
+def _owned_config(
+    path: Path, gateway: Gateway, api_bases: Mapping[str, str], settings: Mapping[str, JsonValue]
+) -> Path:
+    config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    path.write_text(
+        yaml.safe_dump(
+            {
+                **config,
+                "model_list": [
+                    {
+                        "model_name": name,
+                        "litellm_params": {
+                            "model": f"vertex_ai/{_BACKEND}",
+                            "api_base": api_base,
+                            "vertex_project": _PROJECT,
+                            "vertex_location": _LOCATION,
+                            "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
+                        },
+                    }
+                    for name, api_base in api_bases.items()
+                ],
+                "litellm_settings": {**config["litellm_settings"], **settings},
+            }
+        )
+    )
+    return path
+
+
+@pytest.mark.parametrize("system", [_SYSTEM, _SYSTEM_BLOCKS], ids=["string", "blocks"])
+def test_messages_count_tokens_forwards_system_and_tools(gateway: Gateway, system: JsonValue) -> None:
+    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, "litellm_count_source": "provider"}, response.text
+        assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system, "tools": _TOOLS},)
+
+
+def test_messages_count_tokens_without_system_or_tools_sends_bare_body(gateway: Gateway) -> None:
+    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, "litellm_count_source": "provider"}, response.text
+        assert _counted_bodies(wire) == (_PEER_BARE,)
+
+
+def test_utils_token_counter_call_endpoint_forwards_system_and_tools(gateway: Gateway) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request(
+            "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"}
+        )
+        payload: Final = _payload(response)
+        assert (payload["total_tokens"], payload["tokenizer_type"]) == (_PEER_COUNT, "vertex_ai_partner_models")
+        assert (payload["request_model"], payload["model_used"]) == (model, _BACKEND), response.text
+        assert _counted_bodies(wire) == (_PEER_FULL,)
+
+
+def test_utils_token_counter_local_mode_never_calls_the_peer(gateway: Gateway) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        assert _local_count(gateway, _full(model)) > 0
+        assert wire.drain() == ()
+
+
+def test_utils_token_counter_surfaces_the_peer_error_for_openai_tools(gateway: Gateway) -> None:
+    with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        body: Final = {**_bare(model), "tools": _OPENAI_TOOLS}
+        response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"})
+        assert response.status_code == 400, response.text
+        assert "total_tokens" not in response.text, response.text
+        assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": _OPENAI_TOOLS},)
+
+
+def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        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,
+            "litellm_count_source": "provider",
+        }, response.text
+        assert _counted_bodies(wire) == (_PEER_BARE,)
+
+
+def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> None:
+    with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/responses/input_tokens",
+            {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS},
+        )
+        assert response.status_code == 400, response.text
+        assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text
+        (sent,) = _counted_bodies(wire)
+        assert sent["tools"] == _RESPONSES_TOOLS, sent
+
+
+
+
+
+
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+def test_custom_tokenizer_partner_counts_locally_on_every_count_route(gateway: Gateway, tmp_path: Path) -> None:
+    environment: Final = {"HF_HOME": str(tmp_path / "hf-home")}
+    with (
+        wire_server(_tokenizer_hub(_char_tokenizer())) as hub,
+        wire_server(_peer()) as vertex,
+        owned_proxy_process(gateway, tmp_path, {**environment, "HF_ENDPOINT": hub.url}) as owned,
+        owned.gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(
+            owned.gateway,
+            scenario,
+            vertex.url,
+            model=f"vertex_ai/{_GLM_BACKEND}",
+            model_info={"custom_tokenizer": {"identifier": _TOKENIZER_REPO, "revision": "main", "auth_token": None}},
+        )
+        local: Final = _payload(
+            owned.gateway.request("POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "false"})
+        )
+        assert local["tokenizer_type"] == "huggingface_tokenizer", local
+        total: Final = local["total_tokens"]
+        assert isinstance(total, int) and total >= len(_COUNTED_TEXT), local
+        counts: Final = tuple(_route_count(owned.gateway, route, model) for route in _COUNT_ROUTES)
+        assert counts == ((200, total),) * len(_COUNT_ROUTES), counts
+        contents_only: Final = owned.gateway.request(
+            "POST",
+            "/utils/token_counter",
+            {"model": model, "contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]},
+            params={"call_endpoint": "true"},
+        )
+        assert contents_only.status_code == 400, contents_only.text
+        assert _GLM_BACKEND in _error_message(_JSON_OBJECT.validate_json(contents_only.content)), contents_only.text
+        assert vertex.drain() == ()
+        assert {request.target for request in hub.drain()} >= {_TOKENIZER_TARGET}
+
+
+def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
+        assert "totalTokens" in _payload(response), response.text
+        assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},)
+
+
+def test_gemini_count_tokens_route_surfaces_the_peer_error(gateway: Gateway) -> None:
+    with wire_server(_peer(_rejecting(400))) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
+        assert response.status_code == 400, response.text
+        assert "totalTokens" not in response.text, response.text
+        assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},)
+
+
+@pytest.mark.parametrize("status", [400, 500, 503])
+def test_messages_count_tokens_surfaces_the_peer_error_when_peer_errors(gateway: Gateway, status: int) -> None:
+    with wire_server(_peer(_rejecting(status))) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = _count(gateway, _full(model))
+        assert response.status_code == status, response.text
+        assert "input_tokens" not in response.text, response.text
+        assert _counted_bodies(wire) == (_PEER_FULL,)
+
+
+def test_messages_count_tokens_surfaces_the_error_when_token_endpoint_rejects(gateway: Gateway) -> None:
+    def respond(request: Request) -> Reply:
+        if request.target == "/_oauth/token":
+            return Reply(
+                status=400,
+                body=json.dumps({"error": "invalid_grant", "error_description": "scripted refusal"}).encode(),
+            )
+        return _peer()(request)
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(
+            gateway, scenario, wire.url, vertex_credentials=service_account_json(_PROJECT, wire.url)
+        )
+        response: Final = _count(gateway, _full(model))
+        assert 500 <= response.status_code < 600, response.text
+        assert "input_tokens" not in response.text, response.text
+        targets: Final = frozenset(request.target for request in wire.drain())
+        assert targets == {"/_oauth/token"}, targets
+
+
+def test_messages_count_tokens_errors_when_peer_is_unreachable(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, _closed_port_url())
+        response: Final = _count(gateway, _full(model))
+        assert 500 <= response.status_code < 600, response.text
+        assert "input_tokens" not in response.text, response.text
+
+
+@pytest.mark.parametrize(
+    "tools",
+    [5, "", "x" * 5120, ["get_weather"]],
+    ids=["int", "empty_string", "5kb_string", "list_of_strings"],
+)
+def test_messages_count_tokens_rejects_malformed_tools_without_calling_the_peer(
+    gateway: Gateway, tools: JsonValue
+) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        refused: Final = _count(gateway, {**_bare(model), "tools": tools})
+        assert 400 <= refused.status_code < 600, refused.text
+        assert "input_tokens" not in refused.text, refused.text
+        assert wire.drain() == ()
+        assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _PEER_COUNT)
+        assert _counted_bodies(wire) == (_PEER_BARE,)
+
+
+def test_messages_count_tokens_forwards_an_empty_tools_list(gateway: Gateway) -> None:
+    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, "litellm_count_source": "provider"}, response.text
+        assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": []},)
+
+
+def test_messages_count_tokens_forwards_an_empty_system_string(gateway: Gateway) -> None:
+    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, "litellm_count_source": "provider"}, response.text
+        assert _counted_bodies(wire) == ({**_PEER_BARE, "system": ""},)
+
+
+def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway) -> None:
+    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, "litellm_count_source": "provider"}, response.text
+        assert _counted_bodies(wire) == (_PEER_BARE,)
+
+
+def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system(gateway: Gateway) -> None:
+    with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = _count(gateway, {**_bare(model), "system": 5})
+        assert response.status_code == 400, response.text
+        assert "input_tokens" not in response.text, response.text
+        assert _counted_bodies(wire) == ({**_PEER_BARE, "system": 5},)
+
+
+def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway) -> None:
+    system: Final = "Answer in one sentence. " * 214
+    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, "litellm_count_source": "provider"}, response.text
+        assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system},)
+
+
+def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value_each(gateway: Gateway) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}'
+        response: Final = gateway.client.post(
+            "/v1/messages/count_tokens",
+            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, "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
+
+
+def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer(gateway: Gateway) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request(
+            "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued"
+        )
+        assert response.status_code == 401, response.text
+        assert "input_tokens" not in response.text, response.text
+        assert wire.drain() == ()
+
+
+@pytest.mark.parametrize("fields", [{}, {"messages": []}], ids=["missing", "empty"])
+def test_messages_count_tokens_without_messages_is_rejected(gateway: Gateway, fields: dict[str, JsonValue]) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields})
+        assert response.status_code == 400, response.text
+        assert "messages parameter is required" in response.text, response.text
+        assert wire.drain() == ()
+
+
+def test_count_tokens_location_override_targets_the_count_region(gateway: Gateway) -> None:
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(
+            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, "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,)
+
+
+def test_messages_count_tokens_repeated_request_reaches_the_peer_each_time(gateway: Gateway) -> None:
+    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, "litellm_count_source": "provider"},) * 2
+        assert _counted_bodies(wire) == (_PEER_BARE,) * 2
+
+
+@pytest.mark.timeout(240)  # boots an owned two-worker proxy with litellm_settings.disable_token_counter
+def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_locally(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with wire_server(_peer(_rejecting_marked_messages)) as wire:
+        config: Final = _owned_config(
+            tmp_path / "disabled-token-counter.yaml",
+            gateway,
+            {_OWNED_MODEL: wire.url, _OWNED_UNREACHABLE_MODEL: _closed_port_url()},
+            {"disable_token_counter": True},
+        )
+        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, "litellm_count_source": "provider"}, counted.text
+            rejected: Final = _count(
+                owned.gateway, {**_full(_OWNED_MODEL), "messages": [{"role": "user", "content": _REJECT_TEXT}]}
+            )
+            assert rejected.status_code == 400, rejected.text
+            assert _REJECTION in rejected.text and "input_tokens" not in rejected.text, rejected.text
+            unreachable: Final = _count(owned.gateway, _full(_OWNED_UNREACHABLE_MODEL))
+            assert 500 <= unreachable.status_code < 600, unreachable.text
+            assert "input_tokens" not in unreachable.text, unreachable.text
+            local: Final = owned.gateway.request(
+                "POST", "/utils/token_counter", _full(_OWNED_MODEL), params={"call_endpoint": "false"}
+            )
+            assert local.status_code == 503, local.text
+            assert len(_counted_bodies(wire)) == 2
+
+
+def test_peer_outage_between_concurrent_waves_errors_then_recovers(gateway: Gateway) -> None:
+    with ExitStack() as stack:
+        clients: Final = _clients(stack, str(gateway.client.base_url), 8)
+        pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
+        scenario: Final = stack.enter_context(gateway.scenario())
+        with wire_server(_peer()) as wire:
+            port: Final = int(wire.url.rsplit(":", 1)[1])
+            model: Final = _deployment(gateway, scenario, wire.url)
+            body: Final = _full(model)
+
+            def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]:
+                return _generated_then_counted(client, gateway.key, model, body)
+
+            assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients)
+            assert _counted_bodies(wire) == (_PEER_FULL,) * len(clients)
+        outage: Final = tuple(pool.map(lambda client: _counted_on(client, gateway.key, body), clients))
+        for status, tokens in outage:
+            assert 500 <= status < 600 and tokens is None, outage
+        with wire_server(_peer(), port=port) as revived:
+            assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients)
+            assert _counted_bodies(revived) == (_PEER_FULL,) * len(clients)
+
+
+def test_slow_peer_holds_concurrent_counts_without_stalling_the_proxy(gateway: Gateway) -> None:
+    held: Final[SimpleQueue[str]] = SimpleQueue()
+    release: Final = threading.Event()
+
+    def hold(request: Request) -> Reply:
+        held.put(request.target)
+        assert release.wait(timeout=20), "Held count was never released"
+        return _counted(request)
+
+    with ExitStack() as stack:
+        clients: Final = _clients(stack, str(gateway.client.base_url), 6)
+        wire: Final = stack.enter_context(wire_server(_peer(hold)))
+        scenario: Final = stack.enter_context(gateway.scenario())
+        pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
+        model: Final = _deployment(gateway, scenario, wire.url)
+        try:
+            futures: Final = tuple(
+                pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients
+            )
+            eventually(held.qsize, lambda size: size == len(clients), seconds=30)
+            assert gateway.request("GET", "/health/liveliness").status_code == 200
+            assert _local_count(gateway, _full(model)) > 0
+            assert not any(future.done() for future in futures)
+        finally:
+            release.set()
+        assert tuple(future.result(timeout=30) for future in futures) == ((200, 200, _PEER_COUNT),) * len(clients)
+        assert len(_counted_bodies(wire)) == len(clients)
+
+
+@pytest.mark.timeout(300)  # boots an owned two-worker proxy, kills one worker, and waits for its replacement
+def test_worker_sigkill_mid_burst_leaves_the_sibling_counting(gateway: Gateway, tmp_path: Path) -> None:
+    held: Final[SimpleQueue[str]] = SimpleQueue()
+    release: Final = threading.Event()
+
+    def hold(request: Request) -> Reply:
+        held.put(request.target)
+        assert release.wait(timeout=60), "Held count was never released"
+        return _counted(request)
+
+    with ExitStack() as stack:
+        wire: Final = stack.enter_context(wire_server(_peer(hold)))
+        stack.callback(release.set)
+        config: Final = _owned_config(tmp_path / "worker-kill.yaml", gateway, {_OWNED_MODEL: wire.url}, {})
+        owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2))
+        proxy_url: Final = owned.gateway.client.base_url
+        workers: Final = eventually(
+            lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
+            lambda pids: len(pids) == 2,
+            seconds=30,
+        )
+        clients: Final = _clients(stack, str(proxy_url), 12)
+        pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
+        stack.callback(release.set)
+        ports: Final = tuple(_local_port(client) for client in clients)
+        futures: Final = tuple(
+            pool.submit(_counted_or_dropped, client, gateway.key, _full(_OWNED_MODEL)) for client in clients
+        )
+        eventually(held.qsize, lambda size: size == len(clients), seconds=30)
+        shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers}
+        assert sum(map(len, shares.values())) == len(clients), shares
+        victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid]))
+        psutil.Process(victim).send_signal(signal.SIGKILL)
+        release.set()
+        results: Final = tuple(future.result(timeout=60) for future in futures)
+        for port, result in zip(ports, results, strict=True):
+            assert result == (None if port in shares[victim] else (200, _PEER_COUNT)), (port, result, shares)
+        second_wave: Final = _clients(stack, str(proxy_url), 6)
+        assert tuple(_counted_on(client, gateway.key, _full(_OWNED_MODEL)) for client in second_wave) == (
+            (200, _PEER_COUNT),
+        ) * len(second_wave)
+        assert len(_counted_bodies(wire)) == len(clients) + len(second_wave)
+        eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120)
+        assert owned.process.poll() is None
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
+def test_chat_completions_on_the_same_deployment_still_generate(gateway: Gateway, stream: bool) -> None:
+    marker: Final = f"chat control {uuid.uuid4().hex}"
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": marker}],
+                "stream": stream,
+                "cache": {"no-cache": True},
+            },
+        )
+        assert response.status_code == 200, response.text
+        assert _REPLY_TEXT in response.text, response.text
+        assert not stream or response.text.rstrip().endswith("data: [DONE]"), response.text
+        (sent,) = wire.drain()
+        assert sent.target == (_STREAM_TARGET if stream else _MESSAGE_TARGET), sent.target
+        assert marker in sent.body.decode(), sent.body
+
+
+def test_messages_endpoint_on_the_same_deployment_still_generates(gateway: Gateway) -> None:
+    marker: Final = f"messages control {uuid.uuid4().hex}"
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/messages",
+            {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]},
+        )
+        assert response.status_code == 200, response.text
+        assert _REPLY_TEXT in response.text, response.text
+        (sent,) = wire.drain()
+        assert sent.target == _MESSAGE_TARGET, sent.target
+        assert marker in sent.body.decode(), sent.body
+
+
+def test_responses_endpoint_on_the_same_deployment_still_generates(gateway: Gateway) -> None:
+    marker: Final = f"responses control {uuid.uuid4().hex}"
+    with wire_server(_peer()) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(gateway, scenario, wire.url)
+        response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": marker})
+        assert response.status_code == 200, response.text
+        assert _REPLY_TEXT in response.text, response.text
+        (sent,) = wire.drain()
+        assert sent.target == _MESSAGE_TARGET, sent.target
+        assert marker in sent.body.decode(), sent.body

tests/integration/proxy_config.yaml

@@ -1,4 +1,11 @@
 model_list:
+  - model_name: anthropic/claude-opus-5-5
+    litellm_params:
+      model: anthropic/claude-opus-5-5
+      api_base: http://127.0.0.1:8191
+      api_key: synthetic-anthropic-key
+    model_info:
+      token_count_framing: 0
   - model_name: perplexity/pplx-decider-v1-27b
     litellm_params:
       model: perplexity/pplx-decider-v1-27b

tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py

@@ -0,0 +1,90 @@
+import json
+from collections.abc import Callable
+from typing import Final
+
+import litellm
+import pytest
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_BACKEND: Final = "claude-sonnet-4-6"
+_PROJECT: Final = "scripted-project"
+_LOCATION: Final = "us-east5"
+_COUNT_TARGET: Final = (
+    f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/count-tokens:rawPredict"
+)
+_PEER_COUNT: Final = 4242
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_MESSAGES: Final[list[dict[str, str]]] = [{"role": "user", "content": "Count this message"}]
+_SYSTEM: Final = "You are a terse assistant that answers in one sentence"
+_TOOLS: Final[list[dict[str, JsonValue]]] = [
+    {
+        "name": "get_weather",
+        "description": "Look up the current weather for a city",
+        "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+    }
+]
+
+
+def _peer(status: int) -> Callable[[Request], Reply]:
+    def respond(request: Request) -> Reply:
+        if request.target == "/_oauth/token":
+            token: Final = {"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}
+            return Reply(body=json.dumps(token).encode())
+        if status == 200:
+            return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode())
+        rejection: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted"}}
+        return Reply(status=status, body=json.dumps(rejection).encode())
+
+    return respond
+
+
+def _count_requests(requests: tuple[Request, ...]) -> tuple[Request, ...]:
+    return tuple(request for request in requests if "count-tokens" in request.target)
+
+
+@pytest.fixture
+def vertex_environment(monkeypatch: pytest.MonkeyPatch) -> Callable[[str], None]:
+    def configure(token_url: str) -> None:
+        monkeypatch.setenv("VERTEXAI_PROJECT", _PROJECT)
+        monkeypatch.setenv("VERTEXAI_LOCATION", _LOCATION)
+        monkeypatch.setenv("VERTEXAI_CREDENTIALS", service_account_json(_PROJECT, token_url))
+
+    return configure
+
+
+async def test_acount_tokens_forwards_system_and_tools_to_the_partner_peer(
+    vertex_environment: Callable[[str], None],
+) -> None:
+    with wire_server(_peer(200)) as wire:
+        vertex_environment(wire.url)
+        counted: Final = await litellm.acount_tokens(
+            model=f"vertex_ai/{_BACKEND}", messages=_MESSAGES, tools=_TOOLS, system=_SYSTEM, api_base=wire.url
+        )
+        (sent,) = _count_requests(wire.drain())
+        assert (sent.method, sent.target, sent.headers["authorization"]) == (
+            "POST",
+            _COUNT_TARGET,
+            "Bearer scripted-token",
+        )
+        assert _JSON_OBJECT.validate_json(sent.body) == {
+            "model": _BACKEND,
+            "messages": _MESSAGES,
+            "system": _SYSTEM,
+            "tools": _TOOLS,
+        }
+        assert (counted.total_tokens, counted.tokenizer_type) == (_PEER_COUNT, "vertex_ai_partner_models"), counted
+
+
+async def test_acount_tokens_falls_back_to_the_local_tokenizer_when_the_peer_rejects(
+    vertex_environment: Callable[[str], None],
+) -> None:
+    with wire_server(_peer(400)) as wire:
+        vertex_environment(wire.url)
+        counted: Final = await litellm.acount_tokens(
+            model=f"vertex_ai/{_BACKEND}", messages=_MESSAGES, tools=_TOOLS, system=_SYSTEM, api_base=wire.url
+        )
+        assert len(_count_requests(wire.drain())) == 1
+        assert counted.tokenizer_type == "local_tokenizer", counted
+        assert counted.total_tokens > 0 and counted.total_tokens != _PEER_COUNT, counted

tests/integration/spend/test_cache_and_quota.py

@@ -399,9 +399,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/test_litellm/proxy/anthropic_endpoints/test_endpoints.py

@@ -284,7 +284,9 @@ class TestHttpExceptionDictDetail:
                 "base_process_llm_request",
                 new=AsyncMock(side_effect=exc),
             ),
-            patch.object(proxy_server, "proxy_logging_obj") as mock_logging,  # test-quality-ok: module global imported at call time; no injection seam
+            patch.object(
+                proxy_server, "proxy_logging_obj"
+            ) as mock_logging,  # test-quality-ok: module global imported at call time; no injection seam
         ):
             mock_logging.post_call_failure_hook = AsyncMock()
             response = await ep.anthropic_response(
@@ -424,8 +426,6 @@ class TestErrorLogCarriesCallId:
 
     @pytest.mark.asyncio
     async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture):
-        from fastapi import HTTPException
-
         import litellm.proxy.anthropic_endpoints.endpoints as ep
         import litellm.proxy.proxy_server as proxy_server
         from litellm.proxy._types import UserAPIKeyAuth
@@ -440,13 +440,14 @@ class TestErrorLogCarriesCallId:
                 "_read_request_body",
                 new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}),
             ),
-            patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))),  # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
+            patch.object(
+                proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))
+            ),  # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
             caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
-            pytest.raises(HTTPException) as raised,
         ):
-            await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
+            response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
 
-        assert raised.value.status_code == 500
+        assert response.status_code == 500
         record = self._error_record(caplog)
         assert record.litellm_call_id == call_id
         assert call_id in record.getMessage()

tests/test_litellm/proxy/auth/test_auth_utils.py

@@ -30,6 +30,7 @@ from litellm.proxy.auth.auth_utils import (
     log_once_if_budget_reservation_disabled,
     warn_once_if_custom_auth_skips_common_checks,
 )
+from litellm.router import Router
 
 
 class TestCustomAuthCommonChecksWarning:
@@ -1167,6 +1168,124 @@ async def test_managed_batch_routes_pass_team_model_access_check(route, request_
         )
 
 
+_TEAM_DEPLOYMENT_ID: Final = "team-a-deployment-id"
+_TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_6d1f3c2e-0000-4000-8000-000000000000"
+_TEAM_PUBLIC_MODEL_NAME: Final = "team-batch-model"
+_TEAM_MANAGED_BATCH_ID: Final = _encode_managed_id(
+    f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};llm_batch_id:provider-batch-789"
+)
+
+
+def _team_scoped_router(public_model_name: str | None = _TEAM_PUBLIC_MODEL_NAME) -> Router:
+    return Router(
+        model_list=[
+            {
+                "model_name": _TEAM_INTERNAL_MODEL_NAME,
+                "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
+                "model_info": {
+                    "id": _TEAM_DEPLOYMENT_ID,
+                    "team_id": "team-a",
+                    **({"team_public_model_name": public_model_name} if public_model_name else {}),
+                },
+            },
+        ]
+    )
+
+
+def _team_scoped_character_id() -> str:
+    from litellm.types.videos.utils import encode_character_id_with_provider
+
+    return encode_character_id_with_provider(character_id="char_team", provider="openai", model_id=_TEAM_DEPLOYMENT_ID)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "route, request_data",
+    [
+        pytest.param("/v1/batches/{batch_id}", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-retrieve"),
+        pytest.param("/v1/batches/{batch_id}/cancel", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-cancel"),
+        pytest.param(
+            "/v1/fine_tuning/jobs/{fine_tuning_job_id}",
+            {
+                "fine_tuning_job_id": _encode_managed_id(
+                    f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};generic_response_id:ftjob-1"
+                )
+            },
+            id="fine-tuning-job",
+        ),
+        pytest.param(
+            "/v1/vector_stores/{vector_store_id}",
+            {
+                "vector_store_id": _encode_managed_id(
+                    "litellm_proxy:vector_store;unified_id,vs-unified;"
+                    f"target_model_names,{_TEAM_PUBLIC_MODEL_NAME};resource_id,vs_1;model_id,{_TEAM_DEPLOYMENT_ID}"
+                )
+            },
+            id="vector-store",
+        ),
+        pytest.param(
+            "/v1/videos/characters/{character_id}",
+            {"character_id": _team_scoped_character_id()},
+            id="character",
+        ),
+    ],
+)
+async def test_team_scoped_managed_ids_authorize_only_the_owning_team_by_public_name(
+    route: str, request_data: dict[str, str]
+) -> None:
+    from litellm.proxy._types import ProxyException
+    from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
+
+    llm_router: Final = _team_scoped_router()
+
+    async def check(team_id: str) -> None:
+        await can_key_call_resolved_model(
+            model=get_model_from_request(
+                request_data=request_data, route=route, llm_router=llm_router, team_id=team_id
+            ),
+            llm_model_list=None,
+            valid_token=UserAPIKeyAuth(api_key="sk-team", team_id=team_id, team_models=[_TEAM_PUBLIC_MODEL_NAME]),
+            llm_router=llm_router,
+        )
+
+    await check("team-a")
+    with pytest.raises(ProxyException) as exc_info:
+        await check("team-b")
+    assert exc_info.value.code == "403"
+
+
+@pytest.mark.asyncio
+async def test_legacy_team_alias_authorizes_a_team_batch_only_for_keys_granted_the_alias() -> None:
+    from litellm.proxy._types import ProxyException
+    from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
+
+    llm_router: Final = _team_scoped_router(public_model_name=None)
+
+    async def check(key_models: list[str]) -> None:
+        await can_key_call_resolved_model(
+            model=get_model_from_request(
+                request_data={"batch_id": _TEAM_MANAGED_BATCH_ID},
+                route="/v1/batches/{batch_id}",
+                llm_router=llm_router,
+                team_id="team-a",
+            ),
+            llm_model_list=None,
+            valid_token=UserAPIKeyAuth(
+                api_key="sk-team",
+                team_id="team-a",
+                models=key_models,
+                team_models=["legacy-team-model"],
+                team_model_aliases={"legacy-team-model": _TEAM_INTERNAL_MODEL_NAME},
+            ),
+            llm_router=llm_router,
+        )
+
+    await check(["legacy-team-model"])
+    with pytest.raises(ProxyException) as exc_info:
+        await check(["some-other-model"])
+    assert exc_info.value.code == "403"
+
+
 def test_get_model_from_request_resolves_character_id_model_with_router():
     from litellm.types.videos.utils import encode_character_id_with_provider
 

tests/test_litellm/proxy/batches_endpoints/test_endpoints.py

@@ -35,7 +35,7 @@ import logging
 from contextlib import ExitStack
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Final, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -3258,6 +3258,129 @@ async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_h
     cancel_harness.router_acancel.assert_not_called()
 
 
+TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_00000000-0000-0000-0000-000000000000"
+TEAM_PUBLIC_MODEL_NAME: Final = "team-gpt4o-mini"
+
+
+def _team_scoped_deployment() -> Deployment:
+    return Deployment(
+        model_name=TEAM_INTERNAL_MODEL_NAME,
+        litellm_params=LiteLLM_Params(model="gpt-4o-mini"),
+        model_info=ModelInfo(id="dep-team-a-1", team_id="team-a", team_public_model_name=TEAM_PUBLIC_MODEL_NAME),
+    )
+
+
+def _route_unified_batch_to_team_deployment(
+    harness: RetrieveHarness | CancelHarness, deployment: Deployment | None = None
+) -> None:
+    target: Final = deployment or _team_scoped_deployment()
+    harness.router.resolve_model_name_from_model_id = MagicMock(return_value=target.model_name)
+    harness.router.get_deployment_by_model_group_name = MagicMock(return_value=target)
+
+
+@pytest.mark.asyncio
+async def test_retrieve__team_scoped_deployment_authorizes_owning_team_by_public_name(
+    retrieve_harness: RetrieveHarness,
+) -> None:
+    _route_unified_batch_to_team_deployment(retrieve_harness)
+
+    resp: Final = await call_retrieve(
+        retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME)
+    )
+
+    assert retrieve_harness.router_aretrieve.call_count == 1
+    assert resp._hidden_params["model_id"] == "gpt-4o-mini"
+
+
+@pytest.mark.asyncio
+async def test_cancel__team_scoped_deployment_authorizes_owning_team_by_public_name(
+    cancel_harness: CancelHarness,
+) -> None:
+    _route_unified_batch_to_team_deployment(cancel_harness)
+
+    await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME))
+
+    assert cancel_harness.router_acancel.call_count == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "deployment, user",
+    [
+        pytest.param(
+            _team_scoped_deployment(), _key_restricted_to("some-other-model"), id="owning-team-key-without-the-model"
+        ),
+        pytest.param(
+            _team_scoped_deployment(),
+            UserAPIKeyAuth(
+                api_key="sk-team-b",
+                team_id="team-b",
+                team_models=[TEAM_PUBLIC_MODEL_NAME],
+                models=[TEAM_PUBLIC_MODEL_NAME],
+            ),
+            id="other-team-key-granted-the-same-public-name",
+        ),
+        pytest.param(
+            Deployment(
+                model_name="model_name_team-a_ml_00000000-0000-0000-0000-000000000000",
+                litellm_params=LiteLLM_Params(model="gpt-4o-mini"),
+                model_info=ModelInfo(
+                    id="dep-team-a-ml-1", team_id="team-a_ml", team_public_model_name=TEAM_PUBLIC_MODEL_NAME
+                ),
+            ),
+            _key_restricted_to(TEAM_PUBLIC_MODEL_NAME),
+            id="team-whose-id-prefixes-the-owning-team-id",
+        ),
+    ],
+)
+async def test_retrieve__team_scoped_deployment_rejects_keys_outside_its_grant(
+    retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth
+) -> None:
+    _route_unified_batch_to_team_deployment(retrieve_harness, deployment)
+
+    with pytest.raises(ProxyException) as exc_info:
+        await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user)
+
+    assert exc_info.value.code == "403"
+    retrieve_harness.router_aretrieve.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "deployment, user",
+    [
+        pytest.param(
+            _team_scoped_deployment(),
+            _key_restricted_to(TEAM_INTERNAL_MODEL_NAME),
+            id="owning-team-key-granted-the-internal-name",
+        ),
+        pytest.param(
+            Deployment(
+                model_name=TEAM_INTERNAL_MODEL_NAME,
+                litellm_params=LiteLLM_Params(model="gpt-4o-mini"),
+                model_info=ModelInfo(id="dep-team-a-1", team_id="team-a"),
+            ),
+            UserAPIKeyAuth(
+                api_key="sk-legacy",
+                team_id="team-a",
+                team_models=["legacy-team-model"],
+                models=["legacy-team-model"],
+                team_model_aliases={"legacy-team-model": TEAM_INTERNAL_MODEL_NAME},
+            ),
+            id="legacy-team-model-reached-through-a-team-alias",
+        ),
+    ],
+)
+async def test_retrieve__team_scoped_deployment_authorizes_internal_name_and_legacy_alias_grants(
+    retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth
+) -> None:
+    _route_unified_batch_to_team_deployment(retrieve_harness, deployment)
+
+    await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user)
+
+    assert retrieve_harness.router_aretrieve.call_count == 1
+
+
 @pytest.mark.asyncio
 async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner):
     runner, factory = executed_runner

tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py

@@ -3,10 +3,12 @@
 Test to verify the Google GenAI proxy API endpoints
 """
 
-from unittest.mock import AsyncMock, MagicMock, patch
+from typing import Final
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
+from litellm.types.utils import TokenCountResponse
 
 
 def _build_test_client():
@@ -87,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():
@@ -122,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"]
@@ -159,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",
@@ -179,3 +182,61 @@ 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():
+    try:
+        client = _build_test_client()
+    except ImportError as e:
+        pytest.skip(f"Skipping test due to missing dependency: {e}")
+
+    with patch(
+        "litellm.proxy.proxy_server.token_counter",
+        new_callable=AsyncMock,
+        side_effect=RuntimeError("tokenizer unavailable"),
+    ):
+        response = client.post(
+            "/v1beta/models/test-model:countTokens",
+            json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
+        )
+
+    assert response.status_code == 500
+    error = response.json()["error"]
+    assert error["code"] == 500
+    assert error["status"] == "INTERNAL"
+    assert "tokenizer unavailable" not in error["message"]
+
+
+def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs():
+    """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field."""
+    try:
+        client: Final = _build_test_client()
+    except ImportError as e:
+        pytest.skip(f"Skipping test due to missing dependency: {e}")
+
+    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",
+        new_callable=AsyncMock,
+        return_value=fake_response,
+    ):
+        response: Final = client.post(
+            "/v1beta/models/claude-opus-4-8:countTokens",
+            json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
+        )
+
+        assert response.status_code == 200
+        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/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py

@@ -4239,6 +4239,54 @@ class TestModelInfoCostMapEchoFilter:
         assert "litellm_provider" not in info
         assert "supported_openai_params" not in info
 
+    @pytest.fixture
+    def bedrock_sonnet_5_5_mapped_as_on_main(self) -> Iterator[None]:
+        import litellm
+
+        entry: Final = {
+            "litellm_provider": "bedrock_converse",
+            "mode": "chat",
+            "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"},
+        }
+        with pytest.MonkeyPatch.context() as patched:
+            patched.setattr(litellm, "model_cost", {**litellm.model_cost, "us.anthropic.claude-sonnet-5-5": entry})
+            litellm.utils._invalidate_model_cost_lowercase_map()
+            yield
+        litellm.utils._invalidate_model_cost_lowercase_map()
+
+    @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, bedrock_sonnet_5_5_mapped_as_on_main: None
+    ) -> 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/test_litellm/proxy/proxy_server/test_routes_utils.py

@@ -52,6 +52,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/litellm_core_utils/test_token_counter.py

@@ -11,6 +11,7 @@ import threading
 import time
 from collections.abc import Mapping
 from concurrent.futures import Future, wait
+from itertools import accumulate, chain
 from pathlib import Path
 from typing import Final
 from unittest.mock import MagicMock
@@ -1362,14 +1363,51 @@ def _count_user_content(content: list[dict]) -> int:
     )
 
 
+def _pdf_base64(pages: tuple[str, ...], width: int = 612, height: int = 792) -> str:
+    def page_objects(index: int, text: str) -> tuple[bytes, bytes]:
+        escaped: Final = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)")
+        stream: Final = f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1")
+        content: Final = f"<< /Length {len(stream)} >>\nstream\n".encode() + stream + b"\nendstream"
+        page: Final = (
+            f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width} {height}] "
+            f"/Resources << /Font << /F1 3 0 R >> >> /Contents {4 + 2 * index} 0 R >>"
+        ).encode()
+        return content, page
+
+    kids: Final = " ".join(f"{5 + 2 * index} 0 R" for index in range(len(pages)))
+    bodies: Final = (
+        b"<< /Type /Catalog /Pages 2 0 R >>",
+        f"<< /Type /Pages /Count {len(pages)} /Kids [ {kids} ] >>".encode(),
+        b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
+        *chain.from_iterable(page_objects(index, text) for index, text in enumerate(pages)),
+    )
+    header: Final = b"%PDF-1.4\n"
+    objects: Final = tuple(
+        f"{number} 0 obj\n".encode() + body + b"\nendobj\n" for number, body in enumerate(bodies, start=1)
+    )
+    offsets: Final = accumulate((len(header), *(len(obj) for obj in objects[:-1])))
+    xref: Final = f"xref\n0 {len(objects) + 1}\n0000000000 65535 f \n".encode() + b"".join(
+        f"{offset:010d} 00000 n \n".encode() for offset in offsets
+    )
+    trailer: Final = (
+        f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n"
+        f"startxref\n{len(header) + sum(len(obj) for obj in objects)}\n%%EOF\n"
+    ).encode()
+    return base64.b64encode(header + b"".join(objects) + xref + trailer).decode()
+
+
+_READABLE_PDF_BASE64: Final = _pdf_base64(("Count this document without parsing its text locally.",))
+
+
 @pytest.mark.parametrize(
     "source",
     [
         {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
         {"type": "url", "url": "https://example.com/report.pdf"},
         {"type": "file", "file_id": "file-abc123"},
+        {"type": "base64", "media_type": "application/pdf", "data": _READABLE_PDF_BASE64},
     ],
-    ids=["base64", "url", "file"],
+    ids=["base64", "url", "file", "readable_pdf"],
 )
 def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]):
     """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on."""
@@ -1411,7 +1449,8 @@ def test_anthropic_document_title_and_context_add_their_tokens():
     )
 
 
-def test_openai_file_block_prices_like_the_equivalent_anthropic_document():
+@pytest.mark.parametrize("pdf_data", ("JVBERi0xLjQK", _READABLE_PDF_BASE64), ids=("opaque", "readable"))
+def test_openai_file_block_prices_like_the_equivalent_anthropic_document(pdf_data: str) -> None:
     """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise.
 
     Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject`
@@ -1421,17 +1460,21 @@ def test_openai_file_block_prices_like_the_equivalent_anthropic_document():
     prompt = {"type": "text", "text": "Summarize this file."}
     inline_file = {
         "type": "file",
-        "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"},
+        "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64," + pdf_data},
     }
     document = {
         "type": "document",
         "title": "report.pdf",
-        "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
+        "source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data},
     }
 
     assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document])
     assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt])
 
+    assert _count_user_content([prompt, inline_file]) == _count_user_content(
+        [prompt, {"type": "text", "text": "report.pdf"}, {"type": "image", "source": document["source"]}]
+    )
+
 
 def test_openai_file_block_without_inline_bytes_counts_what_it_carries():
     """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens."""
@@ -1441,9 +1484,7 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries():
     assert _count_user_content([prompt, by_id]) == _count_user_content([prompt])
 
     named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}}
-    assert _count_user_content([prompt, named]) == _count_user_content(
-        [prompt, {"type": "text", "text": "report.pdf"}]
-    )
+    assert _count_user_content([prompt, named]) == _count_user_content([prompt, {"type": "text", "text": "report.pdf"}])
 
 
 def _png_data_url(width: int, height: int) -> str:

tests/unit/llms/anthropic/experimental_pass_through/test_convertible_blocks.py

@@ -0,0 +1,192 @@
+import asyncio
+import json
+from collections.abc import Mapping
+from typing import Final
+
+import pytest
+from pydantic import JsonValue, TypeAdapter
+
+from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
+from litellm.llms.anthropic.experimental_pass_through.convertible_blocks import (
+    unconvertible_anthropic_blocks,
+    unconvertible_blocks_error,
+)
+from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import (
+    LiteLLMAnthropicToResponsesAPIAdapter,
+)
+from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents
+from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues
+
+_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]])
+
+_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}
+_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}
+
+_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = (
+    ("user", {"type": "text", "text": "hi"}),
+    ("user", {"type": "image", "source": _PNG}),
+    ("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}),
+    ("user", {"type": "document", "source": _PDF}),
+    ("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}),
+    ("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}),
+    ("assistant", {"type": "text", "text": "hi"}),
+    ("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}),
+    ("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}),
+)
+
+_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = (
+    ("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"),
+    ("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"),
+    ("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"),
+    ("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"),
+    ("user", {"type": "document"}, "document without a source"),
+    ("user", {"text": "untyped"}, "a block without a type"),
+    ("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"),
+    (
+        "assistant",
+        {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}},
+        "server_tool_use",
+    ),
+)
+
+
+def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]:
+    return [{"role": role, "content": [dict(block)]}]  # pyright: ignore[reportReturnType]  # sends block shapes the TypedDicts leave out on purpose
+
+
+def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]:
+    messages: Final = _adapter_messages(role, block)
+    chat: Final = _JSON_OBJECTS.validate_python(
+        LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages)
+    )
+    responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages)
+    chat_kept: Final = any(
+        message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat
+    )
+    return chat_kept, bool(responses)
+
+
+@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS)
+def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None:
+    assert _converted_by_both_adapters(role, block) == (True, True)
+    assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == ()
+
+
+@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS)
+def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None:
+    assert _converted_by_both_adapters(role, block) == (False, False)
+    assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,)
+
+
+def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None:
+    messages: Final = [
+        {"role": "user", "content": "plain text"},
+        {"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]},
+        {"role": "assistant", "content": [{"type": "search_result"}]},
+    ]
+    assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result")
+
+
+def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None:
+    messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}]
+    assert unconvertible_anthropic_blocks(messages) == ()
+
+
+def test_messages_that_are_not_json_are_left_to_the_count() -> None:
+    messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}]
+    assert unconvertible_anthropic_blocks(messages) == ()
+
+
+def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
+    return (
+        {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]},
+        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]},
+    )
+
+
+def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]:
+    messages: Final = _tool_turn(inner)
+    gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {}))
+    openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(
+        [dict(message) for message in messages]  # pyright: ignore[reportArgumentType]  # sends block shapes the TypedDicts leave out on purpose
+    )
+    return json.dumps(gemini, default=str), json.dumps(openai, default=str)
+
+
+_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = (
+    ({"type": "text", "text": "forty-two"}, "forty-two"),
+    ({"type": "image", "source": _PNG}, _PNG["data"]),
+    ({"type": "document", "source": _PDF}, _PDF["data"]),
+)
+
+_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = (
+    (
+        {
+            "type": "search_result",
+            "source": "https://example.com",
+            "title": "t",
+            "content": [{"type": "text", "text": "search-snippet"}],
+        },
+        "search-snippet",
+        "search_result inside a tool_result",
+    ),
+    (
+        {"type": "document", "source": {"type": "text", "data": "text-notes"}},
+        "text-notes",
+        "document with a text source inside a tool_result",
+    ),
+    ({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"),
+)
+
+
+@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT)
+def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None:
+    gemini, openai = _sent_by_each_counter(inner)
+
+    assert marker in gemini
+    assert marker in openai
+    assert unconvertible_anthropic_blocks(_tool_turn(inner)) == ()
+
+
+@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT)
+def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None:
+    gemini, openai = _sent_by_each_counter(inner)
+
+    assert marker not in gemini
+    assert marker not in openai
+    assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,)
+
+
+def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]:
+    return (
+        {"role": "user", "content": "What is new?"},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
+                {
+                    "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": encrypted_content,
+                            "snippet": "LiteLLM ships a weekly release candidate.",
+                        }
+                    ],
+                },
+            ],
+        },
+    )
+
+
+def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None:
+    assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None
+
+
+def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None:
+    error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic")
+    assert error is not None
+    assert error.endswith("server_tool_use, web_search_tool_result")

tests/unit/llms/anthropic/test_count_tokens_oauth.py

@@ -14,6 +14,12 @@ sys.path.insert(
     0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
 )
 
+import httpx
+import pytest
+import respx
+
+import litellm
+from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter
 from litellm.llms.anthropic.count_tokens.transformation import (
     AnthropicCountTokensConfig,
 )
@@ -84,3 +90,48 @@ class TestCountTokensOAuthHeaders:
         assert (
             "oauth-2025-04-20" in beta_value
         ), f"oauth beta missing from OAuth headers: {beta_value}"
+
+
+@pytest.mark.asyncio
+async def test_a_global_anthropic_key_counts_like_chat_would_authenticate(monkeypatch):
+    for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY)
+
+    with respx.mock(assert_all_called=True) as router:
+        route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
+            return_value=httpx.Response(200, json={"input_tokens": 13})
+        )
+        result = await AnthropicTokenCounter().count_tokens(
+            model_to_use="claude-sonnet-4-5",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}},
+            request_model="claude-sonnet-4-5",
+        )
+
+    assert result is not None
+    assert result.total_tokens == 13
+    assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY
+
+
+@pytest.mark.asyncio
+async def test_no_anthropic_key_returns_a_401_count_error(monkeypatch):
+    for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.setattr(litellm, "anthropic_key", None)
+    monkeypatch.setattr(litellm, "api_key", None)
+
+    result = await AnthropicTokenCounter().count_tokens(
+        model_to_use="claude-sonnet-4-5",
+        messages=[{"role": "user", "content": "hi"}],
+        contents=None,
+        deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}},
+        request_model="claude-sonnet-4-5",
+    )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 401
+    assert result.total_tokens == 0

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

@@ -0,0 +1,217 @@
+import json
+import operator
+from typing import Final
+
+import httpx
+import pytest
+import respx
+from pydantic import ValidationError
+
+import litellm
+from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter
+from litellm.types.utils import TokenCountResponse
+
+API_BASE: Final = "https://my-resource.services.ai.azure.com"
+COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens"
+
+
+async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    with respx.mock:
+        respx.post(COUNT_URL).mock(return_value=response)
+        return await AzureAIAnthropicTokenCounter().count_tokens(
+            model_to_use="claude-sonnet-4-5",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}},
+            request_model="azure_ai/claude-sonnet-4-5",
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("payload", "expected_total"),
+    [
+        ({"input_tokens": 15}, 15),
+        ({"input_tokens": "15"}, 15),
+        ({"input_tokens": 15.0}, 15),
+        ({"input_tokens": True}, 1),
+        ({"input_tokens": -3}, -3),
+        ({}, 0),
+        ({"type": "error", "error": {"message": "429 rate limit"}}, 0),
+        ({"input_tokens": 7, "context_management": {"original_input_tokens": 9}, "nested": {"a": [1, None]}}, 7),
+    ],
+)
+async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
+    monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int
+) -> None:
+    result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": expected_total,
+            "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,
+            "status_code": None,
+        }
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "rejected_count",
+    [
+        None,
+        15.5,
+        "Request timed out",
+        "429 Too Many Requests",
+        [15],
+        [*range(403), "Request timed out"],
+        {"value": 429},
+    ],
+)
+async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message(
+    monkeypatch: pytest.MonkeyPatch, rejected_count: object
+) -> None:
+    with pytest.raises(ValidationError) as rejection:
+        TokenCountResponse.model_validate(
+            {
+                "total_tokens": rejected_count,
+                "request_model": "azure_ai/claude-sonnet-4-5",
+                "model_used": "claude-sonnet-4-5",
+                "tokenizer_type": "azure_ai_anthropic_api",
+            }
+        )
+
+    result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count}))
+
+    assert result is not None
+    assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)]
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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),
+            "status_code": 500,
+        }
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "payload",
+    [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False],
+)
+async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup(
+    monkeypatch: pytest.MonkeyPatch, payload: object
+) -> None:
+    with pytest.raises(AttributeError) as failed_lookup:
+        operator.attrgetter("get")(payload)
+
+    result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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),
+            "status_code": 500,
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None:
+    assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("upstream_status", "upstream_text"),
+    [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")],
+)
+async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text(
+    monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str
+) -> None:
+    result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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,
+            "status_code": upstream_status,
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    with pytest.raises(json.JSONDecodeError) as undecodable:
+        json.loads("not json")
+
+    result: Final = await _count(monkeypatch, httpx.Response(200, text="not json"))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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}",
+            "status_code": 500,
+        }
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base"))
+async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.delenv("AZURE_AI_API_BASE", raising=False)
+    monkeypatch.delenv("AZURE_API_BASE", raising=False)
+    if set_base == "env":
+        monkeypatch.setenv("AZURE_API_BASE", API_BASE)
+    else:
+        monkeypatch.setattr(litellm, "api_base", API_BASE)
+    with respx.mock:
+        route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6}))
+        result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
+            model_to_use="claude-sonnet-4-5",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {"api_key": "azure-key"}},
+            request_model="azure_ai/claude-sonnet-4-5",
+        )
+    assert result is not None
+    assert result.total_tokens == 6
+    assert route.called

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/gemini/test_gemini_common_utils.py

@@ -1,8 +1,200 @@
+from typing import Final
 from unittest.mock import AsyncMock, patch
 
+import httpx
 import pytest
+from pydantic import JsonValue, TypeAdapter
 
+import litellm
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter
+from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
+from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter as GeminiCountTokensHandler
+
+_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue])
+
+
+class _GeminiCountTokensPeer:
+    def __init__(self) -> None:
+        self.request_body: dict[str, JsonValue] | None = None
+
+    def __call__(self, request: httpx.Request) -> httpx.Response:
+        self.request_body = _JSON_BODY.validate_json(request.content)
+        return httpx.Response(200, json={"totalTokens": 11}, request=request)
+
+
+@pytest.mark.asyncio
+async def test_google_ai_studio_count_tokens_forwards_system_and_tools() -> None:
+    peer: Final = _GeminiCountTokensPeer()
+    client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
+
+    with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
+        result: Final = await GoogleAIStudioTokenCounter().count_tokens(
+            model_to_use="gemini-2.5-flash",
+            messages=None,
+            contents=[{"role": "user", "parts": [{"text": "hello"}]}],
+            deployment={
+                "litellm_params": {
+                    "api_key": "test-key",
+                    "api_base": "https://count-tokens.test",
+                    "system_instruction": {"parts": [{"text": "deployment-level"}]},
+                }
+            },
+            request_model="gemini/gemini-2.5-flash",
+            tools=[
+                {
+                    "name": "lookup",
+                    "description": "Find a record",
+                    "input_schema": {"type": "object", "properties": {}},
+                }
+            ],
+            system="You are terse.",
+        )
+
+    assert result is not None
+    assert result.error is False
+    assert result.total_tokens == 11
+    assert peer.request_body == {
+        "generateContentRequest": {
+            "model": "models/gemini-2.5-flash",
+            "contents": [{"role": "user", "parts": [{"text": "hello"}]}],
+            "systemInstruction": {"parts": [{"text": "You are terse."}]},
+            "tools": [
+                {
+                    "functionDeclarations": [
+                        {
+                            "name": "lookup",
+                            "description": "Find a record",
+                            "parameters": {"type": "object"},
+                        }
+                    ]
+                }
+            ],
+        }
+    }
+
+
+@pytest.mark.asyncio
+async def test_google_ai_studio_count_tokens_rejects_image_system_without_request() -> None:
+    peer: Final = _GeminiCountTokensPeer()
+    client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
+
+    with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
+        result: Final = await GoogleAIStudioTokenCounter().count_tokens(
+            model_to_use="gemini-2.5-flash",
+            messages=None,
+            contents=[{"role": "user", "parts": [{"text": "hello"}]}],
+            deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}},
+            request_model="gemini/gemini-2.5-flash",
+            system=[
+                {"type": "text", "text": "You are terse."},
+                {
+                    "type": "image",
+                    "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="},
+                },
+            ],
+        )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 400
+    assert peer.request_body is None
+
+
+@pytest.mark.asyncio
+async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "user_url_validation", True)
+    peer: Final = _GeminiCountTokensPeer()
+    client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
+
+    with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
+        result: Final = await GoogleAIStudioTokenCounter().count_tokens(
+            model_to_use="gemini-2.5-flash",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}],
+                }
+            ],
+            contents=None,
+            deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}},
+            request_model="gemini/gemini-2.5-flash",
+            messages_format="anthropic",
+        )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 400
+    assert "http://127.0.0.1:9/image.png" in (result.error_message or "")
+    assert peer.request_body is None
+
+
+@pytest.mark.asyncio
+async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None:
+    rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}'
+    client: Final = AsyncHTTPHandler(
+        transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request))
+    )
+
+    with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
+        result: Final = await GoogleAIStudioTokenCounter().count_tokens(
+            model_to_use="gemini-2.5-flash",
+            messages=None,
+            contents=[{"role": "user", "parts": [{"text": "hello"}]}],
+            deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}},
+            request_model="gemini/gemini-2.5-flash",
+        )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 400
+    assert "Invalid function name" in (result.error_message or "")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("system_instruction", "tools", "expected_body"),
+    (
+        (
+            {"parts": [{"text": "Follow the system rule."}]},
+            [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}],
+            {
+                "generateContentRequest": {
+                    "model": "models/gemini-4-argon",
+                    "contents": [{"role": "user", "parts": [{"text": "hello"}]}],
+                    "systemInstruction": {"parts": [{"text": "Follow the system rule."}]},
+                    "tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}],
+                }
+            },
+        ),
+        (
+            None,
+            [],
+            {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]},
+        ),
+    ),
+)
+async def test_gemini_count_tokens_sends_optional_request_fields(
+    system_instruction: dict[str, JsonValue] | None,
+    tools: list[dict[str, JsonValue]],
+    expected_body: dict[str, JsonValue],
+) -> None:
+    peer: Final = _GeminiCountTokensPeer()
+    token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)))
+
+    result: Final = await token_counter.acount_tokens(
+        contents=[{"role": "user", "parts": [{"text": "hello"}]}],
+        model="gemini-4-argon",
+        api_key="test-key",
+        api_base="https://count-tokens.test",
+        system_instruction=system_instruction,
+        tools=tools,
+    )
+
+    assert result == {"totalTokens": 11}
+    assert peer.request_body == expected_body
 
 
 class TestGeminiModelInfo:
@@ -36,18 +228,10 @@ class TestGeminiModelInfo:
         # Test edge cases where model names end with characters from "models/"
         # These would be incorrectly processed if using strip("models/") instead of replace("models/", "")
         models = [
-            {
-                "name": "models/gemini-1.5-pro"
-            },  # ends with 'o' - would become "gemini-1.5-pr" with strip()
-            {
-                "name": "models/test-model"
-            },  # ends with 'l' - would become "gemini/test-mode" with strip()
-            {
-                "name": "models/custom-models"
-            },  # ends with 's' - would become "gemini/custom-model" with strip()
-            {
-                "name": "models/demo"
-            },  # ends with 'o' - would become "gemini/dem" with strip()
+            {"name": "models/gemini-1.5-pro"},  # ends with 'o' - would become "gemini-1.5-pr" with strip()
+            {"name": "models/test-model"},  # ends with 'l' - would become "gemini/test-mode" with strip()
+            {"name": "models/custom-models"},  # ends with 's' - would become "gemini/custom-model" with strip()
+            {"name": "models/demo"},  # ends with 'o' - would become "gemini/dem" with strip()
         ]
 
         result = gemini_model_info.process_model_name(models)
@@ -98,16 +282,10 @@ class TestGoogleAIStudioTokenCounter:
         token_counter = GoogleAIStudioTokenCounter()
 
         # Test with gemini provider - should return True
-        assert (
-            token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value)
-            is True
-        )
+        assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True
 
         # Test with other providers - should return False
-        assert (
-            token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value)
-            is False
-        )
+        assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False
         assert token_counter.should_use_token_counting_api("anthropic") is False
         assert token_counter.should_use_token_counting_api("vertex_ai") is False
 
@@ -158,7 +336,10 @@ class TestGoogleAIStudioTokenCounter:
 
             # Verify the mock was called correctly
             mock_acount_tokens.assert_called_once_with(
-                model=model_to_use, contents=contents
+                model=model_to_use,
+                contents=contents,
+                system_instruction=None,
+                tools=None,
             )
 
     def test_clean_contents_for_gemini_api_removes_id_field(self):
@@ -176,9 +357,7 @@ class TestGoogleAIStudioTokenCounter:
                         "functionResponse": {
                             "id": "read_many_files-1757526647518-730a691aac11c",  # This should be removed
                             "name": "read_many_files",
-                            "response": {
-                                "output": "No files matching the criteria were found or all were skipped."
-                            },
+                            "response": {"output": "No files matching the criteria were found or all were skipped."},
                         }
                     }
                 ],
@@ -187,9 +366,7 @@ class TestGoogleAIStudioTokenCounter:
         ]
 
         # Clean the contents
-        cleaned_contents = token_counter._clean_contents_for_gemini_api(
-            contents_with_id
-        )
+        cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id)
 
         # Verify the 'id' field was removed
         function_response = cleaned_contents[1]["parts"][0]["functionResponse"]
@@ -198,8 +375,7 @@ class TestGoogleAIStudioTokenCounter:
         assert "response" in function_response
         assert function_response["name"] == "read_many_files"
         assert (
-            function_response["response"]["output"]
-            == "No files matching the criteria were found or all were skipped."
+            function_response["response"]["output"] == "No files matching the criteria were found or all were skipped."
         )
 
     def test_clean_contents_for_gemini_api_preserves_other_fields(self):
@@ -215,9 +391,7 @@ class TestGoogleAIStudioTokenCounter:
         ]
 
         # Clean the contents
-        cleaned_contents = token_counter._clean_contents_for_gemini_api(
-            contents_without_function_response
-        )
+        cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response)
 
         # Verify the contents are unchanged
         assert cleaned_contents == contents_without_function_response

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

@@ -0,0 +1,436 @@
+import json
+import operator
+from typing import Final
+
+import httpx
+import pytest
+import respx
+from pydantic import JsonValue, ValidationError
+
+import litellm
+from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
+from litellm.responses.utils import ResponsesAPIRequestUtils
+from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
+
+COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens"
+
+
+async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    with respx.mock:
+        respx.post(COUNT_URL).mock(return_value=response)
+        return await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {"api_key": "sk-test"}},
+            request_model="openai/gpt-4o",
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("payload", "expected_total"),
+    [
+        ({"input_tokens": 15}, 15),
+        ({"input_tokens": "15"}, 15),
+        ({"input_tokens": 15.0}, 15),
+        ({"input_tokens": True}, 1),
+        ({"input_tokens": -3}, -3),
+        ({}, 0),
+        ({"object": "response.input_tokens", "error": {"message": "429 rate limit"}}, 0),
+        ({"input_tokens": 7, "object": "response.input_tokens", "nested": {"a": [1, None]}}, 7),
+    ],
+)
+async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
+    monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int
+) -> None:
+    result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": expected_total,
+            "request_model": "openai/gpt-4o",
+            "model_used": "gpt-4o",
+            "tokenizer_type": "openai_api",
+            "count_source": None,
+            "original_response": payload,
+            "error": False,
+            "error_message": None,
+            "status_code": None,
+        }
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "rejected_count",
+    [
+        None,
+        15.5,
+        "Request timed out",
+        "429 Too Many Requests",
+        [15],
+        [*range(403), "Request timed out"],
+        {"value": 429},
+    ],
+)
+async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message(
+    monkeypatch: pytest.MonkeyPatch, rejected_count: object
+) -> None:
+    with pytest.raises(ValidationError) as rejection:
+        TokenCountResponse.model_validate(
+            {
+                "total_tokens": rejected_count,
+                "request_model": "openai/gpt-4o",
+                "model_used": "gpt-4o",
+                "tokenizer_type": "openai_api",
+            }
+        )
+
+    result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count}))
+
+    assert result is not None
+    assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)]
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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),
+            "status_code": 500,
+        }
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "payload",
+    [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False],
+)
+async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup(
+    monkeypatch: pytest.MonkeyPatch, payload: object
+) -> None:
+    with pytest.raises(AttributeError) as failed_lookup:
+        operator.attrgetter("get")(payload)
+
+    result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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),
+            "status_code": 500,
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None:
+    assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("upstream_status", "upstream_text"),
+    [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")],
+)
+async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text(
+    monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str
+) -> None:
+    result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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,
+            "status_code": upstream_status,
+        }
+    )
+
+
+async def _sent_count_body(
+    monkeypatch: pytest.MonkeyPatch,
+    messages: list[dict[str, object]],
+    system: object,
+    messages_format: TokenCountMessagesFormat,
+    output_format: JsonValue = None,
+) -> object:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    with respx.mock:
+        route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1}))
+        await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=messages,
+            contents=None,
+            deployment={"litellm_params": {"api_key": "sk-test"}},
+            request_model="openai/gpt-4o",
+            system=system,
+            messages_format=messages_format,
+            output_format=output_format,
+        )
+        return json.loads(route.calls.last.request.content)
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    messages: Final[list[dict[str, object]]] = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Summarize."},
+                {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}},
+                {"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}},
+            ],
+        },
+        {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]},
+        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]},
+    ]
+
+    sent: Final = await _sent_count_body(
+        monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic"
+    )
+
+    assert sent == {
+        "model": "gpt-4o",
+        "input": [
+            {
+                "type": "message",
+                "role": "user",
+                "content": [
+                    {"type": "input_text", "text": "Summarize."},
+                    {"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"},
+                    {"type": "input_image", "image_url": "https://example.com/page.png"},
+                ],
+            },
+            {"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"},
+            {"type": "function_call_output", "call_id": "toolu_01", "output": "found"},
+        ],
+        "instructions": "Be brief.",
+    }
+
+
+@pytest.mark.asyncio
+async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None:
+    messages: Final[list[dict[str, object]]] = [
+        {"role": "system", "content": "Be brief."},
+        {"role": "user", "content": [{"type": "text", "text": "Summarize."}]},
+        {"role": "tool", "tool_call_id": "call_1", "content": "found"},
+    ]
+
+    sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai")
+
+    assert sent == {
+        "model": "gpt-4o",
+        "input": [
+            {"role": "user", "content": "Summarize."},
+            {"type": "function_call_output", "call_id": "call_1", "output": "found"},
+        ],
+        "instructions": "Be brief.",
+    }
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    wrapped: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1")  # pyright: ignore[reportPrivateUsage]  # the wrap inference applies
+    messages: Final[list[dict[str, object]]] = [
+        {"role": "user", "content": "hi"},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"},
+                {"type": "text", "text": "OK"},
+            ],
+        },
+        {"role": "user", "content": "again"},
+    ]
+
+    sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")
+
+    assert sent == {
+        "model": "gpt-4o",
+        "input": [
+            {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
+            {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"},
+            {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]},
+            {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]},
+        ],
+    }
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None:
+    schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
+
+    sent: Final = await _sent_count_body(
+        monkeypatch,
+        [{"role": "user", "content": "Answer in JSON."}],
+        system=None,
+        messages_format="anthropic",
+        output_format={"type": "json_schema", "schema": schema},
+    )
+
+    assert sent == {
+        "model": "gpt-4o",
+        "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}],
+        "text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}},
+    }
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None:
+    snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
+    messages: Final[list[dict[str, object]]] = [
+        {"role": "user", "content": "What is new?"},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
+                {
+                    "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": "",
+                            "snippet": snippet,
+                        }
+                    ],
+                },
+            ],
+        },
+        {"role": "user", "content": "Thanks."},
+    ]
+
+    sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic"))
+
+    assert snippet in sent
+    assert "https://docs.litellm.ai/release_notes" in sent
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    with respx.mock:
+        route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1}))
+        result: Final = await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=[{"role": "user", "content": "hi"}, {"content": "no role"}],
+            contents=None,
+            deployment={"litellm_params": {"api_key": "sk-test"}},
+            request_model="openai/gpt-4o",
+            messages_format="anthropic",
+        )
+
+    assert route.called is False
+    assert result is not None
+    assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required")
+
+
+@pytest.mark.asyncio
+async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None:
+    sent: Final = await _sent_count_body(
+        monkeypatch,
+        [{"role": "user", "content": "hi"}],
+        system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}],
+        messages_format="openai",
+    )
+
+    assert sent == {
+        "model": "gpt-4o",
+        "input": [{"role": "user", "content": "hi"}],
+        "instructions": "Be brief.\nAnswer in English.",
+    }
+
+
+@pytest.mark.asyncio
+async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    with pytest.raises(json.JSONDecodeError) as undecodable:
+        json.loads("not json")
+
+    result: Final = await _count(monkeypatch, httpx.Response(200, text="not json"))
+
+    assert result is not None
+    assert repr(result.model_dump()) == repr(
+        {
+            "total_tokens": 0,
+            "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}",
+            "status_code": 500,
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings")
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    with respx.mock:
+        route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9}))
+        result: Final = await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {}},
+            request_model="openai/gpt-4o",
+        )
+    assert result is not None
+    assert result.total_tokens == 9
+    assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings"
+
+
+@pytest.mark.asyncio
+async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway")
+    monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1")
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    with respx.mock:
+        route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock(
+            return_value=httpx.Response(200, json={"input_tokens": 4})
+        )
+        result: Final = await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {}},
+            request_model="openai/gpt-4o",
+        )
+    assert result is not None
+    assert result.total_tokens == 4
+    assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway"

tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py

@@ -1,6 +1,10 @@
+from typing import Final
 
 import pytest
 
+from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import (
+    LiteLLMAnthropicToResponsesAPIAdapter,
+)
 from litellm.llms.openai.responses.count_tokens.transformation import (
     OpenAICountTokensConfig,
 )
@@ -446,3 +450,108 @@ def test_get_required_headers():
 
     assert headers["Authorization"] == "Bearer sk-test-key"
     assert headers["Content-Type"] == "application/json"
+
+
+PDF_BASE64: Final = "JVBERi0xLjcK"
+PNG_BASE64: Final = "iVBORw0KGgo="
+
+ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "Summarize the indemnity terms."},
+            {
+                "type": "document",
+                "title": "contract.pdf",
+                "source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64},
+            },
+            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}},
+        ],
+    },
+    {
+        "role": "assistant",
+        "content": [
+            {"type": "text", "text": "Let me pull clause 14."},
+            {"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}},
+        ],
+    },
+    {
+        "role": "user",
+        "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}],
+    },
+]
+
+ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [
+    {
+        "name": "lookup_clause",
+        "description": "Fetch a clause by id",
+        "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}},
+    }
+]
+
+
+@pytest.mark.parametrize(
+    "system",
+    [None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]],
+)
+def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None:
+    inference_request: Final[dict[str, object]] = {
+        "model": "gpt-4o",
+        "messages": ANTHROPIC_AGENT_TURN,
+        "tools": ANTHROPIC_TOOLS,
+        **({"system": system} if system is not None else {}),
+    }
+    inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request(
+        inference_request, include_encrypted_reasoning=False
+    )
+
+    counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
+        model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS
+    )
+
+    assert counted == (
+        tuple(inference["input"]),
+        inference.get("instructions"),
+        tuple(inference["tools"]),
+        inference.get("text"),
+    )
+
+
+def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None:
+    input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
+        model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None
+    )
+
+    assert input_items == (
+        {
+            "type": "message",
+            "role": "user",
+            "content": [
+                {"type": "input_text", "text": "Summarize the indemnity terms."},
+                {
+                    "type": "input_file",
+                    "filename": "contract.pdf",
+                    "file_data": f"data:application/pdf;base64,{PDF_BASE64}",
+                },
+                {"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}"},
+            ],
+        },
+        {"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'},
+        {
+            "type": "message",
+            "role": "assistant",
+            "content": [{"type": "output_text", "text": "Let me pull clause 14."}],
+        },
+        {"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."},
+    )
+
+
+def test_anthropic_list_system_becomes_joined_instruction_text() -> None:
+    _, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "hi"}],
+        system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}],
+        tools=None,
+    )
+
+    assert instructions == "First rule.\nSecond rule."

tests/unit/llms/openai/test_openai_workload_identity.py

@@ -279,3 +279,56 @@ class TestResponsesValidateEnvironment:
             headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams()
         )
         assert headers["Authorization"] == "Bearer None"
+
+
+class TestCountTokens:
+    @pytest.mark.asyncio
+    @respx.mock
+    async def test_count_authenticates_with_the_workload_identity_bearer(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
+
+        monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+        monkeypatch.setattr(litellm, "api_key", None)
+        monkeypatch.setattr(litellm, "openai_key", None)
+        mock_token_exchange()
+        count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock(
+            return_value=httpx.Response(200, json={"input_tokens": 12})
+        )
+        result: Final = await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {}},
+            request_model="openai/gpt-4o",
+        )
+        assert result is not None
+        assert result.error is False
+        assert result.total_tokens == 12
+        assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token"
+
+    @pytest.mark.asyncio
+    @respx.mock
+    async def test_failed_mint_becomes_a_count_error_response(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
+
+        monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+        monkeypatch.setattr(litellm, "api_key", None)
+        monkeypatch.setattr(litellm, "openai_key", None)
+        respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"}))
+        count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens")
+        result: Final = await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={"litellm_params": {}},
+            request_model="openai/gpt-4o",
+        )
+        assert result is not None
+        assert result.error is True
+        assert result.status_code == 401
+        assert result.total_tokens == 0
+        assert not count_route.called

tests/unit/llms/vertex_ai/count_tokens/test_context.py

@@ -0,0 +1,507 @@
+import asyncio
+import base64
+import time
+import uuid
+from collections.abc import Generator
+from contextlib import contextmanager
+from copy import deepcopy
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from threading import Thread
+from typing import Final, Literal
+from unittest.mock import patch
+
+import pytest
+from pydantic import JsonValue, TypeAdapter
+
+import litellm
+from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
+    LiteLLMAnthropicMessagesAdapter,
+    truncate_tool_name,
+)
+from litellm.llms.vertex_ai.count_tokens.context import (
+    VertexGeminiCountContext,
+    VertexGeminiCountContextError,
+    gemini_count_contents,
+    inference_request_view,
+    vertex_gemini_count_context,
+)
+from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
+
+_SCHEMA: Final[JsonValue] = {
+    "type": "object",
+    "properties": {},
+    "additionalProperties": False,
+}
+_FUNCTION_DECLARATION: Final = {
+    "functionDeclarations": [
+        {
+            "name": "lookup",
+            "description": "Look up a value",
+            "parameters": {"type": "object"},
+        }
+    ]
+}
+
+
+@pytest.mark.parametrize(
+    "tool",
+    (
+        {
+            "name": "lookup",
+            "description": "Look up a value",
+            "input_schema": _SCHEMA,
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "lookup",
+                "description": "Look up a value",
+                "parameters": _SCHEMA,
+            },
+        },
+        {
+            "type": "function",
+            "name": "lookup",
+            "description": "Look up a value",
+            "parameters": _SCHEMA,
+        },
+    ),
+    ids=("anthropic", "openai-chat", "responses"),
+)
+def test_function_tool_shapes_map_to_vertex_declarations(tool: JsonValue) -> None:
+    context: Final = vertex_gemini_count_context(system=None, tools=[tool])
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.tools == (_FUNCTION_DECLARATION,)
+    assert _SCHEMA["additionalProperties"] is False
+
+
+def test_long_anthropic_tool_name_is_truncated_for_vertex() -> None:
+    name: Final = "mcp__docs__" + "a" * 70
+    context: Final = vertex_gemini_count_context(
+        system=None,
+        tools=[{"name": name, "description": "Look up a value", "input_schema": _SCHEMA}],
+    )
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.tools == (
+        {
+            "functionDeclarations": [
+                {
+                    "name": truncate_tool_name(name),
+                    "description": "Look up a value",
+                    "parameters": {"type": "object"},
+                }
+            ]
+        },
+    )
+    assert len(truncate_tool_name(name)) == 64
+
+
+def test_64_character_anthropic_tool_name_is_unchanged_for_vertex() -> None:
+    name: Final = "mcp__docs__" + "a" * 53
+    context: Final = vertex_gemini_count_context(
+        system=None,
+        tools=[{"name": name, "input_schema": _SCHEMA}],
+    )
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},)
+
+
+def test_long_openai_chat_tool_name_is_unchanged_for_vertex() -> None:
+    name: Final = "mcp__docs__" + "a" * 70
+    context: Final = vertex_gemini_count_context(
+        system=None,
+        tools=[{"type": "function", "function": {"name": name, "parameters": _SCHEMA}}],
+    )
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},)
+
+
+def test_unknown_tools_are_omitted() -> None:
+    context: Final = vertex_gemini_count_context(
+        system=None,
+        tools=[{"type": "web_search_preview"}, {"type": "computer_use"}],
+    )
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.tools is None
+
+
+def test_text_block_system_instruction_keeps_one_part_per_block() -> None:
+    context: Final = vertex_gemini_count_context(
+        system=[{"type": "text", "text": "First instruction."}, {"type": "text", "text": "Second instruction."}],
+        tools=None,
+    )
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.system_instruction == {"parts": ({"text": "First instruction."}, {"text": "Second instruction."})}
+
+
+def test_string_system_instruction_maps_to_one_part() -> None:
+    context: Final = vertex_gemini_count_context(system="Follow this instruction.", tools=None)
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.system_instruction == {"parts": ({"text": "Follow this instruction."},)}
+
+
+def test_non_text_system_block_returns_an_error_value() -> None:
+    context: Final = vertex_gemini_count_context(
+        system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}],
+        tools=None,
+    )
+
+    assert isinstance(context, VertexGeminiCountContextError)
+    assert "text blocks" in context.error_message
+
+
+@pytest.mark.parametrize(("system", "tools"), ((None, None), ("", None), ([], [])))
+def test_empty_context_is_omitted(system: JsonValue | None, tools: JsonValue | None) -> None:
+    context: Final = vertex_gemini_count_context(system=system, tools=tools)
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.system_instruction is None
+    assert context.tools is None
+
+
+@pytest.mark.parametrize(
+    "tool",
+    (
+        {"description": "Look up a value", "input_schema": _SCHEMA},
+        {"type": "bash_20250124", "name": "bash"},
+    ),
+    ids=("unnamed", "anthropic-hosted"),
+)
+def test_anthropic_tools_count_the_same_declarations_inference_sends(tool: dict[str, JsonValue]) -> None:
+    translated_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(tool)])  # pyright: ignore[reportArgumentType]  # raw Anthropic tool dict
+    inference_tools: Final = VertexGeminiConfig()._map_function(  # pyright: ignore[reportPrivateUsage]  # the tool mapping inference uses
+        value=[dict(translated) for translated in translated_tools],
+        optional_params={},
+    )
+
+    context: Final = vertex_gemini_count_context(system=None, tools=[tool])
+
+    assert isinstance(context, VertexGeminiCountContext)
+    assert context.tools is not None
+    assert [tool_value["functionDeclarations"] for tool_value in context.tools] == [
+        inference_tool["function_declarations"] for inference_tool in inference_tools
+    ]
+
+
+@pytest.mark.parametrize(
+    "tools",
+    (
+        "lookup",
+        [1],
+        [{"name": "lookup", "input_schema": "not-a-schema"}],
+        [{"type": "function", "function": {"name": "lookup", "parameters": [1]}}],
+        [{"name": "lookup", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}],
+    ),
+    ids=("not-a-list", "non-object-tool", "string-schema", "list-parameters", "malformed-property"),
+)
+def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None:
+    context: Final = vertex_gemini_count_context(system=None, tools=tools)
+
+    assert isinstance(context, VertexGeminiCountContextError)
+    assert "tools" in context.error_message
+
+
+_ANTHROPIC_CONVERSATION: Final = (
+    {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "Summarize both"},
+            {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}},
+            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}},
+        ],
+    },
+    {
+        "role": "assistant",
+        "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}],
+    },
+    {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]},
+)
+
+
+@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
+async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts(
+    custom_llm_provider: Literal["vertex_ai", "gemini"],
+) -> None:
+    contents: Final = await gemini_count_contents(
+        _ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}
+    )
+
+    assert contents == [
+        {
+            "role": "user",
+            "parts": [
+                {"text": "Summarize both"},
+                {"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}},
+                {"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}},
+            ],
+        },
+        {"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]},
+        {"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]},
+    ]
+
+
+@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
+async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None:
+    contents: Final = await gemini_count_contents(
+        ({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}),
+        "openai",
+        "gemini-2.5-flash",
+        custom_llm_provider,
+        {},
+    )
+
+    assert contents == [
+        {"role": "user", "parts": [{"text": "hello"}]},
+        {"role": "model", "parts": [{"text": "hi"}]},
+    ]
+
+
+_PNG_BYTES: Final = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+)
+_DOWNLOAD_SECONDS: Final = 0.5
+
+
+class _SlowPngHandler(BaseHTTPRequestHandler):
+    def do_GET(self) -> None:
+        time.sleep(_DOWNLOAD_SECONDS)
+        self.send_response(200)
+        self.send_header("Content-Type", "image/png")
+        self.send_header("Content-Length", str(len(_PNG_BYTES)))
+        self.end_headers()
+        self.wfile.write(_PNG_BYTES)
+
+    def log_message(self, format: str, *args: object) -> None:
+        return None
+
+
+@contextmanager
+def _slow_png_server() -> Generator[str, None, None]:
+    with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server:
+        thread: Final = Thread(target=server.serve_forever, daemon=True)
+        thread.start()
+        try:
+            yield f"http://127.0.0.1:{server.server_address[1]}"
+        finally:
+            server.shutdown()
+
+
+async def _event_loop_wake_delay(nap_seconds: float) -> float:
+    started: Final = time.monotonic()
+    await asyncio.sleep(nap_seconds)
+    return time.monotonic() - started
+
+
+@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
+async def test_remote_image_downloads_without_blocking_the_event_loop(
+    custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "user_url_validation", False)
+    with _slow_png_server() as base_url:
+        image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png"
+        wake_delay, contents = await asyncio.gather(
+            _event_loop_wake_delay(0.05),
+            gemini_count_contents(
+                (
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": "Describe it"},
+                            {"type": "image", "source": {"type": "url", "url": image_url}},
+                        ],
+                    },
+                ),
+                "anthropic",
+                "gemini-2.5-flash",
+                custom_llm_provider,
+                {},
+            ),
+        )
+
+    assert contents == [
+        {
+            "role": "user",
+            "parts": [
+                {"text": "Describe it"},
+                {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}},
+            ],
+        }
+    ]
+    assert wake_delay < _DOWNLOAD_SECONDS / 2
+
+
+class _MissingImageHandler(BaseHTTPRequestHandler):
+    def do_GET(self) -> None:
+        self.send_response(404)
+        self.send_header("Content-Length", "0")
+        self.end_headers()
+
+    def log_message(self, format: str, *args: object) -> None:
+        return None
+
+
+@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
+async def test_image_that_cannot_be_fetched_returns_an_error_value(
+    custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "user_url_validation", False)
+    with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server:
+        Thread(target=server.serve_forever, daemon=True).start()
+        try:
+            contents = await gemini_count_contents(
+                (
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "image",
+                                "source": {
+                                    "type": "url",
+                                    "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png",
+                                },
+                            }
+                        ],
+                    },
+                ),
+                "anthropic",
+                "gemini-2.5-flash",
+                custom_llm_provider,
+                {},
+            )
+        finally:
+            server.shutdown()
+
+    assert isinstance(contents, VertexGeminiCountContextError)
+
+
+@pytest.mark.parametrize(
+    "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system")
+)
+def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None:
+    view: Final = inference_request_view(
+        (
+            {"role": "user", "content": "first question " * 50},
+            {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]},
+            {"role": "user", "content": "latest question"},
+        ),
+        system,
+        "anthropic",
+    )
+
+    assert view.messages == [{"role": "user", "content": "latest question"}]
+    assert "Summary of the earlier turns." in str(view.system)
+    assert "Be brief." in str(view.system)
+
+
+def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None:
+    snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
+    view: Final = inference_request_view(
+        (
+            {"role": "user", "content": "What is new?"},
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
+                    {
+                        "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": "",
+                                "snippet": snippet,
+                            }
+                        ],
+                    },
+                ],
+            },
+        ),
+        None,
+        "anthropic",
+    )
+
+    assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"])
+    assert [block["type"] for block in assistant_blocks] == ["text"]
+    assert snippet in assistant_blocks[0]["text"]
+
+
+def test_openai_view_is_the_request_unchanged() -> None:
+    messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},)
+
+    view: Final = inference_request_view(messages, "Be brief.", "openai")
+
+    assert view.messages == messages
+    assert view.system == "Be brief."
+
+
+async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None:
+    with patch(
+        "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
+        return_value="image/png",
+    ) as lookup:
+        contents: Final = await gemini_count_contents(
+            (
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}}
+                    ],
+                },
+            ),
+            "anthropic",
+            "gemini-2.5-flash",
+            "vertex_ai",
+            {"vertex_project": "deployment-project"},
+        )
+
+    assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project"
+    assert not isinstance(contents, VertexGeminiCountContextError)
+    assert contents[0]["parts"][0] == {
+        "file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"}
+    }
+
+
+@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
+def test_anthropic_replayed_web_search_result_is_counted_as_its_text(
+    custom_llm_provider: Literal["vertex_ai", "gemini"],
+) -> None:
+    snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
+    messages: Final = (
+        {"role": "user", "content": "What is new?"},
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
+                {
+                    "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": "",
+                            "snippet": snippet,
+                        }
+                    ],
+                },
+            ],
+        },
+        {"role": "user", "content": "Thanks."},
+    )
+
+    view: Final = inference_request_view(messages, None, "anthropic")
+    contents: Final = asyncio.run(
+        gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {})
+    )
+
+    assert snippet in str(contents)

tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py

@@ -1,11 +1,18 @@
-from unittest.mock import patch
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Final
+from unittest.mock import AsyncMock, patch
 
+import httpx
 import pytest
+from pydantic import JsonValue, TypeAdapter
 
+import litellm
 from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
-
-
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
 from litellm.llms.vertex_ai.common_utils import (
+    VertexAITokenCounter,
     _get_vertex_url,
     convert_anyof_null_to_nullable,
     get_vertex_location_from_url,
@@ -16,6 +23,110 @@ from litellm.llms.vertex_ai.common_utils import (
     validate_vertex_location,
     vertex_request_labels_from_litellm_params,
 )
+from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler
+from litellm.types.utils import TokenCountMessagesFormat
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("messages_format", ("openai", "anthropic"))
+async def test_vertex_gemini_rejects_non_text_system_blocks_before_request(
+    messages_format: TokenCountMessagesFormat,
+) -> None:
+    result: Final = await VertexAITokenCounter().count_tokens(
+        model_to_use="gemini-4-argon",
+        messages=[{"role": "user", "content": "Count this"}],
+        contents=None,
+        request_model="vertex_ai/gemini-4-argon",
+        system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}],
+        messages_format=messages_format,
+    )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 400
+    assert "text blocks" in (result.error_message or "")
+
+
+class _VertexCountTokensPeer:
+    def __init__(self, status_code: int) -> None:
+        self.status_code: Final = status_code
+        self.requests: tuple[httpx.Request, ...] = ()
+
+    def __call__(self, request: httpx.Request) -> httpx.Response:
+        self.requests = (*self.requests, request)
+        if self.status_code == 200:
+            return httpx.Response(200, json={"totalTokens": 37}, request=request)
+        return httpx.Response(
+            self.status_code,
+            json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}},
+            request=request,
+        )
+
+
+@contextmanager
+def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]:
+    peer: Final = _VertexCountTokensPeer(status_code)
+    with (
+        patch.object(
+            gemini_count_tokens_handler,
+            "get_async_httpx_client",
+            return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
+        ),
+        patch.object(
+            vertex_count_tokens_handler.VertexAITokenCounter,
+            "_ensure_access_token_async",
+            new=AsyncMock(return_value=("fake-token", "test-project")),
+        ),
+    ):
+        yield peer
+
+
+_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}}
+
+
+@pytest.mark.asyncio
+async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None:
+    with _vertex_count_tokens_peer(400) as peer:
+        result: Final = await VertexAITokenCounter().count_tokens(
+            model_to_use="gemini-2.5-flash",
+            messages=[{"role": "user", "content": "Hello"}],
+            contents=None,
+            deployment=_VERTEX_DEPLOYMENT,
+            request_model="vertex_ai/gemini-2.5-flash",
+        )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 400
+    assert "Invalid function name" in (result.error_message or "")
+    assert len(peer.requests) == 1
+
+
+@pytest.mark.asyncio
+async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "user_url_validation", True)
+    with _vertex_count_tokens_peer(200) as peer:
+        result: Final = await VertexAITokenCounter().count_tokens(
+            model_to_use="gemini-2.5-flash",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}],
+                }
+            ],
+            contents=None,
+            deployment=_VERTEX_DEPLOYMENT,
+            request_model="vertex_ai/gemini-2.5-flash",
+            messages_format="anthropic",
+        )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 400
+    assert "http://127.0.0.1:9/image.png" in (result.error_message or "")
+    assert peer.requests == ()
 
 
 @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"])
@@ -149,15 +260,11 @@ async def test_get_supports_system_message():
     from litellm.llms.vertex_ai.common_utils import get_supports_system_message
 
     # fine-tuned vertex gemini models will specifiy they are in the /gemini spec format
-    result = get_supports_system_message(
-        model="gemini/1234567890", custom_llm_provider="vertex_ai"
-    )
+    result = get_supports_system_message(model="gemini/1234567890", custom_llm_provider="vertex_ai")
     assert result == True
 
     # non-fine-tuned vertex gemini models will not specifiy they are in the /gemini spec format
-    result = get_supports_system_message(
-        model="random-model-name", custom_llm_provider="vertex_ai"
-    )
+    result = get_supports_system_message(model="random-model-name", custom_llm_provider="vertex_ai")
     assert result == False
 
 
@@ -230,13 +337,9 @@ def test_build_vertex_schema():
                 "properties": {
                     "tags": {"items": {"type": "string"}, "type": "array"},
                     "metadata": {"type": "object"},
-                    "callbacks": {
-                        "anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]
-                    },
+                    "callbacks": {"anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]},
                     "run_name": {"type": "string"},
-                    "max_concurrency": {
-                        "anyOf": [{"type": "integer"}, {"type": "null"}]
-                    },
+                    "max_concurrency": {"anyOf": [{"type": "integer"}, {"type": "null"}]},
                     "recursion_limit": {"type": "integer"},
                     "configurable": {"type": "object"},
                     "run_id": {
@@ -280,9 +383,7 @@ def test_build_vertex_schema():
                         ]
                     },
                     "run_name": {"type": "string"},
-                    "max_concurrency": {
-                        "anyOf": [{"type": "integer", "nullable": True}]
-                    },
+                    "max_concurrency": {"anyOf": [{"type": "integer", "nullable": True}]},
                     "recursion_limit": {"type": "integer"},
                     "configurable": {"type": "object"},
                     "run_id": {"anyOf": [{"type": "string", "nullable": True}]},
@@ -689,9 +790,7 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix):
         ({}, "us-east1", "us-east1"),
     ],
 )
-def test_get_vertex_region_global_only_model(
-    model_cost_entry, vertex_region, expected_region
-):
+def test_get_vertex_region_global_only_model(model_cost_entry, vertex_region, expected_region):
     """Test get_vertex_region resolves region from model_cost supported_regions"""
     import litellm
     from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
@@ -703,9 +802,7 @@ def test_get_vertex_region_global_only_model(
         {"vertex_ai/test-model": model_cost_entry},
         clear=False,
     ):
-        result = vertex_base.get_vertex_region(
-            vertex_region=vertex_region, model="test-model"
-        )
+        result = vertex_base.get_vertex_region(vertex_region=vertex_region, model="test-model")
 
         assert result == expected_region
 
@@ -823,9 +920,7 @@ def test_convert_schema_types_type_array_conversion():
     assert anyof_types[1]["type"] == "number"
 
     # 4. Other properties preserved
-    assert (
-        input_schema["properties"]["studio"]["description"] == "The studio ID or name"
-    )
+    assert input_schema["properties"]["studio"]["description"] == "The studio ID or name"
     assert input_schema["required"] == ["studio"]
 
 
@@ -992,7 +1087,9 @@ def test_construct_target_url_with_version_prefix():
         ),
     ],
 )
-def test_construct_target_url_versionless_project_route_gets_api_version(requested_route: str, expected_url: str) -> None:
+def test_construct_target_url_versionless_project_route_gets_api_version(
+    requested_route: str, expected_url: str
+) -> None:
     from litellm.llms.vertex_ai.common_utils import construct_target_url
 
     target_url = construct_target_url(
@@ -1125,10 +1222,7 @@ def test_fix_enum_types():
     # 2. Non-string enums are removed
     assert "enum" not in input_schema["properties"]["maxLength"]
     assert "enum" not in input_schema["properties"]["enabled"]
-    assert (
-        "enum"
-        not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"]
-    )
+    assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"]
 
     # 3. anyOf with string type keeps enum, non-string removes it
     assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0]
@@ -1182,8 +1276,6 @@ def test_get_token_url():
     assert "v1beta1" not in url
     assert "/v1/" in url
 
-    pass
-
 
 @pytest.mark.asyncio
 async def test_vertex_ai_token_counter_routes_partner_models():
@@ -1229,6 +1321,149 @@ async def test_vertex_ai_token_counter_routes_partner_models():
         assert result.tokenizer_type == "vertex_ai_partner_models"
 
 
+@pytest.mark.asyncio
+async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_request():
+    from typing import Final
+    from unittest.mock import AsyncMock, patch
+
+    from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter
+    from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler
+
+    class FakeResponse:
+        status_code = 200
+
+        def json(self) -> dict[str, int]:
+            return {"input_tokens": 37}
+
+    class FakeHttpClient:
+        posted_bodies: tuple[dict[str, object], ...] = ()
+
+        async def post(
+            self,
+            url: str,
+            headers: dict[str, str],
+            json: dict[str, object],
+            timeout: float,
+        ) -> FakeResponse:
+            self.posted_bodies = (*self.posted_bodies, json)
+            return FakeResponse()
+
+    fake_http_client: Final = FakeHttpClient()
+    counter: Final = VertexAITokenCounter()
+    model: Final = "claude-opus-5-5"
+    messages: Final = [{"role": "user", "content": "Hello"}]
+    system: Final = "Follow the system instructions"
+    tools: Final = [
+        {
+            "name": "lookup",
+            "description": "Look up a value",
+            "input_schema": {"type": "object", "properties": {}},
+        }
+    ]
+    deployment: Final = {
+        "litellm_params": {
+            "vertex_project": "test-project",
+            "vertex_location": "us-east5",
+        }
+    }
+
+    with (
+        patch.object(handler, "get_async_httpx_client", return_value=fake_http_client),
+        patch.object(
+            handler.VertexAIPartnerModelsTokenCounter,
+            "_ensure_access_token_async",
+            new=AsyncMock(return_value=("fake-token", "test-project")),
+        ),
+    ):
+        with_optional_fields: Final = await counter.count_tokens(
+            model_to_use=model,
+            messages=messages,
+            contents=None,
+            deployment=deployment,
+            system=system,
+            tools=tools,
+        )
+        without_optional_fields: Final = await counter.count_tokens(
+            model_to_use=model,
+            messages=messages,
+            contents=None,
+            deployment=deployment,
+        )
+
+    assert fake_http_client.posted_bodies == (
+        {"model": model, "messages": messages, "system": system, "tools": tools},
+        {"model": model, "messages": messages},
+    )
+    assert with_optional_fields is not None
+    assert with_optional_fields.total_tokens == 37
+    assert without_optional_fields is not None
+    assert without_optional_fields.total_tokens == 37
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("provider_failure", "expected_status", "expected_message"),
+    [
+        ("http_400", 400, 'tools.0: Input tag "function" does not match any of the expected tags'),
+        ("credentials", 500, "could not resolve credentials"),
+    ],
+)
+async def test_vertex_ai_token_counter_returns_partner_provider_error_as_value(
+    provider_failure: str, expected_status: int, expected_message: str
+):
+    from typing import Final
+    from unittest.mock import AsyncMock, patch
+
+    import httpx
+
+    from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
+    from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter
+    from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler
+
+    class RejectingHttpClient:
+        async def post(
+            self,
+            url: str,
+            headers: dict[str, str],
+            json: dict[str, object],
+            timeout: float,
+        ) -> None:
+            request: Final = httpx.Request("POST", url)
+            response: Final = httpx.Response(400, text=expected_message, request=request)
+            raise MaskedHTTPStatusError(
+                httpx.HTTPStatusError("Client error '400 Bad Request'", request=request, response=response),
+                message=expected_message,
+                text=expected_message,
+            )
+
+    access_token: Final = (
+        AsyncMock(side_effect=ValueError(expected_message))
+        if provider_failure == "credentials"
+        else AsyncMock(return_value=("fake-token", "test-project"))
+    )
+    with (
+        patch.object(handler, "get_async_httpx_client", return_value=RejectingHttpClient()),
+        patch.object(handler.VertexAIPartnerModelsTokenCounter, "_ensure_access_token_async", new=access_token),
+    ):
+        result: Final = await VertexAITokenCounter().count_tokens(
+            model_to_use="claude-opus-5-5",
+            messages=[{"role": "user", "content": "Hello"}],
+            contents=None,
+            deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}},
+            request_model="vertex-claude",
+            tools=[{"type": "function", "function": {"name": "lookup", "parameters": {}}}],
+        )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == expected_status
+    assert result.error_message is not None
+    assert expected_message in result.error_message
+    assert result.total_tokens == 0
+    assert result.request_model == "vertex-claude"
+    assert result.tokenizer_type == "vertex_ai_partner_models"
+
+
 @pytest.mark.asyncio
 async def test_vertex_ai_token_counter_uses_count_tokens_location():
     """
@@ -1318,6 +1553,49 @@ async def test_vertex_ai_token_counter_routes_gemini_models():
         assert result.total_tokens == 50
 
 
+@pytest.mark.asyncio
+async def test_vertex_gemini_count_tokens_uses_request_system_over_deployment_system_instruction() -> None:
+    class VertexCountTokensPeer:
+        def __init__(self) -> None:
+            self.request_body: dict[str, JsonValue] | None = None
+
+        def __call__(self, request: httpx.Request) -> httpx.Response:
+            self.request_body = TypeAdapter(dict[str, JsonValue]).validate_json(request.content)
+            return httpx.Response(200, json={"totalTokens": 37}, request=request)
+
+    peer: Final = VertexCountTokensPeer()
+    client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
+
+    with (
+        patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client),
+        patch.object(
+            vertex_count_tokens_handler.VertexAITokenCounter,
+            "_ensure_access_token_async",
+            new=AsyncMock(return_value=("fake-token", "test-project")),
+        ),
+    ):
+        result: Final = await VertexAITokenCounter().count_tokens(
+            model_to_use="gemini-2.5-flash",
+            messages=[{"role": "user", "content": "Hello"}],
+            contents=None,
+            deployment={
+                "litellm_params": {
+                    "vertex_project": "test-project",
+                    "vertex_location": "us-central1",
+                    "system_instruction": {"parts": [{"text": "deployment-level"}]},
+                }
+            },
+            request_model="vertex_ai/gemini-2.5-flash",
+            system="You are terse.",
+        )
+
+    assert result is not None
+    assert result.error is False
+    assert result.total_tokens == 37
+    assert peer.request_body is not None
+    assert peer.request_body["systemInstruction"] == {"parts": [{"text": "You are terse."}]}
+
+
 @pytest.mark.asyncio
 async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini():
     """
@@ -1332,9 +1610,7 @@ async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini(
 
     token_counter = VertexAITokenCounter()
 
-    with patch(
-        "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens"
-    ) as mock_acount_tokens:
+    with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens:
         mock_acount_tokens.return_value = {
             "totalTokens": 42,
             "tokenizer_used": "gemini",
@@ -1376,9 +1652,7 @@ async def test_vertex_ai_token_counter_returns_none_when_api_omits_total_tokens(
 
     token_counter = VertexAITokenCounter()
 
-    with patch(
-        "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens"
-    ) as mock_acount_tokens:
+    with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens:
         mock_acount_tokens.return_value = {"tokenizer_used": "gemini"}
 
         result = await token_counter.count_tokens(
@@ -1421,9 +1695,7 @@ async def test_vertex_ai_partner_model_detection():
     # Test Minimax models
     assert VertexAIPartnerModels.is_vertex_partner_model("minimaxai/minimax-m2-maas")
     # Test Moonshot models
-    assert VertexAIPartnerModels.is_vertex_partner_model(
-        "moonshotai/kimi-k2-thinking-maas"
-    )
+    assert VertexAIPartnerModels.is_vertex_partner_model("moonshotai/kimi-k2-thinking-maas")
 
     # Test Gemini models (should NOT be detected as partner model)
     assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-1.5-pro")
@@ -1454,9 +1726,7 @@ def test_vertex_ai_moonshot_uses_openai_handler():
         VertexAIPartnerModels,
     )
 
-    assert VertexAIPartnerModels.should_use_openai_handler(
-        "moonshotai/kimi-k2-thinking-maas"
-    )
+    assert VertexAIPartnerModels.should_use_openai_handler("moonshotai/kimi-k2-thinking-maas")
 
 
 def test_vertex_ai_zai_uses_openai_handler():
@@ -1491,9 +1761,7 @@ def test_vertex_ai_gemma_maas_is_partner_model():
         VertexAIPartnerModels,
     )
 
-    assert VertexAIPartnerModels.is_vertex_partner_model(
-        "google/gemma-4-26b-a4b-it-maas"
-    )
+    assert VertexAIPartnerModels.is_vertex_partner_model("google/gemma-4-26b-a4b-it-maas")
 
 
 def test_vertex_ai_gemma_maas_uses_openai_handler():
@@ -1504,9 +1772,7 @@ def test_vertex_ai_gemma_maas_uses_openai_handler():
         VertexAIPartnerModels,
     )
 
-    assert VertexAIPartnerModels.should_use_openai_handler(
-        "google/gemma-4-26b-a4b-it-maas"
-    )
+    assert VertexAIPartnerModels.should_use_openai_handler("google/gemma-4-26b-a4b-it-maas")
 
 
 def test_vertex_ai_gemma_maas_routes_to_partner_models():
@@ -1708,12 +1974,8 @@ def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata():
 
 def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent():
     optional: dict = {}
-    litellm_params = {
-        "litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}
-    }
-    assert pop_vertex_request_labels(optional, litellm_params) == {
-        "team": "from_litellm_meta"
-    }
+    litellm_params = {"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}}
+    assert pop_vertex_request_labels(optional, litellm_params) == {"team": "from_litellm_meta"}
 
 
 def test_vertex_text_embedding_request_includes_labels_from_metadata():
@@ -1723,9 +1985,7 @@ def test_vertex_text_embedding_request_includes_labels_from_metadata():
         input="hi",
         optional_params={},
         model="text-embedding-004",
-        litellm_params={
-            "metadata": {"requester_metadata": {"project_id": "cost-center-1"}}
-        },
+        litellm_params={"metadata": {"requester_metadata": {"project_id": "cost-center-1"}}},
     )
     assert req.get("labels") == {"project_id": "cost-center-1"}
 

tests/unit/proxy/common_utils/test_token_count_settings.py

@@ -0,0 +1,190 @@
+import json
+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"},
+}
+
+
+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
+
+
+_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_gemini_proxy_request.py

@@ -9,11 +9,12 @@ The request payload is correctly processed and forwarded to the httpx client.
 
 import json
 import unittest.mock
-from typing import Optional
+from typing import Final, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
+from pydantic import JsonValue, ValidationError
 
 # Add the parent directory to the system path
 
@@ -493,3 +494,39 @@ async def test_gemini_proxy_config_with_custom_api_base():
 if __name__ == "__main__":
     # Run the tests
     pytest.main([__file__, "-v"])
+
+
+@pytest.mark.parametrize("system_field", ("systemInstruction", "system_instruction"))
+@pytest.mark.parametrize("parts", ({"text": "Follow these rules"}, [{"text": "Follow these rules"}]))
+@pytest.mark.parametrize("tools", (None, []))
+def test_google_count_tokens_accepts_equivalent_system_parts_and_no_tools(
+    system_field: str, parts: JsonValue, tools: JsonValue
+) -> None:
+    from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
+
+    contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}]
+    request: Final = _google_count_tokens_request(
+        "counted", {"contents": contents, system_field: {"parts": parts}, "tools": tools}
+    )
+
+    assert request.messages == [{"role": "user", "content": "Count this request"}]
+    assert request.system == [{"type": "text", "text": "Follow these rules"}]
+    assert request.tools is None
+    assert request.contents == contents
+
+
+@pytest.mark.parametrize(
+    "parts",
+    (
+        {"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}},
+        [{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}],
+    ),
+)
+def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None:
+    from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
+
+    with pytest.raises(ValidationError):
+        _google_count_tokens_request(
+            "counted",
+            {"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}},
+        )

tests/unit/test_count_tokens_public_api.py

@@ -97,7 +97,7 @@ def test_acount_tokens_with_tools():
         assert result.total_tokens == 30
         mock_handler.assert_called_once()
         call_kwargs = mock_handler.call_args
-        assert call_kwargs.kwargs.get("tools") == tools
+        assert list(call_kwargs.kwargs.get("tools") or ()) == tools
 
 
 def test_acount_tokens_with_system():

tests/unit/test_model_prices_schema.py

@@ -176,6 +176,52 @@ def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict):
     assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}})
 
 
+TOKEN_COUNT_ENTRY: Final = MappingProxyType(
+    {
+        "litellm_provider": "bedrock_converse",
+        "mode": "chat",
+        "token_count_framing": 78,
+        "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"},
+    }
+)
+
+
+@pytest.mark.parametrize(
+    "setting",
+    [
+        {"token_count_framing": 0},
+        {"token_count_framing": 78},
+        {"token_count_model": {"model": "us.anthropic.claude-opus-5-5"}},
+    ],
+)
+def test_generator_classifies_token_count_settings(setting: dict):
+    generator = load_generator()
+    schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)})))
+    assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}})
+
+
+@pytest.mark.parametrize(
+    "setting",
+    [
+        {"token_count_framing": -1},
+        {"token_count_framing": "78"},
+        {"token_count_framing": 1.5},
+        {"token_count_model": {"model": ""}},
+        {"token_count_model": {"model": "   "}},
+        {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": ""}},
+        {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "   "}},
+        {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": 5}},
+        {"token_count_model": {"region": "us-east-1"}},
+        {"token_count_model": "us.anthropic.claude-opus-5-5"},
+        {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "zone": "us-east-1"}},
+    ],
+)
+def test_generated_token_count_schema_rejects_malformed_settings(setting: dict):
+    generator = load_generator()
+    schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)})))
+    assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}})
+
+
 def find_duplicate_keys(path: Path) -> list[str]:
     duplicates: list[str] = []
 

tests/unit/test_router_model_cost_isolation.py

@@ -513,6 +513,40 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing():
     )
 
 
+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.moonshotai.kimi-k3"
+    deploy_id: Final = "kimi-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": "kimi-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

tests/unit/test_utils.py

@@ -936,6 +936,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
                 "supports_video_input": {"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_url_context": {"type": "boolean"},
                 "supports_multimodal": {"type": "boolean"},
                 "uses_embed_content": {"type": "boolean"},

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

@@ -43384,9 +43384,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 */
@@ -43415,6 +43423,8 @@ export interface components {
         };
         /** TokenCountResponse */
         TokenCountResponse: {
+            /** Count Source */
+            count_source?: ("provider" | "token_count_model" | "local_estimate") | null;
             /**
              * Error
              * @default false

uv.lock

@@ -4500,7 +4500,7 @@ wheels = [
 
 [[package]]
 name = "litellm"
-version = "1.104.3"
+version = "1.104.4"
 source = { editable = "." }
 dependencies = [
     { name = "aiohttp" },