fix(count_tokens): count system and tools on Vertex Gemini, fall back to Mantle on Bedrock 404 (#46430)

* fix: align token counting with inference context

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

* fix: preserve Vertex Gemini count message shapes

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

* fix: forward Google countTokens tools in shapes every provider counter accepts

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

* test: add package marker for Vertex count-token unit tests

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

* fix(vertex_ai): ignore deployment system_instruction when counting tokens

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

* fix(count_tokens): forward system and tools on Google AI Studio, keep strict on Anthropic tools

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

* style(count_tokens): ruff format gemini count handler

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

* refactor(count_tokens): return Mapping from count request body builders

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

* fix(count_tokens): truncate long Anthropic tool names on Vertex count, update stale strict test expectations

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

* refactor(count_tokens): simplify Gemini count context and tests

Put the Vertex Gemini count branch back in the same shape as main and the partner model branch, keeping only the system and tools context, the provider error mapping and the gemini tokenizer default

Share one error response helper between the Vertex and AI Studio callers, drop the unread kind fields and the extra validation of map_function output, and say Gemini instead of Vertex AI in the error that AI Studio users also see

Tighten the tests by removing a stray parametrize on a helper, sharing the SDK count args, asserting the exact forwarded system and tools, and folding the Anthropic tool shape tests into one table

* fix(count_tokens): reuse inference tool translation for Vertex Gemini count

Vertex and AI Studio count now convert tools with the same functions inference uses, translate_anthropic_tools_to_openai and the Responses tool converter, before VertexGeminiConfig.map_function. Unnamed and Anthropic hosted tools now count the same declarations inference sends instead of being dropped

Tools that cannot be converted, such as a non list tools value, a non object tool or a schema that is not a JSON object, now return a 400 before any request instead of being counted without them

---------

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

litellm/llms/bedrock/count_tokens/bedrock_token_counter.py

@@ -20,6 +20,7 @@ from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCo
 
 RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
 MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
+_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404})
 
 
 class _CountTokensReply(LiteLLMPydanticObjectBase):
@@ -46,14 +47,14 @@ CountTokensOutcome = CountedTokens | CountTokensFailure
 def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
     return (
         isinstance(outcome, CountTokensFailure)
-        and outcome.status_code == 400
+        and outcome.status_code in _RUNTIME_FALLBACK_STATUSES
         and is_mantle_claude_model(resolved_model)
     )
 
 
 class BedrockTokenCounter(BaseTokenCounter):
     """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens
-    on bedrock-mantle for the Claude models bedrock-runtime answers 400 for"""
+    on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for"""
 
     def __init__(
         self,

litellm/llms/gemini/common_utils.py

@@ -499,12 +499,23 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
         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,
+            vertex_gemini_count_context,
+        )
+
+        count_context: Final = vertex_gemini_count_context(system=system, tools=tools)
+        if isinstance(count_context, VertexGeminiCountContextError):
+            return count_context_error_response(count_context, request_model, model_to_use)
 
         deployment = deployment or {}
         count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {}))
         count_tokens_params: Final = {
             "model": model_to_use,
             "contents": 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(

litellm/llms/gemini/count_tokens/handler.py

@@ -6,7 +6,7 @@ from pydantic import ConfigDict, RootModel, with_config
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 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:
@@ -41,12 +41,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:
@@ -67,6 +67,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,
@@ -97,6 +119,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]:
         """
@@ -143,14 +167,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/transformation.py

@@ -157,9 +157,10 @@ class OpenAICountTokensConfig:
         tools: list[dict[str, Any]],
     ) -> list[dict[str, object]]:
         """
-        Transform OpenAI chat tools format to Responses API tools format.
+        Transform OpenAI chat and Anthropic tools to Responses API tools format.
 
         Chat format:  {"type": "function", "function": {"name": "...", "parameters": {...}}}
+        Anthropic format:  {"name": "...", "input_schema": {...}}
         Responses format: {"type": "function", "name": "...", "parameters": {...}}
         """
         transformed: Final[list[dict[str, object]]] = []
@@ -175,6 +176,16 @@ class OpenAICountTokensConfig:
                 if "strict" in func:
                     item["strict"] = func["strict"]
                 transformed.append(item)
+            elif "input_schema" in tool and "name" in tool and tool.get("type") in (None, "custom"):
+                transformed.append(
+                    {
+                        "type": "function",
+                        "name": tool["name"],
+                        "description": tool.get("description", ""),
+                        "parameters": tool["input_schema"],
+                        "strict": bool(tool.get("strict")),
+                    }
+                )
             else:
                 # Pass through non-function tools (e.g., web_search, file_search)
                 transformed.append(tool)

litellm/llms/vertex_ai/common_utils.py

@@ -1343,30 +1343,56 @@ class VertexAITokenCounter(BaseTokenCounter):
                     original_response=result,
                 )
         else:
+            from litellm.llms.vertex_ai.count_tokens.context import (
+                VertexGeminiCountContextError,
+                count_context_error_response,
+                vertex_gemini_count_context,
+            )
             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
             )
 
+            count_context: Final = vertex_gemini_count_context(system=system, tools=tools)
+            if isinstance(count_context, VertexGeminiCountContextError):
+                return count_context_error_response(count_context, request_model, model_to_use)
+
             resolved_contents: Final = (
-                contents if contents is not None else gemini_convert_messages_with_history(messages=messages or [])
+                contents
+                if contents is not None
+                else gemini_convert_messages_with_history(
+                    messages=messages or []  # pyright: ignore[reportArgumentType]  # preserve direct caller message handling
+                )
             )
 
             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 TokenCountResponse(
+                    total_tokens=0,
+                    request_model=request_model,
+                    model_used=model_to_use,
+                    tokenizer_type="gemini",
+                    error=True,
+                    error_message=error.message,
+                    status_code=error.status_code,
+                )
 
-            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,144 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Final
+
+from pydantic import JsonValue, TypeAdapter, ValidationError
+
+from litellm.types.llms.anthropic import AnthropicMessagesTextParam
+from litellm.types.utils import TokenCountResponse
+
+
+@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 _is_responses_function_tool(tool: Mapping[str, JsonValue]) -> bool:
+    return tool.get("type") == "function" and "function" not in tool and "input_schema" not in tool
+
+
+def _chat_tool(tool: Mapping[str, JsonValue]) -> tuple[Mapping[str, object], ...]:
+    from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
+    from litellm.responses.litellm_completion_transformation.transformation import (
+        LiteLLMCompletionResponsesConfig,
+    )
+
+    if _is_responses_function_tool(tool):
+        chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
+            [dict(tool)]  # pyright: ignore[reportArgumentType]  # validated JSON dict in the Responses function tool shape
+        )
+        return tuple(chat_tools)
+    anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(
+        [dict(tool)]  # pyright: ignore[reportArgumentType]  # validated JSON dict in an Anthropic or chat tool shape
+    )
+    return tuple(anthropic_tools)
+
+
+def _map_function_tools(
+    tools: object | None,
+) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError:
+    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_tool(tool)
+    )  # comprehension-ok: flattens each tool's translation
+    try:
+        mapped_tools: Final = VertexGeminiConfig().map_function(
+            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 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,
+    )

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/proxy/google_endpoints/endpoints.py

@@ -1,9 +1,14 @@
+from collections.abc import Mapping
+from itertools import chain
 from typing import Final
 
-from fastapi import APIRouter, Depends, Request, Response
+from fastapi import APIRouter, Depends, HTTPException, Request, Response
 from fastapi.responses import ORJSONResponse
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
 
+from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types
 from litellm.proxy._types import *
+from litellm.proxy._types import 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_utils.http_parsing_utils import (  # noqa: F401  # legacy module exports
@@ -12,11 +17,83 @@ from litellm.proxy.common_utils.http_parsing_utils import (  # noqa: F401  # leg
 )
 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, ...]
+
+
+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, ...] = ()
+
+
 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": {}}
+
+
+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)
+    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)],
@@ -168,23 +245,14 @@ 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
-
-    # 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
-    )
+    try:
+        token_request: Final = _google_count_tokens_request(model_name, data)
+    except ValidationError as error:
+        raise HTTPException(status_code=400, detail="Invalid Google countTokens request body") from error
 
     # Call the internal token counter function with direct request flag set to False
     token_response: Final = await internal_token_counter(

tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py

@@ -456,6 +456,140 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock
         assert mantle.drain() == ()
 
 
+def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
+    gateway: Final = counting_proxy.gateway
+    system: Final = f"{_SYSTEM} status=404"
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url)
+        response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
+
+
+def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: OwnedProxy, mantle_port: int) -> None:
+    gateway: Final = counting_proxy.gateway
+    system: Final = f"{_SYSTEM} status=404"
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url)
+        response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
+        assert response.status_code == 404, response.text
+        assert _REJECTION in response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
+
+
+def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url, _NOVA)
+        body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]}
+        response: Final = _count(gateway, body)
+        assert response.status_code == 404, response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert mantle.drain() == ()
+
+
+def test_google_count_tokens_forwards_system_and_tools_to_mantle(
+    counting_proxy: OwnedProxy, mantle_port: int
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    system_text: Final = "Use the provided tool"
+    tool_schema: Final = {"type": "object", "properties": {"location": {"type": "string"}}}
+    body: Final[dict[str, JsonValue]] = {
+        "contents": [{"role": "user", "parts": [{"text": "status=404"}]}],
+        "systemInstruction": {"parts": [{"text": system_text}]},
+        "tools": [
+            {
+                "functionDeclarations": [
+                    {"name": "get_weather", "description": "Look up weather", "parameters": tool_schema}
+                ]
+            }
+        ],
+    }
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url)
+        response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
+        assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert _mantle_bodies(mantle) == (
+            _mantle_body(
+                messages=[{"role": "user", "content": "status=404"}],
+                system=[{"type": "text", "text": system_text}],
+                tools=[
+                    {
+                        "name": "get_weather",
+                        "description": "Look up weather",
+                        "input_schema": tool_schema,
+                    }
+                ],
+            ),
+        )
+
+
+def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle(
+    counting_proxy: OwnedProxy, mantle_port: int
+) -> None:
+    gateway: Final = counting_proxy.gateway
+    body: Final[dict[str, JsonValue]] = {
+        "contents": [{"role": "user", "parts": [{"text": "status=404"}]}],
+        "tools": [
+            {
+                "functionDeclarations": [
+                    {
+                        "name": "get_weather",
+                        "parameters": {
+                            "type": "OBJECT",
+                            "properties": {"location": {"type": "STRING"}},
+                            "required": ["location"],
+                        },
+                    }
+                ]
+            }
+        ],
+    }
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url)
+        response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
+        assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert _mantle_bodies(mantle) == (
+            _mantle_body(
+                messages=[{"role": "user", "content": "status=404"}],
+                tools=[
+                    {
+                        "name": "get_weather",
+                        "input_schema": {
+                            "type": "object",
+                            "properties": {"location": {"type": "string"}},
+                            "required": ["location"],
+                        },
+                    }
+                ],
+            ),
+        )
+
+
 def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
     gateway: Final = counting_proxy.gateway
     with (

tests/integration/providers/test_openai_count_tokens_tools_wire.py

@@ -0,0 +1,161 @@
+import json
+from collections.abc import Iterator
+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())
+
+
+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,
+            }
+        ]

tests/integration/providers/test_vertex_gemini_count_tokens_wire.py

@@ -0,0 +1,462 @@
+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)
+
+
+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
+        )

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

@@ -93,8 +93,10 @@ async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None:
 @pytest.mark.parametrize(
     ("model_to_use", "runtime"),
     (
+        ("global.anthropic.claude-opus-4-8", (401, {"Message": "not authorized to perform: bedrock:CountTokens"})),
         ("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})),
         ("amazon.nova-pro-v1:0", (400, UNSUPPORTED)),
+        ("amazon.nova-pro-v1:0", (404, UNSUPPORTED)),
     ),
 )
 async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle(
@@ -120,7 +122,7 @@ async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle(
 @pytest.mark.asyncio
 async def test_mantle_failure_is_reported_with_its_status() -> None:
     bedrock: Final = _Bedrock(
-        runtime=(400, UNSUPPORTED),
+        runtime=(404, UNSUPPORTED),
         mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}),
     )
 
@@ -138,3 +140,37 @@ async def test_mantle_failure_is_reported_with_its_status() -> None:
     assert result.tokenizer_type == "bedrock_mantle_api"
     assert "does not exist" in (result.error_message or "")
     assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
+
+
+@pytest.mark.asyncio
+async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None:
+    bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}))
+    tools: Final = [
+        {
+            "name": "lookup",
+            "description": "Find a record",
+            "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
+        }
+    ]
+
+    result: Final = await _counter(bedrock).count_tokens(
+        model_to_use="global.anthropic.claude-opus-4-8",
+        messages=MESSAGES,
+        contents=None,
+        deployment=DEPLOYMENT,
+        request_model="claude-opus-4-8",
+        system="You are a terse assistant.",
+        tools=tools,
+    )
+
+    assert result is not None
+    assert result.error is False
+    assert result.total_tokens == 2177
+    assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
+    mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content)
+    assert mantle_body == {
+        "model": "anthropic.claude-opus-4-8",
+        "messages": MESSAGES,
+        "system": "You are a terse assistant.",
+        "tools": tools,
+    }

tests/unit/llms/gemini/test_gemini_common_utils.py

@@ -1,8 +1,149 @@
+from typing import Final
 from unittest.mock import AsyncMock, patch
 
+import httpx
 import pytest
+from pydantic import JsonValue, TypeAdapter
 
+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
+@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:
@@ -158,7 +299,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):

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

@@ -1,3 +1,4 @@
+from typing import Final
 
 import pytest
 
@@ -79,6 +80,54 @@ def test_transform_includes_tools():
     assert result["tools"] == tools
 
 
+_SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}}  # mutable-ok: matches the tools param type
+
+
+@pytest.mark.parametrize(
+    ("tool", "expected"),
+    (
+        (
+            {"name": "lookup", "input_schema": _SCHEMA},
+            {"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": False},
+        ),
+        (
+            {"name": "lookup", "input_schema": _SCHEMA, "strict": True},
+            {"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": True},
+        ),
+        (
+            {"type": "custom", "name": "lookup", "description": "Find a record", "input_schema": _SCHEMA},
+            {
+                "type": "function",
+                "name": "lookup",
+                "description": "Find a record",
+                "parameters": _SCHEMA,
+                "strict": False,
+            },
+        ),
+    ),
+    ids=("anthropic", "anthropic-strict", "anthropic-custom-type"),
+)
+def test_transform_anthropic_tool_shapes(tool: dict[str, object], expected: dict[str, object]) -> None:
+    result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens(
+        model="gpt-4o", input="Find a record", tools=[tool]
+    )
+
+    assert result["tools"] == [expected]
+
+
+def test_transform_builtin_tool_passes_through():
+    config: Final = OpenAICountTokensConfig()
+    tools: Final = [{"type": "web_search", "search_context_size": "low"}]
+
+    result: Final = config.transform_request_to_count_tokens(
+        model="gpt-4o",
+        input="Find a record",
+        tools=tools,
+    )
+
+    assert result["tools"] == tools
+
+
 def test_transform_no_instructions_no_tools():
     """Test that None values are not included."""
     config = OpenAICountTokensConfig()
@@ -102,9 +151,7 @@ def test_messages_to_responses_input_basic():
         {"role": "user", "content": "How are you?"},
     ]
 
-    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
-        messages
-    )
+    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
 
     assert len(input_items) == 3
     assert input_items[0] == {"role": "user", "content": "Hello"}
@@ -120,9 +167,7 @@ def test_messages_to_responses_input_with_system():
         {"role": "user", "content": "Hello"},
     ]
 
-    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
-        messages
-    )
+    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
 
     assert len(input_items) == 1
     assert input_items[0] == {"role": "user", "content": "Hello"}
@@ -136,9 +181,7 @@ def test_messages_to_responses_input_with_developer():
         {"role": "user", "content": "Hello"},
     ]
 
-    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
-        messages
-    )
+    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
 
     assert len(input_items) == 1
     assert instructions == "Be concise."
@@ -151,9 +194,7 @@ def test_messages_to_responses_input_with_tool():
         {"role": "tool", "content": "72°F", "tool_call_id": "call_123"},
     ]
 
-    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
-        messages
-    )
+    input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
 
     assert len(input_items) == 2
     assert input_items[1] == {
@@ -424,10 +465,7 @@ def test_validate_request_missing_input():
 def test_get_endpoint_default():
     """Test default endpoint URL."""
     config = OpenAICountTokensConfig()
-    assert (
-        config.get_openai_count_tokens_endpoint()
-        == "https://api.openai.com/v1/responses/input_tokens"
-    )
+    assert config.get_openai_count_tokens_endpoint() == "https://api.openai.com/v1/responses/input_tokens"
 
 
 def test_get_endpoint_custom_base():

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

@@ -0,0 +1,195 @@
+from typing import Final
+
+import pytest
+from pydantic import JsonValue
+
+from litellm.llms.anthropic.pass_through.adapters.transformation import (
+    LiteLLMAnthropicMessagesAdapter,
+    truncate_tool_name,
+)
+from litellm.llms.vertex_ai.count_tokens.context import (
+    VertexGeminiCountContext,
+    VertexGeminiCountContextError,
+    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([tool])  # pyright: ignore[reportArgumentType]  # raw Anthropic tool dict
+    inference_tools: Final = VertexGeminiConfig().map_function(
+        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

tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py

@@ -1,11 +1,17 @@
 import time
 from collections.abc import Iterator
-from unittest.mock import patch
+from typing import Final
+from unittest.mock import AsyncMock, patch
 
+import httpx
 import pytest
+from pydantic import JsonValue, TypeAdapter
 
 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,
     convert_anyof_null_to_nullable,
     convert_vertex_datetime_to_openai_datetime,
     get_vertex_location_from_url,
@@ -17,6 +23,23 @@ 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
+
+
+@pytest.mark.asyncio
+async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> 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"}}],
+    )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 400
+    assert "text blocks" in (result.error_message or "")
 
 
 @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"])
@@ -150,15 +173,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
 
 
@@ -231,13 +250,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": {
@@ -281,9 +296,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}]},
@@ -688,9 +701,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
@@ -702,9 +713,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
 
@@ -822,9 +831,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"]
 
 
@@ -991,7 +998,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(
@@ -1124,10 +1133,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]
@@ -1181,8 +1187,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():
@@ -1460,6 +1464,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():
     """
@@ -1474,9 +1521,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",
@@ -1518,9 +1563,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(
@@ -1563,9 +1606,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")
@@ -1596,9 +1637,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():
@@ -1633,9 +1672,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():
@@ -1646,9 +1683,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():
@@ -1838,12 +1873,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():
@@ -1853,9 +1884,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/google_endpoints/test_endpoints.py

@@ -2,15 +2,85 @@
 Test for google_endpoints/endpoints.py
 """
 
+from typing import Final
+
 import pytest
-import sys, os
 from dotenv import load_dotenv
-
-
-from litellm.proxy.google_endpoints.endpoints import google_count_tokens
-from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
+from pydantic import JsonValue
 from starlette.requests import Request
 
+from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens
+from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
+
+
+@pytest.mark.parametrize(
+    ("system_key", "schema_key"),
+    (("systemInstruction", "parameters"), ("system_instruction", "parametersJsonSchema")),
+)
+def test_google_count_tokens_translates_system_and_function_declarations(
+    system_key: str, schema_key: str
+) -> None:
+    schema: Final = {"type": "object", "properties": {"query": {"type": "string"}}}
+    request_body: Final[dict[str, JsonValue]] = {
+        "contents": [{"role": "user", "parts": [{"text": "Look this up"}]}],
+        system_key: {"parts": [{"text": "Use the provided tool."}]},
+        "tools": [
+            {
+                "functionDeclarations": [
+                    {
+                        "name": "lookup",
+                        "description": "Find a record",
+                        schema_key: schema,
+                    }
+                ]
+            }
+        ],
+    }
+    request: Final = _google_count_tokens_request("claude-sonnet", request_body)
+
+    assert request.system == [{"type": "text", "text": "Use the provided tool."}]
+    assert request.tools == [
+        {
+            "name": "lookup",
+            "description": "Find a record",
+            "input_schema": schema,
+        }
+    ]
+    assert request.messages == [{"role": "user", "content": "Look this up"}]
+
+
+@pytest.mark.parametrize(
+    ("schema_key", "expected_schema"),
+    (
+        ("parameters", {"type": "object", "properties": {"query": {"type": "string"}}}),
+        ("parametersJsonSchema", {"type": "OBJECT", "properties": {"query": {"type": "STRING"}}}),
+    ),
+)
+def test_google_count_tokens_normalizes_parameters_only(
+    schema_key: str, expected_schema: dict[str, JsonValue]
+) -> None:
+    request: Final = _google_count_tokens_request(
+        "claude-sonnet",
+        {
+            "contents": [{"role": "user", "parts": [{"text": "Look this up"}]}],
+            "tools": [
+                {
+                    "functionDeclarations": [
+                        {
+                            "name": "lookup",
+                            schema_key: {
+                                "type": "OBJECT",
+                                "properties": {"query": {"type": "STRING"}},
+                            },
+                        }
+                    ]
+                }
+            ],
+        },
+    )
+
+    assert request.tools == [{"name": "lookup", "input_schema": expected_schema}]
+
 load_dotenv()