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():