From bcafb0a6151ac8d1052e3759d0bce437ab9a9b25 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:56:25 -0700 Subject: [PATCH] fix(context_management): count unreadable blocks and check tool_result text before clearing clear_tool_uses now maps blocks the token counter cannot read (server tool, compaction) to JSON text, so an input_tokens trigger no longer 500s on that history. The Responses API path checks malformed tool_result text on the original messages before clearing, as the chat completions path does. Also removes comments that repeated the code and makes the test fixtures immutable. --- .../editors/clear_tool_uses.py | 57 ++++++--- .../messages/responses_adapters/handler.py | 7 +- .../responses_adapters/streaming_iterator.py | 1 - .../test_clear_tool_uses.py | 41 +++++++ .../test_responses_adapters_handler.py | 109 +++++++++++------- 5 files changed, 152 insertions(+), 63 deletions(-) diff --git a/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py b/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py index d4430ca5ab9..54011b228f3 100644 --- a/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py +++ b/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py @@ -10,6 +10,7 @@ the way the API's 400 does, so the dispatcher can reject the edit, or drop only runs. """ +import json from collections.abc import Mapping, Sequence, Set from dataclasses import dataclass from functools import reduce @@ -38,6 +39,21 @@ if TYPE_CHECKING: T = TypeVar("T") # api.anthropic.com, measured 2026-10-19: text results of 1 to 46 bytes are never cleared or counted toward keep _TINY_RESULT_BYTES: Final = 46 +_TOKEN_COUNTER_BLOCK_TYPES: Final = frozenset( + { + "text", + "image_url", + "image", + "document", + "file", + "tool_use", + "tool_result", + "thinking", + "redacted_thinking", + "tool_reference", + "search_result", + } +) ToolNames = Literal["all"] | frozenset[str] @@ -198,6 +214,29 @@ def _select_in_batches(candidates: Sequence[str], tokens: Mapping[str, int], cle return tuple(candidates[:batch_end]) +def _countable_block(block: Mapping[str, object]) -> Mapping[str, object]: + if block.get("type") in _TOKEN_COUNTER_BLOCK_TYPES: + return block + return {"type": "text", "text": json.dumps(block, default=str)} + + +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": [_countable_block(block) for block in blocks]} + + +def _count_tokens( + model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None +) -> int: + return litellm.token_counter( + model=model, + messages=[_countable_message(message) for message in messages], + tools=cast("list[ChatCompletionToolParam] | None", tools), + ) + + def _trigger_met( trigger: dict[str, object], model: str, @@ -216,11 +255,7 @@ def _trigger_met( if not isinstance(threshold, int): threshold = DEFAULT_INPUT_TOKENS_TRIGGER - current_tokens: Final = litellm.token_counter( - model=model, - messages=messages, - tools=cast("list[ChatCompletionToolParam] | None", tools), - ) + current_tokens: Final = _count_tokens(model, messages, tools) verbose_logger.debug("context_management polyfill: current_tokens: %s", current_tokens) verbose_logger.debug("context_management polyfill: threshold: %s", threshold) return current_tokens > threshold, current_tokens @@ -366,16 +401,8 @@ def apply_clear_tool_uses_20250919( [_message_without_inputs(message, ids_to_clear, input_tools) for message in edited] if input_tools else edited ) - tokens_before: Final = ( - counted_tokens - if counted_tokens is not None - else litellm.token_counter( - model=model, messages=messages, tools=cast("list[ChatCompletionToolParam] | None", tools) - ) - ) - tokens_after: Final = litellm.token_counter( - model=model, messages=final_messages, tools=cast("list[ChatCompletionToolParam] | None", tools) - ) + tokens_before: Final = counted_tokens if counted_tokens is not None else _count_tokens(model, messages, tools) + tokens_after: Final = _count_tokens(model, final_messages, tools) cleared_input_tokens: Final = max(tokens_before - tokens_after, 0) return final_messages, { diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index b8cdb7dfd17..ed51d8bc98c 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -49,7 +49,6 @@ 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]]) -# (native spec for the Responses API, edits for the gateway polyfill) _SplitEdits: TypeAlias = tuple[dict[str, object] | None, Sequence[Mapping[str, object]]] _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter() @@ -86,6 +85,7 @@ async def _polyfilled( """Run the edits the Responses API can't run natively; raise like the chat completions path does.""" if not edits: return None + reject_malformed_tool_result_text(messages, model, kwargs) drop_invalid_edits: Final = drops_unsupported_params(kwargs.get("drop_params")) try: return await apply_context_management( @@ -97,8 +97,6 @@ async def _polyfilled( drop_invalid_edits=drop_invalid_edits, ) except AnthropicContextManagementError: - if drop_invalid_edits: - return None raise except Exception as e: verbose_logger.exception("context_management polyfill failed on the Responses API path: %s", e) @@ -355,7 +353,6 @@ class LiteLLMMessagesToResponsesAPIHandler: ) responses_kwargs: Final = _build_responses_kwargs( max_tokens=max_tokens, - # clear_tool_uses edits only rewrite messages; system stays as sent messages=polyfill_result.messages if polyfill_result is not None else messages, model=model, context_management=native_context_management, @@ -457,11 +454,9 @@ class LiteLLMMessagesToResponsesAPIHandler: if polyfilled_edits else None ) - # run_async_function is untyped: check what it returned rather than assume it polyfill_result: Final = ran if isinstance(ran, PolyfillResult) else None responses_kwargs: Final = _build_responses_kwargs( max_tokens=max_tokens, - # clear_tool_uses edits only rewrite messages; system stays as sent messages=polyfill_result.messages if polyfill_result is not None else messages, model=model, context_management=native_context_management, diff --git a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py index c29094b25de..8b99a21353d 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py @@ -131,7 +131,6 @@ class AnthropicResponsesStreamWrapper: applied_edits: Sequence[AppliedEdit] | None = None, ) -> None: self.responses_stream = responses_stream - # The gateway polyfill's edits, reported on the final message_delta as on the chat completions path self._applied_edits: Final = tuple(applied_edits or ()) self.model = model self.safeguards_evaluator: Final = safeguards_evaluator diff --git a/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py b/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py index 2567d1b76db..e29583327af 100644 --- a/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py +++ b/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py @@ -217,6 +217,47 @@ def test_cleared_input_tokens_is_nonnegative(): assert applied["cleared_input_tokens"] >= 0 +SERVER_TOOL_TURN: Final = { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "Mumbai rain"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [{"type": "web_search_result", "url": "https://example.com", "title": "rain " * 300}], + }, + ], +} +COMPACTION_TURN: Final = {"role": "assistant", "content": [{"type": "compaction", "content": "summary " * 300}]} + + +@pytest.mark.parametrize( + "turn", [pytest.param(SERVER_TOOL_TURN, id="server_tool"), pytest.param(COMPACTION_TURN, id="compaction")] +) +def test_an_input_tokens_trigger_counts_blocks_the_token_counter_cannot_read(turn: dict[str, object]) -> None: + history: Final = _make_history(n_pairs=3) + threshold: Final = litellm.token_counter(model=MODEL, messages=history) + 100 + messages: Final = [history[0], turn, *history[1:]] + + new_messages, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": threshold}, + "keep": {"type": "tool_uses", "value": 1}, + }, + ) + + assert applied is not None + assert applied["cleared_tool_uses"] == 2 + assert applied["cleared_input_tokens"] > 0 + assert _cleared_ids(new_messages) == {"toolu_00", "toolu_01"} + assert new_messages[1] == turn + + 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)) diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py index 06728ab7ca3..33a19358d3c 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py @@ -3,6 +3,7 @@ import json import os import sys from collections.abc import AsyncIterator, Mapping +from itertools import chain from typing import Final, cast from unittest.mock import AsyncMock, patch @@ -769,26 +770,30 @@ CLEAR_EDIT: Final = { COMPACT_EDIT: Final = {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}} +def _read(i: int) -> tuple[dict[str, object], dict[str, object]]: + return ( + { + "role": "assistant", + "content": [{"type": "tool_use", "id": f"toolu_{i}", "name": "Bash", "input": {"command": f"cat {i}"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{i}", "content": f"log {i} " * 40}], + }, + ) + + def _three_reads() -> list[dict[str, object]]: - messages: Final[list[dict[str, object]]] = [{"role": "user", "content": "Read the logs."}] - for i in range(3): - messages.append( - { - "role": "assistant", - "content": [{"type": "tool_use", "id": f"toolu_{i}", "name": "Bash", "input": {"command": f"cat {i}"}}], - } - ) - messages.append( - { - "role": "user", - "content": [{"type": "tool_result", "tool_use_id": f"toolu_{i}", "content": f"log {i} " * 40}], - } - ) - return messages + return [{"role": "user", "content": "Read the logs."}, *chain.from_iterable(_read(i) for i in range(3))] async def _cleared_bridge( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, *, stream: bool, edits: list[dict[str, object]] + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + *, + stream: bool, + edits: list[dict[str, object]], + drop_params: bool = False, ) -> tuple[object, Mapping[str, object]]: monkeypatch.setenv("OPENAI_API_KEY", "sk-clear-test") monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") @@ -806,6 +811,7 @@ async def _cleared_bridge( stream=stream, custom_llm_provider="openai", context_management={"edits": edits}, + drop_params=drop_params, ) body: Final = json.loads(upstream.calls.last.request.content) if not stream: @@ -817,7 +823,7 @@ async def _cleared_bridge( async def test_clear_tool_uses_clears_old_results_before_the_responses_api_call( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - from litellm.llms.anthropic.pass_through.context_management.constants import CLEARED_TOOL_RESULT_PLACEHOLDER + from litellm.llms.anthropic.messages.context_management.constants import CLEARED_TOOL_RESULT_PLACEHOLDER served, body = await _cleared_bridge(respx_mock, monkeypatch, stream=False, edits=[dict(CLEAR_EDIT)]) @@ -862,7 +868,7 @@ async def test_compact_stays_native_while_clear_tool_uses_runs_in_the_gateway( async def test_a_malformed_clear_tool_uses_edit_is_a_400_before_the_responses_api_call( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError + from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").respond( status_code=200, json=TEXT_RESPONSE @@ -880,8 +886,35 @@ async def test_a_malformed_clear_tool_uses_edit_is_a_400_before_the_responses_ap assert not upstream.called +@pytest.mark.asyncio +async def test_a_malformed_text_in_a_result_clearing_would_drop_is_still_a_400( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").respond( + status_code=200, json=TEXT_RESPONSE + ) + monkeypatch.setenv("OPENAI_API_KEY", "sk-clear-test") + reads: Final = _three_reads() + malformed_oldest_result: Final = { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_0", "content": [{"type": "text", "text": None}]}], + } + with pytest.raises(litellm.BadRequestError) as rejected: + await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=1024, + messages=[*reads[:2], malformed_oldest_result, *reads[3:]], + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + custom_llm_provider="openai", + context_management={"edits": [dict(CLEAR_EDIT)]}, + ) + assert rejected.value.status_code == 400 + assert "tool_result toolu_0 content.0.text: Input should be a valid string" in rejected.value.message + assert not upstream.called + + def test_the_sync_path_clears_old_results_too(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.llms.anthropic.pass_through.context_management.constants import CLEARED_TOOL_RESULT_PLACEHOLDER + from litellm.llms.anthropic.messages.context_management.constants import CLEARED_TOOL_RESULT_PLACEHOLDER monkeypatch.setenv("OPENAI_API_KEY", "sk-clear-test") monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") @@ -916,8 +949,8 @@ def test_the_sync_path_clears_old_results_too(respx_mock: respx.MockRouter, monk async def test_a_polyfill_crash_is_a_500_before_the_responses_api_call( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError - from litellm.llms.anthropic.pass_through.responses_adapters import handler + from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError + from litellm.llms.anthropic.messages.responses_adapters import handler upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").respond( status_code=200, json=TEXT_RESPONSE @@ -937,31 +970,25 @@ async def test_a_polyfill_crash_is_a_500_before_the_responses_api_call( @pytest.mark.asyncio -async def test_under_drop_params_a_rejected_edit_is_skipped_and_the_request_goes_out( +async def test_under_drop_params_a_malformed_edit_still_clears_without_its_bad_field( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError - from litellm.llms.anthropic.pass_through.responses_adapters import handler + from litellm.llms.anthropic.messages.context_management.constants import CLEARED_TOOL_RESULT_PLACEHOLDER - monkeypatch.setenv("OPENAI_API_KEY", "sk-clear-test") - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - litellm.in_memory_llm_clients_cache.flush_cache() - upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").respond( - status_code=200, json=TEXT_RESPONSE + served, body = await _cleared_bridge( + respx_mock, monkeypatch, stream=False, edits=[{**CLEAR_EDIT, "exclude_tools": "Bash"}], drop_params=True ) - rejected: Final = AnthropicContextManagementError(status_code=400, message="bad edit") - monkeypatch.setattr(handler, "apply_context_management", AsyncMock(side_effect=rejected)) - served: Final = await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( - max_tokens=1024, - messages=_three_reads(), - model="openai/gpt-5.6-luna", - custom_llm_provider="openai", - context_management={"edits": [dict(CLEAR_EDIT)]}, - drop_params=True, - ) - assert upstream.called - assert isinstance(served, dict) and "context_management" not in served + outputs: Final = { + item["call_id"]: item["output"] for item in body["input"] if item.get("type") == "function_call_output" + } + assert [CLEARED_TOOL_RESULT_PLACEHOLDER in json.dumps(outputs[f"toolu_{i}"]) for i in range(3)] == [ + True, + True, + False, + ] + assert isinstance(served, dict) + assert served["context_management"]["applied_edits"][0]["cleared_tool_uses"] == 2 @pytest.mark.asyncio