From 46e4891c644d71c28ed836aec0e43860a7dd7576 Mon Sep 17 00:00:00 2001 From: Jan-Stefan Janetzky Date: Tue, 14 Jul 2026 13:17:31 +0000 Subject: [PATCH] fix(compression): close post-dispatch lock scope --- agent/conversation_compression.py | 129 ++++++++++-------- .../agent/test_compression_concurrent_fork.py | 125 +++++++++++++++++ 2 files changed, 194 insertions(+), 60 deletions(-) diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 2cc764740251..78784d4c930f 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -935,10 +935,19 @@ def compress_context( if not _existing_sp: _existing_sp = agent._build_system_prompt(system_message) return messages, _existing_sp + _lock_released = False + def _release_lock() -> None: """Release the lock keyed on the OLD session_id (before rotation).""" + nonlocal _lock_released + if _lock_released: + return + _lock_released = True if _lock_refresher is not None: - _lock_refresher.stop() + try: + _lock_refresher.stop() + except Exception as _stop_err: + logger.debug("compression lock refresher stop failed: %s", _stop_err) if _lock_db is not None and _lock_sid and _lock_holder: try: _lock_db.release_compression_lock(_lock_sid, _lock_holder) @@ -1054,73 +1063,73 @@ def compress_context( _release_lock() raise - # Capture boundary quality before session-rotation callbacks run. Built-in - # and plugin lifecycle hooks may reset per-session compressor fields while - # rebinding to the child id; the completed attempt's verdict must survive - # that rebind and be recorded only after the full boundary commits. - _compression_made_progress = bool( - getattr(agent.context_compressor, "_last_compression_made_progress", False) - ) - _compression_used_fallback = bool( - getattr(agent.context_compressor, "_last_summary_fallback_used", False) - ) + try: + # Capture boundary quality before session-rotation callbacks run. Built-in + # and plugin lifecycle hooks may reset per-session compressor fields while + # rebinding to the child id; the completed attempt's verdict must survive + # that rebind and be recorded only after the full boundary commits. + _compression_made_progress = bool( + getattr(agent.context_compressor, "_last_compression_made_progress", False) + ) + _compression_used_fallback = bool( + getattr(agent.context_compressor, "_last_summary_fallback_used", False) + ) - # If compression aborted (aux LLM failed to produce a usable summary) - # the compressor returns the input messages unchanged. Surface the - # error to the user, skip the session-rotation work entirely (no - # session has logically ended), and let auto-compress callers detect - # the no-op via len(returned) == len(input). - if getattr(agent.context_compressor, "_last_compress_aborted", False): - try: - _err = getattr(agent.context_compressor, "_last_summary_error", None) or "unknown error" - if getattr(agent, "_last_compression_summary_warning", None) != _err: - agent._last_compression_summary_warning = _err - agent._emit_warning( - f"⚠ Compression aborted: {_err}. " - "No messages were dropped — conversation continues unchanged. " - "Run /compress to retry, or /new to start a fresh session." - ) + # If compression aborted (aux LLM failed to produce a usable summary) + # the compressor returns the input messages unchanged. Surface the + # error to the user, skip the session-rotation work entirely (no + # session has logically ended), and let auto-compress callers detect + # the no-op via len(returned) == len(input). + if getattr(agent.context_compressor, "_last_compress_aborted", False): + try: + _err = getattr(agent.context_compressor, "_last_summary_error", None) or "unknown error" + if getattr(agent, "_last_compression_summary_warning", None) != _err: + agent._last_compression_summary_warning = _err + agent._emit_warning( + f"⚠ Compression aborted: {_err}. " + "No messages were dropped — conversation continues unchanged. " + "Run /compress to retry, or /new to start a fresh session." + ) + _existing_sp = getattr(agent, "_cached_system_prompt", None) + if not _existing_sp: + _existing_sp = agent._build_system_prompt(system_message) + return messages, _existing_sp + finally: + _release_lock() + + # A compressor that returns the exact input object made no structural + # progress. Do not rotate/rewrite the session or arm post-compression + # deferral in that case; its own anti-thrash counter records the no-op. + if compressed is messages: + logger.info( + "Compression made no progress (session=%s) — skipping boundary rewrite.", + agent.session_id or "none", + ) _existing_sp = getattr(agent, "_cached_system_prompt", None) if not _existing_sp: _existing_sp = agent._build_system_prompt(system_message) - return messages, _existing_sp - finally: _release_lock() + return messages, _existing_sp - # A compressor that returns the exact input object made no structural - # progress. Do not rotate/rewrite the session or arm post-compression - # deferral in that case; its own anti-thrash counter records the no-op. - if compressed is messages: - logger.info( - "Compression made no progress (session=%s) — skipping boundary rewrite.", - agent.session_id or "none", - ) - _existing_sp = getattr(agent, "_cached_system_prompt", None) - if not _existing_sp: - _existing_sp = agent._build_system_prompt(system_message) - _release_lock() - return messages, _existing_sp - - if not compressed: - logger.error( - "context compression returned an empty transcript; refusing to " - "rotate session=%s so the parent remains resumable", - agent.session_id or "none", - ) - try: - agent._emit_warning( - "⚠ Compression returned an empty transcript. " - "No session split was performed; conversation continues unchanged." + if not compressed: + logger.error( + "context compression returned an empty transcript; refusing to " + "rotate session=%s so the parent remains resumable", + agent.session_id or "none", ) - except Exception: - pass - _existing_sp = getattr(agent, "_cached_system_prompt", None) - if not _existing_sp: - _existing_sp = agent._build_system_prompt(system_message) - _release_lock() - return messages, _existing_sp + try: + agent._emit_warning( + "⚠ Compression returned an empty transcript. " + "No session split was performed; conversation continues unchanged." + ) + except Exception: + pass + _existing_sp = getattr(agent, "_cached_system_prompt", None) + if not _existing_sp: + _existing_sp = agent._build_system_prompt(system_message) + _release_lock() + return messages, _existing_sp - try: summary_error = getattr(agent.context_compressor, "_last_summary_error", None) if summary_error: if getattr(agent, "_last_compression_summary_warning", None) != summary_error: diff --git a/tests/agent/test_compression_concurrent_fork.py b/tests/agent/test_compression_concurrent_fork.py index 2c66c210583d..3405bcdfd086 100644 --- a/tests/agent/test_compression_concurrent_fork.py +++ b/tests/agent/test_compression_concurrent_fork.py @@ -608,6 +608,131 @@ def test_signature_introspection_exception_releases_lock_and_refresher( assert not refreshers[0]._thread.is_alive() +def test_noop_prompt_exception_releases_lock_and_refresher( + tmp_path: Path, monkeypatch +) -> None: + """No-op prompt rebuild failures must not escape the lock cleanup scope.""" + from agent.conversation_compression import ( + _CompressionLockLeaseRefresher as RealLeaseRefresher, + ) + + refreshers = [] + + class RecordingLeaseRefresher(RealLeaseRefresher): + def start(self): + refreshers.append(self) + return super().start() + + monkeypatch.setattr( + "agent.conversation_compression._CompressionLockLeaseRefresher", + RecordingLeaseRefresher, + ) + + db = SessionDB(db_path=tmp_path / "state.db") + parent_sid = "NOOP_PROMPT_EXCEPTION_TEST" + db.create_session(parent_sid, source="discord") + agent = _build_agent_with_db(db, parent_sid) + agent._compression_lock_refresh_interval = 0.1 + messages = [{"role": "user", "content": f"m{i}"} for i in range(20)] + agent.context_compressor.compress.side_effect = lambda *_a, **_kw: messages + agent._cached_system_prompt = None + agent._build_system_prompt = lambda *_a, **_kw: (_ for _ in ()).throw( + RuntimeError("prompt rebuild boom") + ) + + with pytest.raises(RuntimeError, match="prompt rebuild boom"): + agent._compress_context(messages, "sys", approx_tokens=120_000) + + assert db.get_compression_lock_holder(parent_sid) is None + assert len(refreshers) == 1 + assert not refreshers[0]._thread.is_alive() + + +def test_post_dispatch_attribute_exception_releases_lock_and_refresher( + tmp_path: Path, monkeypatch +) -> None: + """Plugin state lookup failures after dispatch must release the lock.""" + from agent.conversation_compression import ( + _CompressionLockLeaseRefresher as RealLeaseRefresher, + ) + + refreshers = [] + + class RecordingLeaseRefresher(RealLeaseRefresher): + def start(self): + refreshers.append(self) + return super().start() + + class AttributeBombEngine: + name = "attribute-bomb" + + def compress(self, messages, **_kwargs): + return [messages[0], messages[-1]] + + def __getattribute__(self, name): + if name == "_last_compression_made_progress": + raise RuntimeError("post-dispatch attribute boom") + return object.__getattribute__(self, name) + + monkeypatch.setattr( + "agent.conversation_compression._CompressionLockLeaseRefresher", + RecordingLeaseRefresher, + ) + + db = SessionDB(db_path=tmp_path / "state.db") + parent_sid = "POST_DISPATCH_ATTRIBUTE_EXCEPTION_TEST" + db.create_session(parent_sid, source="discord") + agent = _build_agent_with_db(db, parent_sid) + agent._compression_lock_refresh_interval = 0.1 + agent.context_compressor = AttributeBombEngine() + messages = [{"role": "user", "content": f"m{i}"} for i in range(20)] + + with pytest.raises(RuntimeError, match="post-dispatch attribute boom"): + agent._compress_context(messages, "sys", approx_tokens=120_000) + + assert db.get_compression_lock_holder(parent_sid) is None + assert len(refreshers) == 1 + assert not refreshers[0]._thread.is_alive() + + +def test_refresher_stop_exception_does_not_block_lock_release( + tmp_path: Path, monkeypatch +) -> None: + """Refresher cleanup failure must not prevent holder-qualified DB release.""" + refreshers = [] + + class StopFailingLeaseRefresher: + def __init__(self, *_args, **_kwargs): + self.stop_calls = 0 + refreshers.append(self) + + def start(self): + return self + + def stop(self): + self.stop_calls += 1 + raise RuntimeError("refresher stop boom") + + monkeypatch.setattr( + "agent.conversation_compression._CompressionLockLeaseRefresher", + StopFailingLeaseRefresher, + ) + + db = SessionDB(db_path=tmp_path / "state.db") + parent_sid = "REFRESHER_STOP_EXCEPTION_TEST" + db.create_session(parent_sid, source="discord") + agent = _build_agent_with_db(db, parent_sid) + agent.context_compressor.compress.side_effect = RuntimeError("engine boom") + messages = [{"role": "user", "content": f"m{i}"} for i in range(20)] + + with pytest.raises(RuntimeError, match="engine boom"): + agent._compress_context(messages, "sys", approx_tokens=120_000) + + assert db.get_compression_lock_holder(parent_sid) is None + assert len(refreshers) == 1 + assert refreshers[0].stop_calls == 1 + + def _make_legacy_session_db_class() -> type: """Model the class retained in ``sys.modules`` before the lock API existed.