fix(context_management)!: count the system prompt toward the clear_tool_uses input_tokens trigger (#47857)

* fix(context_management): count the system prompt toward the clear_tool_uses input_tokens trigger

Anthropic's API compares an input_tokens trigger with the whole prompt, system
included; the polyfill counted messages and tools only, so it cleared late by the
system prompt's size. compact_20260112 already counts it: share its system_to_text.

* fix(context_management): take read-only types in the system prompt count

The type-discipline gate (LIT001) rejects new list and dict annotations.

* chore: rerun CI (proxy-build-from-pip timed out)

* fix(context_management): count the system prompt through compact's shared counter

Finishes the main merge's conflict in clear_tool_uses.py. The input_tokens trigger now adds the system prompt through count_system_tokens, the same counter compact uses, while the count reused for cleared_input_tokens stays messages and tools only, so the reported savings are just the cleared tool results

* perf(context_management): skip the token count when nothing can be cleared

The keep check runs before the trigger count, the system prompt reader accepts any mapping block, and the system prompt test gets short parameter ids

* test(context_management): keep the off-loop count test counting after the keep check moved first

* test(context_management): cover a client compaction summary counting toward the trigger

The trigger also skips the system prompt count when messages and tools alone already pass it

* fix(context_management): move a client compaction summary into the system prompt when clearing runs before compact

* fix(context_management): clear only the history after a client compaction block so compact still replaces its summary

* fix(context_management): keep tool calls beside a client compaction block clearable

---------

Co-authored-by: AlonSoko <25619863+AlonSoko@users.noreply.github.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Alon Sokolovski 2026-10-29 06:35:32 +02:00 • committed by GitHub
parent de42b67252
commit 764a7a6536
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 158 additions and 21 deletions

View file

@ -31,6 +31,7 @@ from ..constants import (
)
from ..errors import AnthropicContextManagementError
from ..placeholders import build_cleared_tool_result_content
from .compact import count_system_tokens, history_after_client_compaction
if TYPE_CHECKING:
from litellm.types.llms.openai import ChatCompletionToolParam
@ -203,9 +204,10 @@ def _trigger_met(
model: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
system: str | Sequence[Mapping[str, object]] | None,
counted_tool_uses: int,
) -> tuple[bool, int | None]:
"""Return (trigger_met, input_tokens if counted for reuse)."""
"""Return (trigger_met, message and tool tokens for reuse); like the API, input_tokens adds the system prompt."""
trigger_type: Final = trigger.get("type", "input_tokens")
threshold = trigger.get("value")
@ -216,14 +218,17 @@ def _trigger_met(
if not isinstance(threshold, int):
threshold = DEFAULT_INPUT_TOKENS_TRIGGER
current_tokens: Final = litellm.token_counter(
message_tokens: Final = litellm.token_counter(
model=model,
messages=messages,
tools=cast("list[ChatCompletionToolParam] | None", tools),
)
if message_tokens > threshold:
return True, message_tokens
current_tokens: Final = message_tokens + count_system_tokens(model, system)
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
return current_tokens > threshold, message_tokens
def _resolve_keep_count(keep: dict[str, object]) -> int:
@ -237,7 +242,7 @@ def _resolve_keep_count(keep: dict[str, object]) -> int:
def _last_completed_tool_use_id(
messages: list[dict[str, object]],
messages: Sequence[Mapping[str, object]],
) -> str | None:
"""Latest completed tool_result id; never cleared."""
last_id: str | None = None
@ -253,7 +258,7 @@ def _last_completed_tool_use_id(
def _clear_tool_results(
messages: list[dict[str, object]], ids_to_clear: set[str]
messages: list[dict[str, object]], ids_to_clear: Set[str]
) -> tuple[list[dict[str, object]], int]:
"""Clear matching tool_result content; return (messages, cleared_count)."""
cleared = 0
@ -311,7 +316,11 @@ 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, it reads only the history after a client compaction block and counts the block's summary as system.
"""
kept, counted_system = history_after_client_compaction(messages, system)
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 ())
@ -326,7 +335,7 @@ def apply_clear_tool_uses_20250919(
"value": DEFAULT_KEEP_TOOL_USES,
}
blocks: Final = _all_blocks(messages)
blocks: Final = _all_blocks(kept)
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)
@ -335,15 +344,15 @@ def apply_clear_tool_uses_20250919(
for tool_use_id, name in _tool_uses_in_order(blocks)
if name not in exclude_tools and (tool_use_id not in tiny or _clears_input(name, input_tools))
)
met, counted_tokens = _trigger_met(trigger, model, messages, tools, len(clearable))
if not met:
return messages, None
keep_count: Final = _resolve_keep_count(keep)
if len(clearable) <= keep_count:
return messages, None
last_completed_id: Final = _last_completed_tool_use_id(messages)
met, counted_tokens = _trigger_met(trigger, model, kept, tools, counted_system, len(clearable))
if not met:
return messages, None
last_completed_id: Final = _last_completed_tool_use_id(kept)
candidates: Final = tuple(
(tool_use_id, name)
for tool_use_id, name in clearable[: len(clearable) - keep_count]
@ -358,7 +367,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(messages, ids_to_clear)
verbose_logger.debug("context_management polyfill: edited: %s", edited)
if cleared_count == 0:
return messages, None
@ -370,11 +379,12 @@ 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=kept, tools=cast("list[ChatCompletionToolParam] | None", tools)
)
)
kept_after, _ = history_after_client_compaction(final_messages, system)
tokens_after: Final = litellm.token_counter(
model=model, messages=final_messages, tools=cast("list[ChatCompletionToolParam] | None", tools)
model=model, messages=kept_after, tools=cast("list[ChatCompletionToolParam] | None", tools)
)
cleared_input_tokens: Final = max(tokens_before - tokens_after, 0)

View file

@ -189,6 +189,21 @@ def _strip_compaction_blocks(
return cleaned
def history_after_client_compaction(
messages: Sequence[dict[str, object]],
system: str | Sequence[Mapping[str, object]] | None,
) -> tuple[list[dict[str, object]], str | Sequence[Mapping[str, object]] | None]:
"""Return the history the API keeps after the latest compaction block, without it, and the system with its summary."""
sliced, block = _slice_around_compaction_block(messages)
if block is None:
return list(messages), system
tail: Final = _strip_compaction_blocks(sliced)
summary: Final = block.get("content")
if not isinstance(summary, str) or not summary:
return tail, system
return tail, _augment_system_with_summary(_system_to_text(system) or None, summary)
def _augment_system_with_summary(
system: str | list[dict[str, object]] | None,
summary_text: str,
@ -320,14 +335,16 @@ def _count_effective_tokens(
content: Final = compaction_block.get("content") or ""
if content:
total += litellm.token_counter(model=model, text=content)
return total + count_system_tokens(model, system)
def count_system_tokens(model: str, system: str | Sequence[Mapping[str, object]] | None) -> int:
system_text: Final = _system_to_text(system)
if system_text:
total += litellm.token_counter(model=model, text=system_text)
return total
return litellm.token_counter(model=model, text=system_text) if system_text else 0
def _system_to_text(
system: str | list[dict[str, object]] | None,
system: str | Sequence[Mapping[str, object]] | None,
) -> str:
"""Flatten an Anthropic-style ``system`` value into a single string for
token counting. Returns ``""`` when ``system`` carries no text."""
@ -338,7 +355,7 @@ def _system_to_text(
return "\n".join(
text
for block in system
if isinstance(block, dict)
if isinstance(block, Mapping)
and block.get("type") == "text"
and isinstance(text := block.get("text"), str)
and text

View file

@ -10,7 +10,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm
from litellm.llms.anthropic.messages.adapters.handler import _run_polyfill_if_enabled
from litellm.llms.anthropic.messages.adapters.handler import _run_polyfill_if_enabled, prepare_context_managed_request
from litellm.llms.anthropic.messages.context_management import (
AnthropicContextManagementError,
PolyfillResult,
@ -690,6 +690,115 @@ def test_tool_uses_trigger_does_not_count_excluded_tools(trigger_value: int, tri
assert (applied is not None) == triggers
@pytest.mark.parametrize(
"system",
["Reference notes. " * 400, [{"type": "text", "text": "Reference notes. " * 400}]],
ids=["string", "blocks"],
)
def test_input_tokens_trigger_counts_the_system_prompt(system: str | list[dict[str, object]]) -> None:
"""Like the API, the trigger compares the whole prompt: a system prompt can push it over."""
messages = _make_history(n_pairs=3)
history_tokens: Final = litellm.token_counter(model=MODEL, messages=messages)
edit_spec: Final = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": history_tokens + 100},
"keep": {"type": "tool_uses", "value": 1},
}
_, without_system = apply_clear_tool_uses_20250919(
model=MODEL, messages=deepcopy(messages), tools=None, system=None, edit_spec=edit_spec
)
_, with_system = apply_clear_tool_uses_20250919(
model=MODEL, messages=deepcopy(messages), tools=None, system=system, edit_spec=edit_spec
)
_, met_without_system = apply_clear_tool_uses_20250919(
model=MODEL,
messages=deepcopy(messages),
tools=None,
system=None,
edit_spec={**edit_spec, "trigger": {"type": "input_tokens", "value": 0}},
)
assert without_system is None
assert with_system is not None
assert met_without_system is not None
assert with_system["cleared_tool_uses"] == 2
assert with_system["cleared_input_tokens"] == met_without_system["cleared_input_tokens"] > 0
@pytest.mark.parametrize("compact_position", [None, "after", "before"], ids=["alone", "before_compact", "after_compact"])
async def test_input_tokens_trigger_counts_a_client_compaction_summary(compact_position: str | None) -> None:
"""A client compaction summary counts toward the trigger and calls beside it stay clearable, in any edit order."""
tail: Final = _make_history(n_pairs=3)
edit_spec: Final = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": litellm.token_counter(model=MODEL, messages=tail) + 100},
"keep": {"type": "tool_uses", "value": 1},
}
compaction: Final = {"type": "compaction", "content": "Earlier work. " * 400}
compacted_history: Final = [tail[0], {**tail[1], "content": [compaction, *tail[1]["content"]]}, *tail[2:]]
compact_spec: Final = {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}}
edits: Final = {None: [edit_spec], "after": [edit_spec, compact_spec], "before": [compact_spec, edit_spec]}[
compact_position
]
async def prepare(messages: list[dict[str, object]]) -> PolyfillResult | None:
return await prepare_context_managed_request(
model=MODEL,
messages=messages,
tools=None,
system=None,
context_management_spec={"edits": edits},
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
)
def cleared(result: PolyfillResult | None) -> list[object]:
applied: Final = result.applied_edits if result is not None else []
return [edit.get("cleared_tool_uses") for edit in applied if edit.get("type") == "clear_tool_uses_20250919"]
without_summary: Final = await prepare(deepcopy(tail))
with_summary: Final = await prepare(deepcopy(compacted_history))
assert cleared(without_summary) == []
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="<summary>Fresh summary</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

View file

@ -151,6 +151,7 @@ async def test_sync_editor_counts_tokens_off_the_event_loop():
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 10_000_000},
"keep": {"type": "tool_uses", "value": 0},
}
]
},