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.
This commit is contained in:
AlonSoko 2026-10-21 21:52:59 +03:00
parent b6f91a221b
commit c6c4541700
3 changed files with 43 additions and 6 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 system_to_text
T = TypeVar("T")
# api.anthropic.com, measured 2026-10-19: text results of 1 to 46 bytes are never cleared or counted toward keep
@ -195,11 +196,20 @@ def _select_in_batches(candidates: Sequence[str], tokens: Mapping[str, int], cle
return tuple(candidates[:batch_end])
def _with_system(
messages: list[dict[str, object]], system: str | list[dict[str, object]] | None
) -> list[dict[str, object]]:
"""The messages to count: like the API, an ``input_tokens`` trigger counts the system prompt too."""
system_text: Final = system_to_text(system)
return [{"role": "system", "content": system_text}, *messages] if system_text else messages
def _trigger_met(
trigger: dict[str, object],
model: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
system: str | list[dict[str, object]] | None,
counted_tool_uses: int,
) -> tuple[bool, int | None]:
"""Return (trigger_met, input_tokens if counted for reuse)."""
@ -215,7 +225,7 @@ def _trigger_met(
threshold = DEFAULT_INPUT_TOKENS_TRIGGER
current_tokens: Final = litellm.token_counter(
model=model,
messages=messages,
messages=_with_system(messages, system),
tools=cast(Any, tools),
)
verbose_logger.debug("context_management polyfill: current_tokens: %s", current_tokens)
@ -332,7 +342,7 @@ 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))
met, counted_tokens = _trigger_met(trigger, model, messages, tools, system, len(clearable))
if not met:
return messages, None
@ -366,9 +376,11 @@ def apply_clear_tool_uses_20250919(
tokens_before: Final = (
counted_tokens
if counted_tokens is not None
else litellm.token_counter(model=model, messages=messages, tools=cast(Any, tools))
else litellm.token_counter(model=model, messages=_with_system(messages, system), tools=cast(Any, tools))
)
tokens_after: Final = litellm.token_counter(
model=model, messages=_with_system(final_messages, system), tools=cast(Any, tools)
)
tokens_after: Final = litellm.token_counter(model=model, messages=final_messages, tools=cast(Any, tools))
cleared_input_tokens: Final = max(tokens_before - tokens_after, 0)
return final_messages, {

View file

@ -320,13 +320,13 @@ def _count_effective_tokens(
content: Final = compaction_block.get("content") or ""
if content:
total += litellm.token_counter(model=model, text=content)
system_text: Final = _system_to_text(system)
system_text: Final = system_to_text(system)
if system_text:
total += litellm.token_counter(model=model, text=system_text)
return total
def _system_to_text(
def system_to_text(
system: str | list[dict[str, object]] | None,
) -> str:
"""Flatten an Anthropic-style ``system`` value into a single string for

View file

@ -690,6 +690,31 @@ 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}]],
)
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
)
assert without_system is None
assert with_system is not None
assert with_system["cleared_tool_uses"] == 2
assert with_system["cleared_input_tokens"] < history_tokens
@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