hermes-agent/tests/agent/test_memory_session_switch.py

"""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
    # Mode + API state used by _resolve_retain_target; stub the resolver
    # so tests don't actually probe the API. Real probe behavior is
    # exercised by tests in tests/plugins/memory/test_hindsight_provider.py.
    provider._mode = "cloud"
    provider._api_url = ""
    provider._api_key = ""
    provider._client = None
    provider._resolve_retain_target = lambda fb: (fb, 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"