From 928bcdde246484fd2bd2c46d31f98833ced41aeb Mon Sep 17 00:00:00 2001 From: WeiYusc <3115763429@qq.com> Date: Mon, 6 Jul 2026 00:59:53 +0800 Subject: [PATCH] fix(compression): refresh gateway activity during compaction Refresh the agent activity tracker while context compression is blocked in the auxiliary summarizer so gateway watchdogs do not report inactivity during long compactions. Add regression coverage for successful heartbeats, exception cleanup, touch failures, and strict-signature compressor fallback. (cherry picked from commit c09e58b7709cc60c5b454701f0ecf840e759222f) --- agent/conversation_compression.py | 69 ++++++++- .../agent/test_compression_concurrent_fork.py | 134 ++++++++++++++++++ .../test_codex_app_server_compaction.py | 57 ++++++++ 3 files changed, 259 insertions(+), 1 deletion(-) diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index cd6d93bc7cd2..1135f23f682d 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -31,6 +31,7 @@ from __future__ import annotations import copy import inspect import logging +import math import os import tempfile import uuid @@ -231,6 +232,51 @@ def _supported_compression_kwargs( return {name: value for name, value in candidates.items() if name in parameters} +class _CompressionActivityHeartbeat: + """Refresh the agent inactivity tracker while compression blocks in an aux call.""" + + def __init__(self, agent: Any, interval_seconds: float | None = None) -> None: + self._agent = agent + if interval_seconds is None: + interval_seconds = getattr(agent, "_compression_activity_heartbeat_interval", 60.0) + try: + interval_seconds = float(interval_seconds or 60.0) + except (TypeError, ValueError): + interval_seconds = 60.0 + if not math.isfinite(interval_seconds): + interval_seconds = 60.0 + self._interval_seconds = max(0.1, interval_seconds) + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._run, + name="compression-activity-heartbeat", + daemon=True, + ) + + def start(self) -> "_CompressionActivityHeartbeat": + self._touch("context compression started") + self._thread.start() + return self + + def stop(self, desc: str = "context compression completed") -> None: + self._stop.set() + if self._thread.is_alive() and threading.current_thread() is not self._thread: + self._thread.join(timeout=1.0) + self._touch(desc) + + def _touch(self, desc: str) -> None: + try: + touch = getattr(self._agent, "_touch_activity", None) + if callable(touch): + touch(desc) + except Exception: + logger.debug("compression activity heartbeat touch failed", exc_info=True) + + def _run(self) -> None: + while not self._stop.wait(self._interval_seconds): + self._touch("context compression in progress") + + class _CompressionLockLeaseRefresher: def __init__( self, @@ -1020,6 +1066,7 @@ def compress_context( existing_prompt = agent._build_system_prompt(system_message) return messages, existing_prompt + _activity_heartbeat: Optional[_CompressionActivityHeartbeat] = None try: if _lock_holder is not None: _lock_refresher = _CompressionLockLeaseRefresher( @@ -1070,13 +1117,20 @@ def compress_context( ) messages_before_compression = copy.deepcopy(messages) + _activity_heartbeat = _CompressionActivityHeartbeat(agent).start() compressed = compress_fn(messages, **compress_kwargs) except BaseException: # ANY exception after lock acquisition — memory hook, capability # inspection, engine lookup, or compress() — must release the lock so # the session isn't permanently blocked from future compression. + if _activity_heartbeat is not None: + _activity_heartbeat.stop("context compression failed") + _activity_heartbeat = None _release_lock() raise + finally: + if _activity_heartbeat is not None: + _activity_heartbeat.stop("context compression completed") try: # Capture boundary quality before session-rotation callbacks run. Built-in @@ -1594,7 +1648,20 @@ def _compress_context_via_codex_app_server( except Exception: pass - result = codex_session.compact_thread() + _activity_heartbeat: Optional[_CompressionActivityHeartbeat] = None + try: + _activity_heartbeat = _CompressionActivityHeartbeat(agent).start() + result = codex_session.compact_thread() + except BaseException: + if _activity_heartbeat is not None: + _activity_heartbeat.stop("context compression failed") + raise + + if getattr(result, "interrupted", False) or getattr(result, "error", None): + _activity_heartbeat.stop("context compression failed") + else: + _activity_heartbeat.stop("context compression completed") + if getattr(result, "should_retire", False): try: codex_session.close() diff --git a/tests/agent/test_compression_concurrent_fork.py b/tests/agent/test_compression_concurrent_fork.py index 5296c67ec323..3b8b1957a289 100644 --- a/tests/agent/test_compression_concurrent_fork.py +++ b/tests/agent/test_compression_concurrent_fork.py @@ -94,6 +94,140 @@ def _count_children(db: SessionDB, parent_sid: str) -> int: return len(rows) +def _wait_for_touch(touch_calls: list[str], value: str, timeout: float = 1.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if value in touch_calls: + return + time.sleep(0.01) + pytest.fail(f"Timed out waiting for touch activity {value!r}; calls={touch_calls!r}") + + +def test_compression_activity_heartbeat_touches_agent_during_long_compress(tmp_path: Path) -> None: + """Long compression must refresh agent activity so gateway watchdogs do not fire.""" + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "HEARTBEAT_TEST" + db.create_session(session_id, source="test") + + agent = _build_agent_with_db(db, session_id) + agent._compression_activity_heartbeat_interval = 0.1 + touch_calls: list[str] = [] + agent._touch_activity = lambda desc: touch_calls.append(desc) + + def _slow_compress(*_a, **_kw): + _wait_for_touch(touch_calls, "context compression in progress") + return [ + {"role": "user", "content": "[CONTEXT COMPACTION] summary"}, + {"role": "user", "content": "tail"}, + ] + + agent.context_compressor.compress.side_effect = _slow_compress + messages = [{"role": "user", "content": f"m{i}"} for i in range(20)] + + agent._compress_context(messages, "sys", approx_tokens=120_000) + + assert touch_calls[0] == "context compression started" + assert "context compression in progress" in touch_calls + assert touch_calls[-1] == "context compression completed" + assert db.get_compression_lock_holder(session_id) is None + + +def test_compression_activity_heartbeat_stops_on_compress_exception(tmp_path: Path) -> None: + """Exception paths must stop the heartbeat and release the compression lock.""" + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "HEARTBEAT_FAIL_TEST" + db.create_session(session_id, source="test") + + agent = _build_agent_with_db(db, session_id) + agent._compression_activity_heartbeat_interval = 0.1 + touch_calls: list[str] = [] + agent._touch_activity = lambda desc: touch_calls.append(desc) + + def _failing_compress(*_a, **_kw): + _wait_for_touch(touch_calls, "context compression in progress") + raise RuntimeError("compress boom") + + agent.context_compressor.compress.side_effect = _failing_compress + messages = [{"role": "user", "content": f"m{i}"} for i in range(20)] + + with pytest.raises(RuntimeError, match="compress boom"): + agent._compress_context(messages, "sys", approx_tokens=120_000) + + assert touch_calls[0] == "context compression started" + assert "context compression in progress" in touch_calls + assert touch_calls[-1] == "context compression failed" + assert db.get_compression_lock_holder(session_id) is None + + +def test_compression_activity_heartbeat_ignores_touch_errors(tmp_path: Path) -> None: + """Activity touch failures must not affect compression success semantics.""" + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "HEARTBEAT_TOUCH_ERROR_TEST" + db.create_session(session_id, source="test") + + agent = _build_agent_with_db(db, session_id) + agent._compression_activity_heartbeat_interval = 0.1 + agent._touch_activity = lambda _desc: (_ for _ in ()).throw(RuntimeError("touch boom")) + messages = [{"role": "user", "content": f"m{i}"} for i in range(20)] + + compressed, _sp = agent._compress_context(messages, "sys", approx_tokens=120_000) + + assert compressed[0]["content"] == "[CONTEXT COMPACTION] summary" + assert db.get_compression_lock_holder(session_id) is None + + +def test_compression_activity_heartbeat_strict_signature_fallback_releases_lock(tmp_path: Path) -> None: + """Strict compressor signatures still fall back while heartbeat cleanup runs.""" + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "HEARTBEAT_TYPEERROR_TEST" + db.create_session(session_id, source="test") + + agent = _build_agent_with_db(db, session_id) + agent._compression_activity_heartbeat_interval = "not-a-number" + touch_calls: list[str] = [] + agent._touch_activity = lambda desc: touch_calls.append(desc) + messages = [{"role": "user", "content": f"m{i}"} for i in range(20)] + + def _strict_compress(*_args, **kwargs): + if "focus_topic" in kwargs or "force" in kwargs: + raise TypeError("unexpected keyword") + return [ + {"role": "user", "content": "[CONTEXT COMPACTION] strict summary"}, + {"role": "user", "content": "tail"}, + ] + + agent.context_compressor.compress.side_effect = _strict_compress + + compressed, _sp = agent._compress_context(messages, "sys", approx_tokens=120_000) + + assert compressed[0]["content"] == "[CONTEXT COMPACTION] strict summary" + assert touch_calls[0] == "context compression started" + assert touch_calls[-1] == "context compression completed" + assert db.get_compression_lock_holder(session_id) is None + assert agent.context_compressor.compress.call_count == 2 + + +def test_compression_activity_heartbeat_nonfinite_interval_falls_back(tmp_path: Path) -> None: + """Non-finite heartbeat intervals must not reach Event.wait().""" + from agent.conversation_compression import _CompressionActivityHeartbeat + + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "HEARTBEAT_NONFINITE_INTERVAL_TEST" + db.create_session(session_id, source="test") + + agent = _build_agent_with_db(db, session_id) + touch_calls: list[str] = [] + agent._touch_activity = lambda desc: touch_calls.append(desc) + + heartbeat = _CompressionActivityHeartbeat(agent, interval_seconds=float("inf")) + + assert heartbeat._interval_seconds == 60.0 + heartbeat.start() + heartbeat.stop() + assert touch_calls == ["context compression started", "context compression completed"] + + + def test_concurrent_compression_does_not_fork_session(tmp_path: Path) -> None: """Two AIAgents that share a session_id MUST NOT both rotate it. diff --git a/tests/run_agent/test_codex_app_server_compaction.py b/tests/run_agent/test_codex_app_server_compaction.py index 93bbe0eadf22..98f8cb8059a4 100644 --- a/tests/run_agent/test_codex_app_server_compaction.py +++ b/tests/run_agent/test_codex_app_server_compaction.py @@ -1,5 +1,8 @@ +import time from types import SimpleNamespace +import pytest + from agent.codex_runtime import _record_codex_app_server_compaction from agent.conversation_compression import COMPACTION_STATUS, compress_context from agent.transports.codex_app_server_session import TurnResult @@ -19,6 +22,26 @@ class FakeCodexSession: self.closed = True +class SlowCodexSession(FakeCodexSession): + def __init__(self, result, touch_calls): + super().__init__(result) + self.touch_calls = touch_calls + + def compact_thread(self): + self.calls += 1 + _wait_for_touch(self.touch_calls, "context compression in progress") + return self.result + + +def _wait_for_touch(touch_calls, desc, timeout=1.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if desc in touch_calls: + return + time.sleep(0.01) + pytest.fail(f"timed out waiting for touch {desc!r}; saw {touch_calls!r}") + + class DummyAgent: def __init__( self, @@ -43,6 +66,11 @@ class DummyAgent: self.warnings = [] self.events = [] self.built_prompts = [] + self.touch_calls = [] + self._compression_activity_heartbeat_interval = 0.1 + + def _touch_activity(self, desc): + self.touch_calls.append(desc) def _emit_status(self, message): self.statuses.append(message) @@ -79,6 +107,33 @@ def test_codex_app_server_native_auto_mode_leaves_thread_compaction_to_codex(): assert agent.events == [] +def test_codex_app_server_compaction_heartbeat_refreshes_activity_while_waiting(): + agent = DummyAgent( + TurnResult(thread_id="thread-1", turn_id="compact-turn-1") + ) + agent._codex_session = SlowCodexSession( + agent._codex_session.result, + agent.touch_calls, + ) + messages = [{"role": "user", "content": "hi"}] + + returned, prompt = compress_context( + agent, + messages, + "system", + approx_tokens=100000, + task_id="test", + force=True, + ) + + assert returned is messages + assert prompt == "cached prompt" + assert agent._codex_session.calls == 1 + assert "context compression started" in agent.touch_calls + assert "context compression in progress" in agent.touch_calls + assert agent.touch_calls[-1] == "context compression completed" + + def test_codex_app_server_manual_compression_routes_to_codex_thread(): agent = DummyAgent( TurnResult(thread_id="thread-1", turn_id="compact-turn-1") @@ -162,6 +217,8 @@ def test_codex_app_server_compression_failure_preserves_bookkeeping(): assert agent.context_compressor.compression_count == 0 assert agent.context_compressor.last_prompt_tokens == 123 assert agent.warnings + assert agent.touch_calls[0] == "context compression started" + assert agent.touch_calls[-1] == "context compression failed" def test_codex_app_server_native_compaction_notice_emits_status_and_event():