fix(token_counter): count null and SDK-shaped custom tool calls and map them for provider counts

Pick the tool call payload with is_custom_tool_call_dict so a custom call that carries "function": null counts its input, and count a null payload as an empty call instead of failing the request. Tool calls with neither key still raise.

Send custom tool calls, their results, and custom tool definitions to OpenAI's input_tokens endpoint as custom_tool_call, custom_tool_call_output, and flat custom tools, and count a custom tool definition's description and grammar locally.

litellm/litellm_core_utils/token_counter.py

@@ -35,10 +35,14 @@ 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.prompt_templates.chat_template_tools import (
     ChatTemplateToolSection,
     chat_template_tool_section,
 )
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    convert_custom_tool_format_to_responses_shape,
+)
 from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
 
 if TYPE_CHECKING:
@@ -64,7 +68,7 @@ from litellm.types.llms.openai import (
     ChatCompletionToolParam,
     OpenAIMessageContentListBlock,
 )
-from litellm.types.utils import Message, SelectTokenizerResponse
+from litellm.types.utils import Message, SelectTokenizerResponse, is_custom_tool_call_dict
 
 
 def get_modified_max_tokens(
@@ -570,13 +574,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 "function" in tool_call:
-                kind, field = "function", "arguments"
-            elif "custom" in tool_call:
-                kind, field = "custom", "input"
-            else:
+            if (fields := as_str_mapping(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(str(tool_call[kind].get(field, "")))
+            total += count_function(_tool_call_text(fields))
         return total
     if key == "function_call":
         if not isinstance(value, Mapping):
@@ -585,6 +585,14 @@ 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"))
+    if payload is None:
+        return ""
+    return str(payload.get("input" if is_custom else "arguments", ""))
+
+
 def _count_messages(
     params: _MessageCountParams,
     messages: list[AllMessageValues],
@@ -1149,6 +1157,10 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi
     if isinstance(function, Mapping):
         yield function
         return
+    custom: Final = as_str_mapping(tool.get("custom"))
+    if custom is not None:
+        yield _custom_tool_definition(custom)
+        return
     declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations")
     if isinstance(declarations, list):
         for declaration in declarations:
@@ -1173,6 +1185,18 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi
     }
 
 
+def _custom_tool_definition(custom: Mapping[str, object]) -> Mapping[str, object]:
+    format_obj: Final = as_str_mapping(custom.get("format"))
+    grammar: Final = (
+        convert_custom_tool_format_to_responses_shape(format_obj).get("definition") if format_obj is not None else None
+    )
+    description_parts: Final = tuple(str(part) for part in (custom.get("description"), grammar) if part)
+    return {
+        "name": custom.get("name"),
+        **({"description": "\n".join(description_parts)} if description_parts else {}),
+    }
+
+
 def _format_single_function_definition(function: Mapping[str, object]) -> tuple[str, ...]:
     function_name: Final = function.get("name")
     if not function_name:

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

@@ -5,14 +5,17 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp
 """
 
 import copy
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Final, Literal, Protocol, TypeAlias
 
 from openai.types.responses.input_token_count_params import InputTokenCountParams
 from pydantic import JsonValue, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
-from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes
+from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, require_attributes
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    convert_custom_tool_format_to_responses_shape,
+)
 from litellm.litellm_core_utils.tool_fields import DropPolicy
 from litellm.litellm_core_utils.url_utils import join_api_base
 from litellm.llms.anthropic.messages.convertible_blocks import malformed_tool_result_text_error
@@ -23,7 +26,7 @@ from litellm.llms.anthropic.messages.responses_adapters.transformation import (
 from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema
 from litellm.types.llms.openai import ResponsesAPIRequestParams
-from litellm.types.utils import all_litellm_params
+from litellm.types.utils import all_litellm_params, is_custom_tool_call_dict
 
 
 class ResponsesInputTextPart(TypedDict):
@@ -82,6 +85,47 @@ def _function_fields(value: object) -> _FunctionFields:
     return require_attributes(value, _FunctionFields, "get", "__contains__", "__getitem__")
 
 
+def _is_custom_tool_call(tool_call: _FieldReader) -> bool:
+    return is_custom_tool_call_dict({"type": tool_call.get("type"), "custom": tool_call.get("custom")})
+
+
+def _assistant_tool_calls(messages: Sequence[Mapping[str, object]]) -> Iterable[_FieldReader]:
+    for msg in messages:
+        tool_calls = msg.get("tool_calls")
+        if msg.get("role") == "assistant" and tool_calls:
+            yield from (_field_reader(tool_call) for tool_call in iter_items(tool_calls))
+
+
+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"))
+        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", {}))
+    return {
+        "type": "function_call",
+        "call_id": call_id,
+        "name": function.get("name", ""),
+        "arguments": function.get("arguments", ""),
+    }
+
+
+def _responses_custom_tool(custom: Mapping[str, object]) -> Mapping[str, object]:
+    description: Final = custom.get("description")
+    format_obj: Final = as_str_mapping(custom.get("format"))
+    return {
+        "type": "custom",
+        "name": custom.get("name", ""),
+        **({"description": description} if description is not None else {}),
+        **({"format": convert_custom_tool_format_to_responses_shape(format_obj)} if format_obj is not None else {}),
+    }
+
+
 _ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema)
 
 
@@ -239,6 +283,8 @@ 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(
                     {
@@ -267,6 +313,9 @@ class OpenAICountTokensConfig:
         """
         input_items: Final[list[dict[str, object]]] = []
         instructions_parts: Final[list[str]] = []
+        custom_tool_call_ids: Final = frozenset(
+            tool_call.get("id") for tool_call in _assistant_tool_calls(messages) if _is_custom_tool_call(tool_call)
+        )
 
         for msg in messages:
             role = msg.get("role", "")
@@ -297,26 +346,20 @@ class OpenAICountTokensConfig:
                 if content:
                     input_items.append({"role": "assistant", "content": content})
                 if tool_calls:
-                    for tc in iter_items(tool_calls):
-                        tool_call = _field_reader(tc)
-                        function = tool_call.get("function", {})
-                        call_id = tool_call.get("id", "")
-                        function_fields = _field_reader(function)
-                        input_items.append(
-                            {
-                                "type": "function_call",
-                                "call_id": call_id,
-                                "name": function_fields.get("name", ""),
-                                "arguments": function_fields.get("arguments", ""),
-                            }
-                        )
+                    input_items.extend(
+                        dict(_responses_tool_call_item(_field_reader(tool_call)))
+                        for tool_call in iter_items(tool_calls)
+                    )
                 elif not content:
                     input_items.append({"role": "assistant", "content": content})
             elif role == "tool":
+                tool_call_id = msg.get("tool_call_id", "")
                 input_items.append(
                     {
-                        "type": "function_call_output",
-                        "call_id": msg.get("tool_call_id", ""),
+                        "type": "custom_tool_call_output"
+                        if tool_call_id in custom_tool_call_ids
+                        else "function_call_output",
+                        "call_id": tool_call_id,
                         "output": content if isinstance(content, str) else str(content),
                     }
                 )

tests/unit/litellm_core_utils/test_token_counter.py

@@ -377,24 +377,32 @@ 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 test_token_counter_custom_tool_call_counts_input_like_function_arguments():
-    code = "print('hello world from a custom tool')\n" * 50
+def _tool_call_turn(tool_call: dict[str, object]) -> list[dict[str, object]]:
+    return [
+        {"role": "user", "content": "run the code"},
+        {"role": "assistant", "content": None, "tool_calls": [tool_call]},
+        {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+    ]
 
-    def assistant_turn(tool_call: dict[str, object]) -> list[dict[str, object]]:
-        return [
-            {"role": "user", "content": "run the code"},
-            {"role": "assistant", "content": None, "tool_calls": [tool_call]},
-            {"role": "tool", "tool_call_id": "call_1", "content": "done"},
-        ]
 
-    custom_tokens = token_counter(
-        model="gpt-4o",
-        messages=assistant_turn({"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": code}}),
-    )
+_CUSTOM_CODE = "print('hello world from a custom tool')\n" * 50
+
+
+@pytest.mark.parametrize(
+    "custom_call",
+    (
+        {"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
+        {"id": "call_1", "type": "custom", "function": None, "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
+        {"id": "call_1", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
+    ),
+    ids=("typed", "sdk-shaped-null-function", "typeless"),
+)
+def test_token_counter_custom_tool_call_counts_input_like_function_arguments(custom_call: dict[str, object]):
+    custom_tokens = token_counter(model="gpt-4o", messages=_tool_call_turn(custom_call))
     function_tokens = token_counter(
         model="gpt-4o",
-        messages=assistant_turn(
-            {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": code}}
+        messages=_tool_call_turn(
+            {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}}
         ),
     )
 
@@ -404,6 +412,57 @@ def test_token_counter_custom_tool_call_counts_input_like_function_arguments():
     assert custom_tokens > 300, f"custom tool call input must contribute real tokens, got {custom_tokens}"
 
 
+@pytest.mark.parametrize(
+    ("null_call", "empty_call"),
+    (
+        (
+            {"id": "call_1", "type": "custom", "custom": None},
+            {"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": ""}},
+        ),
+        (
+            {"id": "call_1", "type": "function", "function": None},
+            {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": ""}},
+        ),
+    ),
+    ids=("custom", "function"),
+)
+def test_token_counter_null_tool_call_payload_counts_as_empty_call(
+    null_call: dict[str, object], empty_call: dict[str, object]
+):
+    assert token_counter(model="gpt-4o", messages=_tool_call_turn(null_call)) == token_counter(
+        model="gpt-4o", messages=_tool_call_turn(empty_call)
+    )
+
+
+def test_token_counter_rejects_tool_call_without_function_or_custom_payload():
+    with pytest.raises(ValueError, match="must contain a function or custom key"):
+        token_counter(model="gpt-4o", messages=_tool_call_turn({"id": "call_1", "type": "function"}))
+
+
+def test_token_counter_counts_custom_tool_definition_description_and_grammar():
+    messages = [{"role": "user", "content": "run the code"}]
+    grammar = "start: " + " | ".join(f'"command_{index}"' for index in range(100))
+
+    def custom_tool(**fields: object) -> dict[str, object]:
+        return {"type": "custom", "custom": {"name": "code_exec", **fields}}
+
+    no_tools = token_counter(model="gpt-4o", messages=messages)
+    bare = token_counter(model="gpt-4o", messages=messages, tools=[custom_tool()])
+    described = token_counter(
+        model="gpt-4o",
+        messages=messages,
+        tools=[
+            custom_tool(
+                description="Runs one shell command",
+                format={"type": "grammar", "grammar": {"syntax": "lark", "definition": grammar}},
+            )
+        ],
+    )
+
+    assert bare > no_tools, f"a custom tool's name must count, got {bare} vs {no_tools} without tools"
+    assert described - bare > 300, f"description and grammar must count, got {described} vs {bare} bare"
+
+
 @pytest.mark.parametrize(
     "message_count_pair",
     MESSAGES_TEXT,

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

@@ -225,6 +225,55 @@ def test_messages_to_responses_input_reads_a_tool_call_that_forwards_its_fields(
     }
 
 
+def test_messages_to_responses_input_maps_custom_tool_calls_and_their_outputs():
+    messages: Final = [
+        {"role": "user", "content": "Run it"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {"id": "call_custom", "type": "custom", "custom": {"name": "code_exec", "input": "print(2+3)"}},
+                {"id": "call_fn", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}},
+            ],
+        },
+        {"role": "tool", "tool_call_id": "call_custom", "content": "5"},
+        {"role": "tool", "tool_call_id": "call_fn", "content": "sunny"},
+    ]
+
+    input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages)
+
+    assert input_items[1:] == [
+        {"type": "custom_tool_call", "call_id": "call_custom", "name": "code_exec", "input": "print(2+3)"},
+        {"type": "function_call", "call_id": "call_fn", "name": "get_weather", "arguments": "{}"},
+        {"type": "custom_tool_call_output", "call_id": "call_custom", "output": "5"},
+        {"type": "function_call_output", "call_id": "call_fn", "output": "sunny"},
+    ]
+
+
+def test_transform_custom_tool_uses_the_flat_responses_shape():
+    tool: Final = {
+        "type": "custom",
+        "custom": {
+            "name": "code_exec",
+            "description": "Run python",
+            "format": {"type": "grammar", "grammar": {"syntax": "lark", "definition": "start: /.+/"}},
+        },
+    }
+
+    result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens(
+        model="gpt-4o", input="Run it", tools=[tool]
+    )
+
+    assert result["tools"] == [
+        {
+            "type": "custom",
+            "name": "code_exec",
+            "description": "Run python",
+            "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"},
+        }
+    ]
+
+
 def test_messages_to_responses_input_preserves_images():
     """An image block must survive the round trip, or OpenAI counts only the text.