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)
This commit is contained in:
WeiYusc 2026-07-06 00:59:53 +08:00 committed by Teknium
parent ea0fd393db
commit 928bcdde24
3 changed files with 259 additions and 1 deletions

View file

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

View file

@ -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.

View file

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