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:
kshitijk4poor
2026-08-20 11:04:39 +05:30
committed by kshitij
parent 5c03fbedc6
commit 0596ccdeb3
4 changed files with 127 additions and 16 deletions
+57 -12
View File
@@ -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"
+3 -1
View File
@@ -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:
+64 -2
View File
@@ -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}]
+3 -1
View File
@@ -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