fix(compression): durable-sync the prune runway on model switch + fast no-op for incapable stores
Three review follow-ups on the salvaged #79286 commit: - update_model() zeroed the in-memory prune runway but left the durable model_config copy stale, breaking the method's own durable-sync discipline (the strike reset three lines above keeps its durable copy in sync). A restart after a model switch resurrected a runway computed under the old model's trigger sizes. New _clear_durable_proactive_prune_rearm() removes the persisted key via patch_session_model_config() without touching the transcript. - The archive_and_compact capability check ran AFTER the expensive 3-pass prune scan, so a duck-typed session store lacking the method paid the full scan on every eligible iteration forever with pruning permanently no-opping. Hoist it above the scan (all in-tree stores pass a real SessionDB; this only affects third-party stores). - _load_proactive_prune_rearm_tokens now uses the shared get_session_model_config_value() accessor instead of inlining a 5th copy of the model_config JSON parse, matching its sibling loaders' typed-accessor pattern. Also documents why the rotation-publish-failure branch restores only the runway field rather than the full attempt snapshot. Tests: model-switch durable clear, patch_session_model_config merge/delete/no-op, and a guard proving incapable stores skip the scan.
This commit is contained in:
kshitij
kshitij
parent
565b2c42eb
commit
241605d1ea
|
|
@ -1730,19 +1730,11 @@ class ContextCompressor(ContextEngine):
|
|||
"""Restore the cache-boundary runway for a resumed durable session."""
|
||||
session_db = getattr(self, "_session_db", None)
|
||||
session_id = getattr(self, "_session_id", "")
|
||||
getter = getattr(session_db, "get_session", None)
|
||||
getter = getattr(session_db, "get_session_model_config_value", None)
|
||||
if not session_id or not callable(getter):
|
||||
return
|
||||
try:
|
||||
session = getter(session_id) or {}
|
||||
raw = session.get("model_config")
|
||||
if isinstance(raw, str):
|
||||
raw = json.loads(raw) if raw.strip() else {}
|
||||
value = (
|
||||
raw.get(PROACTIVE_PRUNE_REARM_MODEL_CONFIG_KEY, 0)
|
||||
if isinstance(raw, dict)
|
||||
else 0
|
||||
)
|
||||
value = getter(session_id, PROACTIVE_PRUNE_REARM_MODEL_CONFIG_KEY, 0)
|
||||
self._proactive_prune_rearm_tokens = max(
|
||||
0,
|
||||
int(value) if isinstance(value, (int, float, str)) else 0,
|
||||
|
|
@ -1752,6 +1744,23 @@ class ContextCompressor(ContextEngine):
|
|||
except Exception as exc:
|
||||
logger.debug("proactive prune runway lookup failed (non-sqlite): %s", exc)
|
||||
|
||||
def _clear_durable_proactive_prune_rearm(self) -> None:
|
||||
"""Remove the persisted runway key without touching the transcript.
|
||||
|
||||
Best-effort companion to zeroing the in-memory mirror at sites that
|
||||
void the runway (model switch): without it a restart would reload a
|
||||
runway computed under thresholds that no longer apply.
|
||||
"""
|
||||
session_db = getattr(self, "_session_db", None)
|
||||
session_id = getattr(self, "_session_id", "")
|
||||
patcher = getattr(session_db, "patch_session_model_config", None)
|
||||
if not session_id or not callable(patcher):
|
||||
return
|
||||
try:
|
||||
patcher(session_id, {PROACTIVE_PRUNE_REARM_MODEL_CONFIG_KEY: None})
|
||||
except Exception as exc:
|
||||
logger.debug("proactive prune runway clear failed: %s", exc)
|
||||
|
||||
def _persist_fallback_compression_streak(self) -> None:
|
||||
session_db = getattr(self, "_session_db", None)
|
||||
session_id = getattr(self, "_session_id", "")
|
||||
|
|
@ -2111,7 +2120,12 @@ class ContextCompressor(ContextEngine):
|
|||
self._clear_compression_failure_cooldown()
|
||||
self._verify_compaction_cleared_threshold = False
|
||||
self._last_compression_made_progress = False
|
||||
# The prune runway was computed against the PREVIOUS model's trigger
|
||||
# sizes. Same durable-sync discipline as the strike reset above: clear
|
||||
# the model_config copy too, so a restart doesn't resurrect a runway
|
||||
# this recalibration just voided.
|
||||
self._proactive_prune_rearm_tokens = 0
|
||||
self._clear_durable_proactive_prune_rearm()
|
||||
|
||||
# When the MINIMUM_CONTEXT_LENGTH floor meets/exceeds a small context
|
||||
# window, compacting at the percentage (50% → 32K of a 64K window) wastes
|
||||
|
|
@ -3101,6 +3115,18 @@ class ContextCompressor(ContextEngine):
|
|||
before = sum(_estimate_msg_budget_tokens(m) for m in messages)
|
||||
if before < self._proactive_prune_rearm_tokens:
|
||||
return messages, 0
|
||||
# Capability gate BEFORE the expensive 3-pass scan: a bound store that
|
||||
# can't persist the prune atomically (duck-typed/plugin session store
|
||||
# without archive_and_compact) makes every prune a permanent no-op, so
|
||||
# don't pay the scan for it on every eligible iteration.
|
||||
session_db = getattr(self, "_session_db", None)
|
||||
session_id = getattr(self, "_session_id", "")
|
||||
if (
|
||||
session_db
|
||||
and session_id
|
||||
and not callable(getattr(session_db, "archive_and_compact", None))
|
||||
):
|
||||
return messages, 0
|
||||
pruned_msgs, pruned_count = self._prune_old_tool_results(
|
||||
messages,
|
||||
protect_tail_count=self.protect_last_n,
|
||||
|
|
@ -3128,14 +3154,10 @@ class ContextCompressor(ContextEngine):
|
|||
self.proactive_prune_min_reclaim_tokens,
|
||||
)
|
||||
next_rearm_tokens = after + runway
|
||||
session_db = getattr(self, "_session_db", None)
|
||||
session_id = getattr(self, "_session_id", "")
|
||||
if session_db and session_id:
|
||||
archive_and_compact = getattr(session_db, "archive_and_compact", None)
|
||||
if not callable(archive_and_compact):
|
||||
return messages, 0
|
||||
# The capability gate above guarantees archive_and_compact exists.
|
||||
try:
|
||||
archive_and_compact(
|
||||
session_db.archive_and_compact(
|
||||
session_id,
|
||||
pruned_msgs,
|
||||
model_config_patch={
|
||||
|
|
|
|||
|
|
@ -3340,6 +3340,18 @@ def compress_context(
|
|||
messages[:] = copy.deepcopy(messages_before_compression)
|
||||
compressed = messages
|
||||
_compression_made_progress = False
|
||||
# Restore ONLY the prune runway, not the full attempt
|
||||
# snapshot: _restore_compressor_attempt_state is reserved
|
||||
# for pre-commit cancels (fence deny / explicit cancel),
|
||||
# while this branch is post-attempt — the other snapshot
|
||||
# fields (telemetry, aborted flags) must keep the failed
|
||||
# attempt's values. The runway is a property of transcript
|
||||
# state, and the transcript was just rolled back to its
|
||||
# pre-compression copy, so the runway rolls back with it.
|
||||
# (compress() zeroed it in-memory on summary success; the
|
||||
# durable copy was never cleared — that clear only rides
|
||||
# the atomic archive_and_compact / child-row publication
|
||||
# that just failed.)
|
||||
if "_proactive_prune_rearm_tokens" in _compressor_attempt_snapshot:
|
||||
agent.context_compressor._proactive_prune_rearm_tokens = (
|
||||
_compressor_attempt_snapshot[
|
||||
|
|
|
|||
|
|
@ -209,3 +209,65 @@ def test_archive_model_config_patch_rolls_back_with_transcript(tmp_path: Path) -
|
|||
|
||||
assert db.get_messages_as_conversation(session_id)[0]["content"] == "original"
|
||||
assert _model_config(db, session_id) == {"keep": "value", _REARM_KEY: 120_000}
|
||||
|
||||
|
||||
def test_model_switch_clears_durable_runway(tmp_path: Path) -> None:
|
||||
"""update_model must clear BOTH the in-memory and the durable runway."""
|
||||
db = SessionDB(db_path=tmp_path / "state.db")
|
||||
session_id = "MODEL_SWITCH_CLEARS_RUNWAY"
|
||||
db.create_session(
|
||||
session_id,
|
||||
source="telegram",
|
||||
model_config={"keep": "value", _REARM_KEY: 120_000},
|
||||
)
|
||||
agent = _build_agent(db, session_id)
|
||||
compressor = agent.context_compressor
|
||||
assert compressor._proactive_prune_rearm_tokens == 120_000
|
||||
|
||||
compressor.update_model("other/model", 200_000)
|
||||
|
||||
assert compressor._proactive_prune_rearm_tokens == 0
|
||||
assert _REARM_KEY not in _model_config(db, session_id)
|
||||
assert _model_config(db, session_id)["keep"] == "value"
|
||||
|
||||
|
||||
def test_patch_session_model_config_merge_and_delete(tmp_path: Path) -> None:
|
||||
db = SessionDB(db_path=tmp_path / "state.db")
|
||||
session_id = "PATCH_MODEL_CONFIG"
|
||||
db.create_session(
|
||||
session_id, source="cli", model_config={"keep": "value", "drop": 1},
|
||||
)
|
||||
|
||||
db.patch_session_model_config(session_id, {"drop": None, "added": 7})
|
||||
assert _model_config(db, session_id) == {"keep": "value", "added": 7}
|
||||
|
||||
# Missing rows and empty patches are no-ops, never errors.
|
||||
db.patch_session_model_config("NO_SUCH_SESSION", {"x": 1})
|
||||
db.patch_session_model_config(session_id, {})
|
||||
|
||||
|
||||
def test_incapable_store_short_circuits_before_prune_scan(tmp_path: Path) -> None:
|
||||
"""A bound store without archive_and_compact must not pay the prune scan."""
|
||||
db = SessionDB(db_path=tmp_path / "state.db")
|
||||
session_id = "INCAPABLE_STORE_FAST_NOOP"
|
||||
db.create_session(session_id, source="telegram")
|
||||
db.append_messages_batch(session_id, _history())
|
||||
agent = _build_agent(db, session_id)
|
||||
_configure_pruning(agent)
|
||||
compressor = agent.context_compressor
|
||||
|
||||
class _NoArchiveStore:
|
||||
pass
|
||||
|
||||
compressor.bind_session_state(_NoArchiveStore(), session_id)
|
||||
messages = db.get_messages_as_conversation(session_id)
|
||||
with patch.object(
|
||||
type(compressor), "_prune_old_tool_results",
|
||||
side_effect=AssertionError("scan must not run for incapable stores"),
|
||||
):
|
||||
result, count = compressor.prune_tool_results_only(
|
||||
messages, current_tokens=120_000,
|
||||
)
|
||||
|
||||
assert result is messages
|
||||
assert count == 0
|
||||
|
|
|
|||
Loading…
Reference in New Issue