fix(compression): correct prune boundary + exempt native compaction checkpoints
Two corrections on top of the #71077 base (the whole bug class): 1. Turn boundary = last USER message, not last assistant message. A Codex turn spans several assistant messages (assistant+tool_calls -> tool -> ... -> final assistant) whose reasoning items must replay together; the last-assistant boundary would strip reasoning mid-chain from the active turn (the gap flagged in PR #71077 review). 2. type="compaction" checkpoints (native server-side compaction, PR #81747) are exempt: they carry already-pruned history, not per-turn reasoning. Pruning filters items instead of popping the sidecar key. Sibling site fixed in the same class: the Codex incomplete-continuation dedup path blind-overwrote codex_reasoning_items on visually-duplicate interim messages, which would drop the only copy of a checkpoint captured on the earlier response. Extracted merge_interim_reasoning_items() into agent/native_compaction.py; newer reasoning wins, prior checkpoints are preserved unless the newer payload carries its own.
This commit is contained in:
Teknium
+44
-17
@@ -247,8 +247,8 @@ def _strip_persistence_markers(messages: List[Dict[str, Any]]) -> None:
|
||||
|
||||
|
||||
def _prune_stale_reasoning_replay(messages: List[Dict[str, Any]]) -> int:
|
||||
"""Strip stale replay fields (``codex_reasoning_items``) from retained
|
||||
assistant messages older than the most recent assistant turn.
|
||||
"""Strip stale per-turn replay items (``codex_reasoning_items``) from
|
||||
assistant messages that belong to turns older than the active one.
|
||||
|
||||
During Codex/Responses sessions, every retained assistant message carries
|
||||
encrypted reasoning blobs (``codex_reasoning_items``) that are only needed
|
||||
@@ -260,32 +260,59 @@ def _prune_stale_reasoning_replay(messages: List[Dict[str, Any]]) -> int:
|
||||
Operates in place on the fully assembled compacted message list. Returns
|
||||
the number of messages that were pruned (for diagnostics). #71058.
|
||||
|
||||
The pruning rule is conservative: everything up to (but not including) the
|
||||
*last* assistant message in the list gets its stale replay fields stripped.
|
||||
The final assistant message retains its items because it is the most recent
|
||||
turn and its replay chain may still be active. When there is no assistant
|
||||
message at all (shouldn't happen in practice, but defensive) nothing is
|
||||
stripped.
|
||||
Two safety rules define the prune:
|
||||
|
||||
* **Turn boundary is the last user message, not the last assistant
|
||||
message.** A single Codex turn spans several assistant messages
|
||||
(assistant+tool_calls -> tool -> assistant+tool_calls -> ... -> final
|
||||
assistant), and the Responses API requires the reasoning items that
|
||||
bridge those function calls to be replayed together. Everything after
|
||||
the last user message is the active turn and keeps its items; only
|
||||
messages at or before that boundary are stale. (An earlier draft used
|
||||
the last assistant message and would have stripped reasoning mid-chain
|
||||
from the in-flight turn.)
|
||||
|
||||
* **Native compaction checkpoints are exempt.** ``type: "compaction"``
|
||||
items in the same sidecar are the server-side stand-in for already
|
||||
pruned history (see ``agent/native_compaction.py``) — cumulative
|
||||
context carriers, not per-turn reasoning. They must survive on every
|
||||
retained message, so pruning filters items instead of popping the key.
|
||||
"""
|
||||
# Find the last assistant message index — everything before it is stale.
|
||||
last_asst_idx = -1
|
||||
# Find the last real user message — everything after it is the active
|
||||
# turn. Synthetic continuation rows and tool results never mark a turn
|
||||
# boundary.
|
||||
last_user_idx = -1
|
||||
for i in range(len(messages) - 1, -1, -1):
|
||||
msg = messages[i]
|
||||
if isinstance(msg, dict) and msg.get("role") == "assistant":
|
||||
last_asst_idx = i
|
||||
if isinstance(msg, dict) and msg.get("role") == "user":
|
||||
last_user_idx = i
|
||||
break
|
||||
if last_asst_idx <= 0:
|
||||
# No assistant message, or only one (nothing to prune).
|
||||
if last_user_idx < 0:
|
||||
# No user boundary found — cannot distinguish the active turn, so
|
||||
# prune nothing (fail open toward correctness, not size).
|
||||
return 0
|
||||
|
||||
pruned = 0
|
||||
for i in range(last_asst_idx):
|
||||
for i in range(last_user_idx):
|
||||
msg = messages[i]
|
||||
if not isinstance(msg, dict) or msg.get("role") != "assistant":
|
||||
continue
|
||||
for key in _STALE_REPLAY_PRUNE_KEYS:
|
||||
if msg.pop(key, None) is not None:
|
||||
pruned += 1
|
||||
items = msg.get(key)
|
||||
if not isinstance(items, list) or not items:
|
||||
continue
|
||||
kept = [
|
||||
item
|
||||
for item in items
|
||||
if isinstance(item, dict) and item.get("type") == "compaction"
|
||||
]
|
||||
if len(kept) == len(items):
|
||||
continue # nothing stale in this sidecar
|
||||
if kept:
|
||||
msg[key] = kept
|
||||
else:
|
||||
msg.pop(key, None)
|
||||
pruned += 1
|
||||
return pruned
|
||||
|
||||
|
||||
|
||||
@@ -6159,7 +6159,20 @@ def run_conversation(
|
||||
"codex_message_items",
|
||||
):
|
||||
if _key in interim_msg:
|
||||
last_msg[_key] = interim_msg[_key]
|
||||
if _key == "codex_reasoning_items":
|
||||
# Merge instead of overwrite: a native
|
||||
# compaction checkpoint captured on the
|
||||
# earlier incomplete response is the only
|
||||
# copy — the continuation won't re-emit
|
||||
# it. See merge_interim_reasoning_items.
|
||||
from agent.native_compaction import (
|
||||
merge_interim_reasoning_items,
|
||||
)
|
||||
last_msg[_key] = merge_interim_reasoning_items(
|
||||
last_msg.get(_key), interim_msg[_key]
|
||||
)
|
||||
else:
|
||||
last_msg[_key] = interim_msg[_key]
|
||||
else:
|
||||
messages.append(interim_msg)
|
||||
agent._emit_interim_assistant_message(interim_msg)
|
||||
|
||||
@@ -154,3 +154,33 @@ def is_native_compaction_rejection(error: Any) -> bool:
|
||||
"""
|
||||
text = str(error or "").lower()
|
||||
return "context_management" in text or "compact_threshold" in text
|
||||
|
||||
|
||||
def merge_interim_reasoning_items(
|
||||
prior_items: Any,
|
||||
new_items: Any,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Merge ``codex_reasoning_items`` across Codex incomplete-continuation
|
||||
dedup, preserving native compaction checkpoints.
|
||||
|
||||
The incomplete-retry path updates a visually-duplicate interim assistant
|
||||
message in place with the newer response's replay payload. A checkpoint
|
||||
captured on the EARLIER response is a cumulative context carrier the
|
||||
continuation won't re-emit (the replayed checkpoint keeps the server
|
||||
render under threshold), so a blind overwrite drops the only copy and the
|
||||
next request balloons back to full history. Rule: newer items win, but
|
||||
prior checkpoints are prepended unless the newer payload carries its own.
|
||||
"""
|
||||
kept_checkpoints = [
|
||||
item
|
||||
for item in (prior_items if isinstance(prior_items, list) else [])
|
||||
if isinstance(item, dict) and item.get("type") == "compaction"
|
||||
]
|
||||
new_list = list(new_items) if isinstance(new_items, list) else []
|
||||
new_has_checkpoint = any(
|
||||
isinstance(item, dict) and item.get("type") == "compaction"
|
||||
for item in new_list
|
||||
)
|
||||
if new_has_checkpoint or not kept_checkpoints:
|
||||
return new_list
|
||||
return kept_checkpoints + new_list
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
"""Tests for stale codex_reasoning_items pruning during compaction (#71058).
|
||||
|
||||
Salvaged from PR #71077 (@webtecnica) with two correctness fixes:
|
||||
the prune boundary is the last USER message (a Codex turn spans multiple
|
||||
assistant messages whose reasoning items must replay together), and native
|
||||
compaction checkpoints (type="compaction") are exempt because they carry
|
||||
already-pruned history, not per-turn reasoning.
|
||||
"""
|
||||
|
||||
from agent.context_compressor import (
|
||||
_STALE_REPLAY_PRUNE_KEYS,
|
||||
_prune_stale_reasoning_replay,
|
||||
)
|
||||
|
||||
|
||||
def _reasoning(item_id="rs_1"):
|
||||
return {"type": "reasoning", "encrypted_content": "blob-" + item_id, "id": item_id}
|
||||
|
||||
|
||||
def _compaction():
|
||||
return {"type": "compaction", "encrypted_content": "checkpoint-blob"}
|
||||
|
||||
|
||||
def test_prior_turn_reasoning_items_are_pruned():
|
||||
messages = [
|
||||
{"role": "user", "content": "turn 1"},
|
||||
{"role": "assistant", "content": "a1", "codex_reasoning_items": [_reasoning("rs_a")]},
|
||||
{"role": "user", "content": "turn 2"},
|
||||
{"role": "assistant", "content": "a2", "codex_reasoning_items": [_reasoning("rs_b")]},
|
||||
]
|
||||
pruned = _prune_stale_reasoning_replay(messages)
|
||||
assert pruned == 1
|
||||
assert "codex_reasoning_items" not in messages[1]
|
||||
# Active turn (after last user message) keeps its items.
|
||||
assert messages[3]["codex_reasoning_items"] == [_reasoning("rs_b")]
|
||||
|
||||
|
||||
def test_multi_message_active_turn_chain_is_never_pruned():
|
||||
"""A Codex turn spans assistant+tool_calls -> tool -> assistant; ALL of
|
||||
the active turn's reasoning items must survive (the #71077 review gap)."""
|
||||
messages = [
|
||||
{"role": "user", "content": "old turn"},
|
||||
{"role": "assistant", "content": "old", "codex_reasoning_items": [_reasoning("rs_old")]},
|
||||
{"role": "user", "content": "active turn"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "t", "arguments": "{}"}}],
|
||||
"codex_reasoning_items": [_reasoning("rs_chain1")],
|
||||
},
|
||||
{"role": "tool", "content": "result", "tool_call_id": "c1"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [{"id": "c2", "type": "function", "function": {"name": "t", "arguments": "{}"}}],
|
||||
"codex_reasoning_items": [_reasoning("rs_chain2")],
|
||||
},
|
||||
{"role": "tool", "content": "result", "tool_call_id": "c2"},
|
||||
{"role": "assistant", "content": "done", "codex_reasoning_items": [_reasoning("rs_final")]},
|
||||
]
|
||||
pruned = _prune_stale_reasoning_replay(messages)
|
||||
assert pruned == 1 # only the old turn
|
||||
assert "codex_reasoning_items" not in messages[1]
|
||||
for idx in (3, 5, 7):
|
||||
assert messages[idx].get("codex_reasoning_items"), f"active-chain msg {idx} lost its items"
|
||||
|
||||
|
||||
def test_native_compaction_checkpoints_survive_pruning():
|
||||
"""type="compaction" items are cumulative context carriers — they must
|
||||
survive on stale messages even when reasoning items are stripped."""
|
||||
messages = [
|
||||
{"role": "user", "content": "turn 1"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "a1",
|
||||
"codex_reasoning_items": [_compaction(), _reasoning("rs_a")],
|
||||
},
|
||||
{"role": "user", "content": "turn 2"},
|
||||
{"role": "assistant", "content": "a2"},
|
||||
]
|
||||
pruned = _prune_stale_reasoning_replay(messages)
|
||||
assert pruned == 1
|
||||
# Reasoning stripped, checkpoint kept.
|
||||
assert messages[1]["codex_reasoning_items"] == [_compaction()]
|
||||
|
||||
|
||||
def test_checkpoint_only_sidecar_untouched_and_uncounted():
|
||||
messages = [
|
||||
{"role": "user", "content": "turn 1"},
|
||||
{"role": "assistant", "content": "a1", "codex_reasoning_items": [_compaction()]},
|
||||
{"role": "user", "content": "turn 2"},
|
||||
{"role": "assistant", "content": "a2"},
|
||||
]
|
||||
pruned = _prune_stale_reasoning_replay(messages)
|
||||
assert pruned == 0
|
||||
assert messages[1]["codex_reasoning_items"] == [_compaction()]
|
||||
|
||||
|
||||
def test_no_user_boundary_prunes_nothing():
|
||||
messages = [
|
||||
{"role": "assistant", "content": "a1", "codex_reasoning_items": [_reasoning("rs_a")]},
|
||||
{"role": "assistant", "content": "a2", "codex_reasoning_items": [_reasoning("rs_b")]},
|
||||
]
|
||||
assert _prune_stale_reasoning_replay(messages) == 0
|
||||
assert messages[0]["codex_reasoning_items"]
|
||||
assert messages[1]["codex_reasoning_items"]
|
||||
|
||||
|
||||
def test_non_codex_messages_untouched():
|
||||
messages = [
|
||||
{"role": "user", "content": "u1"},
|
||||
{"role": "assistant", "content": "plain"},
|
||||
{"role": "user", "content": "u2"},
|
||||
{"role": "assistant", "content": "plain2"},
|
||||
]
|
||||
assert _prune_stale_reasoning_replay(messages) == 0
|
||||
assert messages == [
|
||||
{"role": "user", "content": "u1"},
|
||||
{"role": "assistant", "content": "plain"},
|
||||
{"role": "user", "content": "u2"},
|
||||
{"role": "assistant", "content": "plain2"},
|
||||
]
|
||||
|
||||
|
||||
def test_prune_keys_contract():
|
||||
"""codex_message_items are replayed for prefix-cache continuity and must
|
||||
NOT be in the prune set; the prune targets reasoning blobs only."""
|
||||
assert "codex_reasoning_items" in _STALE_REPLAY_PRUNE_KEYS
|
||||
assert "codex_message_items" not in _STALE_REPLAY_PRUNE_KEYS
|
||||
|
||||
|
||||
class TestInterimMergePreservesCheckpoints:
|
||||
"""Sibling site: the Codex incomplete-continuation dedup path must not
|
||||
drop checkpoints when overwriting a visually-duplicate interim message."""
|
||||
|
||||
def test_prior_checkpoint_survives_overwrite(self):
|
||||
from agent.native_compaction import merge_interim_reasoning_items
|
||||
|
||||
prior = [_compaction(), _reasoning("rs_old")]
|
||||
newer = [_reasoning("rs_new")]
|
||||
merged = merge_interim_reasoning_items(prior, newer)
|
||||
assert _compaction() in merged
|
||||
assert _reasoning("rs_new") in merged
|
||||
assert _reasoning("rs_old") not in merged # newer reasoning wins
|
||||
|
||||
def test_newer_checkpoint_wins_outright(self):
|
||||
from agent.native_compaction import merge_interim_reasoning_items
|
||||
|
||||
prior = [{"type": "compaction", "encrypted_content": "old-ckpt"}]
|
||||
newer = [{"type": "compaction", "encrypted_content": "new-ckpt"}, _reasoning("rs_new")]
|
||||
merged = merge_interim_reasoning_items(prior, newer)
|
||||
assert merged == newer
|
||||
|
||||
def test_no_prior_checkpoint_is_plain_overwrite(self):
|
||||
from agent.native_compaction import merge_interim_reasoning_items
|
||||
|
||||
assert merge_interim_reasoning_items(
|
||||
[_reasoning("rs_old")], [_reasoning("rs_new")]
|
||||
) == [_reasoning("rs_new")]
|
||||
|
||||
def test_non_list_inputs_are_safe(self):
|
||||
from agent.native_compaction import merge_interim_reasoning_items
|
||||
|
||||
assert merge_interim_reasoning_items(None, None) == []
|
||||
assert merge_interim_reasoning_items(None, [_reasoning("r")]) == [_reasoning("r")]
|
||||
assert merge_interim_reasoning_items([_compaction()], None) == [_compaction()]
|
||||
Reference in New Issue
Block a user