From f9d11ccbceee127e8259abb159cb53b357c85220 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:15:39 -0700 Subject: [PATCH] fix(context_management): count only blocks the model receives and keep list specs unsent clear_tool_uses now leaves blocks token_counter cannot read, like server tool results and compaction, out of its count instead of counting them as JSON text. No adapter sends them to the model, so an input_tokens trigger no longer fires on tokens the model never sees. On the Responses API path, an OpenAI-style list context_management stays unsent, as it was before, instead of reaching OpenAI's native compaction, where a threshold under 1000 is a 400. On the chat completions path, a clear edit listed before compact_20260112 now runs on the turns after the latest compaction block, as Anthropic does, instead of counting and clearing turns the compaction block drops. The crash test that replaced apply_context_management with a mock is removed. --- .../anthropic/messages/adapters/handler.py | 3 +- .../editors/clear_tool_uses.py | 9 +----- .../context_management/editors/compact.py | 6 ++-- .../messages/responses_adapters/handler.py | 12 ++++--- .../test_clear_tool_uses.py | 32 ++++++++++++------- .../context_management/test_compact.py | 6 ++-- .../test_responses_adapters_handler.py | 30 ++++++++--------- .../vertex_ai/count_tokens/test_context.py | 2 +- 8 files changed, 53 insertions(+), 47 deletions(-) diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index 71ccfbaf915..704bec9056f 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -160,6 +160,7 @@ async def prepare_context_managed_request( """Apply client compaction history, then optional context_management polyfill.""" from litellm.llms.anthropic.messages.context_management.editors.compact import ( apply_client_compaction_block_history, + slice_around_compaction_block, ) # Skip the client-history pre-processing when a ``compact_20260112`` @@ -175,7 +176,7 @@ async def prepare_context_managed_request( if polyfill_will_run: history_result: PolyfillResult | None = None - working_messages: _AnthropicMessages = messages + working_messages: _AnthropicMessages = list(slice_around_compaction_block(messages)[0]) working_system: _AnthropicSystem = system else: history_result = apply_client_compaction_block_history( 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 54011b228f3..08b00a23819 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,7 +10,6 @@ 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 @@ -214,17 +213,11 @@ 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]} + return {**message, "content": [block for block in blocks if block.get("type") in _TOKEN_COUNTER_BLOCK_TYPES]} def _count_tokens( diff --git a/litellm/llms/anthropic/messages/context_management/editors/compact.py b/litellm/llms/anthropic/messages/context_management/editors/compact.py index 254d1c4c058..fac23817927 100644 --- a/litellm/llms/anthropic/messages/context_management/editors/compact.py +++ b/litellm/llms/anthropic/messages/context_management/editors/compact.py @@ -137,7 +137,7 @@ def _find_latest_compaction_index( return None, None -def _slice_around_compaction_block( +def slice_around_compaction_block( messages: Sequence[_MsgT], ) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]: """Apply Anthropic's "drop everything before the compaction block" rule. @@ -598,7 +598,7 @@ def apply_client_compaction_block_history( tail is forwarded unchanged (with compaction blocks stripped) so recent turns the summary does not cover are preserved. """ - effective_messages, prior_compaction_block = _slice_around_compaction_block(messages) + effective_messages, prior_compaction_block = slice_around_compaction_block(messages) if prior_compaction_block is None: return None @@ -669,7 +669,7 @@ async def apply_compact_20260112( # opt-in gate below so that even when summarization is disabled we still # strip Anthropic-only ``compaction`` blocks from messages going to # non-Anthropic backends (which would reject them). - effective_messages, prior_compaction_block = _slice_around_compaction_block(messages) + effective_messages, prior_compaction_block = slice_around_compaction_block(messages) prior_summary_text: Final = prior_compaction_block.get("content") if prior_compaction_block else None augmented_system: str | list[dict[str, object]] | None = system if isinstance(prior_summary_text, str) and prior_summary_text: diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index ed51d8bc98c..9e71c83fec0 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -50,20 +50,24 @@ AnthropicRequestMessages: TypeAlias = list[AllAnthropicMessageValues] | list[dic BlockRewrite: TypeAlias = Callable[[Mapping[str, object]], Mapping[str, object]] _BLOCK_LIST: Final = TypeAdapter(list[dict[str, object]]) _SplitEdits: TypeAlias = tuple[dict[str, object] | None, Sequence[Mapping[str, object]]] +AnthropicContextManagementSpec: TypeAlias = dict[str, object] | list[dict[str, object]] | None _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter() def _split_context_management( - spec: dict[str, object] | None, # mutable-ok: the request's spec, handed to normalize_spec + spec: AnthropicContextManagementSpec, kwargs: Mapping[str, object], ) -> _SplitEdits: """Return (the edits the Responses API runs natively, the edits the gateway polyfill runs). The Responses API compacts natively (``compact_20260112`` becomes its ``compaction``); any other edit, such as ``clear_tool_uses_20250919``, has no Responses API form and runs in the polyfill, as - it does on the chat completions path. + it does on the chat completions path. An OpenAI-style list spec is not Anthropic's shape and stays + unsent, as it was before the polyfill ran here. """ + if not isinstance(spec, dict): + return None, () if not spec or context_management_explicitly_dropped(kwargs.get("additional_drop_params")): return spec, () edits: Final = normalize_spec(spec) or [] @@ -332,7 +336,7 @@ class LiteLLMMessagesToResponsesAPIHandler: max_tokens: int, messages: AnthropicRequestMessages, model: str, - context_management: dict[str, object] | None = None, + context_management: AnthropicContextManagementSpec = None, metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, @@ -398,7 +402,7 @@ class LiteLLMMessagesToResponsesAPIHandler: max_tokens: int, messages: AnthropicRequestMessages, model: str, - context_management: dict[str, object] | None = None, + context_management: AnthropicContextManagementSpec = None, metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, 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 e29583327af..9d721cfb5f7 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 @@ -231,15 +231,10 @@ SERVER_TOOL_TURN: Final = { 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( +def _clear_above( + messages: list[dict[str, object]], threshold: int +) -> tuple[list[dict[str, object]], AppliedEdit | None]: + return apply_clear_tool_uses_20250919( model=MODEL, messages=messages, tools=None, @@ -251,11 +246,24 @@ def test_an_input_tokens_trigger_counts_blocks_the_token_counter_cannot_read(tur }, ) + +@pytest.mark.parametrize( + "turn", [pytest.param(SERVER_TOOL_TURN, id="server_tool"), pytest.param(COMPACTION_TURN, id="compaction")] +) +def test_an_input_tokens_trigger_skips_blocks_no_adapter_sends(turn: dict[str, object]) -> None: + history: Final = _make_history(n_pairs=3) + sent_tokens: Final = litellm.token_counter(model=MODEL, messages=history) + messages: Final = [history[0], turn, *history[1:]] + + 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 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 + assert _cleared_ids(cleared) == {"toolu_00", "toolu_01"} + assert cleared[1] == turn def _blocks_of(messages: list[dict[str, object]], block_type: str) -> tuple[dict[str, object], ...]: diff --git a/tests/unit/llms/anthropic/messages/context_management/test_compact.py b/tests/unit/llms/anthropic/messages/context_management/test_compact.py index 159c3586c53..27d4fd996d3 100644 --- a/tests/unit/llms/anthropic/messages/context_management/test_compact.py +++ b/tests/unit/llms/anthropic/messages/context_management/test_compact.py @@ -29,7 +29,7 @@ from litellm.llms.anthropic.messages.context_management.editors.compact import ( _augment_system_with_summary, _extract_summary_text, _select_last_user_question, - _slice_around_compaction_block, + slice_around_compaction_block, _strip_compaction_blocks, apply_client_compaction_block_history, apply_compact_20260112, @@ -153,7 +153,7 @@ def test_applied_edits_for_response_includes_compact_when_block_present(): def test_slice_around_compaction_block_found(): messages = _messages_with_compaction("my summary") - sliced, block = _slice_around_compaction_block(messages) + sliced, block = slice_around_compaction_block(messages) assert block is not None assert block["type"] == "compaction" assert block["content"] == "my summary" @@ -164,7 +164,7 @@ def test_slice_around_compaction_block_found(): def test_slice_around_compaction_block_not_found(): messages = _simple_messages() - sliced, block = _slice_around_compaction_block(messages) + sliced, block = slice_around_compaction_block(messages) assert block is None assert sliced is messages # same object, no copy 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 33a19358d3c..40007533e0b 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 @@ -946,27 +946,27 @@ def test_the_sync_path_clears_old_results_too(respx_mock: respx.MockRouter, monk @pytest.mark.asyncio -async def test_a_polyfill_crash_is_a_500_before_the_responses_api_call( +async def test_an_openai_style_context_management_list_is_still_not_sent_to_the_responses_api( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - 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 ) monkeypatch.setenv("OPENAI_API_KEY", "sk-clear-test") - monkeypatch.setattr(handler, "apply_context_management", AsyncMock(side_effect=RuntimeError("editor bug"))) - with pytest.raises(AnthropicContextManagementError) as raised: - 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)]}, - ) - assert raised.value.status_code == 500 - assert not upstream.called + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + + 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=[{"type": "compaction", "compact_threshold": 500}], + ) + + assert "context_management" not in json.loads(upstream.calls.last.request.content) + assert isinstance(served, dict) + assert "context_management" not in served @pytest.mark.asyncio diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index b2a0b735b0f..38eb87debfd 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -649,7 +649,7 @@ async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when assert not any(record.levelno >= logging.ERROR for record in caplog.records) -async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_before_compact_fails() -> None: +async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None: clear_before_compact: Final[JsonValue] = { "edits": [ {