fix(compression): salvage follow-up — todo snapshot last-resort, reuse prune helpers
Review follow-up on the salvaged #90353:
- Todo snapshot (+ coupled pruned-skill reload notice, 7a16840add) is now
reduced only as a LAST resort after reasoning/tool/summary shrink ops,
and the reload notice survives even then.
- Reuse existing helpers/constants instead of re-hardcoding:
_PRUNED_TOOL_PLACEHOLDER, _PRUNE_MIN_CHARS, _NEWEST_TURN_ONLY_BUDGET_KEYS,
and _prune_stale_reasoning_replay (codex sidecar shrink, #71058 boundary).
- Assistant-role messages without the summary metadata key are no longer
truncatable by the summary-cap heuristic.
- Caller passes budget so the estimator runs 3x, not 5x, per would-grow pass.
This commit is contained in:
+57
-12
@@ -350,10 +350,8 @@ _SUMMARY_END_MARKER = (
|
||||
_MERGED_PRIOR_CONTEXT_HEADER = "[PRIOR CONTEXT — for reference only; not a new message]"
|
||||
_MERGED_SUMMARY_DELIMITER = "[END OF PRIOR CONTEXT — COMPACTION SUMMARY BELOW]"
|
||||
|
||||
_SALVAGE_TOOL_PLACEHOLDER = "[Old tool output cleared to save context space]"
|
||||
_SALVAGE_SUMMARY_MAX_CHARS = 8_000
|
||||
_SALVAGE_KEEP_RECENT_TOOLS = 2
|
||||
_SALVAGE_REASONING_KEYS = ("reasoning", "reasoning_content", "reasoning_details")
|
||||
|
||||
|
||||
def _looks_like_compaction_summary(msg: Dict[str, Any], content: str) -> bool:
|
||||
@@ -363,11 +361,15 @@ def _looks_like_compaction_summary(msg: Dict[str, Any], content: str) -> bool:
|
||||
return False
|
||||
if content.startswith(_MERGED_PRIOR_CONTEXT_HEADER):
|
||||
return False
|
||||
# Content heuristics alone must never authorize mutating a live user turn.
|
||||
# Compressor-generated user-role summaries carry this private marker;
|
||||
# ordinary user input does not.
|
||||
# Content heuristics alone must never authorize mutating a live turn.
|
||||
# Compressor-generated summaries carry this private marker; ordinary
|
||||
# user input — and live assistant replies or kept tool bodies that
|
||||
# merely quote a summary header/marker — do not. Tool messages are
|
||||
# handled exclusively by the stub/keep-recent pass, never the cap.
|
||||
if msg.get("role") == "tool":
|
||||
return False
|
||||
if (
|
||||
msg.get("role") == "user"
|
||||
msg.get("role") in ("user", "assistant")
|
||||
and not msg.get(COMPRESSED_SUMMARY_METADATA_KEY)
|
||||
):
|
||||
return False
|
||||
@@ -380,9 +382,40 @@ def _looks_like_compaction_summary(msg: Dict[str, Any], content: str) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def _salvage_reduce_todo_snapshot(out: List[Dict[str, Any]]) -> None:
|
||||
"""Last-resort shrink: reduce or drop the synthetic todo snapshot.
|
||||
|
||||
The snapshot is the only in-transcript todo re-injection at a compaction
|
||||
boundary, and since 7a16840add the pruned-skill reload notice is coupled
|
||||
into the same string — so it is only touched when the cheaper shrink ops
|
||||
could not get under budget. When the snapshot carries a reload notice,
|
||||
keep just the notice (the coupling must survive salvage); otherwise drop
|
||||
the row entirely.
|
||||
"""
|
||||
from agent.conversation_compression import _PRUNED_SKILL_RELOAD_NOTICE_HEADER
|
||||
|
||||
for i in range(len(out) - 1, -1, -1):
|
||||
msg = out[i]
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
if msg.get("_todo_snapshot_synthetic") and msg.get("role") == "user":
|
||||
content = msg.get("content")
|
||||
notice_idx = (
|
||||
content.find(_PRUNED_SKILL_RELOAD_NOTICE_HEADER)
|
||||
if isinstance(content, str)
|
||||
else -1
|
||||
)
|
||||
if isinstance(content, str) and notice_idx >= 0:
|
||||
msg["content"] = content[notice_idx:]
|
||||
else:
|
||||
del out[i]
|
||||
return
|
||||
|
||||
|
||||
def salvage_grown_transcript(
|
||||
original: List[Dict[str, Any]],
|
||||
candidate: List[Dict[str, Any]],
|
||||
budget: Optional[int] = None,
|
||||
) -> Optional[List[Dict[str, Any]]]:
|
||||
"""Mechanically shrink a compression candidate, or return ``None``.
|
||||
|
||||
@@ -390,10 +423,17 @@ def salvage_grown_transcript(
|
||||
tool bodies, stale reasoning, or a synthetic todo snapshot tip the final
|
||||
candidate over the input size. Work on copies and admit the salvage only
|
||||
when the same rough estimator proves it is strictly smaller than the input.
|
||||
|
||||
Shrink order is cheapest-information-loss first: stale reasoning keys and
|
||||
codex replay sidecars, then old tool bodies, then an oversized summary cap.
|
||||
The synthetic todo snapshot (which carries the pruned-skill reload notice,
|
||||
see ``_salvage_reduce_todo_snapshot``) is only reduced as a LAST resort
|
||||
when everything else still leaves the candidate at or over budget.
|
||||
"""
|
||||
if not candidate or not original:
|
||||
return None
|
||||
budget = estimate_messages_tokens_rough(original)
|
||||
if budget is None:
|
||||
budget = estimate_messages_tokens_rough(original)
|
||||
if budget <= 0:
|
||||
return None
|
||||
|
||||
@@ -404,8 +444,6 @@ def salvage_grown_transcript(
|
||||
if not isinstance(msg, dict):
|
||||
out.append(msg)
|
||||
continue
|
||||
if msg.get("_todo_snapshot_synthetic") and msg.get("role") == "user":
|
||||
continue
|
||||
copied = dict(msg)
|
||||
out.append(copied)
|
||||
role = copied.get("role")
|
||||
@@ -414,17 +452,18 @@ def salvage_grown_transcript(
|
||||
elif role == "assistant":
|
||||
last_assistant_idx = len(out) - 1
|
||||
|
||||
salvage_reasoning_keys = _NEWEST_TURN_ONLY_BUDGET_KEYS + ("reasoning_details",)
|
||||
keep_tools = set(tool_indices[-_SALVAGE_KEEP_RECENT_TOOLS:])
|
||||
for index, msg in enumerate(out):
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
if msg.get("role") == "assistant" and index != last_assistant_idx:
|
||||
for key in _SALVAGE_REASONING_KEYS:
|
||||
for key in salvage_reasoning_keys:
|
||||
msg.pop(key, None)
|
||||
if msg.get("role") == "tool" and index not in keep_tools:
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str) and len(content) > 200:
|
||||
msg["content"] = _SALVAGE_TOOL_PLACEHOLDER
|
||||
if isinstance(content, str) and len(content) > _PRUNE_MIN_CHARS:
|
||||
msg["content"] = _PRUNED_TOOL_PLACEHOLDER
|
||||
content = msg.get("content")
|
||||
if (
|
||||
isinstance(content, str)
|
||||
@@ -436,6 +475,12 @@ def salvage_grown_transcript(
|
||||
+ "\n…[summary truncated so compaction can shrink]\n\n"
|
||||
+ _SUMMARY_END_MARKER
|
||||
)
|
||||
# Heavier codex replay sidecars (encrypted reasoning blobs) — reuse the
|
||||
# proven prune with its last-user-turn safety boundary (#71058).
|
||||
_prune_stale_reasoning_replay(out)
|
||||
|
||||
if estimate_messages_tokens_rough(out) >= budget:
|
||||
_salvage_reduce_todo_snapshot(out)
|
||||
|
||||
if not any(
|
||||
isinstance(message, dict) and message.get("role") == "user"
|
||||
|
||||
@@ -3392,7 +3392,9 @@ def compress_context(
|
||||
# candidate over. Give it one mechanical salvage pass.
|
||||
from agent.context_compressor import salvage_grown_transcript
|
||||
|
||||
_salvaged = salvage_grown_transcript(messages, compressed)
|
||||
_salvaged = salvage_grown_transcript(
|
||||
messages, compressed, budget=_rough_in
|
||||
)
|
||||
if _salvaged is not None:
|
||||
_salv_est = estimate_messages_tokens_rough(_salvaged)
|
||||
if _salv_est < _rough_in:
|
||||
|
||||
@@ -8,7 +8,12 @@ from agent.context_compressor import (
|
||||
from agent.model_metadata import estimate_messages_tokens_rough
|
||||
|
||||
|
||||
def test_salvage_stubs_old_tools_and_drops_todo():
|
||||
def test_salvage_stubs_old_tools_and_keeps_todo_when_stubbing_suffices():
|
||||
"""Tool stubbing alone gets under budget → the todo snapshot survives.
|
||||
|
||||
The snapshot is the only in-transcript todo re-injection at the boundary
|
||||
(and may carry the pruned-skill reload notice), so it is last-resort only.
|
||||
"""
|
||||
original = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
@@ -29,12 +34,69 @@ def test_salvage_stubs_old_tools_and_drops_todo():
|
||||
|
||||
assert out is not None
|
||||
assert estimate_messages_tokens_rough(out) < estimate_messages_tokens_rough(original)
|
||||
assert not any(m.get("_todo_snapshot_synthetic") for m in out)
|
||||
assert any(m.get("_todo_snapshot_synthetic") for m in out)
|
||||
tools = [m["content"] for m in out if m.get("role") == "tool"]
|
||||
assert tools[-1] == "keep-latest"
|
||||
assert any("cleared to save context space" in t for t in tools)
|
||||
|
||||
|
||||
def test_salvage_drops_todo_only_as_last_resort():
|
||||
"""When cheaper ops cannot get under budget, the snapshot is dropped."""
|
||||
original = [
|
||||
{"role": "user", "content": "please do the thing " + ("o" * 600)},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
grown = [
|
||||
{"role": "user", "content": "summary of the ask"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Current todos:\n- [ ] " + ("t" * 800),
|
||||
"_todo_snapshot_synthetic": True,
|
||||
},
|
||||
]
|
||||
assert estimate_messages_tokens_rough(grown) > estimate_messages_tokens_rough(original)
|
||||
|
||||
out = salvage_grown_transcript(original, grown)
|
||||
|
||||
assert out is not None
|
||||
assert estimate_messages_tokens_rough(out) < estimate_messages_tokens_rough(original)
|
||||
assert not any(m.get("_todo_snapshot_synthetic") for m in out)
|
||||
|
||||
|
||||
def test_salvage_last_resort_preserves_pruned_skill_reload_notice():
|
||||
"""7a16840add couples the reload notice into the snapshot — it survives."""
|
||||
from agent.conversation_compression import _PRUNED_SKILL_RELOAD_NOTICE_HEADER
|
||||
|
||||
notice = (
|
||||
f"{_PRUNED_SKILL_RELOAD_NOTICE_HEADER}\n"
|
||||
"Reload with skill_view(name='example-skill') before acting."
|
||||
)
|
||||
original = [
|
||||
{"role": "user", "content": "please do the thing " + ("o" * 3000)},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
grown = [
|
||||
{"role": "user", "content": "summary of the ask"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Current todos:\n- [ ] " + ("t" * 4000) + f"\n\n{notice}",
|
||||
"_todo_snapshot_synthetic": True,
|
||||
},
|
||||
]
|
||||
assert estimate_messages_tokens_rough(grown) > estimate_messages_tokens_rough(original)
|
||||
|
||||
out = salvage_grown_transcript(original, grown)
|
||||
|
||||
assert out is not None
|
||||
assert estimate_messages_tokens_rough(out) < estimate_messages_tokens_rough(original)
|
||||
snapshot_rows = [m for m in out if m.get("_todo_snapshot_synthetic")]
|
||||
assert len(snapshot_rows) == 1
|
||||
assert snapshot_rows[0]["content"].startswith(_PRUNED_SKILL_RELOAD_NOTICE_HEADER)
|
||||
assert "Current todos" not in snapshot_rows[0]["content"]
|
||||
|
||||
|
||||
def test_salvage_returns_none_when_nothing_can_shrink():
|
||||
original = [{"role": "user", "content": "tiny"}]
|
||||
huge = [{"role": "user", "content": "X" * 200_000}]
|
||||
|
||||
@@ -345,7 +345,9 @@ class TestInPlaceAntiGrowthGuard:
|
||||
tool_bodies = [m.get("content") for m in compressed if m.get("role") == "tool"]
|
||||
assert any(isinstance(body, str) and body.startswith("keep-me") for body in tool_bodies)
|
||||
assert any("cleared to save context space" in (body or "") for body in tool_bodies)
|
||||
assert not any(m.get("_todo_snapshot_synthetic") for m in compressed)
|
||||
# Tool stubbing alone got under budget, so the todo snapshot (the
|
||||
# only in-transcript todo re-injection) survives the salvage.
|
||||
assert any(m.get("_todo_snapshot_synthetic") for m in compressed)
|
||||
|
||||
def test_in_place_still_commits_shrinking_compression(self):
|
||||
"""The guard must not block legitimate compressions — a result SMALLER
|
||||
|
||||
Reference in New Issue
Block a user