"""Host-side ``AIAgent._compress_context`` wrapper. Publishes the commit fence ``hard_interrupt()`` reads, runs the compressor on a snapshot under the progress timeout, mirrors ``_DB_PERSISTED_MARKER`` stamps back onto the live lists and rebinds the session context. Extracted from ``run_agent.py``; every method resolves through ``AIAgent``'s MRO unchanged. """ import contextlib import copy import functools import logging import threading from agent.session_activity import ActivityProvenance # Same logger name as the origin module so log records / caplog filters are unchanged. logger = logging.getLogger("run_agent") class _CommitFenceRegistration: """One ``_compress_context`` attempt's entry on the agent's commit-fence stack. The fence slot is a stack, not a save/restore cell: overlapping attempts register in order but may complete in any order, so each attempt removes only its own entry and the published slot always tracks the newest *live* attempt. The stall-fallback retry swaps ``fence`` in place so the slot update stays tied to the owning registration. """ __slots__ = ("fence",) def __init__(self, fence) -> None: self.fence = fence 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.""" if cached := getattr(agent, "_cached_system_prompt", None): 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. record = getattr(getattr(agent, "context_compressor", None), "record_timeout_failure", None) if callable(record): try: if total_exhausted: record("host compress_context total ceiling exhausted", failure_kind="ceiling_exhausted") else: record("host compress_context timeout (no summary progress)", failure_kind="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( f"⚠ Context compression timed out 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( f"⚠ Context compression commit is taking unusually 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) def _run_under_progress_timeout( agent, run, messages, system_message, *, active_fence, registration, fence_registration_lock, idle_timeout, total_ceiling, approx_tokens=None, ): """Run ``run(fence, target_messages=snapshot)`` on the pool under the progress-aware timeout. The pooled worker must NEVER share the caller's live transcript — a late engine after a host timeout could rewrite it. It deep-snapshots on the worker and publishes only via an ADMITTED commit; a no-op/abort returns the snapshot unchanged, so the ORIGINAL list is handed back to keep identity semantics.""" from agent.conversation_compression import ( CompressionCommitFence, request_exceeds_model_window, run_compress_context_with_progress_timeout, ) def _snapshot_worker(fence=None, *, same_turn_fallback_recovery=False): # #76354 review F3: the pooled worker must NEVER share the caller's live transcript. Plugin/legacy # context engines are allowed to mutate their input list in place; after a host timeout the worker # stays alive, so a shared list would let a late engine rewrite the live conversation (roles, # ordering, persisted content) behind the caller's back. Deep-snapshot here, on the worker thread, # so the caller's list object is never touched by pooled code. Results are published to # caller-visible state only via the returned value of an ADMITTED commit (the host discards results # on timeout/cancel); durable SessionDB mutation is already gated behind the commit fence inside # compress_context. snapshot = copy.deepcopy(messages) result_msgs, result_prompt = run( fence, target_messages=snapshot, same_turn_fallback_recovery=same_turn_fallback_recovery ) return (messages if result_msgs is snapshot else result_msgs), result_prompt # The stall-fallback retry is the same recovery attempt as the stalled primary, but the cancelled primary # worker records its stall_interrupted cooldown while unwinding — racing the retry's automatic gate # (#112387). The retry therefore bypasses ONLY the summary-failure cooldown (never clears it; the # structural breakers stay in force), exactly like provider-proven overflow recovery. _same_turn_fallback_worker = functools.partial(_snapshot_worker, same_turn_fallback_recovery=True) timeout_cause = {"total_exhausted": False, "progress_observed": False} def _on_timeout_cause(total_exhausted, progress_observed): timeout_cause.update(total_exhausted=total_exhausted, progress_observed=progress_observed) def _on_timeout(idle, waited, since_progress): _report_compression_timeout( agent, idle=idle, waited=waited, since_progress=since_progress, total_ceiling=total_ceiling, **timeout_cause ) def _publish_new_fence(): # The stall-fallback retry needs a fence the aborted attempt cannot veto; publish # it on the slot hard_interrupt() reads, but only while this attempt still owns the # top registration. A later attempt's live fence must not be clobbered. retry_fence = CompressionCommitFence() with fence_registration_lock: registration.fence = retry_fence stack = vars(agent).get("_compression_commit_fence_stack") or () if stack and stack[-1] is registration: agent._active_compression_commit_fence = retry_fence return retry_fence return run_compress_context_with_progress_timeout( worker=_snapshot_worker, messages=messages, system_prompt_fallback=lambda: _timeout_fallback_prompt(agent, 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=lambda waited, ceiling: _warn_commit_overrun(agent, waited, ceiling), fence=active_fence, telemetry_agent=agent, new_fence=_publish_new_fence, fallback_worker=_same_turn_fallback_worker, request_exceeds_window=request_exceeds_model_window(agent, approx_tokens) is True, ) def _mirror_result_onto_live_lists(agent, result, messages, *, direct_path: bool) -> None: """Mirror persisted-marker stamps from the result list onto the live list(s).""" if not (isinstance(result, tuple) and result and isinstance(result[0], list)): return result_messages = result[0] # Direct-path callers bypass the snapshot worker but still need the post-publish mirror. if direct_path or result_messages is not messages: _sync_persisted_markers(messages, result_messages) session_messages = getattr(agent, "_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_persisted_markers(session_messages, result_messages) def _rebind_caller_session_context(agent) -> None: """Propagate a rotated session id to the CALLER's thread/ContextVar (idempotent otherwise). The worker thread rotated hermes_logging's thread-local id; post-compression tools must resolve HERMES_SESSION_ID to the child id.""" with contextlib.suppress(Exception): from hermes_logging import set_session_context set_session_context(agent.session_id) try: from gateway.session_context import set_current_session_id if agent.session_id: set_current_session_id(agent.session_id) except Exception: logger.debug("post-compression session ContextVar rebind failed", exc_info=True) class CompressionFacadeMixin: """``_compress_context`` (see module docstring).""" def _compress_context( self, messages: list, system_message: str, *, approx_tokens: int = None, task_id: str = "default", focus_topic: str = None, force: bool = False, bypass_cooldown: bool = False, defer_context_engine_notification: bool = False, commit_fence=None, verbatim_tail: list = None, trigger: str = None, ) -> tuple: """Forwarder — see ``agent.conversation_compression.compress_context``. ``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. ``force=True`` is passed by the manual ``/compress`` slash command so users can bypass the summary-failure cooldown after an auto-compress abort. Auto-compress callers use the default ``force=False``. See #100661. """ # 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. # A stalled compression must not be mistaken for a structural no-op and followed by the oversized # provider request it was meant to prevent. The typed helper upgrades the simple attribute to # thread-local state guarded by a per-agent lock so overlapping automatic/manual entrypoints cannot # clobber each other's outcome (#98741). from agent.conversation_compression import ( CompressionCommitFence, compress_context, reset_context_compression_timeout_outcome, resolve_context_compression_timeouts, ) reset_context_compression_timeout_outcome(self) from agent.portal_tags import ( get_affinity_scope, get_conversation_context, reset_affinity_scope, reset_conversation_context, set_affinity_scope, set_conversation_context, ) 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. 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) # Initialized alongside `token`: the turn-lease timeout/interrupt early returns leave the try block # before set_affinity_scope() runs, and the finally reads this name unconditionally # (UnboundLocalError otherwise — the 4 red cross-process lease tests on PR #97158). 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(). Publication is serialized so overlapping automatic/manual # entrypoints cannot replace the fence of the attempt currently committing. active_fence = commit_fence or CompressionCommitFence() fence_registration_lock = vars(self).setdefault("_compression_commit_fence_lock", threading.RLock()) registration = _CommitFenceRegistration(active_fence) try: with fence_registration_lock: fence_stack = vars(self).setdefault("_compression_commit_fence_stack", []) fence_stack.append(registration) self._active_compression_commit_fence = active_fence def _run(fence=None, target_messages=None, same_turn_fallback_recovery=False): return compress_context( self, target_messages if target_messages is not None else messages, system_message, approx_tokens=approx_tokens, task_id=task_id, focus_topic=focus_topic, force=force, bypass_cooldown=bypass_cooldown or same_turn_fallback_recovery, defer_context_engine_notification=(defer_context_engine_notification), commit_fence=fence, verbatim_tail=verbatim_tail, trigger=trigger, ) # Callers that already own a progress-aware wait (gateway session # hygiene) pass commit_fence and must not be double-wrapped. direct_path = commit_fence is not None if not direct_path: idle_timeout, total_ceiling = resolve_context_compression_timeouts() direct_path = idle_timeout <= 0 if direct_path: result = _run(active_fence) else: result = _run_under_progress_timeout( self, _run, messages, system_message, active_fence=active_fence, registration=registration, fence_registration_lock=fence_registration_lock, idle_timeout=idle_timeout, total_ceiling=total_ceiling, approx_tokens=approx_tokens, ) _mirror_result_onto_live_lists(self, result, messages, direct_path=direct_path) _rebind_caller_session_context(self) return result finally: # Remove only THIS attempt's registration. Completion order is not registration # order: restoring a saved "previous" fence would resurrect a dead fence over a # live newer attempt (and popping the slot outright would delete that attempt's # fence mid-run). The slot always tracks the newest live registration. with fence_registration_lock: # Re-read rather than reuse the try-block local: an early exception before # the append would otherwise raise NameError here and mask the real failure. fence_stack = vars(self).get("_compression_commit_fence_stack") or () for idx in range(len(fence_stack) - 1, -1, -1): if fence_stack[idx] is registration: del fence_stack[idx] break if fence_stack: self._active_compression_commit_fence = fence_stack[-1].fence else: vars(self).pop("_active_compression_commit_fence", None) # Restore whatever the caller had, so a compaction never leaks its tag into the surrounding scope. if token is not None: reset_conversation_context(token) if affinity_token is not None: reset_affinity_scope(affinity_token)