refactor(agent/compression): unify session rebind, scoped-twin stamping; lift facade timeout callbacks to module helpers

This commit is contained in:
Teknium
2026-09-02 18:17:52 -07:00
parent 113f04616b
commit 7524e9ec8a
3 changed files with 219 additions and 271 deletions
+155 -201
View File
@@ -14,6 +14,133 @@ from agent.session_activity import ActivityProvenance
logger = logging.getLogger("run_agent")
def _timeout_fallback_prompt(agent, system_message: str) -> str:
"""Cached prompt, else a fresh build, else the raw ``system_message`` (never raises).
Resolved lazily by the timeout wrapper: an eager rebuild would raise before
compress_context runs when ``_cached_system_prompt`` is unset and the builder fails.
"""
cached = getattr(agent, "_cached_system_prompt", None)
if cached:
return cached
try:
return agent._build_system_prompt(system_message)
except Exception:
logger.debug(
"compress_context timeout fallback prompt rebuild "
"failed; using raw system_message",
exc_info=True,
)
return system_message or ""
def _report_compression_timeout(
agent,
*,
idle: float,
waited: float,
since_progress: float,
total_ceiling: float,
total_exhausted: bool,
progress_observed: bool,
) -> None:
"""Host-side timeout bookkeeping: log, activity stamp, cooldown ladder, user warning."""
from agent.conversation_compression import mark_context_compression_timed_out
mark_context_compression_timed_out(agent)
if total_exhausted:
logger.warning(
"Context compression reached its total ceiling "
"after %.1fs (progress observed=%s); continuing "
"without compression",
waited,
progress_observed,
)
else:
logger.warning(
"Context compression made no progress for %.1fs "
"(total wait %.1fs, ceiling %.1fs); continuing "
"without compression",
since_progress,
waited,
total_ceiling,
)
touch = getattr(agent, "_touch_activity", None)
if callable(touch):
try:
touch(
"context compression timed out",
provenance=ActivityProvenance.AGENT_COMPRESSION_TIMEOUT,
)
except Exception:
logger.debug("compress_context timeout activity touch failed", exc_info=True)
# Same timeout cooldown ladder as summary-LLM timeouts: avoid re-burning the
# full idle budget every turn.
compressor = getattr(agent, "context_compressor", None)
record = getattr(compressor, "record_timeout_failure", None) if compressor is not None else None
if callable(record):
try:
record(
(
"host compress_context total ceiling exhausted"
if total_exhausted
else "host compress_context timeout (no summary progress)"
),
failure_kind="ceiling_exhausted" if total_exhausted else "stalled",
)
except Exception:
logger.debug("failed to record compress_context timeout cooldown", exc_info=True)
emit = getattr(agent, "_emit_warning", None)
if not callable(emit):
return
if total_exhausted:
progress = " after summary output was observed" if progress_observed else ""
emit(
"⚠ Context compression reached its total ceiling "
f"after {waited:.1f}s{progress}. No messages were "
"dropped — continuing without compression. Run "
"/compress to retry or /new for a clean session."
)
else:
emit(
"⚠ Context compression timed out "
f"after {idle:.1f}s with no output from the summary "
"model. No messages were dropped — continuing "
"without compression. Run /compress to retry, /new "
"for a clean session, or check "
"auxiliary.compression."
)
def _warn_commit_overrun(agent, waited: float, ceiling: float) -> None:
"""Commit-phase ceiling breach: the SessionDB mutation must complete, so only surface it."""
emit = getattr(agent, "_emit_warning", None)
if callable(emit):
emit(
"⚠ Context compression commit is taking unusually "
f"long ({waited:.0f}s, ceiling {ceiling:.0f}s). "
"Waiting for it to finish safely — if this persists, "
"check SessionDB health (disk / lock contention)."
)
def _sync_persisted_markers(target_messages, source_messages) -> None:
"""Mirror ``_DB_PERSISTED_MARKER`` stamps from the worker's snapshot onto a live list.
Matched by scoped identity; timestamp-less repeated content is ambiguous, so every
scoped match is stamped. Imported UNCONDITIONALLY: a silent fallback literal would
split the stamping key from the flush's and resurrect the duplicate-row bug.
"""
from agent.context_compressor import _DB_PERSISTED_MARKER
from agent.conversation_compression import _stamp_scoped_twins
if not isinstance(target_messages, list) or not isinstance(source_messages, list):
return
for source_message in source_messages:
if isinstance(source_message, dict) and source_message.get(_DB_PERSISTED_MARKER):
_stamp_scoped_twins(target_messages, source_message)
class CompressionFacadeMixin:
"""``_compress_context`` (see module docstring)."""
@@ -35,12 +162,11 @@ class CompressionFacadeMixin:
``force=True`` (manual /compress) bypasses the summary-failure cooldown; ``bypass_cooldown=True``
(provider-proven overflow recovery) runs one real attempt while the cooldown stays armed.
"""
# Per-attempt timeout signal for turn-start preflight and in-loop consumers (#98424): a stalled
# compression must not be mistaken for a structural no-op. Thread-local + per-agent lock (#98741).
# Per-attempt timeout signal for turn-start preflight and in-loop consumers: a stalled
# compression must not be mistaken for a structural no-op. Thread-local + per-agent lock.
from agent.conversation_compression import (
CompressionCommitFence,
compress_context,
mark_context_compression_timed_out,
reset_context_compression_timeout_outcome,
resolve_context_compression_timeouts,
run_compress_context_with_progress_timeout,
@@ -57,32 +183,27 @@ class CompressionFacadeMixin:
from agent.prompt_cache_scope import declared_conversation_scope_safe
# Out-of-turn compaction (/compact, gateway /compress, partial head compression) runs outside
# run_conversation's ambient scope; publish the root as a fallback so the summarizer's call carries
# the conversation tag. No-op for in-turn callers.
# the conversation tag. No-op for in-turn callers. Same for the ROUTING scope when declared.
token = None
if get_conversation_context() is None:
root = self._conversation_root_id()
if root:
token = set_conversation_context(root)
# Same fallback for the ROUTING scope, only when the host declared one (pre-#96811 fallback
# otherwise).
affinity_token = None
if get_affinity_scope() is None:
declared = declared_conversation_scope_safe(self)
if declared:
affinity_token = set_affinity_scope(declared)
# Every compression has a fence; hard_interrupt() uses this exact instance to serialize cancel
# admission against begin_commit().
# admission against begin_commit(). Publication is serialized so overlapping automatic/manual
# entrypoints cannot replace the fence of the attempt currently committing.
active_fence = commit_fence or CompressionCommitFence()
# Serialize fence publication so overlapping automatic/manual entrypoints cannot replace the
# fence of the attempt currently committing.
fence_registration_lock = vars(self).setdefault(
"_compression_commit_fence_lock", threading.RLock()
)
with fence_registration_lock:
missing_fence = object()
previous_fence = vars(self).get(
"_active_compression_commit_fence", missing_fence
)
previous_fence = vars(self).get("_active_compression_commit_fence", missing_fence)
self._active_compression_commit_fence = active_fence
try:
def _run(fence=None, target_messages=None):
@@ -113,143 +234,34 @@ class CompressionFacadeMixin:
result = _run(active_fence)
else:
def _snapshot_worker(fence=None):
# #76354 F3: the pooled worker must NEVER share the caller's live transcript — a late
# engine after a host timeout could rewrite it. Deep-snapshot on the worker; results
# publish only via an ADMITTED commit.
# The pooled worker must NEVER share the caller's live transcript — a late
# engine after a host timeout could rewrite it. Deep-snapshot on the worker;
# results publish only via an ADMITTED commit. A no-op/abort returns the
# snapshot unchanged: hand back the ORIGINAL list so identity semantics hold.
snapshot = copy.deepcopy(messages)
result_msgs, result_prompt = _run(
fence, target_messages=snapshot
)
result_msgs, result_prompt = _run(fence, target_messages=snapshot)
if result_msgs is snapshot:
# No-op/abort returned the snapshot unchanged: hand back the ORIGINAL list so
# identity-based semantics keep working.
return messages, result_prompt
return result_msgs, result_prompt
# Resolve the fallback prompt lazily: an eager rebuild would raise before compress_context
# runs when _cached_system_prompt is unset and _build_system_prompt fails.
def _fallback_prompt():
cached = getattr(self, "_cached_system_prompt", None)
if cached:
return cached
try:
return self._build_system_prompt(system_message)
except Exception:
logger.debug(
"compress_context timeout fallback prompt rebuild "
"failed; using raw system_message",
exc_info=True,
)
return system_message or ""
timeout_cause = {
"total_exhausted": False,
"progress_observed": False,
}
timeout_cause = {"total_exhausted": False, "progress_observed": False}
def _on_timeout_cause(total_exhausted, progress_observed):
timeout_cause["total_exhausted"] = total_exhausted
timeout_cause["progress_observed"] = progress_observed
def _on_timeout(idle, waited, since_progress):
mark_context_compression_timed_out(self)
total_exhausted = timeout_cause["total_exhausted"]
progress_observed = timeout_cause["progress_observed"]
if total_exhausted:
logger.warning(
"Context compression reached its total ceiling "
"after %.1fs (progress observed=%s); continuing "
"without compression",
waited,
progress_observed,
)
else:
logger.warning(
"Context compression made no progress for %.1fs "
"(total wait %.1fs, ceiling %.1fs); continuing "
"without compression",
since_progress,
waited,
total_ceiling,
)
touch = getattr(self, "_touch_activity", None)
if callable(touch):
try:
touch(
"context compression timed out",
provenance=ActivityProvenance.AGENT_COMPRESSION_TIMEOUT,
)
except Exception:
logger.debug(
"compress_context timeout activity touch failed",
exc_info=True,
)
# Same timeout cooldown ladder as summary-LLM timeouts
# (#62452): avoid re-burning the full idle budget every turn.
compressor = getattr(self, "context_compressor", None)
if compressor is not None:
record = getattr(compressor, "record_timeout_failure", None)
if callable(record):
try:
reason = (
"host compress_context total ceiling "
"exhausted"
if total_exhausted
else "host compress_context timeout "
"(no summary progress)"
)
record(
reason,
failure_kind=(
"ceiling_exhausted"
if total_exhausted
else "stalled"
),
)
except Exception:
logger.debug(
"failed to record compress_context timeout "
"cooldown",
exc_info=True,
)
emit = getattr(self, "_emit_warning", None)
if callable(emit):
if total_exhausted:
progress = (
" after summary output was observed"
if progress_observed
else ""
)
emit(
"⚠ Context compression reached its total ceiling "
f"after {waited:.1f}s{progress}. No messages were "
"dropped — continuing without compression. Run "
"/compress to retry or /new for a clean session."
)
else:
emit(
"⚠ Context compression timed out "
f"after {idle:.1f}s with no output from the summary "
"model. No messages were dropped — continuing "
"without compression. Run /compress to retry, /new "
"for a clean session, or check "
"auxiliary.compression."
)
def _on_commit_overrun(waited, ceiling):
# Commit-phase ceiling breach: the SessionDB mutation must complete, so this only surfaces
# the overrun.
emit = getattr(self, "_emit_warning", None)
if callable(emit):
emit(
"⚠ Context compression commit is taking unusually "
f"long ({waited:.0f}s, ceiling {ceiling:.0f}s). "
"Waiting for it to finish safely — if this persists, "
"check SessionDB health (disk / lock contention)."
)
_report_compression_timeout(
self,
idle=idle,
waited=waited,
since_progress=since_progress,
total_ceiling=total_ceiling,
**timeout_cause,
)
def _publish_new_fence():
# The stall-fallback retry (#78981) needs a fence the aborted attempt cannot veto; publish
# The stall-fallback retry needs a fence the aborted attempt cannot veto; publish
# it on the slot hard_interrupt() reads. The finally restores the caller's fence.
retry_fence = CompressionCommitFence()
with fence_registration_lock:
@@ -259,66 +271,16 @@ class CompressionFacadeMixin:
result = run_compress_context_with_progress_timeout(
worker=_snapshot_worker,
messages=messages,
system_prompt_fallback=_fallback_prompt,
system_prompt_fallback=lambda: _timeout_fallback_prompt(self, system_message),
idle_timeout_seconds=idle_timeout,
total_ceiling_seconds=total_ceiling,
on_timeout=_on_timeout,
on_timeout_cause=_on_timeout_cause,
on_commit_overrun=_on_commit_overrun,
on_commit_overrun=lambda waited, ceiling: _warn_commit_overrun(self, waited, ceiling),
fence=active_fence,
telemetry_agent=self,
new_fence=_publish_new_fence,
)
# Imported UNCONDITIONALLY: a silent fallback literal would split the stamping key from the
# flush's and resurrect the duplicate-row bug.
from agent.context_compressor import _DB_PERSISTED_MARKER
from agent.conversation_compression import (
_messages_match_scoped_identity,
)
def _sync_persisted_markers(target_messages, source_messages):
if not isinstance(target_messages, list) or not isinstance(
source_messages, list
):
return
# Stamps land on the worker's snapshot first; mirror them onto the live lists by scoped
# identity. Timestamp-less repeated content is ambiguous, so every scoped match is stamped.
for source_message in source_messages:
if not (
isinstance(source_message, dict)
and source_message.get(_DB_PERSISTED_MARKER)
):
continue
source_timestamp = source_message.get("timestamp")
matched_exact_timestamp = False
if source_timestamp is not None:
for target_message in target_messages:
if not isinstance(target_message, dict):
continue
if target_message.get(_DB_PERSISTED_MARKER):
continue
if not _messages_match_scoped_identity(
target_message, source_message
):
continue
if target_message.get("timestamp") != source_timestamp:
continue
target_message[_DB_PERSISTED_MARKER] = True
matched_exact_timestamp = True
if matched_exact_timestamp:
continue
for target_message in target_messages:
if not isinstance(target_message, dict):
continue
if target_message.get(_DB_PERSISTED_MARKER):
continue
if not _messages_match_scoped_identity(
target_message, source_message
):
continue
target_message[_DB_PERSISTED_MARKER] = True
if isinstance(result, tuple) and result:
result_messages = result[0]
if isinstance(result_messages, list):
@@ -326,31 +288,23 @@ class CompressionFacadeMixin:
if direct_path or result_messages is not messages:
_sync_persisted_markers(messages, result_messages)
session_messages = getattr(self, "_session_messages", None)
if (
isinstance(session_messages, list)
and session_messages is not messages
):
# Durable-parent adoption can leave `_session_messages` on the pre-adoption list; sync
# both.
if isinstance(session_messages, list) and session_messages is not messages:
# Durable-parent adoption can leave `_session_messages` on the pre-adoption list.
_sync_persisted_markers(session_messages, result_messages)
# The worker thread rotated hermes_logging's thread-local session id; propagate to this thread
# (#34089).
# The worker thread rotated hermes_logging's thread-local session id; propagate to this thread.
try:
from hermes_logging import set_session_context
set_session_context(self.session_id)
except Exception:
pass
# #76354 F5: rebind the session ContextVar in the CALLER's context so post-compression tools
# Rebind the session ContextVar in the CALLER's context so post-compression tools
# resolve HERMES_SESSION_ID to the child id (idempotent when no rotation happened).
try:
from gateway.session_context import set_current_session_id
if self.session_id:
set_current_session_id(self.session_id)
except Exception:
logger.debug(
"post-compression session ContextVar rebind failed",
exc_info=True,
)
logger.debug("post-compression session ContextVar rebind failed", exc_info=True)
return result
finally:
with fence_registration_lock:
+59 -65
View File
@@ -1691,6 +1691,22 @@ def _mark_compression_blocked_transient(agent: Any, compressor: Any) -> None:
pass
def _rebind_session_context(session_id: str) -> None:
"""Point the worker thread's session ContextVar and log context at ``session_id``."""
try:
from gateway.session_context import set_current_session_id
set_current_session_id(session_id)
except Exception:
os.environ["HERMES_SESSION_ID"] = session_id
try:
from hermes_logging import set_session_context
set_session_context(session_id)
except Exception:
pass
def _adopt_live_compression_child(
agent: Any,
session_db: Any,
@@ -1724,18 +1740,7 @@ def _adopt_live_compression_child(
return None
agent.session_id = child_session_id
try:
from gateway.session_context import set_current_session_id
set_current_session_id(child_session_id)
except Exception:
os.environ["HERMES_SESSION_ID"] = child_session_id
try:
from hermes_logging import set_session_context
set_session_context(child_session_id)
except Exception:
pass
_rebind_session_context(child_session_id)
agent._session_db_created = True
if child.get("system_prompt"):
@@ -1840,7 +1845,6 @@ def _compression_lock_holder(agent: Any) -> str:
pid+tid tell crashed holders apart in diagnostics; instance id and per-acquire
uuid disambiguate co-resident agents on one thread or pooled compressions.
"""
import threading
return (
f"pid={os.getpid()}"
f":tid={threading.get_ident()}"
@@ -2690,6 +2694,43 @@ def _messages_match_scoped_identity(left: Any, right: Any) -> bool:
return True
def _stamp_scoped_twins(
targets: list, source: dict, *, exact_counts_stamped: bool = False
) -> None:
"""Stamp ``_db_persisted`` on every unstamped scoped twin of ``source`` in ``targets``.
Exact-timestamp twins are preferred: when the source carries a timestamp and any
exact twin was stamped (or, with ``exact_counts_stamped``, merely exists), the
broad scoped pass is skipped so a content-equal old duplicate is left alone.
"""
from agent.context_compressor import _DB_PERSISTED_MARKER
source_timestamp = source.get("timestamp")
exact_hit = False
if source_timestamp is not None:
for target in targets:
if (
not isinstance(target, dict)
or target.get("timestamp") != source_timestamp
or not _messages_match_scoped_identity(target, source)
):
continue
if target.get(_DB_PERSISTED_MARKER):
exact_hit = exact_hit or exact_counts_stamped
continue
target[_DB_PERSISTED_MARKER] = True
exact_hit = True
if exact_hit:
return
for target in targets:
if (
isinstance(target, dict)
and not target.get(_DB_PERSISTED_MARKER)
and _messages_match_scoped_identity(target, source)
):
target[_DB_PERSISTED_MARKER] = True
_PENDING_CONTEXT_ENGINE_NOTIFICATION = (
"_pending_context_engine_compression_notification"
)
@@ -3715,64 +3756,17 @@ def _publish_rotated_compaction(
):
# Adoption may leave _session_messages on the pre-adoption list with an out-of-
# range idx; stamp every scoped twin against the ANCHOR SOURCE, as the wrapper.
_anchor_timestamp = _compressed_anchor_source.get(
"timestamp"
# An already-stamped exact twin still suppresses the broad pass here, or a
# content-equal old duplicate would get stamped.
_stamp_scoped_twins(
_session_messages, _compressed_anchor_source, exact_counts_stamped=True
)
_found_exact_timestamp_candidate = False
if _anchor_timestamp is not None:
for _twin_message in _session_messages:
if (
isinstance(_twin_message, dict)
and _twin_message.get("timestamp")
== _anchor_timestamp
and _messages_match_scoped_identity(
_twin_message,
_compressed_anchor_source,
)
):
# Count an exact scoped twin REGARDLESS of marker: an already-stamped twin must
# still suppress the broad fallback or a content-equal old dup gets stamped.
_found_exact_timestamp_candidate = True
if not _twin_message.get(
_DB_PERSISTED_MARKER
):
_twin_message[
_DB_PERSISTED_MARKER
] = True
if not _found_exact_timestamp_candidate:
# No exact twin anywhere (or timestamp-less anchor): stamp every scoped match.
# An already-stamped exact hit never opens this branch.
for _twin_message in _session_messages:
if (
isinstance(_twin_message, dict)
and not _twin_message.get(
_DB_PERSISTED_MARKER
)
and _messages_match_scoped_identity(
_twin_message,
_compressed_anchor_source,
)
):
_twin_message[
_DB_PERSISTED_MARKER
] = True
for _handoff_message in compressed:
if isinstance(_handoff_message, dict):
_handoff_message[_DB_PERSISTED_MARKER] = True
agent.session_id = new_session_id
agent._db_flush_scan_prefix = None
try:
from gateway.session_context import set_current_session_id
set_current_session_id(agent.session_id)
except Exception:
os.environ["HERMES_SESSION_ID"] = agent.session_id
try:
from hermes_logging import set_session_context
set_session_context(agent.session_id)
except Exception:
pass
_rebind_session_context(agent.session_id)
agent._session_db_created = True
# Carry /goal to the child: load_goal is a flat per-session lookup with no
# parent walk, so the goal would silently die at the boundary.
@@ -1,8 +1,8 @@
"""Item 1 regression — the run_agent._compress_context fallback shim must be loud.
Before the fix, _compress_context wrapped the imports of _DB_PERSISTED_MARKER
(agent.context_compressor) and _messages_match_scoped_identity
(agent.conversation_compression) in a try/except that silently defined local
(agent.context_compressor) and the scoped-identity stamping helper
(agent.conversation_compression._stamp_scoped_twins) in a try/except that silently defined local
fallbacks (a hard-coded ``"_db_persisted"`` literal and a local copy of the
identity helper) with NO logging. If the canonical constant/helper is renamed
or removed upstream, the import raises, the fallback silently keeps stamping
@@ -17,7 +17,7 @@ The ``already_present`` outcome is load-bearing: it keeps compress_context
from touching the deleted module-global name (which would raise NameError on
BOTH pre- and post-fix code and make the test non-discriminating), because the
stamp block at conversation_compression.py:3834-3859 — the ONLY in-module use
of _messages_match_scoped_identity — is skipped for already_present.
of _stamp_scoped_twins — is skipped for already_present.
"""
import os
@@ -50,7 +50,7 @@ def _build_agent_with_db(db: SessionDB, session_id: str, platform: str = "telegr
compressor = MagicMock()
# A real user row in the stub return makes _ensure_compressed_has_user_turn
# return `already_present`, so the in-module stamp block (the only user of
# _messages_match_scoped_identity inside compress_context) is skipped and
# _stamp_scoped_twins inside compress_context) is skipped and
# the deleted name is referenced ONLY by the run_agent shim import.
compressor.compress.return_value = [
{"role": "user", "content": "real user row"},
@@ -91,7 +91,7 @@ class TestCompressContextFallbackShim:
# stamp block). Pre-fix the except branch silently defines a fallback;
# post-fix the unconditional import must raise ImportError.
monkeypatch.delattr(
conversation_compression, "_messages_match_scoped_identity"
conversation_compression, "_stamp_scoped_twins"
)
with pytest.raises(ImportError):
agent._compress_context(