From eed5157c633ae5dc98af35deeab17c28bb6000ad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:11:43 -0700 Subject: [PATCH] fix(context_management): clear only the history after a client compaction block so compact still replaces its summary --- .../anthropic/messages/adapters/handler.py | 20 +++++------ .../editors/clear_tool_uses.py | 21 ++++++----- .../context_management/editors/compact.py | 15 ++++++++ .../test_clear_tool_uses.py | 36 +++++++++++++++++++ 4 files changed, 74 insertions(+), 18 deletions(-) diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index 66b206b528a..71ccfbaf915 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -163,17 +163,17 @@ async def prepare_context_managed_request( ) # Skip the client-history pre-processing when a ``compact_20260112`` - # polyfill edit runs first: that editor already slices around any client-sent + # polyfill spec will run: that editor already slices around any client-sent # compaction block in its Phase A (and uses the full post-compaction tail # for its token-threshold check). Pre-collapsing to just the latest user # question here would starve the polyfill of conversation context and # silently drop intermediate turns. - compact_runs_first: Final = _compact_runs_first( + polyfill_will_run: Final = _polyfill_will_run( context_management_spec=context_management_spec, additional_drop_params=additional_drop_params, ) - if compact_runs_first: + if polyfill_will_run: history_result: PolyfillResult | None = None working_messages: _AnthropicMessages = messages working_system: _AnthropicSystem = system @@ -211,7 +211,7 @@ async def prepare_context_managed_request( # ``_run_polyfill_if_enabled``), apply the slice-only fallback now so # Anthropic-specific ``compaction`` content blocks don't leak through # to non-Anthropic backends that would reject them. - if compact_runs_first and history_result is None: + if polyfill_will_run and history_result is None: history_result = apply_client_compaction_block_history( messages=messages, system=system, @@ -219,18 +219,18 @@ async def prepare_context_managed_request( return history_result -def _compact_runs_first( +def _polyfill_will_run( *, context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: - """Return True when ``compact_20260112`` is the first edit the polyfill dispatcher runs. + """Return True when ``compact_20260112`` will run via the polyfill dispatcher. Mirrors the gating in ``_run_polyfill_if_enabled``: an empty spec or an explicit ``context_management`` entry in ``additional_drop_params`` - short-circuits the polyfill. The pre-processing skip only applies when - ``apply_compact_20260112`` (which has its own compaction-block slicing) - sees the messages before any other editor does. + short-circuits the polyfill. The pre-processing skip only applies when the + dispatcher will actually invoke ``apply_compact_20260112`` (which has its + own compaction-block slicing). """ edits: Final = _normalize_spec_edits( context_management_spec=context_management_spec, @@ -243,7 +243,7 @@ def _compact_runs_first( COMPACT_EDIT_TYPE, ) - return bool(edits) and edits[0].get("type") == COMPACT_EDIT_TYPE + return any(edit.get("type") == COMPACT_EDIT_TYPE for edit in edits) def _spec_has_non_compact_edits( 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 31f0f9e0963..88a4640633c 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 @@ -31,7 +31,7 @@ from ..constants import ( ) from ..errors import AnthropicContextManagementError from ..placeholders import build_cleared_tool_result_content -from .compact import count_system_tokens +from .compact import after_client_compaction, count_system_tokens if TYPE_CHECKING: from litellm.types.llms.openai import ChatCompletionToolParam @@ -316,7 +316,12 @@ def apply_clear_tool_uses_20250919( system: str | list[dict[str, object]] | None, edit_spec: ClearToolUsesEditSpec, ) -> tuple[list[dict[str, object]], AppliedEdit | None]: - """Apply clear_tool_uses to an edit ``check_clear_tool_uses_spec`` returned; return (messages, AppliedEdit or None).""" + """Apply clear_tool_uses to an edit ``check_clear_tool_uses_spec`` returned; return (messages, AppliedEdit or None). + + Like the API, only the history after a client compaction block is edited, with its summary counted as system. + """ + tail_start, counted_system = after_client_compaction(messages, system) or (0, system) + tail: Final = messages[tail_start:] fields: Final = _ClearToolUsesFields.model_validate(edit_spec) clear_at_least: Final = 0 if fields.clear_at_least is None else fields.clear_at_least.value exclude_tools: Final = frozenset(fields.exclude_tools or ()) @@ -331,7 +336,7 @@ def apply_clear_tool_uses_20250919( "value": DEFAULT_KEEP_TOOL_USES, } - blocks: Final = _all_blocks(messages) + blocks: Final = _all_blocks(tail) results: Final = _blocks_by_id(blocks, "tool_result", "tool_use_id") uses: Final = _blocks_by_id(blocks, "tool_use", "id") tiny: Final = _tiny_result_ids(results) @@ -344,11 +349,11 @@ def apply_clear_tool_uses_20250919( if len(clearable) <= keep_count: return messages, None - met, counted_tokens = _trigger_met(trigger, model, messages, tools, system, len(clearable)) + met, counted_tokens = _trigger_met(trigger, model, tail, tools, counted_system, len(clearable)) if not met: return messages, None - last_completed_id: Final = _last_completed_tool_use_id(messages) + last_completed_id: Final = _last_completed_tool_use_id(tail) candidates: Final = tuple( (tool_use_id, name) for tool_use_id, name in clearable[: len(clearable) - keep_count] @@ -363,7 +368,7 @@ def apply_clear_tool_uses_20250919( _select_in_batches(candidate_ids, tokens, clear_at_least) if clear_at_least > 0 else candidate_ids ) - edited, cleared_count = _clear_tool_results(messages, set(ids_to_clear)) + edited, cleared_count = _clear_tool_results(tail, set(ids_to_clear)) verbose_logger.debug("context_management polyfill: edited: %s", edited) if cleared_count == 0: return messages, None @@ -375,7 +380,7 @@ def apply_clear_tool_uses_20250919( counted_tokens if counted_tokens is not None else litellm.token_counter( - model=model, messages=messages, tools=cast("list[ChatCompletionToolParam] | None", tools) + model=model, messages=tail, tools=cast("list[ChatCompletionToolParam] | None", tools) ) ) tokens_after: Final = litellm.token_counter( @@ -383,7 +388,7 @@ def apply_clear_tool_uses_20250919( ) cleared_input_tokens: Final = max(tokens_before - tokens_after, 0) - return final_messages, { + return [*messages[:tail_start], *final_messages], { "type": CLEAR_TOOL_USES_EDIT_TYPE, "cleared_tool_uses": cleared_count, "cleared_input_tokens": cleared_input_tokens, diff --git a/litellm/llms/anthropic/messages/context_management/editors/compact.py b/litellm/llms/anthropic/messages/context_management/editors/compact.py index a4cbb056b8c..0fe9516bcb2 100644 --- a/litellm/llms/anthropic/messages/context_management/editors/compact.py +++ b/litellm/llms/anthropic/messages/context_management/editors/compact.py @@ -167,6 +167,21 @@ def _slice_around_compaction_block( return sliced_messages, compaction_block +def after_client_compaction( + messages: Sequence[Mapping[str, object]], + system: str | Sequence[Mapping[str, object]] | None, +) -> tuple[int, str | Sequence[Mapping[str, object]] | None] | None: + """Return where the history after the latest compaction block starts and the system text with its summary.""" + sliced, block = _slice_around_compaction_block(messages) + if block is None: + return None + tail_start: Final = len(messages) - len(sliced) + 1 + summary: Final = block.get("content") + if not isinstance(summary, str) or not summary: + return tail_start, system + return tail_start, _augment_system_with_summary(_system_to_text(system) or None, summary) + + def _strip_compaction_blocks( messages: Sequence[dict[str, object]], ) -> list[dict[str, object]]: 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 557fd49347a..175f82a6715 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 @@ -762,6 +762,42 @@ async def test_input_tokens_trigger_counts_a_client_compaction_summary(compact_p assert cleared(with_summary) == [2] +async def test_compact_after_clearing_replaces_the_client_compaction_summary(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.proxy_server import general_settings + + monkeypatch.setitem(general_settings, COMPACT_SUMMARY_MODEL_SETTING_KEY, "summary-model") + summary: Final = MagicMock() + summary.choices = [MagicMock(message=MagicMock(content="Fresh summary"))] + summary.usage = MagicMock(prompt_tokens=50, completion_tokens=100) + router: Final = MagicMock() + router.acompletion = AsyncMock(return_value=summary) + tail: Final = _make_history(n_pairs=3) + result: Final = await prepare_context_managed_request( + model=MODEL, + messages=[{"role": "assistant", "content": [{"type": "compaction", "content": "Old summary. " * 20_000}]}, *tail], + tools=None, + system="You are a helpful assistant.", + context_management_spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 1_000}, + "keep": {"type": "tool_uses", "value": 1}, + }, + {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}}, + ] + }, + litellm_metadata=None, + additional_drop_params=None, + llm_router=router, + ) + assert result is not None + assert isinstance(result.system, str) + assert "Fresh summary" in result.system + assert "Old summary" not in result.system + assert result.system.endswith("You are a helpful assistant.") + + @pytest.mark.parametrize(("clear_tool_inputs", "triggers"), [(False, False), (True, True)]) def test_tool_uses_trigger_counts_tiny_results_only_when_their_inputs_are_cleared( clear_tool_inputs: bool, triggers: bool