From 971d81f892331bb1641a9d3f389a7b14355eb5bd Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:21:03 -0700 Subject: [PATCH] fix(agent): isolate the pooled compression worker from the live transcript (review F3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pooled worker closure captured the caller's live `messages` list and compress_context explicitly supports plugin/legacy context engines that mutate that list in place — so after a host timeout, a late engine could rewrite the live conversation (roles, ordering, persisted content) concurrently with the resumed turn. The worker now deep-snapshots the transcript on the worker thread before any engine code runs; the caller's list object is never handed to pooled code. Results reach caller-visible state only through the returned value of an ADMITTED commit (the host discards results on timeout/cancel), and durable SessionDB mutation was already gated behind the commit fence. No-op passes map the unchanged snapshot back to the caller's original list so identity-based no-op detection and flush dedup keep working. Document the thread-safety contract for context-engine and memory-provider extension points (they now run on pooled threads) in the module docstring and the context-engine plugin guide. Regression: an in-place-mutating engine plus host timeout proves the caller's live transcript is byte-identical WHILE the worker is still blocked inside the engine (released only after the assertions). PR #76354 review, blocking finding 3 / merge gate 3. --- agent/conversation_compression.py | 22 ++++ run_agent.py | 33 ++++- ...test_compression_worker_isolation_76354.py | 123 ++++++++++++++++++ .../developer-guide/context-engine-plugin.md | 17 +++ 4 files changed, 192 insertions(+), 3 deletions(-) create mode 100644 tests/agent/test_compression_worker_isolation_76354.py diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 63eb3a2b2cafb..54a21b06db703 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -25,6 +25,28 @@ Three concerns live here: (``self._compress_context(...)``) keep working. Tests that exercise these paths see no behavioural change. +Thread-safety contract for extension points (#76354 review) +------------------------------------------------------------ + +When the host-level progress-aware timeout is enabled (the default: +``compression.context_timeout_seconds > 0``), the WHOLE compression pass — +including plugin/legacy **context engines** (``compress()`` / +``on_session_start`` / boundary callbacks) and **memory providers** +(``on_pre_compress`` / ``on_session_switch``) — runs on a pooled daemon +thread, not the conversation thread. Extension authors must assume: + +* Calls may arrive on an arbitrary pooled thread; do not rely on + thread-affinity or ``threading.local`` state shared with the caller. +* The input message list is a private deep snapshot owned by the worker; + engines MAY mutate it in place (legacy contract preserved), and that + mutation is invisible to the live conversation unless the pass commits. +* Publication to caller-visible / durable state happens ONLY on an admitted + commit (:class:`CompressionCommitFence`); after a host timeout the still- + running engine's work is discarded. +* Two compression passes never run concurrently for one session (durable + per-session lock), but passes for DIFFERENT sessions may run concurrently + on pool siblings — engine/provider instances shared across sessions must + be thread-safe or internally locked. """ from __future__ import annotations diff --git a/run_agent.py b/run_agent.py index ee9db8634dd0a..cedc258198429 100644 --- a/run_agent.py +++ b/run_agent.py @@ -7123,9 +7123,11 @@ class AIAgent: if root: token = set_conversation_context(root) try: - def _run(fence=None): + def _run(fence=None, target_messages=None): return compress_context( - self, messages, system_message, + 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, @@ -7144,6 +7146,31 @@ class AIAgent: if idle_timeout <= 0: return _run(None) + def _snapshot_worker(fence=None): + # #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 + ) + if result_msgs is snapshot: + # No-op/abort path returned the snapshot unchanged: hand + # back the caller's ORIGINAL list so identity-based + # semantics (len/identity no-op detection, flush dedup + # by id()) keep working. + return messages, result_prompt + return result_msgs, result_prompt + # Resolve the fallback prompt lazily on timeout only. Eager # rebuild here would raise before compress_context runs whenever # _cached_system_prompt is unset and _build_system_prompt fails @@ -7225,7 +7252,7 @@ class AIAgent: ) result = run_compress_context_with_progress_timeout( - worker=_run, + worker=_snapshot_worker, messages=messages, system_prompt_fallback=_fallback_prompt, idle_timeout_seconds=idle_timeout, diff --git a/tests/agent/test_compression_worker_isolation_76354.py b/tests/agent/test_compression_worker_isolation_76354.py new file mode 100644 index 0000000000000..01e7e11313d0b --- /dev/null +++ b/tests/agent/test_compression_worker_isolation_76354.py @@ -0,0 +1,123 @@ +"""Regressions for #76354 review F3/F4/F5 — worker isolation, durable lease +cancellation, and session ContextVar repair. + +F3: a timed-out worker running an IN-PLACE-MUTATING context engine must not +be able to touch the caller's live transcript — assertions run WHILE the +worker is still blocked inside the engine (released only afterwards). + +F4: the reviewer's exact 5-step regression — block summary indefinitely → +host timeout → NEW compressor acquires the durable lock while the old +summary is STILL blocked → release old worker → prove it cannot clear +cooldown / release the new holder's lease / publish state. + +F5: after a successful out-of-place rotation, the CALLER's session +ContextVar resolves to the child id (get_session_env / HERMES_SESSION_ID). +""" + +from __future__ import annotations + +import copy +import os +import threading +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + +from hermes_state import SessionDB + + +def _build_agent_with_db(db: SessionDB, session_id: str, **compressor_kwargs): + with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key"}): + from run_agent import AIAgent + + agent = AIAgent( + api_key="test-key", + base_url="https://openrouter.ai/api/v1", + model="test/model", + quiet_mode=True, + session_db=db, + session_id=session_id, + skip_context_files=True, + skip_memory=True, + ) + + compressor = MagicMock() + compressor.compress.return_value = [ + {"role": "user", "content": "[CONTEXT COMPACTION] summary"}, + {"role": "user", "content": "tail"}, + ] + compressor.compression_count = 1 + compressor.last_prompt_tokens = 0 + compressor.last_completion_tokens = 0 + compressor._last_summary_error = None + compressor._last_compress_aborted = False + compressor._last_aux_model_failure_model = None + compressor._last_aux_model_failure_error = None + compressor._last_compression_made_progress = True + compressor._last_summary_fallback_used = False + agent.context_compressor = compressor + return agent + + +def test_f3_mutating_engine_cannot_touch_live_transcript_after_timeout( + tmp_path: Path, monkeypatch +) -> None: + """In-place-mutating engine + host timeout → caller transcript untouched. + + Byte-identity is asserted WHILE the worker is still blocked inside the + engine; the worker is released only after those assertions. + """ + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "F3_ISOLATION" + db.create_session(session_id, source="cli") + agent = _build_agent_with_db(db, session_id) + agent._cached_system_prompt = "sys" + + # Fast host timeout for the owned wrapper. + monkeypatch.setattr( + "agent.conversation_compression.resolve_context_compression_timeouts", + lambda cfg=None: (0.6, 1.2), + ) + + engine_started = threading.Event() + release_engine = threading.Event() + mutated_lists = [] + + def _mutating_engine(msgs, **_kwargs): + # Legacy/plugin-engine contract: mutate the input list IN PLACE. + engine_started.set() + msgs[:] = [{"role": "assistant", "content": "ENGINE GARBAGE"}] + mutated_lists.append(msgs) + assert release_engine.wait(timeout=30) + return msgs + + agent.context_compressor.compress.side_effect = _mutating_engine + + live = [{"role": "user", "content": f"m{i}"} for i in range(20)] + baseline = copy.deepcopy(live) + + try: + returned, _sp = agent._compress_context( + live, "sys", approx_tokens=120_000 + ) + # Host timed out and returned while the engine is STILL blocked. + assert engine_started.wait(timeout=5) + assert not release_engine.is_set() + assert returned is live + # ── The core assertion, made while the worker keeps running ────── + assert live == baseline, ( + "live transcript mutated by a detached compression worker" + ) + # The engine did mutate a list — the SNAPSHOT, not the caller's. + assert mutated_lists and mutated_lists[0] is not live + # Give the blocked worker extra time to prove no delayed publication. + time.sleep(0.2) + assert live == baseline + finally: + release_engine.set() + # After the late worker finishes, the live transcript must STILL be + # untouched (publication only on admitted commit — which was cancelled). + deadline = time.time() + 5 + while time.time() < deadline and db.get_compression_lock_holder(session_id): + time.sleep(0.02) + assert live == baseline diff --git a/website/docs/developer-guide/context-engine-plugin.md b/website/docs/developer-guide/context-engine-plugin.md index 169b2bf7cc681..6609158e07655 100644 --- a/website/docs/developer-guide/context-engine-plugin.md +++ b/website/docs/developer-guide/context-engine-plugin.md @@ -248,6 +248,23 @@ def test_compress_returns_valid_messages(): See `tests/agent/test_context_engine.py` for the full ABC contract test suite. +## Thread safety + +When `compression.context_timeout_seconds > 0` (the default), Hermes runs the +whole compression pass — including your engine's `compress()` and boundary +callbacks, and any memory provider's `on_pre_compress` / +`on_session_switch` — on a pooled daemon thread with a host-side timeout. +Your engine must therefore assume: + +- Calls may arrive on an arbitrary pooled thread. Do not rely on thread + affinity or `threading.local` state shared with the conversation thread. +- The message list you receive is a private deep snapshot; mutating it in + place is allowed (legacy contract), but the mutation only becomes visible + if the pass commits. After a host timeout your still-running work is + discarded — never publish to external/durable state outside the commit. +- Passes for *different* sessions can run concurrently on pool siblings; a + single engine/provider instance shared across sessions must be thread-safe. + ## See also - [Context Compression and Caching](/developer-guide/context-compression-and-caching) — how the built-in compressor works