fix: dedupe persisted compaction handoffs
This commit is contained in:
WXBR
Teknium
parent
020bd1ba0a
commit
2b84ed921c
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue