From c71f81695612f4a9bbe22a4e0a68b94ed75b5b36 Mon Sep 17 00:00:00 2001 From: Ben Barclay Date: Tue, 9 Jun 2026 09:47:01 -0500 Subject: [PATCH] fix(compression): clear all per-session state in on_session_end, not just _previous_summary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The original cross-session contamination fix (#38788) only cleared _previous_summary in on_session_end(), but on_session_reset() clears 14+ per-session variables. When a session ends (cron exit, gateway expiry, session-id rotation) and the compressor instance is reused, the surviving stale state causes: - _ineffective_compression_count surviving → next session skips compression prematurely (anti-thrashing guard misfires) - _summary_failure_cooldown_until surviving → next session blocks summary generation for an unrelated transient error - _last_compress_aborted surviving → callers think compression is still aborted - _last_aux_model_failure_* surviving → stale error warnings shown - _last_summary_dropped_count / _last_summary_fallback_used surviving → misleading user warnings - _context_probed / _context_probe_persistable surviving → stale context-probe state Also fix on_session_reset() which was missing _last_compress_aborted clearing — a /new or /reset would inherit the aborted flag from the prior conversation. Add 6 targeted tests covering the leak vectors and a parity test ensuring on_session_end and on_session_reset always clear the same surface. --- agent/context_compressor.py | 43 +++- ...ext_compressor_session_end_clears_state.py | 242 ++++++++++++++++++ 2 files changed, 274 insertions(+), 11 deletions(-) create mode 100644 tests/agent/test_context_compressor_session_end_clears_state.py diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 0a5574e8e85..dac0cfca1ad 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -640,26 +640,47 @@ class ContextCompressor(ContextEngine): self._ineffective_compression_count = 0 self._summary_failure_cooldown_until = 0.0 # transient errors must not block a fresh session self._last_summary_error = None + self._last_compress_aborted = False self.last_real_prompt_tokens = 0 self.last_compression_rough_tokens = 0 self.last_rough_tokens_when_real_prompt_fit = 0 self.awaiting_real_usage_after_compression = False def on_session_end(self, session_id: str, messages: List[Dict[str, Any]]) -> None: - """Clear per-session compaction state at a real session boundary. + """Clear all per-session compaction state at a real session boundary. - ``_previous_summary`` is per-session iterative-summary state. It is - cleared on ``on_session_reset()`` (/new, /reset), but session *end* - (CLI exit, gateway expiry, session-id rotation) goes through - ``on_session_end()`` instead — which inherited a no-op from - ``ContextEngine``. Without clearing here, a cron/background session's - summary could survive on a reused compressor instance and leak into the - next live session via the ``_generate_summary()`` iterative-update path - (#38788). ``compress()`` already guards the leak at the point of use; - this is defense-in-depth that drops the stale summary the moment the - owning session ends. + Session end (CLI exit, gateway expiry, session-id rotation) goes + through this method rather than ``on_session_reset()`` (/new, /reset). + The original fix (#38788) only cleared ``_previous_summary``, but the + same cross-session contamination risk applies to every per-session + variable that ``on_session_reset()`` clears: stale + ``_ineffective_compression_count`` can suppress compression in a + subsequent live session; ``_summary_failure_cooldown_until`` can block + summary generation; ``_last_compress_aborted`` can make callers think + compression is still aborted; ``_last_aux_model_failure_*`` can surface + stale error warnings; ``_last_summary_dropped_count`` / + ``_last_summary_fallback_used`` can produce misleading user warnings. + + ``compress()`` already guards ``_previous_summary`` leakage at the + point of use; this is defense-in-depth that resets the full per-session + surface the moment the owning session ends. """ self._previous_summary = None + self._last_summary_error = None + self._last_summary_dropped_count = 0 + self._last_summary_fallback_used = False + self._last_aux_model_failure_error = None + self._last_aux_model_failure_model = None + self._last_compression_savings_pct = 100.0 + self._ineffective_compression_count = 0 + self._summary_failure_cooldown_until = 0.0 + self._last_compress_aborted = False + self._context_probed = False + self._context_probe_persistable = False + self.last_real_prompt_tokens = 0 + self.last_compression_rough_tokens = 0 + self.last_rough_tokens_when_real_prompt_fit = 0 + self.awaiting_real_usage_after_compression = False def bind_session_state(self, session_db: Any = None, session_id: str = "") -> None: """Bind the current session row so durable cooldowns can round-trip.""" diff --git a/tests/agent/test_context_compressor_session_end_clears_state.py b/tests/agent/test_context_compressor_session_end_clears_state.py new file mode 100644 index 00000000000..72c7b4d5ebb --- /dev/null +++ b/tests/agent/test_context_compressor_session_end_clears_state.py @@ -0,0 +1,242 @@ +"""Tests for on_session_end() clearing all per-session compressor state. + +Bug: on_session_end() (added in #38788) only cleared _previous_summary, but +on_session_reset() clears 14+ per-session variables. When a session ends +(cron exit, gateway expiry, session-id rotation) and the compressor instance +is reused, these stale values survive: + +- _ineffective_compression_count: can suppress compression in next session +- _summary_failure_cooldown_until: can block summary generation +- _last_compress_aborted: can make callers think compression is aborted +- _last_aux_model_failure_*: can surface stale error warnings +- _last_summary_dropped_count / _last_summary_fallback_used: misleading warnings +- _context_probed / _context_probe_persistable: stale context probe state + +Fix: on_session_end() now clears all per-session state, matching +on_session_reset()'s surface. +""" + +import sys +import types +from pathlib import Path + +# Ensure repo root is importable +sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) + +# Stub out optional heavy dependencies not installed in the test environment +sys.modules.setdefault("fire", types.SimpleNamespace(Fire=lambda *a, **k: None)) +sys.modules.setdefault("firecrawl", types.SimpleNamespace(Firecrawl=object)) +sys.modules.setdefault("fal_client", types.SimpleNamespace()) + +from agent.context_compressor import ContextCompressor + + +def _make_compressor(): + """Build a ContextCompressor with enough state to pass compress() guards.""" + c = ContextCompressor.__new__(ContextCompressor) + c.quiet_mode = True + c.model = "test/model" + c.provider = "test" + c.base_url = "http://test" + c.api_key = "test-key" + c.api_mode = "" + c.context_length = 128000 + c.threshold_tokens = 64000 + c.threshold_percent = 0.50 + c.tail_token_budget = 20000 + c.protect_last_n = 12 + c.summary_model = "" + c.last_prompt_tokens = 100000 + c.last_completion_tokens = 0 + c.last_total_tokens = 100000 + c._summary_failure_cooldown_until = 0.0 + c._max_compaction_summary_tokens = 0 + c.summary_budget_tokens = 0 + c.abort_on_summary_failure = False + c._last_compress_aborted = False + c._summary_model_fallen_back = False + c.compression_count = 0 + c._context_probed = False + c._context_probe_persistable = False + c._last_compression_savings_pct = 100.0 + c._ineffective_compression_count = 0 + c._last_summary_error = None + c._last_summary_dropped_count = 0 + c._last_summary_fallback_used = False + c._last_aux_model_failure_error = None + c._last_aux_model_failure_model = None + c.last_real_prompt_tokens = 0 + c.last_compression_rough_tokens = 0 + c.last_rough_tokens_when_real_prompt_fit = 0 + c.awaiting_real_usage_after_compression = False + c._previous_summary = None + return c + + +def _simulate_cron_session_state(c): + """Simulate per-session state that a cron compaction would leave behind.""" + c._previous_summary = "Cron session summary that must not leak" + c._last_summary_error = "Cron session summary error" + c._last_summary_dropped_count = 5 + c._last_summary_fallback_used = True + c._last_aux_model_failure_error = "Cron aux model error" + c._last_aux_model_failure_model = "cron-model/v1" + c._last_compression_savings_pct = 3.0 + c._ineffective_compression_count = 2 + c._summary_failure_cooldown_until = 9999999999.0 + c._last_compress_aborted = True + c._context_probed = True + c._context_probe_persistable = True + c.last_real_prompt_tokens = 50000 + c.last_compression_rough_tokens = 60000 + c.last_rough_tokens_when_real_prompt_fit = 55000 + c.awaiting_real_usage_after_compression = True + + +def test_on_session_end_clears_all_per_session_state(): + """on_session_end() must clear every per-session variable, not just + _previous_summary. Otherwise stale state from a prior session + (e.g. a cron job) contaminates the next live session.""" + c = _make_compressor() + _simulate_cron_session_state(c) + + c.on_session_end("cron-session-1", []) + + assert c._previous_summary is None, ( + f"_previous_summary must be None after on_session_end, got {c._previous_summary!r}" + ) + assert c._last_summary_error is None, ( + f"_last_summary_error must be None after on_session_end, got {c._last_summary_error!r}" + ) + assert c._last_summary_dropped_count == 0, ( + f"_last_summary_dropped_count must be 0, got {c._last_summary_dropped_count}" + ) + assert c._last_summary_fallback_used is False, ( + f"_last_summary_fallback_used must be False, got {c._last_summary_fallback_used}" + ) + assert c._last_aux_model_failure_error is None, ( + f"_last_aux_model_failure_error must be None, got {c._last_aux_model_failure_error!r}" + ) + assert c._last_aux_model_failure_model is None, ( + f"_last_aux_model_failure_model must be None, got {c._last_aux_model_failure_model!r}" + ) + assert c._last_compression_savings_pct == 100.0, ( + f"_last_compression_savings_pct must be 100.0, got {c._last_compression_savings_pct}" + ) + assert c._ineffective_compression_count == 0, ( + f"_ineffective_compression_count must be 0, got {c._ineffective_compression_count}" + ) + assert c._summary_failure_cooldown_until == 0.0, ( + f"_summary_failure_cooldown_until must be 0.0, got {c._summary_failure_cooldown_until}" + ) + assert c._last_compress_aborted is False, ( + f"_last_compress_aborted must be False, got {c._last_compress_aborted}" + ) + assert c._context_probed is False, ( + f"_context_probed must be False, got {c._context_probed}" + ) + assert c._context_probe_persistable is False, ( + f"_context_probe_persistable must be False, got {c._context_probe_persistable}" + ) + assert c.last_real_prompt_tokens == 0, ( + f"last_real_prompt_tokens must be 0, got {c.last_real_prompt_tokens}" + ) + assert c.last_compression_rough_tokens == 0, ( + f"last_compression_rough_tokens must be 0, got {c.last_compression_rough_tokens}" + ) + assert c.last_rough_tokens_when_real_prompt_fit == 0, ( + f"last_rough_tokens_when_real_prompt_fit must be 0, got {c.last_rough_tokens_when_real_prompt_fit}" + ) + assert c.awaiting_real_usage_after_compression is False, ( + f"awaiting_real_usage_after_compression must be False, got {c.awaiting_real_usage_after_compression}" + ) + + +def test_on_session_end_matches_on_session_reset_surface(): + """Both on_session_end and on_session_reset must clear the same set of + per-session variables. If one is updated and the other isn't, it's a + cross-session contamination bug waiting to happen.""" + c1 = _make_compressor() + c2 = _make_compressor() + _simulate_cron_session_state(c1) + _simulate_cron_session_state(c2) + + c1.on_session_end("session-1", []) + c2.on_session_reset() + + per_session_attrs = [ + "_previous_summary", + "_last_summary_error", + "_last_summary_dropped_count", + "_last_summary_fallback_used", + "_last_aux_model_failure_error", + "_last_aux_model_failure_model", + "_last_compression_savings_pct", + "_ineffective_compression_count", + "_summary_failure_cooldown_until", + "_last_compress_aborted", + "_context_probed", + "_context_probe_persistable", + "last_real_prompt_tokens", + "last_compression_rough_tokens", + "last_rough_tokens_when_real_prompt_fit", + "awaiting_real_usage_after_compression", + ] + + for attr in per_session_attrs: + v_end = getattr(c1, attr) + v_reset = getattr(c2, attr) + assert v_end == v_reset, ( + f"on_session_end and on_session_reset must produce the same " + f"value for {attr}: on_session_end={v_end!r}, " + f"on_session_reset={v_reset!r}" + ) + + +def test_ineffective_compression_count_does_not_leak_across_sessions(): + """A cron session that hit ineffective compression limits must not + suppress compression in a subsequent live session.""" + c = _make_compressor() + c._ineffective_compression_count = 2 # hit the anti-thrashing limit + c._last_compression_savings_pct = 3.0 + + c.on_session_end("cron-session", []) + + # After session end, the next session must start with a clean slate + assert c._ineffective_compression_count == 0 + assert c._last_compression_savings_pct == 100.0 + + +def test_summary_failure_cooldown_does_not_leak_across_sessions(): + """A cron session's summary failure cooldown must not block summary + generation in a subsequent live session.""" + c = _make_compressor() + c._summary_failure_cooldown_until = 9999999999.0 + + c.on_session_end("cron-session", []) + + assert c._summary_failure_cooldown_until == 0.0 + + +def test_compress_aborted_flag_does_not_leak_across_sessions(): + """A cron session's _last_compress_aborted flag must not make callers + think compression is still aborted in a subsequent live session.""" + c = _make_compressor() + c._last_compress_aborted = True + + c.on_session_end("cron-session", []) + + assert c._last_compress_aborted is False + + +def test_aux_model_failure_does_not_leak_across_sessions(): + """Stale aux model failure info from a cron session must not produce + misleading error warnings in a subsequent live session.""" + c = _make_compressor() + c._last_aux_model_failure_error = "cron-model/v1 failed" + c._last_aux_model_failure_model = "cron-model/v1" + + c.on_session_end("cron-session", []) + + assert c._last_aux_model_failure_error is None + assert c._last_aux_model_failure_model is None