From 241605d1ea34aac28b67bb0701294574d2ebf79c Mon Sep 17 00:00:00 2001 From: kshitij Date: Thu, 6 Aug 2026 02:10:07 +0530 Subject: [PATCH] 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. --- agent/context_compressor.py | 54 +++++++++++----- agent/conversation_compression.py | 12 ++++ .../test_proactive_prune_restart_safety.py | 62 +++++++++++++++++++ 3 files changed, 112 insertions(+), 16 deletions(-) diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 6c70898514db5..7dd1e8967316d 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -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={ diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 98734317645ec..f2b3a6a003571 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -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[ diff --git a/tests/agent/test_proactive_prune_restart_safety.py b/tests/agent/test_proactive_prune_restart_safety.py index 8b7a7b9918d9d..83adb437e5b54 100644 --- a/tests/agent/test_proactive_prune_restart_safety.py +++ b/tests/agent/test_proactive_prune_restart_safety.py @@ -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