fix(context_management): apply clear edits in OpenAI count_tokens and skip nested off-spec blocks when counting

litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py

@@ -213,11 +213,23 @@ def _select_in_batches(candidates: Sequence[str], tokens: Mapping[str, int], cle
     return tuple(candidates[:batch_end])
 
 
+def _without_uncountable_blocks(block: Mapping[str, object]) -> Mapping[str, object]:
+    nested: Final = _parse(_BLOCKS, block.get("content"))
+    if nested is None:
+        return block
+    return {**block, "content": [inner for inner in nested if inner.get("type") in _TOKEN_COUNTER_BLOCK_TYPES]}
+
+
 def _countable_message(message: Mapping[str, object]) -> Mapping[str, object]:
     blocks: Final = _parse(_BLOCKS, message.get("content"))
     if blocks is None:
         return message
-    return {**message, "content": [block for block in blocks if block.get("type") in _TOKEN_COUNTER_BLOCK_TYPES]}
+    return {
+        **message,
+        "content": [
+            _without_uncountable_blocks(block) for block in blocks if block.get("type") in _TOKEN_COUNTER_BLOCK_TYPES
+        ],
+    }
 
 
 def _count_tokens(

litellm/llms/anthropic/messages/responses_adapters/handler.py

@@ -34,6 +34,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,
 )
 from litellm.types.llms.openai import ResponsesAPIResponse
+from litellm.types.utils import CountContextManagement
 from litellm.utils import ProviderConfigManager
 
 from ..adapters.handler import context_management_explicitly_dropped, request_safeguards_evaluator
@@ -49,6 +50,7 @@ from .transformation import LiteLLMAnthropicToResponsesAPIAdapter
 AnthropicRequestMessages: TypeAlias = list[AllAnthropicMessageValues] | list[dict[str, object]]
 BlockRewrite: TypeAlias = Callable[[Mapping[str, object]], Mapping[str, object]]
 _BLOCK_LIST: Final = TypeAdapter(list[dict[str, object]])
+_SPEC: Final = TypeAdapter(dict[str, object])
 _SplitEdits: TypeAlias = tuple[dict[str, object] | None, Sequence[Mapping[str, object]]]
 AnthropicContextManagementSpec: TypeAlias = dict[str, object] | list[dict[str, object]] | None
 
@@ -109,6 +111,28 @@ async def _polyfilled(
         ) from e
 
 
+async def count_view_messages(
+    *,
+    model: str,
+    messages: Sequence[Mapping[str, object]],
+    tools: Sequence[Mapping[str, object]] | None,
+    context_management: CountContextManagement | None,
+) -> Sequence[Mapping[str, object]]:
+    """The messages /v1/messages sends the Responses API once the polyfill has run the edits it can't run natively."""
+    if context_management is None or not isinstance(context_management.spec, dict):
+        return messages
+    _, edits = _split_context_management(_SPEC.validate_python(context_management.spec), {})
+    result: Final = await _polyfilled(
+        model=model,
+        messages=_BLOCK_LIST.validate_python(messages),
+        tools=tools,
+        system=None,
+        edits=edits,
+        kwargs={"drop_params": context_management.drop_invalid_edits},
+    )
+    return messages if result is None else result.messages
+
+
 def _with_applied_edits(
     response: AnthropicMessagesResponse, polyfill_result: PolyfillResult | None
 ) -> AnthropicMessagesResponse:

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

@@ -15,7 +15,9 @@ from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.tool_fields import DropPolicy
+from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
 from litellm.llms.anthropic.messages.convertible_blocks import deployment_drops_params, unconvertible_blocks_error
+from litellm.llms.anthropic.messages.responses_adapters.handler import count_view_messages
 from litellm.llms.anthropic.messages.responses_adapters.transformation import (
     LiteLLMAnthropicToResponsesAPIAdapter,
 )
@@ -95,6 +97,27 @@ def _count_error(model_to_use: str, request_model: str, message: str, status_cod
     )
 
 
+async def _sent_messages(
+    model: str,
+    messages: Sequence[Mapping[str, object]],
+    tools: Sequence[Mapping[str, object]] | None,
+    messages_format: TokenCountMessagesFormat,
+    context_management: CountContextManagement | None,
+) -> Sequence[Mapping[str, object]] | AnthropicContextManagementError:
+    match messages_format:
+        case "openai":
+            return messages
+        case "anthropic":
+            try:
+                return await count_view_messages(
+                    model=model, messages=messages, tools=tools, context_management=context_management
+                )
+            except AnthropicContextManagementError as error:
+                return error
+        case _:
+            assert_never(messages_format)
+
+
 class OpenAITokenCounter(BaseTokenCounter):
     """Token counter implementation for OpenAI provider using the Responses API."""
 
@@ -175,13 +198,26 @@ class OpenAITokenCounter(BaseTokenCounter):
                 rejected_before_provider_call=True,
             )
 
+        sent_messages: Final = await _sent_messages(model_to_use, messages, tools, messages_format, context_management)
+        if isinstance(sent_messages, AnthropicContextManagementError):
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="openai_api",
+                error=True,
+                error_message=sent_messages.message,
+                status_code=sent_messages.status_code,
+                rejected_before_provider_call=True,
+            )
+
         resolved: Final = _count_credential(deployment, request_model or model_to_use)
         if isinstance(resolved, OpenAIError):
             return _count_error(model_to_use, request_model, resolved.message, resolved.status_code)
 
         input_items, instructions, counted_tools, text = self._responses_request_parts(
             model_to_use=model_to_use,
-            messages=messages,
+            messages=sent_messages,
             tools=tools,
             system=system,
             messages_format=messages_format,

tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py

@@ -266,6 +266,28 @@ def test_an_input_tokens_trigger_skips_blocks_no_adapter_sends(turn: dict[str, o
     assert cleared[1] == turn
 
 
+def _first_result_with(blocks: list[dict[str, object]]) -> dict[str, object]:
+    return {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_00", "content": blocks}]}
+
+
+def test_an_input_tokens_trigger_skips_blocks_no_adapter_sends_inside_a_tool_result() -> None:
+    history: Final = _make_history(n_pairs=3)
+    result_text: Final = {"type": "text", "text": "Result 0: " + "x" * 200}
+    search_result: Final = {"type": "web_search_result", "url": "https://example.com", "title": "rain " * 300}
+    sent_tokens: Final = litellm.token_counter(
+        model=MODEL, messages=[*history[:2], _first_result_with([result_text]), *history[3:]]
+    )
+    messages: Final = [*history[:2], _first_result_with([result_text, search_result]), *history[3:]]
+
+    untouched, not_applied = _clear_above(messages, sent_tokens + 100)
+    cleared, applied = _clear_above(messages, sent_tokens - 100)
+
+    assert not_applied is None
+    assert untouched == messages
+    assert applied is not None
+    assert _cleared_ids(cleared) == {"toolu_00", "toolu_01"}
+
+
 def _blocks_of(messages: list[dict[str, object]], block_type: str) -> tuple[dict[str, object], ...]:
     contents: Final = (message["content"] for message in messages)
     blocks: Final = chain.from_iterable(content for content in contents if isinstance(content, list))

tests/unit/llms/openai/responses/count_tokens/test_token_counter.py

@@ -2,6 +2,7 @@ import copy
 import json
 import operator
 from collections.abc import Mapping
+from itertools import chain
 from typing import Final
 
 import httpx
@@ -11,9 +12,10 @@ from pydantic import JsonValue, ValidationError
 
 import litellm
 from litellm.litellm_core_utils.tool_fields import DropPolicy
+from litellm.llms.anthropic.messages.context_management.constants import CLEARED_TOOL_RESULT_PLACEHOLDER
 from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
 from litellm.responses.utils import ResponsesAPIRequestUtils
-from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
+from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse
 
 COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens"
 
@@ -484,6 +486,90 @@ async def test_anthropic_format_tool_result_text_block_without_text_still_counts
     assert {"type": "function_call_output", "call_id": "toolu_log", "output": "build ok\n\ndone"} in sent["input"]
 
 
+def _tool_pair(index: int) -> tuple[dict[str, object], dict[str, object]]:
+    return (
+        {"role": "assistant", "content": [{"type": "tool_use", "id": f"toolu_{index}", "name": "log", "input": {}}]},
+        {
+            "role": "user",
+            "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": f"log {index} " * 40}],
+        },
+    )
+
+
+TOOL_PAIRS: Final[list[dict[str, object]]] = [
+    {"role": "user", "content": "Compare the logs."},
+    *chain.from_iterable(_tool_pair(index) for index in range(3)),
+]
+
+
+async def _count_with_edit(
+    monkeypatch: pytest.MonkeyPatch, edit: dict[str, JsonValue], drop_invalid_edits: bool
+) -> tuple[TokenCountResponse | None, respx.Route]:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    with respx.mock:
+        route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9}))
+        result: Final = await OpenAITokenCounter().count_tokens(
+            model_to_use="gpt-4o",
+            messages=copy.deepcopy(TOOL_PAIRS),
+            contents=None,
+            deployment={"litellm_params": {"api_key": "sk-test"}},
+            request_model="openai/gpt-4o",
+            messages_format="anthropic",
+            context_management=CountContextManagement(spec={"edits": [edit]}, drop_invalid_edits=drop_invalid_edits),
+        )
+        return result, route
+
+
+def _sent_outputs(route: respx.Route) -> dict[str, str]:
+    sent: Final = json.loads(route.calls.last.request.content)
+    return {item["call_id"]: item["output"] for item in sent["input"] if item["type"] == "function_call_output"}
+
+
+CLEAR_EDIT: Final[dict[str, JsonValue]] = {
+    "type": "clear_tool_uses_20250919",
+    "trigger": {"type": "tool_uses", "value": 1},
+    "keep": {"type": "tool_uses", "value": 1},
+}
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_counts_the_history_the_clear_edit_leaves_for_inference(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    result, route = await _count_with_edit(monkeypatch, CLEAR_EDIT, drop_invalid_edits=False)
+
+    assert result is not None
+    assert (result.error, result.total_tokens) == (False, 9)
+    assert _sent_outputs(route) == {
+        "toolu_0": CLEARED_TOOL_RESULT_PLACEHOLDER,
+        "toolu_1": CLEARED_TOOL_RESULT_PLACEHOLDER,
+        "toolu_2": "log 2 " * 40,
+    }
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_malformed_clear_edit_answers_400_without_calling_openai(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    result, route = await _count_with_edit(monkeypatch, {**CLEAR_EDIT, "exclude_tools": "log"}, drop_invalid_edits=False)
+
+    assert route.called is False
+    assert result is not None
+    assert (result.error, result.status_code) == (True, 400)
+    assert "exclude_tools" in str(result.error_message)
+
+
+@pytest.mark.asyncio
+async def test_anthropic_format_drop_params_counts_the_clear_without_its_malformed_field(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    result, route = await _count_with_edit(monkeypatch, {**CLEAR_EDIT, "exclude_tools": "log"}, drop_invalid_edits=True)
+
+    assert result is not None
+    assert result.error is False
+    assert _sent_outputs(route)["toolu_0"] == CLEARED_TOOL_RESULT_PLACEHOLDER
+
+
 @pytest.mark.asyncio
 async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)