fix(token-counting): count mid-conversation tool changes (#51447)

* fix(token-counting): count mid-conversation tool changes

* fix(token-counting): narrow tool change mappings before counting

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>

litellm/litellm_core_utils/token_counter.py

@@ -1944,13 +1944,16 @@ def _count_content_list(
                 tool_name = str(c.get("tool_name") or "")
                 if tool_name:
                     num_tokens += count_function(tool_name)
+            elif (change := as_str_mapping(c)) is not None and change.get("type") in ("tool_addition", "tool_removal"):
+                num_tokens += _tool_change_tokens(count_function, as_str_mapping(change.get("tool")))
             else:
                 content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__
                 raise UnsupportedContentTypeError(
                     f"Invalid content item type: {content_type}. "
                     f"Expected str or dict with 'type' field "
                     f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
-                    f"compaction, tool_reference, search_result, or Anthropic server tool blocks)."
+                    f"compaction, tool_reference, tool_addition, tool_removal, search_result, "
+                    f"or Anthropic server tool blocks)."
                 )
         return num_tokens
     except Exception as e:
@@ -1962,6 +1965,24 @@ def _count_content_list(
         )
 
 
+def _tool_change_tokens(count_function: TokenCounterFunction, tool: Mapping[str, object] | None) -> int:
+    match tool:
+        case {"type": "tool_definition", "definition": definition} if (
+            function := as_str_mapping(definition)
+        ) is not None:
+            return count_function(_format_function_definitions((function,)))
+        case {"type": "tool_reference", "name": str() as name}:
+            return count_function(name)
+        case {"type": "mcp_tool_reference", "server_name": str() as server_name, "name": str() as name}:
+            return count_function(server_name) + count_function(name)
+        case {"type": "mcp_toolset_reference", "server_name": str() as server_name}:
+            return count_function(server_name)
+        case _:
+            raise UnsupportedContentTypeError(
+                f"Invalid tool change target type: {None if tool is None else tool.get('type')}"
+            )
+
+
 def _format_function_definitions(tools: Sequence[object]) -> str:
     """Formats tool definitions in the format that OpenAI appears to use.
     Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java

tests/unit/litellm_core_utils/test_token_counter.py

@@ -2,6 +2,7 @@
 #    This tests litellm.token_counter.token_counter() function
 import asyncio
 import base64
+import copy
 import importlib
 import io
 import json
@@ -37,6 +38,8 @@ from litellm.constants import (
 )
 from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.token_counter import (
+    UnsupportedContentTypeError,
+    _count_content_list,
     _encoding_count,
     _get_exact_count_function,
     _get_extrapolating_count_function,
@@ -45,6 +48,7 @@ from litellm.litellm_core_utils.token_counter import (
     estimate_image_tokens,
     high_detail_image_token_upper_bound,
     is_openai_model,
+    messages_reach_token_count,
     offload_token_count,
 )
 from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
@@ -2170,6 +2174,127 @@ def test_count_content_list_rejects_unknown_type():
     assert "tool_reference" in message
 
 
+_CLAUDE: Final = "anthropic/claude-sonnet-5-5"
+_REMOVE_WEATHER: Final = {"type": "tool_removal", "tool": {"type": "tool_reference", "name": "get_weather"}}
+
+
+def _weather_tool(description: str, *extra_properties: str) -> dict[str, object]:
+    return {
+        "name": "get_weather",
+        "description": description,
+        "input_schema": {
+            "type": "object",
+            "properties": {
+                "city": {"type": "string", "description": "City name"},
+                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                **{name: {"type": "string"} for name in extra_properties},
+            },
+            "required": ["city"],
+        },
+    }
+
+
+def _add_definition(definition: Mapping[str, object]) -> dict[str, object]:
+    return {"type": "tool_addition", "tool": {"type": "tool_definition", "definition": definition}}
+
+
+def _conversation(*system_blocks: Mapping[str, object]) -> list[dict[str, object]]:
+    return [
+        {"role": "user", "content": "What is the weather in Paris?"},
+        {"role": "system", "content": [{"type": "text", "text": "Your tools changed."}, *system_blocks]},
+    ]
+
+
+def _then_remove_weather(messages: list[dict[str, object]]) -> list[dict[str, object]]:
+    return [*messages, {"role": "user", "content": "Never mind."}, {"role": "system", "content": [_REMOVE_WEATHER]}]
+
+
+def test_tool_addition_counts_an_inline_definition_like_the_same_tool_declared_in_tools() -> None:
+    brief: Final = _weather_tool("Weather.")
+    detailed: Final = _weather_tool("Current weather, temperature and wind for a city.", "country", "date")
+    add_by_reference: Final = {"type": "tool_addition", "tool": {"type": "tool_reference", "name": "get_weather"}}
+    by_reference: Final = _conversation(add_by_reference)
+
+    inline_delta: Final = token_counter_new(model=_CLAUDE, messages=_conversation(_add_definition(detailed))) - (
+        token_counter_new(model=_CLAUDE, messages=_conversation(_add_definition(brief)))
+    )
+    declared_delta: Final = token_counter_new(model=_CLAUDE, messages=by_reference, tools=[detailed]) - (
+        token_counter_new(model=_CLAUDE, messages=by_reference, tools=[brief])
+    )
+
+    assert declared_delta > 0
+    assert inline_delta == declared_delta
+
+
+@pytest.mark.parametrize("change", ["tool_addition", "tool_removal"])
+@pytest.mark.parametrize(
+    ("target", "referenced_names"),
+    [
+        ({"type": "tool_reference", "name": "get_weather"}, ("get_weather",)),
+        ({"type": "mcp_tool_reference", "server_name": "calendar", "name": "list_events"}, ("calendar", "list_events")),
+        ({"type": "mcp_toolset_reference", "server_name": "calendar"}, ("calendar",)),
+    ],
+    ids=["tool_reference", "mcp_tool_reference", "mcp_toolset_reference"],
+)
+def test_tool_change_by_reference_counts_the_referenced_names(
+    change: str, target: dict[str, object], referenced_names: tuple[str, ...]
+) -> None:
+    counted: Final = _count_content_list(len, [{"type": change, "tool": target}], False, None)
+
+    assert counted == sum(map(len, referenced_names))
+
+
+def test_tool_removal_never_subtracts_an_earlier_inline_definition() -> None:
+    with_definition: Final = _conversation(_add_definition(_weather_tool("Current weather for a city.")))
+    without_definition: Final = _conversation()
+
+    removal_with_definition: Final = token_counter_new(
+        model=_CLAUDE, messages=_then_remove_weather(with_definition)
+    ) - token_counter_new(model=_CLAUDE, messages=with_definition)
+    removal_without_definition: Final = token_counter_new(
+        model=_CLAUDE, messages=_then_remove_weather(without_definition)
+    ) - token_counter_new(model=_CLAUDE, messages=without_definition)
+
+    assert removal_with_definition > 0
+    assert removal_with_definition == removal_without_definition
+
+
+def test_messages_reach_token_count_agrees_with_token_counter_across_tool_changes() -> None:
+    messages: Final = _then_remove_weather(_conversation(_add_definition(_weather_tool("Current weather for a city."))))
+    total: Final = token_counter_new(model=_CLAUDE, messages=messages)
+
+    reached: Final = (
+        messages_reach_token_count(model=_CLAUDE, messages=messages, threshold=total),
+        messages_reach_token_count(model=_CLAUDE, messages=messages, threshold=total + 1),
+    )
+
+    assert reached == (True, False)
+
+
+def test_counting_tool_changes_leaves_the_messages_unchanged() -> None:
+    messages: Final = _then_remove_weather(_conversation(_add_definition(_weather_tool("Current weather for a city."))))
+    sent: Final = copy.deepcopy(messages)
+
+    token_counter_new(model=_CLAUDE, messages=messages)
+    messages_reach_token_count(model=_CLAUDE, messages=messages, threshold=1)
+
+    assert messages == sent
+
+
+@pytest.mark.parametrize(
+    "tool_change",
+    [
+        {"type": "tool_addition", "tool": {"type": "future_reference", "name": "get_weather"}},
+        {"type": "tool_addition", "tool": {"type": "tool_definition", "definition": "get_weather"}},
+        {"type": "tool_removal"},
+    ],
+    ids=["unknown_target_type", "definition_not_an_object", "missing_target"],
+)
+def test_tool_change_with_an_unrecognized_target_stays_unsupported(tool_change: dict[str, object]) -> None:
+    with pytest.raises(UnsupportedContentTypeError, match="Invalid tool change target type"):
+        _count_content_list(len, [tool_change], False, None)
+
+
 @pytest.mark.parametrize(
     "source",
     [

tests/unit/proxy/test_proxy_token_counter.py

@@ -1682,6 +1682,75 @@ async def test_anthropic_count_tokens_on_a_claude_deployment_drops_a_malformed_t
     assert sent["messages"][2]["content"][0]["content"] == [{"type": "text", "text": "build ok"}]
 
 
+def _weather_tool(description: str) -> dict[str, object]:
+    return {
+        "name": "get_weather",
+        "description": description,
+        "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+    }
+
+
+@pytest.mark.asyncio
+async def test_anthropic_count_tokens_falls_back_to_a_local_estimate_that_counts_mid_conversation_tool_changes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    import litellm.proxy.proxy_server as proxy_server
+
+    monkeypatch.setattr(
+        proxy_server,
+        "llm_router",
+        Router(
+            model_list=[
+                {
+                    "model_name": "claude-framed",
+                    "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-ant-test"},
+                    "model_info": {"token_count_framing": 100},
+                }
+            ]
+        ),
+    )
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    brief: Final = _weather_tool("Weather.")
+    detailed: Final = _weather_tool("Current weather, temperature and wind for a city in any country.")
+
+    def conversation(tool: Mapping[str, object]) -> list[dict[str, object]]:
+        return [
+            {"role": "user", "content": "What is the weather in Paris?"},
+            {"role": "system", "content": [{"type": "tool_addition", "tool": tool}]},
+            {"role": "user", "content": "Never mind."},
+            {
+                "role": "system",
+                "content": [{"type": "tool_removal", "tool": {"type": "tool_reference", "name": "get_weather"}}],
+            },
+        ]
+
+    async def estimate(tool: Mapping[str, object], tools: list[dict[str, object]] | None) -> TokenCountResponse:
+        return await count_request_tokens(
+            request=TokenCountRequest(model="claude-framed", messages=conversation(tool), tools=tools),
+            call_endpoint=True,
+            messages_format="anthropic",
+        )
+
+    by_reference: Final = {"type": "tool_reference", "name": "get_weather"}
+    with respx.mock:
+        route: Final = respx.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
+            return_value=httpx.Response(
+                404, json={"type": "error", "error": {"type": "not_found_error", "message": "count unavailable"}}
+            )
+        )
+        inline_brief: Final = await estimate({"type": "tool_definition", "definition": brief}, None)
+        inline_detailed: Final = await estimate({"type": "tool_definition", "definition": detailed}, None)
+        declared_brief: Final = await estimate(by_reference, [brief])
+        declared_detailed: Final = await estimate(by_reference, [detailed])
+
+    inline_delta: Final = inline_detailed.total_tokens - inline_brief.total_tokens
+    declared_delta: Final = declared_detailed.total_tokens - declared_brief.total_tokens
+    assert route.called
+    assert {inline_brief.count_source, inline_detailed.count_source} == {"local_estimate"}
+    assert declared_delta > 0
+    assert inline_delta == declared_delta
+
+
 class _ScriptedCounter(BaseTokenCounter):
     def __init__(self, replies: dict[str, TokenCountResponse], burst: int = 1) -> None:
         self._replies = replies