diff --git a/agent/context_compressor.py b/agent/context_compressor.py index f99e508e5293..705cf0fdfe00 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -1022,38 +1022,51 @@ class ContextCompressor(ContextEngine): self._fallback_compression_streak = 0 self._persist_fallback_compression_streak() - def get_active_compression_failure_cooldown(self) -> Optional[Dict[str, Any]]: + def get_active_compression_failure_cooldown( + self, + *, + refresh: bool = False, + ) -> Optional[Dict[str, Any]]: """Return the live compression-failure cooldown for the bound session.""" now_mono = time.monotonic() + local_state = None if self._summary_failure_cooldown_until > now_mono: - return { + local_state = { "cooldown_until": time.time() + ( self._summary_failure_cooldown_until - now_mono ), "remaining_seconds": self._summary_failure_cooldown_until - now_mono, "error": self._last_summary_error, } + if not refresh: + return local_state session_db = getattr(self, "_session_db", None) session_id = getattr(self, "_session_id", "") if not session_db or not session_id: - return None + return local_state getter = getattr(session_db, "get_compression_failure_cooldown", None) if getter is None: - return None + return local_state try: state = getter(session_id) except sqlite3.Error as exc: logger.debug("compression failure cooldown lookup failed: %s", exc) - return None + return local_state except Exception: - return None + return local_state if not state: + if refresh: + self._summary_failure_cooldown_until = 0.0 + self._last_summary_error = None return None remaining_seconds = float(state.get("remaining_seconds") or 0.0) if remaining_seconds <= 0: + if refresh: + self._summary_failure_cooldown_until = 0.0 + self._last_summary_error = None return None self._summary_failure_cooldown_until = now_mono + remaining_seconds diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index d5099f08444c..bb88ff68abca 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -76,6 +76,35 @@ def _lock_api_is_absent_on_session_db(lock_db: Any) -> bool: return False +def _refresh_persisted_compression_guards(compressor: Any) -> None: + """Refresh durable automatic-compression guards on a built-in compressor.""" + method_calls = ( + ("get_active_compression_failure_cooldown", {"refresh": True}), + ("_load_fallback_compression_streak", {}), + ) + for method_name, kwargs in method_calls: + method = getattr(type(compressor), method_name, None) + if not callable(method): + continue + try: + method(compressor, **kwargs) + except Exception as exc: + logger.debug("compression guard refresh failed (%s): %s", method_name, exc) + + +def _session_was_rotated_by_compression(session_db: Any, session_id: str) -> bool: + """Return whether another path already rotated this compression parent.""" + getter = getattr(type(session_db), "get_session", None) + if not callable(getter): + return False + session = getter(session_db, session_id) + return bool( + session + and session.get("ended_at") is not None + and session.get("end_reason") == "compression" + ) + + def _compression_lock_holder(agent: Any) -> str: """Build a unique holder id for the lock: pid:tid:agent-instance:uuid. @@ -612,6 +641,7 @@ def compress_context( # breaker state. Gateway hygiene constructs a fresh AIAgent, so the # persisted fallback streak is loaded by bind_session_state() before this. if not force: + _refresh_persisted_compression_guards(agent.context_compressor) blocked = getattr( type(agent.context_compressor), "_automatic_compression_blocked", @@ -814,6 +844,58 @@ def compress_context( except Exception as _rel_err: logger.debug("compression lock release failed: %s", _rel_err) + # A delayed contender can acquire the parent lock after the winning path + # has released it and completed rotation. The lock serializes work but does + # not by itself prove that this stale agent still owns a live parent. + if _lock_db is not None and _lock_sid: + try: + _parent_already_rotated = _session_was_rotated_by_compression( + _lock_db, _lock_sid + ) + except Exception as _session_err: + logger.warning( + "compression session ownership lookup failed for session=%s " + "(%s: %s) - skipping compression this cycle", + _lock_sid, + type(_session_err).__name__, + _session_err, + ) + _release_lock() + _existing_sp = getattr(agent, "_cached_system_prompt", None) + if not _existing_sp: + _existing_sp = agent._build_system_prompt(system_message) + return messages, _existing_sp + if _parent_already_rotated: + logger.info( + "compression skipped: session=%s was already rotated by " + "another compression path", + _lock_sid, + ) + _release_lock() + _existing_sp = getattr(agent, "_cached_system_prompt", None) + if not _existing_sp: + _existing_sp = agent._build_system_prompt(system_message) + return messages, _existing_sp + + # The agent may have been constructed before another path completed an + # in-place compaction on the same session. Re-read durable breaker state + # after acquiring the session lock so this final gate cannot act on the + # stale snapshot loaded by bind_session_state(). + if not force: + compressor = agent.context_compressor + _refresh_persisted_compression_guards(compressor) + blocked = getattr( + type(compressor), + "_automatic_compression_blocked", + None, + ) + if callable(blocked) and blocked(compressor): + _release_lock() + existing_prompt = getattr(agent, "_cached_system_prompt", None) + if not existing_prompt: + existing_prompt = agent._build_system_prompt(system_message) + return messages, existing_prompt + # Notify external memory provider before compression discards context if agent._memory_manager: try: diff --git a/tests/agent/test_compression_rotation_state.py b/tests/agent/test_compression_rotation_state.py index a88db5f364b5..07e714aca31c 100644 --- a/tests/agent/test_compression_rotation_state.py +++ b/tests/agent/test_compression_rotation_state.py @@ -18,9 +18,12 @@ These tests drive the real ``compress_context`` path against a real SessionDB. from __future__ import annotations import os +import time from pathlib import Path from unittest.mock import MagicMock, patch +import pytest + from agent.context_compressor import ContextCompressor from hermes_state import SessionDB @@ -65,6 +68,31 @@ def _msgs(n=20): return [{"role": "user", "content": f"m{i}"} for i in range(n)] +def _bound_context_compressor(db: SessionDB, session_id: str) -> ContextCompressor: + with patch( + "agent.context_compressor.get_model_context_length", + return_value=100_000, + ): + compressor = ContextCompressor( + model="test/model", + threshold_percent=0.85, + protect_first_n=2, + protect_last_n=2, + quiet_mode=True, + ) + compressor.bind_session_state(db, session_id) + return compressor + + +@pytest.fixture +def refresh_state_db(tmp_path: Path): + db = SessionDB(db_path=tmp_path / "state.db") + try: + yield db + finally: + db.close() + + class TestGoalMigratesOnRotation: def test_goal_follows_compression_rotation(self, tmp_path: Path): db = SessionDB(db_path=tmp_path / "state.db") @@ -224,3 +252,245 @@ class TestFallbackStreakFollowsRotation: assert child != parent assert compressor._fallback_compression_streak == 1 assert db.get_compression_fallback_streak(child) == 1 + + +class TestAutomaticCompressionStateRefreshAfterLock: + def test_prebound_agent_rejects_parent_rotated_before_lock_acquisition( + self, + refresh_state_db: SessionDB, + ): + db = refresh_state_db + parent_id = "STALE_ROTATED_PARENT" + child_id = "CANONICAL_COMPRESSION_CHILD" + db.create_session(parent_id, source="telegram") + agent = _build_agent_with_db(db, parent_id, platform="telegram") + compressor = _bound_context_compressor(db, parent_id) + + # A competing path completes rotation after this call's initial checks + # but before it acquires the parent lock. + real_acquire = db.try_acquire_compression_lock + + def _acquire_after_rotation(*args, **kwargs): + db.end_session(parent_id, "compression") + db.create_session( + child_id, + source="telegram", + parent_session_id=parent_id, + ) + return real_acquire(*args, **kwargs) + + db.try_acquire_compression_lock = _acquire_after_rotation + agent.context_compressor = compressor + agent.compression_in_place = False + agent._compression_feasibility_checked = True + messages = _msgs() + + with patch.object( + compressor, + "compress", + side_effect=AssertionError("stale parent was compressed again"), + ) as compress: + returned, _ = agent._compress_context( + messages, + "sys", + approx_tokens=120_000, + force=True, + ) + + children = db._conn.execute( + "SELECT id FROM sessions WHERE parent_session_id = ?", + (parent_id,), + ).fetchall() + assert returned is messages + assert agent.session_id == parent_id + assert [row["id"] for row in children] == [child_id] + compress.assert_not_called() + assert db.get_compression_lock_holder(parent_id) is None + + def test_prebound_agent_reloads_persisted_streak_before_compressing( + self, + refresh_state_db: SessionDB, + ): + db = refresh_state_db + session_id = "STALE_FALLBACK_BREAKER" + db.create_session(session_id, source="telegram") + db.set_compression_fallback_streak(session_id, 1) + agent = _build_agent_with_db(db, session_id, platform="telegram") + compressor = _bound_context_compressor(db, session_id) + assert compressor._fallback_compression_streak == 1 + + # A second agent finishes an in-place fallback boundary after this + # call's initial gate but while it is acquiring the session lock. + real_acquire = db.try_acquire_compression_lock + + def _acquire_after_fallback(*args, **kwargs): + db.set_compression_fallback_streak(session_id, 2) + return real_acquire(*args, **kwargs) + + db.try_acquire_compression_lock = _acquire_after_fallback + agent.context_compressor = compressor + agent.compression_in_place = True + agent._compression_feasibility_checked = True + messages = _msgs() + + with patch.object( + compressor, + "compress", + side_effect=AssertionError("stale agent bypassed fallback breaker"), + ) as compress: + returned, _ = agent._compress_context( + messages, + "sys", + approx_tokens=120_000, + ) + + assert returned is messages + assert compressor._fallback_compression_streak == 2 + compress.assert_not_called() + assert db.get_compression_lock_holder(session_id) is None + + def test_prebound_agent_reloads_persisted_cooldown_before_compressing( + self, + refresh_state_db: SessionDB, + ): + db = refresh_state_db + session_id = "STALE_COMPRESSION_COOLDOWN" + db.create_session(session_id, source="telegram") + agent = _build_agent_with_db(db, session_id, platform="telegram") + compressor = _bound_context_compressor(db, session_id) + assert compressor.get_active_compression_failure_cooldown() is None + + # Another agent records a provider cooldown after this call's initial + # gate but while it is acquiring the session lock. + real_acquire = db.try_acquire_compression_lock + + def _acquire_after_cooldown(*args, **kwargs): + db.record_compression_failure_cooldown( + session_id, + time.time() + 60, + "rate limited", + ) + return real_acquire(*args, **kwargs) + + db.try_acquire_compression_lock = _acquire_after_cooldown + agent.context_compressor = compressor + agent.compression_in_place = True + agent._compression_feasibility_checked = True + messages = _msgs() + + with patch.object( + compressor, + "compress", + side_effect=AssertionError("stale agent bypassed compression cooldown"), + ) as compress: + returned, _ = agent._compress_context( + messages, + "sys", + approx_tokens=120_000, + ) + + assert returned is messages + assert compressor.get_active_compression_failure_cooldown() is not None + compress.assert_not_called() + assert db.get_compression_lock_holder(session_id) is None + + def test_prebound_agent_drops_stale_blocker_before_initial_gate( + self, + refresh_state_db: SessionDB, + ): + db = refresh_state_db + session_id = "CLEARED_FALLBACK_BREAKER" + db.create_session(session_id, source="telegram") + db.set_compression_fallback_streak(session_id, 2) + agent = _build_agent_with_db(db, session_id, platform="telegram") + compressor = _bound_context_compressor(db, session_id) + assert compressor._fallback_compression_streak == 2 + + # A healthy boundary on another agent clears the durable breaker after + # this compressor was bound. The initial gate must not remain stuck on + # its stale in-memory snapshot. + db.set_compression_fallback_streak(session_id, 0) + agent.context_compressor = compressor + agent.compression_in_place = True + agent._compression_feasibility_checked = True + messages = _msgs() + + with patch.object(compressor, "compress", return_value=messages) as compress: + returned, _ = agent._compress_context( + messages, + "sys", + approx_tokens=120_000, + ) + + assert returned is messages + assert compressor._fallback_compression_streak == 0 + compress.assert_called_once() + assert db.get_compression_lock_holder(session_id) is None + + def test_prebound_agent_drops_stale_cooldown_before_initial_gate( + self, + refresh_state_db: SessionDB, + ): + db = refresh_state_db + session_id = "CLEARED_COMPRESSION_COOLDOWN" + db.create_session(session_id, source="telegram") + db.record_compression_failure_cooldown( + session_id, + time.time() + 60, + "rate limited", + ) + agent = _build_agent_with_db(db, session_id, platform="telegram") + compressor = _bound_context_compressor(db, session_id) + assert compressor.get_active_compression_failure_cooldown() is not None + + # A successful forced retry on another agent clears the durable row. + # This prebound compressor must not keep honoring its stale local timer. + db.clear_compression_failure_cooldown(session_id) + agent.context_compressor = compressor + agent.compression_in_place = True + agent._compression_feasibility_checked = True + messages = _msgs() + + with patch.object(compressor, "compress", return_value=messages) as compress: + returned, _ = agent._compress_context( + messages, + "sys", + approx_tokens=120_000, + ) + + assert returned is messages + assert compressor.get_active_compression_failure_cooldown() is None + compress.assert_called_once() + assert db.get_compression_lock_holder(session_id) is None + + def test_force_still_bypasses_refreshed_persisted_breaker( + self, + refresh_state_db: SessionDB, + ): + db = refresh_state_db + session_id = "FORCED_FALLBACK_RETRY" + db.create_session(session_id, source="telegram") + db.set_compression_fallback_streak(session_id, 2) + agent = _build_agent_with_db(db, session_id, platform="telegram") + compressor = _bound_context_compressor(db, session_id) + agent.context_compressor = compressor + agent.compression_in_place = True + agent._compression_feasibility_checked = True + messages = _msgs() + + with patch.object(compressor, "compress", return_value=messages) as compress: + returned, _ = agent._compress_context( + messages, + "sys", + approx_tokens=120_000, + force=True, + ) + + assert returned is messages + compress.assert_called_once_with( + messages, + current_tokens=120_000, + focus_topic=None, + force=True, + ) + assert db.get_compression_lock_holder(session_id) is None