feat(compress): preserve recent N user messages during context compression

Add _ensure_last_n_user_messages_in_tail to guarantee the last N user
messages survive compression in the uncompressed tail, with surrounding
assistant/tool context preserved.

- Add min_tail_user_messages parameter (default 3) to ContextCompressor
- New _ensure_last_n_user_messages_in_tail method generalizes single-user protection
- Skip context-summary handoff banners when counting user messages
- User messages are clean boundaries — skip _align_boundary_backward
- Wire through cli.py, agent_init.py, and gateway cache busting keys

Config:
  compression:
    min_tail_user_messages: 3

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Jerry
2026-07-13 22:20:47 +08:00
committed by Teknium
parent 69365109b3
commit a9c868225e
5 changed files with 294 additions and 0 deletions
+2
View File
@@ -1810,6 +1810,7 @@ def init_agent(
compression_enabled = str(_compression_cfg.get("enabled", True)).lower() in {"true", "1", "yes"}
compression_target_ratio = float(_compression_cfg.get("target_ratio", 0.20))
compression_protect_last = int(_compression_cfg.get("protect_last_n", 20))
compression_min_tail_users = int(_compression_cfg.get("min_tail_user_messages", 3))
# Cap on compression retry rounds before a turn gives up with "max
# compression attempts reached" (compression.max_attempts). Hardcoding 3
# strands sessions that legitimately need more rounds — e.g. a restart
@@ -2348,6 +2349,7 @@ def init_agent(
proactive_prune_tokens=compression_proactive_prune_tokens,
proactive_prune_min_result_chars=compression_proactive_prune_min_chars,
proactive_prune_min_reclaim_tokens=compression_proactive_prune_min_reclaim,
min_tail_user_messages=compression_min_tail_users,
)
_bind_session_state = getattr(agent.context_compressor, "bind_session_state", None)
if callable(_bind_session_state):
+67
View File
@@ -1873,6 +1873,7 @@ class ContextCompressor(ContextEngine):
proactive_prune_tokens: int = 0,
proactive_prune_min_result_chars: int = 8000,
proactive_prune_min_reclaim_tokens: int = 4096,
min_tail_user_messages: int = 1,
):
self.model = model
self.base_url = base_url
@@ -1927,6 +1928,7 @@ class ContextCompressor(ContextEngine):
self.proactive_prune_min_reclaim_tokens = max(
0, int(proactive_prune_min_reclaim_tokens or 0)
)
self.min_tail_user_messages = min_tail_user_messages
self.summary_target_ratio = max(0.10, min(summary_target_ratio, 0.80))
self.quiet_mode = quiet_mode
# Output-token reservation: the provider carves max_tokens out of the
@@ -4546,6 +4548,60 @@ This compaction should PRIORITISE preserving all information related to the focu
return max(pair_end, head_end + 1)
return adjusted
def _ensure_last_n_user_messages_in_tail(
self,
messages: List[Dict[str, Any]],
cut_idx: int,
head_end: int,
n: int,
) -> int:
"""Guarantee the last N user messages are in the protected tail.
Generalizes ``_ensure_last_user_message_in_tail`` to preserve an
arbitrary number of recent user messages. This prevents the token-
budget-based tail cut from consuming recent conversation turns
when large tool outputs fill the budget (COMPRESS-01).
When *n* <= 1, delegates directly to the existing single-message
method for byte-identical regression safety (COMPRESS-08).
If the conversation has fewer than *n* user messages, the earliest
available user message is used without error (COMPRESS-07).
A user message is already a clean boundary — there is no
tool_call/result group that spans across it, so
``_align_boundary_backward`` is intentionally NOT called.
Calling it can pull the cut past the user message into the
preceding assistant(tool_calls)→tool group and split it (#22566).
"""
if n <= 1:
return self._ensure_last_user_message_in_tail(messages, cut_idx, head_end)
# Collect real user message indices walking backward from end.
# Skip context-summary handoff banners — they are internal
# continuity state, not real user turns.
user_indices = []
for i in range(len(messages) - 1, head_end - 1, -1):
msg = messages[i]
if msg.get("role") == "user" and not self._is_context_summary_content(
msg.get("content")
):
user_indices.append(i)
if len(user_indices) == 0:
return cut_idx
if len(user_indices) < n:
target_idx = user_indices[-1]
else:
target_idx = user_indices[n - 1]
if target_idx >= cut_idx:
return cut_idx
cut_idx = target_idx
return max(cut_idx, head_end + 1)
def _find_turn_pair_end(
self,
messages: List[Dict[str, Any]],
@@ -4675,6 +4731,17 @@ This compaction should PRIORITISE preserving all information related to the focu
# monotonic — the tail can only grow, never shrink.
cut_idx = self._ensure_last_assistant_message_in_tail(messages, cut_idx, head_end)
# Extend to the last N actionable user messages when configured
# (compression.min_tail_user_messages). This prevents the
# token-budget tail from consuming recent turns when large tool
# outputs fill the budget. The anchor only walks ``cut_idx``
# backward (monotonic — the tail can only grow, never shrink), and
# a user message is a clean boundary, so the forward re-alignment
# below remains a no-op for the anchored index.
cut_idx = self._ensure_last_n_user_messages_in_tail(
messages, cut_idx, head_end, self.min_tail_user_messages,
)
# The floor guarantees forward progress — compression must always claim
# at least one message or the caller's compress_start >= compress_end
# guard turns the pass into a no-op that re-runs forever (the same loop
+1
View File
@@ -468,6 +468,7 @@ def load_cli_config() -> Dict[str, Any]:
"compression": {
"enabled": True, # Auto-compress when approaching context limit
"threshold": 0.50, # Compress at 50% of model's context limit
"min_tail_user_messages": 3, # Min recent user messages to preserve in tail after compression
},
"agent": {
"max_turns": 90, # Default max tool-calling iterations (shared with subagents)
+1
View File
@@ -18125,6 +18125,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
("compression", "proactive_prune_tokens"),
("compression", "proactive_prune_min_result_chars"),
("compression", "proactive_prune_min_reclaim_tokens"),
("compression", "min_tail_user_messages"),
("agent", "disabled_toolsets"),
("memory", "provider"),
("checkpoints", "enabled"),
+223
View File
@@ -4076,3 +4076,226 @@ class TestSummaryPromptBounding:
# handoff prefix either.
marker_only = bounded[bounded.index("\n\n...[summary input truncated"):]
assert ContextCompressor.classify_summary_content(marker_only.lstrip()) is None
class TestMinTailUserMessages:
"""COMPRESS-01,02,07,08: N-user-message tail protection.
Tests the ``_ensure_last_n_user_messages_in_tail`` method and its
integration through ``_find_tail_cut_by_tokens``.
"""
def test_n3_preserves_last_3_user_messages(self):
"""COMPRESS-01: _find_tail_cut_by_tokens with min_tail_user_messages=3
guarantees the last 3 user-role messages are in the tail."""
with patch("agent.context_compressor.get_model_context_length", return_value=200_000):
c = ContextCompressor(
model="test/model",
threshold_percent=0.50,
protect_first_n=2,
quiet_mode=True,
)
c.tail_token_budget = 200
messages = [
{"role": "user", "content": "head msg 1"},
{"role": "assistant", "content": "head reply 1"},
{"role": "user", "content": "middle 1"},
{"role": "assistant", "content": "middle 1 reply"},
{"role": "user", "content": "middle 2"},
{"role": "assistant", "content": "middle 2 reply"},
{"role": "user", "content": "recent 3"},
{"role": "assistant", "content": "recent 3 reply"},
{"role": "user", "content": "recent 2"},
{"role": "assistant", "content": "recent 2 reply"},
{"role": "user", "content": "recent 1"},
{"role": "assistant", "content": "recent 1 reply"},
]
head_end = c.protect_first_n
cut = c._find_tail_cut_by_tokens(messages, head_end)
tail_messages = messages[cut:]
tail_user_contents = [m["content"] for m in tail_messages if m["role"] == "user"]
assert len(tail_user_contents) >= 3, (
f"Expected >= 3 user messages in tail, got {len(tail_user_contents)}"
)
assert "recent 1" in tail_user_contents
assert "recent 2" in tail_user_contents
assert "recent 3" in tail_user_contents
assert cut >= head_end + 1
def test_n3_tool_group_integrity(self):
"""COMPRESS-02: When the 3rd-to-last user message is preceded by
assistant(tool_calls) + tool results, the boundary is set at the
user message (a clean boundary). The preceding tool group stays
together — it is either entirely in the compressed region or
entirely in the tail, never split across the boundary."""
with patch("agent.context_compressor.get_model_context_length", return_value=200_000):
c = ContextCompressor(
model="test/model",
threshold_percent=0.50,
protect_first_n=1,
quiet_mode=True,
)
c.tail_token_budget = 200
messages = [
{"role": "user", "content": "start"},
{"role": "assistant", "content": None,
"tool_calls": [{"function": {"name": "read_file", "arguments": "{}"}}]},
{"role": "tool", "content": "result content",
"tool_call_id": "call_1"},
{"role": "user", "content": "user 3rd last"},
{"role": "assistant", "content": "reply 3rd last"},
{"role": "user", "content": "user 2nd last"},
{"role": "assistant", "content": "reply 2nd last"},
{"role": "user", "content": "user last"},
{"role": "assistant", "content": "reply last"},
]
head_end = c.protect_first_n
cut = c._find_tail_cut_by_tokens(messages, head_end)
compressed = messages[:cut]
tool_positions = [
i for i, m in enumerate(compressed)
if m.get("role") in ("assistant", "tool")
and (m.get("tool_calls") or m.get("tool_call_id"))
]
if len(tool_positions) >= 2:
assert tool_positions[-1] - tool_positions[0] == len(tool_positions) - 1, (
"Tool group must stay contiguous across the boundary"
)
tail_messages = messages[cut:]
tail_users = [m["content"] for m in tail_messages if m["role"] == "user"]
assert "user 3rd last" in tail_users
assert "user 2nd last" in tail_users
assert "user last" in tail_users
assert cut >= head_end + 1
def test_n1_regression_safety(self):
"""COMPRESS-08: N=1 produces identical tail positioning to the existing
_ensure_last_user_message_in_tail method."""
with patch("agent.context_compressor.get_model_context_length", return_value=200_000):
c = ContextCompressor(
model="test/model",
threshold_percent=0.85,
protect_first_n=2,
quiet_mode=True,
)
c.min_tail_user_messages = 1
messages = [
{"role": "user", "content": "head 1"},
{"role": "assistant", "content": "head reply 1"},
{"role": "user", "content": "middle user"},
{"role": "assistant", "content": "middle reply"},
{"role": "user", "content": "last user"},
{"role": "assistant", "content": "last reply"},
]
head_end = c.protect_first_n
cut1 = c._find_tail_cut_by_tokens(messages, head_end)
tail1 = messages[cut1:]
# Verify the last user message is in the tail
tail_users = [m["content"] for m in tail1 if m["role"] == "user"]
assert "last user" in tail_users
assert cut1 >= head_end + 1
def test_fewer_than_n_user_messages(self):
"""COMPRESS-07: When the conversation has fewer than N user messages,
the earliest available user message is used without error."""
with patch("agent.context_compressor.get_model_context_length", return_value=200_000):
c = ContextCompressor(
model="test/model",
threshold_percent=0.50,
protect_first_n=2,
quiet_mode=True,
)
messages = [
{"role": "user", "content": "first"},
{"role": "assistant", "content": "reply 1"},
{"role": "user", "content": "second"},
{"role": "assistant", "content": "reply 2"},
]
# Only 2 user messages, but N=5 — should use earliest found
head_end = c.protect_first_n
result = c._ensure_last_n_user_messages_in_tail(
messages, cut_idx=3, head_end=head_end, n=5
)
# Should not crash, boundary should be before the first user message
# (index 0) or at most cut_idx
assert result <= 3
def test_nth_user_already_in_tail_no_reposition(self):
"""When the Nth user message is already in the tail, cut_idx is unchanged."""
with patch("agent.context_compressor.get_model_context_length", return_value=200_000):
c = ContextCompressor(
model="test/model",
threshold_percent=0.50,
protect_first_n=2,
quiet_mode=True,
)
messages = [
{"role": "user", "content": "first"},
{"role": "assistant", "content": "reply 1"},
{"role": "user", "content": "second"},
{"role": "assistant", "content": "reply 2"},
{"role": "user", "content": "third"},
{"role": "assistant", "content": "reply 3"},
]
head_end = c.protect_first_n
# cut_idx at 2 means all users from index 2 onward are in tail
result = c._ensure_last_n_user_messages_in_tail(
messages, cut_idx=2, head_end=head_end, n=3
)
assert result == 2 # unchanged
def test_n5_preserves_last_5_user_messages(self):
"""COMPRESS-06: min_tail_user_messages=5 protects last 5 user messages."""
with patch("agent.context_compressor.get_model_context_length", return_value=200_000):
c = ContextCompressor(
model="test/model",
threshold_percent=0.50,
protect_first_n=1,
quiet_mode=True,
)
c.min_tail_user_messages = 5
c.tail_token_budget = 500
# protect_first_n=1 → head_end=1, so index 0 is head.
# u1..u5 must all be at indices >= head_end+1 (=2) to survive the clamp.
messages = [
{"role": "user", "content": "head 1"}, # 0 (head)
{"role": "assistant", "content": "head reply"}, # 1 (head_end boundary)
{"role": "user", "content": "u1"}, # 2
{"role": "assistant", "content": "a1"}, # 3
{"role": "user", "content": "u2"}, # 4
{"role": "assistant", "content": "a2"}, # 5
{"role": "user", "content": "u3"}, # 6
{"role": "assistant", "content": "a3"}, # 7
{"role": "user", "content": "u4"}, # 8
{"role": "assistant", "content": "a4"}, # 9
{"role": "user", "content": "u5"}, # 10
{"role": "assistant", "content": "a5"}, # 11
]
head_end = c.protect_first_n # = 1
cut = c._find_tail_cut_by_tokens(messages, head_end)
tail_users = [m["content"] for m in messages[cut:] if m["role"] == "user"]
assert len(tail_users) >= 5, f"Expected >=5 users in tail, got {len(tail_users)}"
for u in ("u1", "u2", "u3", "u4", "u5"):
assert u in tail_users
assert cut >= head_end + 1
def test_no_user_messages_beyond_head(self):
"""When there are no user messages beyond head_end, cut_idx is unchanged."""
with patch("agent.context_compressor.get_model_context_length", return_value=200_000):
c = ContextCompressor(
model="test/model",
threshold_percent=0.50,
protect_first_n=5,
quiet_mode=True,
)
messages = [
{"role": "user", "content": "msg 1"},
{"role": "assistant", "content": "reply 1"},
{"role": "user", "content": "msg 2"},
{"role": "assistant", "content": "reply 2"},
]
head_end = c.protect_first_n # = 5 > len(messages)
result = c._ensure_last_n_user_messages_in_tail(
messages, cut_idx=2, head_end=head_end, n=3
)
assert result == 2 # unchanged