fix(agent): isolate the pooled compression worker from the live transcript (review F3)
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.
This commit is contained in:
Teknium
parent
efdd229884
commit
971d81f892
|
|
@ -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
|
||||
|
|
|
|||
33
run_agent.py
33
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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue