From bcbf2b7a3d5dd397f148d244ca65661068eaa516 Mon Sep 17 00:00:00 2001
From: John Karam <168653085+johnkaram02@users.noreply.github.com>
Date: Tue, 27 Oct 2026 23:53:19 +0000
Subject: [PATCH 1/2] fix(token_counter): count custom tool calls instead of
 raising

token_counter raised "Unsupported tool call ... must contain a function key"
for any assistant tool call without a function key. OpenAI custom tool calls
carry their payload under custom.input instead, and LiteLLM returns them as
ChatCompletionMessageCustomToolCall, so passing a conversation with one back
to token_counter failed

Count custom.input the same way function.arguments is counted
---
 litellm/litellm_core_utils/token_counter.py   | 10 ++++---
 .../litellm_core_utils/test_token_counter.py  | 27 +++++++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py
index edb5ecf0f8d..5784e91f98c 100644
--- a/litellm/litellm_core_utils/token_counter.py
+++ b/litellm/litellm_core_utils/token_counter.py
@@ -570,10 +570,12 @@ 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" not in tool_call:
-                raise ValueError(f"Unsupported tool call {tool_call} must contain a function key")
-            function_arguments = tool_call["function"].get("arguments", "")
-            total += count_function(str(function_arguments))
+            if "function" in tool_call:
+                total += count_function(str(tool_call["function"].get("arguments", "")))
+            elif "custom" in tool_call:
+                total += count_function(str(tool_call["custom"].get("input", "")))
+            else:
+                raise ValueError(f"Unsupported tool call {tool_call} must contain a function or custom key")
         return total
     if key == "function_call":
         if not isinstance(value, Mapping):
diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py
index 9921bbe343a..5d72b1463aa 100644
--- a/tests/unit/litellm_core_utils/test_token_counter.py
+++ b/tests/unit/litellm_core_utils/test_token_counter.py
@@ -377,6 +377,33 @@ 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 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}}),
+    )
+    function_tokens = token_counter(
+        model="gpt-4o",
+        messages=assistant_turn(
+            {"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": code}}
+        ),
+    )
+
+    assert custom_tokens == function_tokens, (
+        f"custom tool call input must count like function arguments, got {custom_tokens} vs {function_tokens}"
+    )
+    assert custom_tokens > 300, f"custom tool call input must contribute real tokens, got {custom_tokens}"
+
+
 @pytest.mark.parametrize(
     "message_count_pair",
     MESSAGES_TEXT,

From fb5414393e63c065ee6b00ed8a4a14fa77c3b5a6 Mon Sep 17 00:00:00 2001
From: John Karam <168653085+johnkaram02@users.noreply.github.com>
Date: Thu, 29 Oct 2026 13:19:34 +0000
Subject: [PATCH 2/2] refactor(token_counter): read function and custom tool
 call text through one lookup

---
 litellm/litellm_core_utils/token_counter.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py
index 5784e91f98c..6c85a6606c6 100644
--- a/litellm/litellm_core_utils/token_counter.py
+++ b/litellm/litellm_core_utils/token_counter.py
@@ -571,11 +571,12 @@ def _count_function_call_tokens(
         total = 0
         for tool_call in value:
             if "function" in tool_call:
-                total += count_function(str(tool_call["function"].get("arguments", "")))
+                kind, field = "function", "arguments"
             elif "custom" in tool_call:
-                total += count_function(str(tool_call["custom"].get("input", "")))
+                kind, field = "custom", "input"
             else:
                 raise ValueError(f"Unsupported tool call {tool_call} must contain a function or custom key")
+            total += count_function(str(tool_call[kind].get(field, "")))
         return total
     if key == "function_call":
         if not isinstance(value, Mapping):