diff --git a/agent/agent_init.py b/agent/agent_init.py index e239c48cfbd9..6271453cd1e0 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -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): diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 1eb10489eb3c..8166398175e8 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -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 diff --git a/cli.py b/cli.py index 7d86deb1d7fc..d0515cb10720 100644 --- a/cli.py +++ b/cli.py @@ -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) diff --git a/gateway/run.py b/gateway/run.py index 98f6b86da786..1a3be8a897a5 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -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"), diff --git a/tests/agent/test_context_compressor.py b/tests/agent/test_context_compressor.py index 926678108fef..26446e266ab9 100644 --- a/tests/agent/test_context_compressor.py +++ b/tests/agent/test_context_compressor.py @@ -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