mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-05-01 08:21:50 +08:00
0a5ee01e48
Follow-up to the cherry-picked PR #17447. The original flush spawned a bare threading.Thread for the buffer-flush path, overwriting self._sync_thread — which is aliased to the long-lived writer thread. Two consequences: 1. No serialization with the writer queue. If old-session retains were still queued in _retain_queue, the flush ran concurrently with the writer and both threads could call aretain_batch against the same document_id. 2. The pre-spawn 'self._sync_thread.join(timeout=5.0)' tried to join the long-lived writer, which never exits, so the join was a no-op that just timed out — never actually serialized anything. Fix: enqueue the flush closure on _retain_queue via _ensure_writer + put(). Natural FIFO ordering behind any pending retains, no new thread, no broken join. Shutdown-aware so it doesn't enqueue after teardown. Tests updated to drain via _retain_queue.join() instead of the stale _sync_thread.join(). Added regression guard test_flush_serializes_behind_pending_retains_via_writer_queue that blocks the writer mid-retain to prove the flush waits in FIFO behind the old retain. Also seeds _retain_queue / _shutting_down / stubbed _ensure_writer on the bare-object test helper in test_memory_session_switch.py so that path doesn't blow up under the new queue-enqueue. tests/plugins/memory/test_hindsight_provider.py + tests/agent/test_memory_session_switch.py: 103/103 passing.
321 lines
11 KiB
Python
321 lines
11 KiB
Python
"""Tests for the on_session_switch hook and session_id propagation.
|
|
|
|
Covers #6672: memory providers must be notified when AIAgent.session_id
|
|
rotates mid-process (via /resume, /branch, /reset, /new, or context
|
|
compression). Without the notification, providers that cache per-session
|
|
state in initialize() (Hindsight, and any plugin that stores session_id
|
|
for scoped writes) keep writing into the old session's record.
|
|
"""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from agent.memory_manager import MemoryManager
|
|
from agent.memory_provider import MemoryProvider
|
|
|
|
|
|
class _RecordingProvider(MemoryProvider):
|
|
"""Provider that records every lifecycle call for assertion."""
|
|
|
|
def __init__(self, name="rec"):
|
|
self._name = name
|
|
self.switch_calls: list[dict] = []
|
|
self.sync_calls: list[dict] = []
|
|
self.queue_calls: list[dict] = []
|
|
self.initialize_calls: list[dict] = []
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return self._name
|
|
|
|
def is_available(self) -> bool: # pragma: no cover - unused
|
|
return True
|
|
|
|
def initialize(self, session_id, **kwargs):
|
|
self.initialize_calls.append({"session_id": session_id, **kwargs})
|
|
|
|
def get_tool_schemas(self):
|
|
return []
|
|
|
|
def sync_turn(self, user_content, assistant_content, *, session_id=""):
|
|
self.sync_calls.append(
|
|
{"user": user_content, "asst": assistant_content, "session_id": session_id}
|
|
)
|
|
|
|
def queue_prefetch(self, query, *, session_id=""):
|
|
self.queue_calls.append({"query": query, "session_id": session_id})
|
|
|
|
def on_session_switch(
|
|
self,
|
|
new_session_id,
|
|
*,
|
|
parent_session_id="",
|
|
reset=False,
|
|
**kwargs,
|
|
):
|
|
self.switch_calls.append(
|
|
{
|
|
"new": new_session_id,
|
|
"parent": parent_session_id,
|
|
"reset": reset,
|
|
"extra": kwargs,
|
|
}
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MemoryProvider ABC — default on_session_switch is a no-op
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _MinimalProvider(MemoryProvider):
|
|
"""Provider that does NOT override on_session_switch — ABC default must no-op."""
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "minimal"
|
|
|
|
def is_available(self) -> bool:
|
|
return True
|
|
|
|
def initialize(self, session_id, **kwargs): # pragma: no cover - unused
|
|
pass
|
|
|
|
def get_tool_schemas(self):
|
|
return []
|
|
|
|
|
|
def test_abc_default_on_session_switch_is_noop():
|
|
"""Providers that don't override the hook must not raise."""
|
|
p = _MinimalProvider()
|
|
# All three call styles must be accepted without raising
|
|
p.on_session_switch("new-id")
|
|
p.on_session_switch("new-id", parent_session_id="old-id")
|
|
p.on_session_switch("new-id", parent_session_id="old-id", reset=True)
|
|
p.on_session_switch("new-id", parent_session_id="old-id", reset=True, reason="new_session")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MemoryManager.on_session_switch — fan-out
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_manager_fans_out_to_all_providers():
|
|
mm = MemoryManager()
|
|
# Only one external provider is allowed; use the builtin slot for p1.
|
|
p1 = _RecordingProvider(name="builtin")
|
|
p2 = _RecordingProvider(name="hindsight")
|
|
mm.add_provider(p1)
|
|
mm.add_provider(p2)
|
|
|
|
mm.on_session_switch("new-sid", parent_session_id="old-sid", reset=False, reason="resume")
|
|
|
|
assert len(p1.switch_calls) == 1
|
|
assert len(p2.switch_calls) == 1
|
|
for call in (p1.switch_calls[0], p2.switch_calls[0]):
|
|
assert call["new"] == "new-sid"
|
|
assert call["parent"] == "old-sid"
|
|
assert call["reset"] is False
|
|
assert call["extra"] == {"reason": "resume"}
|
|
|
|
|
|
def test_manager_ignores_empty_session_id():
|
|
"""Empty string session_id must not trigger provider hooks.
|
|
|
|
Prevents accidental fires during shutdown when self.session_id may be
|
|
cleared. Providers expect a meaningful id to switch TO.
|
|
"""
|
|
mm = MemoryManager()
|
|
p = _RecordingProvider()
|
|
mm.add_provider(p)
|
|
mm.on_session_switch("")
|
|
mm.on_session_switch(None) # type: ignore[arg-type]
|
|
assert p.switch_calls == []
|
|
|
|
|
|
def test_manager_isolates_provider_failures():
|
|
"""A provider that raises must not block other providers."""
|
|
|
|
class _Broken(_RecordingProvider):
|
|
def on_session_switch(self, *args, **kwargs): # type: ignore[override]
|
|
raise RuntimeError("boom")
|
|
|
|
mm = MemoryManager()
|
|
# MemoryManager rejects a second external provider, so pair broken
|
|
# (builtin slot) with a good external one.
|
|
broken = _Broken(name="builtin")
|
|
good = _RecordingProvider(name="good")
|
|
mm.add_provider(broken)
|
|
mm.add_provider(good)
|
|
|
|
# Must not raise — exceptions in one provider are swallowed + logged
|
|
mm.on_session_switch("new-sid", parent_session_id="old-sid")
|
|
assert len(good.switch_calls) == 1
|
|
assert good.switch_calls[0]["new"] == "new-sid"
|
|
|
|
|
|
def test_manager_reset_flag_preserved():
|
|
mm = MemoryManager()
|
|
p = _RecordingProvider()
|
|
mm.add_provider(p)
|
|
mm.on_session_switch("new-sid", reset=True, reason="new_session")
|
|
assert p.switch_calls[0]["reset"] is True
|
|
assert p.switch_calls[0]["extra"] == {"reason": "new_session"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MemoryManager.sync_all / queue_prefetch_all — session_id propagation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_sync_all_propagates_session_id_to_providers():
|
|
"""run_agent.py's sync_all call must pass session_id through to providers.
|
|
|
|
Without this, a provider that updates _session_id defensively in
|
|
sync_turn (as Hindsight does at hindsight/__init__.py:1199) never
|
|
sees the new id and keeps writing under the old one.
|
|
"""
|
|
mm = MemoryManager()
|
|
p = _RecordingProvider()
|
|
mm.add_provider(p)
|
|
mm.sync_all("hello", "world", session_id="sess-42")
|
|
assert p.sync_calls == [
|
|
{"user": "hello", "asst": "world", "session_id": "sess-42"}
|
|
]
|
|
|
|
|
|
def test_queue_prefetch_all_propagates_session_id_to_providers():
|
|
mm = MemoryManager()
|
|
p = _RecordingProvider()
|
|
mm.add_provider(p)
|
|
mm.queue_prefetch_all("next query", session_id="sess-42")
|
|
assert p.queue_calls == [{"query": "next query", "session_id": "sess-42"}]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Hindsight reference implementation — state-flush semantics
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_hindsight_provider():
|
|
"""Build a bare HindsightMemoryProvider that skips network setup.
|
|
|
|
We instantiate without importing optional deps at class-level by
|
|
bypassing __init__ and seeding the attributes on_session_switch
|
|
reads/writes. This keeps the test hermetic.
|
|
"""
|
|
import threading
|
|
hindsight_mod = pytest.importorskip("plugins.memory.hindsight")
|
|
provider = object.__new__(hindsight_mod.HindsightMemoryProvider)
|
|
provider._session_id = "old-sid"
|
|
provider._parent_session_id = ""
|
|
provider._document_id = "old-sid-20260101_000000_000000"
|
|
provider._session_turns = ["turn-1", "turn-2"]
|
|
provider._turn_counter = 2
|
|
provider._turn_index = 2
|
|
# Attrs read by _build_metadata / _build_retain_kwargs when the
|
|
# buffer-flush path on session switch fires. Empty strings keep the
|
|
# metadata minimal but well-formed.
|
|
provider._retain_source = ""
|
|
provider._platform = ""
|
|
provider._user_id = ""
|
|
provider._user_name = ""
|
|
provider._chat_id = ""
|
|
provider._chat_name = ""
|
|
provider._chat_type = ""
|
|
provider._thread_id = ""
|
|
provider._agent_identity = ""
|
|
provider._agent_workspace = ""
|
|
provider._retain_tags = []
|
|
provider._retain_context = "test-context"
|
|
provider._retain_async = False
|
|
provider._bank_id = "test-bank"
|
|
# Prefetch state the switch path drains/clears.
|
|
provider._prefetch_thread = None
|
|
provider._prefetch_lock = threading.Lock()
|
|
provider._prefetch_result = ""
|
|
# Sync thread tracking (legacy alias at the writer).
|
|
provider._sync_thread = None
|
|
# Writer queue infra the flush-on-switch path enqueues onto. We stub
|
|
# _ensure_writer / _register_atexit so no real thread is spawned;
|
|
# tests exercising flush delivery live in
|
|
# tests/plugins/memory/test_hindsight_provider.py where the full
|
|
# writer-queue wiring is in place.
|
|
import queue as _queue
|
|
provider._retain_queue = _queue.Queue()
|
|
provider._shutting_down = threading.Event()
|
|
provider._atexit_registered = True
|
|
provider._ensure_writer = lambda: None
|
|
provider._register_atexit = lambda: None
|
|
# Stub the network-touching helper so any enqueued flush closure is
|
|
# a no-op if ever drained in a unit test.
|
|
provider._run_hindsight_operation = lambda _op: None
|
|
return provider
|
|
|
|
|
|
def test_hindsight_on_session_switch_updates_session_id_and_mints_fresh_doc():
|
|
provider = _make_hindsight_provider()
|
|
old_doc = provider._document_id
|
|
|
|
provider.on_session_switch(
|
|
"new-sid", parent_session_id="old-sid", reset=False, reason="resume"
|
|
)
|
|
|
|
assert provider._session_id == "new-sid"
|
|
assert provider._parent_session_id == "old-sid"
|
|
# Document id MUST be fresh — else next retain overwrites old session doc
|
|
assert provider._document_id != old_doc
|
|
assert provider._document_id.startswith("new-sid-")
|
|
|
|
|
|
def test_hindsight_on_session_switch_clears_turn_buffers():
|
|
"""Accumulated _session_turns must not leak into the next session.
|
|
|
|
Hindsight batches turns under a single _document_id. If the buffer
|
|
isn't cleared on switch, the next retain under the new _document_id
|
|
flushes turns that belong to the previous session.
|
|
"""
|
|
provider = _make_hindsight_provider()
|
|
provider.on_session_switch("new-sid", parent_session_id="old-sid")
|
|
assert provider._session_turns == []
|
|
assert provider._turn_counter == 0
|
|
assert provider._turn_index == 0
|
|
|
|
|
|
def test_hindsight_on_session_switch_clears_on_reset_true():
|
|
"""reset=True (from /new, /reset) must also flush buffers."""
|
|
provider = _make_hindsight_provider()
|
|
provider.on_session_switch("new-sid", reset=True, reason="new_session")
|
|
assert provider._session_id == "new-sid"
|
|
assert provider._session_turns == []
|
|
assert provider._turn_counter == 0
|
|
|
|
|
|
def test_hindsight_on_session_switch_ignores_empty_id():
|
|
"""Empty new_session_id must be a no-op to avoid corrupting state."""
|
|
provider = _make_hindsight_provider()
|
|
before = (
|
|
provider._session_id,
|
|
provider._document_id,
|
|
list(provider._session_turns),
|
|
provider._turn_counter,
|
|
)
|
|
provider.on_session_switch("")
|
|
provider.on_session_switch(None) # type: ignore[arg-type]
|
|
after = (
|
|
provider._session_id,
|
|
provider._document_id,
|
|
list(provider._session_turns),
|
|
provider._turn_counter,
|
|
)
|
|
assert before == after
|
|
|
|
|
|
def test_hindsight_preserves_parent_across_empty_parent_arg():
|
|
"""Omitting parent_session_id must NOT overwrite an existing one."""
|
|
provider = _make_hindsight_provider()
|
|
provider._parent_session_id = "original-parent"
|
|
provider.on_session_switch("new-sid") # no parent passed
|
|
assert provider._parent_session_id == "original-parent"
|