fix(token_counter): count SDK tool call objects and flat custom tools, keep Anthropic options tools as functions

Reused ChatCompletionMessageToolCall and custom tool call objects are read
through their get and contains methods, a null custom or function payload maps
to an empty call on the provider count, flat Responses-shaped custom tools are
counted like the chat shape, and Anthropic tools that carry a custom options
object stay function tools.

litellm/litellm_core_utils/token_counter.py

@@ -47,7 +47,7 @@ from litellm.constants import (
     TOKEN_COUNTER_MAX_EXACT_CHARS,
 )
 from litellm.litellm_core_utils.asyncify import asyncify
-from litellm.litellm_core_utils.core_helpers import as_str_mapping
+from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes
 from litellm.litellm_core_utils.prompt_templates.chat_template_tools import (
     ChatTemplateToolSection,
     chat_template_tool_section,
@@ -893,7 +893,9 @@ def _count_function_call_tokens(
             raise ValueError(f"Unsupported type {type(value)} for key tool_calls in message {message}")
         total = 0
         for tool_call in value:
-            if (fields := as_str_mapping(tool_call)) is None or ("function" not in fields and "custom" not in fields):
+            if (fields := _tool_call_fields(tool_call)) is None or (
+                "function" not in fields and "custom" not in fields
+            ):
                 raise ValueError(f"Unsupported tool call {tool_call} must contain a function or custom key")
             total += count_function(_tool_call_text(fields))
         return total
@@ -904,9 +906,19 @@ def _count_function_call_tokens(
     raise ValueError(f"Unexpected key {key!r}; expected 'tool_calls' or 'function_call'")
 
 
-def _tool_call_text(tool_call: Mapping[str, object]) -> str:
-    is_custom: Final = is_custom_tool_call_dict(tool_call)
-    payload: Final = as_str_mapping(tool_call.get("custom" if is_custom else "function"))
+class _ToolCallFields(Protocol):
+    def get(self, key: str, default: object = None, /) -> object: ...
+
+    def __contains__(self, key: object, /) -> bool: ...
+
+
+def _tool_call_fields(value: object) -> _ToolCallFields | None:
+    return value if has_attributes(value, _ToolCallFields, "get", "__contains__") else None
+
+
+def _tool_call_text(tool_call: _ToolCallFields) -> str:
+    is_custom: Final = is_custom_tool_call_dict({"type": tool_call.get("type"), "custom": tool_call.get("custom")})
+    payload: Final = _tool_call_fields(tool_call.get("custom" if is_custom else "function"))
     if payload is None:
         return ""
     return str(payload.get("input" if is_custom else "arguments", ""))
@@ -1812,8 +1824,8 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi
         yield function
         return
     custom: Final = as_str_mapping(tool.get("custom"))
-    if custom is not None:
-        yield _custom_tool_definition(custom)
+    if "input_schema" not in tool and (custom is not None or tool.get("type") == "custom"):
+        yield _custom_tool_definition(custom if custom is not None else tool)
         return
     declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations")
     if isinstance(declarations, list):

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

@@ -99,14 +99,14 @@ def _assistant_tool_calls(messages: Sequence[Mapping[str, object]]) -> Iterable[
 def _responses_tool_call_item(tool_call: _FieldReader) -> Mapping[str, object]:
     call_id: Final = tool_call.get("id", "")
     if _is_custom_tool_call(tool_call):
-        custom: Final = _field_reader(tool_call.get("custom"))
+        custom: Final = _field_reader(tool_call.get("custom") or {})
         return {
             "type": "custom_tool_call",
             "call_id": call_id,
             "name": custom.get("name", ""),
             "input": custom.get("input", ""),
         }
-    function: Final = _field_reader(tool_call.get("function", {}))
+    function: Final = _field_reader(tool_call.get("function") or {})
     return {
         "type": "function_call",
         "call_id": call_id,
@@ -283,8 +283,6 @@ class OpenAICountTokensConfig:
                 if "strict" in func:
                     item["strict"] = func["strict"]
                 transformed.append(item)
-            elif tool.get("type") == "custom" and (custom := as_str_mapping(tool.get("custom"))) is not None:
-                transformed.append(dict(_responses_custom_tool(custom)))
             elif "input_schema" in tool and "name" in tool and tool.get("type") in (None, "custom"):
                 transformed.append(
                     {
@@ -295,6 +293,8 @@ class OpenAICountTokensConfig:
                         "strict": bool(tool.get("strict")),
                     }
                 )
+            elif tool.get("type") == "custom" and (custom := as_str_mapping(tool.get("custom"))) is not None:
+                transformed.append(dict(_responses_custom_tool(custom)))
             else:
                 # Pass through non-function tools (e.g., web_search, file_search)
                 transformed.append(tool)

tests/unit/litellm_core_utils/test_token_counter.py

@@ -50,7 +50,13 @@ from litellm.litellm_core_utils.token_counter import (
 from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
 from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
-from litellm.types.utils import SelectTokenizerResponse
+from litellm.types.utils import (
+    ChatCompletionCustomToolCallPayload,
+    ChatCompletionMessageCustomToolCall,
+    ChatCompletionMessageToolCall,
+    Function,
+    SelectTokenizerResponse,
+)
 from tests.large_text import text
 from tests.unit.litellm_core_utils.event_loop_lag import (
     assert_loop_stayed_free,
@@ -385,7 +391,7 @@ def test_token_counter_legacy_function_call_counts_arguments():
     assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}"
 
 
-def _tool_call_turn(tool_call: dict[str, object]) -> list[dict[str, object]]:
+def _tool_call_turn(tool_call: object) -> list[dict[str, object]]:
     return [
         {"role": "user", "content": "run the code"},
         {"role": "assistant", "content": None, "tool_calls": [tool_call]},
@@ -420,6 +426,37 @@ def test_token_counter_custom_tool_call_counts_input_like_function_arguments(cus
     assert custom_tokens > 300, f"custom tool call input must contribute real tokens, got {custom_tokens}"
 
 
+@pytest.mark.parametrize(
+    ("sdk_call", "dict_call"),
+    (
+        (
+            ChatCompletionMessageToolCall(
+                id="call_1", type="function", function=Function(name="code_exec", arguments=_CUSTOM_CODE)
+            ),
+            {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}},
+        ),
+        (
+            ChatCompletionMessageCustomToolCall(
+                id="call_1", custom=ChatCompletionCustomToolCallPayload(name="code_exec", input=_CUSTOM_CODE)
+            ),
+            {"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
+        ),
+        (
+            {"id": "call_1", "type": "function", "function": Function(name="code_exec", arguments=_CUSTOM_CODE)},
+            {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}},
+        ),
+    ),
+    ids=("function", "custom", "dict-holding-function-object"),
+)
+def test_token_counter_counts_reused_sdk_tool_call_objects_like_their_dicts(
+    sdk_call: object, dict_call: dict[str, object]
+):
+    sdk_tokens = token_counter(model="gpt-4o", messages=_tool_call_turn(sdk_call))
+
+    assert sdk_tokens == token_counter(model="gpt-4o", messages=_tool_call_turn(dict_call))
+    assert sdk_tokens > 300, f"a reused SDK tool call must contribute its payload tokens, got {sdk_tokens}"
+
+
 @pytest.mark.parametrize(
     ("null_call", "empty_call"),
     (
@@ -471,6 +508,29 @@ def test_token_counter_counts_custom_tool_definition_description_and_grammar():
     assert described - bare > 300, f"description and grammar must count, got {described} vs {bare} bare"
 
 
+def test_token_counter_counts_responses_shaped_custom_tool_like_the_chat_shape():
+    messages = [{"role": "user", "content": "run the code"}]
+    grammar = "start: " + " | ".join(f'"command_{index}"' for index in range(100))
+    chat_tool = {
+        "type": "custom",
+        "custom": {
+            "name": "code_exec",
+            "description": "Runs one shell command",
+            "format": {"type": "grammar", "grammar": {"syntax": "lark", "definition": grammar}},
+        },
+    }
+    responses_tool = {
+        "type": "custom",
+        "name": "code_exec",
+        "description": "Runs one shell command",
+        "format": {"type": "grammar", "syntax": "lark", "definition": grammar},
+    }
+
+    assert token_counter(model="gpt-4o", messages=messages, tools=[responses_tool]) == token_counter(
+        model="gpt-4o", messages=messages, tools=[chat_tool]
+    )
+
+
 @pytest.mark.parametrize(
     "message_count_pair",
     MESSAGES_TEXT,
@@ -830,6 +890,7 @@ def test_token_counter_counts_anthropic_tool_choices_like_openai_choices(
         {**_anthropic_tool(_WEATHER_FUNCTION), "strict": True},
         {**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]},
         {**_anthropic_tool(_WEATHER_FUNCTION), "custom": {"eager_input_streaming": True}},
+        {**_anthropic_tool(_WEATHER_FUNCTION), "type": "custom", "custom": {"eager_input_streaming": True}},
         _anthropic_tool(_WEB_SEARCH_FUNCTION),
     ],
 )

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

@@ -110,8 +110,24 @@ _SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}}  # muta
                 "strict": False,
             },
         ),
+        (
+            {
+                "type": "custom",
+                "name": "lookup",
+                "description": "Find a record",
+                "input_schema": _SCHEMA,
+                "custom": {"eager_input_streaming": True},
+            },
+            {
+                "type": "function",
+                "name": "lookup",
+                "description": "Find a record",
+                "parameters": _SCHEMA,
+                "strict": False,
+            },
+        ),
     ),
-    ids=("anthropic", "anthropic-strict", "anthropic-custom-type"),
+    ids=("anthropic", "anthropic-strict", "anthropic-custom-type", "anthropic-custom-options"),
 )
 def test_transform_anthropic_tool_shapes(tool: dict[str, object], expected: dict[str, object]) -> None:
     result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens(
@@ -252,6 +268,30 @@ def test_messages_to_responses_input_maps_custom_tool_calls_and_their_outputs():
     ]
 
 
+@pytest.mark.parametrize(
+    ("tool_call", "expected"),
+    (
+        (
+            {"id": "call_1", "type": "custom", "custom": None},
+            {"type": "custom_tool_call", "call_id": "call_1", "name": "", "input": ""},
+        ),
+        (
+            {"id": "call_1", "type": "function", "function": None},
+            {"type": "function_call", "call_id": "call_1", "name": "", "arguments": ""},
+        ),
+    ),
+    ids=("custom", "function"),
+)
+def test_messages_to_responses_input_maps_null_tool_call_payloads_to_empty_calls(
+    tool_call: dict[str, object], expected: dict[str, object]
+) -> None:
+    input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(
+        [{"role": "user", "content": "Run it"}, {"role": "assistant", "content": None, "tool_calls": [tool_call]}]
+    )
+
+    assert input_items[1:] == [expected]
+
+
 def test_transform_custom_tool_uses_the_flat_responses_shape():
     tool: Final = {
         "type": "custom",