diff --git a/agent/context_compressor.py b/agent/context_compressor.py index a0ee3068d8ea2..8b5e983e6e336 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -3108,6 +3108,125 @@ This compaction should PRIORITISE preserving all information related to the focu return idx, cls._strip_summary_prefix(_content_text_for_contains(content)) return None, "" + @classmethod + def _strip_context_summary_handoff_message( + cls, + message: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + """Drop stale handoff data while preserving merged prior-tail content.""" + if not isinstance(message, dict): + return message + + content = message.get("content") + is_summary = ( + cls._is_context_summary_content(content) + or cls._has_compressed_summary_metadata(message) + ) + if not is_summary: + return message.copy() + + if isinstance(content, str): + if _MERGED_SUMMARY_DELIMITER in content: + prior = content.split(_MERGED_SUMMARY_DELIMITER, 1)[0].strip() + if prior.startswith(_MERGED_PRIOR_CONTEXT_HEADER): + prior = prior[len(_MERGED_PRIOR_CONTEXT_HEADER):].lstrip() + if prior: + unwrapped = message.copy() + unwrapped["content"] = prior + unwrapped.pop(COMPRESSED_SUMMARY_METADATA_KEY, None) + return unwrapped + else: + marker_idx = content.find(_SUMMARY_END_MARKER) + if marker_idx >= 0: + remainder = content[marker_idx + len(_SUMMARY_END_MARKER):].lstrip() + if remainder: + unwrapped = message.copy() + unwrapped["content"] = remainder + unwrapped.pop(COMPRESSED_SUMMARY_METADATA_KEY, None) + return unwrapped + + if isinstance(content, list): + prior_blocks: list[Any] = [] + found_delimiter = False + for item in content: + if isinstance(item, str): + if _MERGED_SUMMARY_DELIMITER in item: + before = item.split(_MERGED_SUMMARY_DELIMITER, 1)[0] + if before.strip(): + prior_blocks.append(before) + found_delimiter = True + break + prior_blocks.append(item) + continue + if isinstance(item, dict): + text = item.get("text") + if isinstance(text, str) and _MERGED_SUMMARY_DELIMITER in text: + before = text.split(_MERGED_SUMMARY_DELIMITER, 1)[0] + if before.strip(): + copied = item.copy() + copied["text"] = before + prior_blocks.append(copied) + found_delimiter = True + break + prior_blocks.append(item.copy()) + continue + prior_blocks.append(item) + + if not found_delimiter: + legacy_blocks: list[Any] = [] + found_marker = False + for index, item in enumerate(content): + text = item if isinstance(item, str) else item.get("text") if isinstance(item, dict) else None + if not isinstance(text, str) or _SUMMARY_END_MARKER not in text: + continue + remainder = text.split(_SUMMARY_END_MARKER, 1)[1].lstrip() + if remainder: + if isinstance(item, dict): + copied = item.copy() + copied["text"] = remainder + legacy_blocks.append(copied) + else: + legacy_blocks.append(remainder) + for later in content[index + 1:]: + legacy_blocks.append(later.copy() if isinstance(later, dict) else later) + found_marker = True + break + if found_marker and legacy_blocks: + unwrapped = message.copy() + unwrapped["content"] = legacy_blocks + unwrapped.pop(COMPRESSED_SUMMARY_METADATA_KEY, None) + return unwrapped + + if found_delimiter: + for index, item in enumerate(prior_blocks): + if isinstance(item, str): + if item.lstrip().startswith(_MERGED_PRIOR_CONTEXT_HEADER): + leading = item.lstrip()[len(_MERGED_PRIOR_CONTEXT_HEADER):].lstrip() + if leading: + prior_blocks[index] = leading + else: + prior_blocks.pop(index) + break + elif isinstance(item, dict) and isinstance(item.get("text"), str): + text = item["text"] + if text.lstrip().startswith(_MERGED_PRIOR_CONTEXT_HEADER): + leading = text.lstrip()[len(_MERGED_PRIOR_CONTEXT_HEADER):].lstrip() + if leading: + copied = item.copy() + copied["text"] = leading + prior_blocks[index] = copied + else: + prior_blocks.pop(index) + break + + if prior_blocks: + unwrapped = message.copy() + unwrapped["content"] = prior_blocks + unwrapped.pop(COMPRESSED_SUMMARY_METADATA_KEY, None) + return unwrapped + + return None + # ------------------------------------------------------------------ # Tool-call / tool-result pair integrity helpers # ------------------------------------------------------------------ @@ -3943,7 +4062,9 @@ This compaction should PRIORITISE preserving all information related to the focu existing, "\n\n" + _compression_note if isinstance(existing, str) and existing else _compression_note, ) - compressed.append(msg) + stripped = self._strip_context_summary_handoff_message(msg) + if stripped is not None: + compressed.append(stripped) # If LLM summary failed, insert a deterministic fallback so the model # gets at least locally recoverable continuity anchors instead of a @@ -3959,9 +4080,19 @@ This compaction should PRIORITISE preserving all information related to the focu reason=self._last_summary_error, ) + tail_messages: List[Dict[str, Any]] = [] + for i in range(compress_end, n_messages): + msg = _fresh_compaction_message_copy(messages[i]) + stripped = self._strip_context_summary_handoff_message(msg) + if stripped is not None: + tail_messages.append(stripped) + _merge_summary_into_tail = False last_head_role = compressed[-1].get("role", "user") if compressed else "user" - first_tail_role = messages[compress_end].get("role", "user") if compress_end < n_messages else "user" + # NOTE: derive the tail's leading role from tail_messages (post + # handoff-strip), not messages[compress_end] — a stripped stale + # handoff must not influence alternation-safe role selection. + first_tail_role = tail_messages[0].get("role", "user") if tail_messages else None # When the only protected head message is the system prompt, the # summary becomes the first *visible* message in the API request # (most adapters — Anthropic, Bedrock — send the system prompt as @@ -3989,11 +4120,9 @@ This compaction should PRIORITISE preserving all information related to the focu # always has at least one user turn. if not _force_user_leading: _user_survives = any( - messages[i].get("role") == "user" - for i in range(0, compress_start) + message.get("role") == "user" for message in compressed ) or any( - messages[i].get("role") == "user" - for i in range(compress_end, n_messages) + message.get("role") == "user" for message in tail_messages ) if not _user_survives: _force_user_leading = True @@ -4005,7 +4134,7 @@ This compaction should PRIORITISE preserving all information related to the focu summary_role = "assistant" # If the chosen role collides with the tail AND flipping wouldn't # collide with the head, flip it. - if summary_role == first_tail_role: + if first_tail_role and summary_role == first_tail_role: flipped = "assistant" if summary_role == "user" else "user" if flipped != last_head_role and not _force_user_leading: summary_role = flipped @@ -4014,7 +4143,7 @@ This compaction should PRIORITISE preserving all information related to the focu # (e.g. head=assistant, tail=user — neither role works). # Merge the summary into the first tail message instead # of inserting a standalone message that breaks alternation. - _merge_summary_into_tail = True + _merge_summary_into_tail = bool(tail_messages) # When the summary lands as a standalone role="user" message, # weak models read the verbatim "## Active Task" quote of a past @@ -4036,9 +4165,8 @@ This compaction should PRIORITISE preserving all information related to the focu ), }) - for i in range(compress_end, n_messages): - msg = _fresh_compaction_message_copy(messages[i]) - if _merge_summary_into_tail and i == compress_end: + for tail_idx, msg in enumerate(tail_messages): + if _merge_summary_into_tail and tail_idx == 0: # Merge the summary into the first tail message, but place # the END MARKER at the very end so the model sees an # unambiguous boundary. Old tail content is preserved as diff --git a/tests/agent/test_context_compressor_summary_continuity.py b/tests/agent/test_context_compressor_summary_continuity.py index b3c8d7471f9c9..445a16a78e54a 100644 --- a/tests/agent/test_context_compressor_summary_continuity.py +++ b/tests/agent/test_context_compressor_summary_continuity.py @@ -2,7 +2,14 @@ from unittest.mock import MagicMock, patch -from agent.context_compressor import ContextCompressor, SUMMARY_PREFIX +from agent.context_compressor import ( + COMPRESSED_SUMMARY_METADATA_KEY, + ContextCompressor, + SUMMARY_PREFIX, + _MERGED_PRIOR_CONTEXT_HEADER, + _MERGED_SUMMARY_DELIMITER, + _SUMMARY_END_MARKER, +) def _compressor() -> ContextCompressor: @@ -36,6 +43,21 @@ def _messages_with_handoff(summary_body: str): ] +def _messages_with_merged_handoff(summary_body: str, prior_tail: str): + merged = { + "role": "user", + "content": ( + f"{_MERGED_PRIOR_CONTEXT_HEADER}\n{prior_tail}\n\n" + f"{_MERGED_SUMMARY_DELIMITER}\n\n" + f"{SUMMARY_PREFIX}\n{summary_body}\n\n{_SUMMARY_END_MARKER}" + ), + COMPRESSED_SUMMARY_METADATA_KEY: True, + } + messages = _messages_with_handoff(summary_body) + messages[1] = merged + return messages + + def test_existing_previous_summary_is_not_serialized_again_as_new_turn(): """Same-process iterative compression should not feed the old handoff twice.""" compressor = _compressor() @@ -112,3 +134,127 @@ def test_handoff_in_protected_head_is_replaced_not_duplicated(): assert "UPDATED summary body" in str(summary_messages[0]["content"]) assert old_summary not in str(summary_messages[0]["content"]) assert old_summary not in "\n".join(str(msg.get("content") or "") for msg in compressed) + + +def test_recompression_drops_prior_protected_handoff_from_output(): + """Repeated compression must not preserve stale handoff bubbles forever.""" + compressor = _compressor() + old_summary = "DUPLICATE-HANDOFF-BODY unique old facts" + + with patch.object( + compressor, + "_generate_summary", + return_value=ContextCompressor._with_summary_prefix( + "updated summary with old facts folded in" + ), + ): + result = compressor.compress(_messages_with_handoff(old_summary)) + + joined = "\n".join(str(message.get("content", "")) for message in result) + assert old_summary not in joined + assert joined.count(SUMMARY_PREFIX) == 1 + assert "updated summary with old facts folded in" in joined + + +def test_legacy_string_merged_handoff_preserves_real_tail_text(): + """Pre-delimiter string handoffs still unwrap content after the end marker.""" + message = { + "role": "user", + "content": ( + f"{SUMMARY_PREFIX}\nold summary\n\n" + f"{_SUMMARY_END_MARKER}\n\nreal tail message" + ), + COMPRESSED_SUMMARY_METADATA_KEY: True, + } + + result = ContextCompressor._strip_context_summary_handoff_message(message) + + assert result == {"role": "user", "content": "real tail message"} + + +def test_recompression_of_current_merged_handoff_preserves_prior_tail_once(): + """Current merged handoffs lose only stale summary data on recompression.""" + compressor = _compressor() + old_summary = "CURRENT-MERGED-OLD-SUMMARY unique continuity facts" + prior_tail = "PRESERVED-PRIOR-TAIL real user content" + + with patch.object( + compressor, + "_generate_summary", + return_value=ContextCompressor._with_summary_prefix( + "fresh replacement summary" + ), + ): + result = compressor.compress( + _messages_with_merged_handoff(old_summary, prior_tail) + ) + + joined = "\n".join(str(message.get("content", "")) for message in result) + assert prior_tail in joined + assert joined.count(prior_tail) == 1 + assert old_summary not in joined + assert joined.count(SUMMARY_PREFIX) == 1 + assert "fresh replacement summary" in joined + + +def test_current_multimodal_merged_handoff_preserves_original_blocks(): + """Unwrapping current list content must retain text and image blocks.""" + prior_text = {"type": "text", "text": "real multimodal tail"} + prior_image = { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + } + message = { + "role": "user", + "content": [ + {"type": "text", "text": f"{_MERGED_PRIOR_CONTEXT_HEADER}\n"}, + prior_text, + prior_image, + { + "type": "text", + "text": ( + f"\n\n{_MERGED_SUMMARY_DELIMITER}\n\n" + f"{SUMMARY_PREFIX}\nstale summary\n\n{_SUMMARY_END_MARKER}" + ), + }, + ], + COMPRESSED_SUMMARY_METADATA_KEY: True, + } + + result = ContextCompressor._strip_context_summary_handoff_message(message) + + assert result == { + "role": "user", + "content": [prior_text, prior_image], + } + + +def test_legacy_multimodal_merged_handoff_preserves_original_blocks(): + """Persisted pre-delimiter list handoffs must not lose their real tail.""" + prior_text = {"type": "text", "text": "legacy real tail"} + prior_image = { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,BBBB"}, + } + message = { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + f"{SUMMARY_PREFIX}\nlegacy stale summary\n\n" + f"{_SUMMARY_END_MARKER}\n\n" + ), + }, + prior_text, + prior_image, + ], + COMPRESSED_SUMMARY_METADATA_KEY: True, + } + + result = ContextCompressor._strip_context_summary_handoff_message(message) + + assert result == { + "role": "user", + "content": [prior_text, prior_image], + }