hermes-agent/tests/agent/test_system_prompt_restore.py

"""Tests for ``agent.conversation_loop._restore_or_build_system_prompt``.

Validates the gateway DB-roundtrip path that keeps the system prompt
byte-stable across turns (fresh AIAgent → must restore from session DB
instead of rebuilding).  Covers:

  * Successful restore from a stored prompt (present row).
  * Legitimate first-turn build (no history).
  * Silent-failure recovery paths:
      - DB read raises → WARNING + fresh build
      - Row has system_prompt=NULL → WARNING + fresh build
      - Row has system_prompt="" → WARNING + fresh build
      - DB write fails → WARNING (subsequent turns will miss cache)
"""

from __future__ import annotations

import logging
from unittest.mock import MagicMock

import pytest

from agent.conversation_loop import _restore_or_build_system_prompt
from agent.surface_switch import _SURFACE_NAME_END, _SURFACE_SWITCH_NOTE_PREFIX, identity_line_value


def _make_agent(session_db=None, prebuilt_prompt: str = "BUILT_PROMPT"):
    """Construct the minimal agent fake the helper needs."""
    agent = MagicMock()
    agent._cached_system_prompt = None
    agent.session_id = "test-session-id"
    agent.model = "test-model"
    agent.provider = "openrouter"
    agent.platform = "cli"
    agent._session_db = session_db
    # MagicMock attributes are truthy by default; the static-prefix
    # reconstruction is gated on _use_prompt_caching, so default it off
    # for the legacy restore tests (the reconstruction tests enable it).
    agent._use_prompt_caching = False
    agent._build_system_prompt = MagicMock(return_value=prebuilt_prompt)
    return agent


# ---------------------------------------------------------------------------
# Surface switch (#104414)
# ---------------------------------------------------------------------------


class TestSurfaceSwitch:
    """A desktop <-> TUI switch must not rebuild the system prompt.

    The rebuild it used to trigger changed the first blocks of a 200K+ token request, so
    the whole conversation re-prefilled at a ~1% cache hit. The stored bytes are now reused
    and the new surface's guidance is delivered as a per-turn note behind the cached prefix.
    """

    @staticmethod
    def _stored(platform: str) -> str:
        return (
            "SYSTEM PROMPT BODY\n\nConversation started: Monday, January 05, 2026\n"
            "Model: test-model\nProvider: openrouter\n"
            f"Platform: {platform}"
        )

    @staticmethod
    def _announced(platform: str) -> list:
        """A transcript whose newest surface note says the model is on ``platform``."""
        return [
            {"role": "user", "content": "hi",
             "api_content": f"hi\n\n{_SURFACE_SWITCH_NOTE_PREFIX}{platform}{_SURFACE_NAME_END} superseded]"},
            {"role": "assistant", "content": "hello"},
        ]

    @staticmethod
    def _tool(name: str) -> dict:
        return {"type": "function", "function": {"name": name, "parameters": {}}}

    def _restore(self, *, stored: str, current: str, history=None, tool_names=None, tools=None):
        db = MagicMock()
        row = {"system_prompt": self._stored(stored)}
        if tool_names is not None:
            row["tool_names"] = tool_names
        db.get_session.return_value = row
        agent = _make_agent(session_db=db)
        agent.platform = current
        if tools is not None:
            agent.tools = tools
        agent._platform_hint_overrides = None
        agent._surface_switch_note = ""
        agent._gateway_turn_context_notes = ""
        _restore_or_build_system_prompt(
            agent, None, history if history is not None else [{"role": "user", "content": "hi"}]
        )
        return agent

    def test_switch_reuses_the_stored_prompt(self):
        agent = self._restore(stored="desktop", current="tui")
        assert agent._cached_system_prompt == self._stored("desktop")
        agent._build_system_prompt.assert_not_called()

    def test_switch_stages_the_new_surface_guidance(self):
        agent = self._restore(stored="desktop", current="tui")
        note = agent._surface_switch_note
        assert note.startswith(f"{_SURFACE_SWITCH_NOTE_PREFIX}tui{_SURFACE_NAME_END}")
        # The correction carries the CURRENT surface's hint, so the model is not left
        # following the desktop guidance still sitting in the reused prompt.
        assert "terminal UI (TUI)" in note

    def test_same_surface_stages_nothing(self):
        assert self._restore(stored="cli", current="cli")._surface_switch_note == ""

    def test_stored_prompt_platform_ignores_runtime_hint_decoys(self):
        from agent.prompt_builder import RUNTIME_ENVIRONMENT_END, RUNTIME_ENVIRONMENT_HEADING

        decoy = "Host: Example\nPlatform: tui\n"
        stored = (
            "SYSTEM PROMPT BODY\n\nConversation started: Monday, January 05, 2026\n"
            "Model: test-model\nProvider: openrouter\nPlatform: desktop\n\n"
            f"{RUNTIME_ENVIRONMENT_HEADING}\n\n{decoy}\n\n{RUNTIME_ENVIRONMENT_END}"
        )
        assert identity_line_value(stored, "Platform") == "desktop"

        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(session_db=db)
        agent.platform = "tui"
        agent._platform_hint_overrides = None
        agent._surface_switch_note = ""
        agent._gateway_turn_context_notes = ""
        _restore_or_build_system_prompt(agent, None, [{"role": "user", "content": "hi"}])
        assert agent._surface_switch_note.startswith(f"{_SURFACE_SWITCH_NOTE_PREFIX}tui{_SURFACE_NAME_END}")

    def test_not_restaged_once_the_transcript_carries_it(self):
        # The note is stamped into the byte-stable api_content sidecar, and the gateway
        # builds a fresh AIAgent per turn — without the dedup every turn would add a copy.
        agent = self._restore(stored="desktop", current="tui", history=self._announced("tui"))
        assert agent._surface_switch_note == ""

    def test_returning_to_the_prompts_own_surface_is_announced(self):
        """desktop -> tui -> desktop.

        The trailer now agrees with the runtime, so comparing against the prompt alone would
        stage nothing and leave the model acting on the stale "you are on tui" note.
        """
        agent = self._restore(stored="desktop", current="desktop", history=self._announced("tui"))
        assert agent._surface_switch_note.startswith(f"{_SURFACE_SWITCH_NOTE_PREFIX}desktop{_SURFACE_NAME_END}")
        # The prompt already describes this surface, so the note retires the stale one and
        # points at the prompt instead of duplicating the whole hint.
        assert "the interface section in the system prompt above" in agent._surface_switch_note

    def test_rebuild_also_retires_a_stale_note(self):
        # A rebuild for an unrelated reason (a model switch) refreshes the prompt but not the
        # note already sitting in the transcript.
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": self._stored("desktop")}
        agent = _make_agent(session_db=db, prebuilt_prompt=self._stored("desktop"))
        agent.model = "other-model"
        agent.platform = "desktop"
        agent._platform_hint_overrides = None
        agent._surface_switch_note = ""

        _restore_or_build_system_prompt(agent, None, self._announced("tui"))

        agent._build_system_prompt.assert_called_once()
        assert agent._surface_switch_note.startswith(f"{_SURFACE_SWITCH_NOTE_PREFIX}desktop{_SURFACE_NAME_END}")

    def test_tool_prefix_stays_pinned_on_the_turn_that_announces_a_switch(self):
        """tools[] is serialized ahead of the prompt this branch went out of its way to keep.

        Rebuilding the array for the new surface would move token 0 and re-prefill the whole
        request — the cost #104414 is about — on the very turn the fix exists to make cheap.
        """
        from unittest.mock import patch

        with (
            patch("tools.mcp_tool_agent.restore_agent_tool_prefix") as pin,
            patch("tools.mcp_tool_agent.persist_agent_tool_names") as persist,
        ):
            self._restore(stored="desktop", current="tui", tool_names='["desktop_ui_tool"]')
        pin.assert_called_once()
        persist.assert_not_called()

    def test_the_note_names_the_tools_the_pin_carried_forward(self):
        """A pinned tool this surface did not build is inert here — say so.

        Keeping it on the wire is what preserves the prefix, so the model has to learn from the
        note that calling it only returns ``tool_error("desktop only")``.
        """
        from unittest.mock import patch

        def _carry_desktop_tool(agent, saved_names):
            agent.tools = [self._tool("read_file"), self._tool("focus_pane")]
            return True

        with patch("tools.mcp_tool_agent.restore_agent_tool_prefix", _carry_desktop_tool):
            agent = self._restore(stored="desktop", current="tui", tool_names='["focus_pane"]',
                                  tools=[self._tool("read_file")])
        assert "focus_pane" in agent._surface_switch_note
        assert "were not loaded for this interface" in agent._surface_switch_note
        # Only the carried-over name: a tool this surface built is not inert.
        assert "read_file" not in agent._surface_switch_note.split("were not loaded for this interface")[1]

    def test_note_rides_the_user_message_channel_once(self):
        from agent.turn_context import _merge_gateway_notes, consume_surface_switch_note

        agent = self._restore(stored="desktop", current="tui")
        staged = agent._surface_switch_note
        assert _merge_gateway_notes(agent, [{"role": "user", "content": "hi"}], 0, "") == staged
        assert consume_surface_switch_note(agent) == ""


# ---------------------------------------------------------------------------
# Happy paths
# ---------------------------------------------------------------------------


class TestStoredPromptReuse:
    def test_present_row_is_reused_verbatim(self, caplog):
        """Continuing session with a stored prompt → reuse byte-for-byte."""
        stored = "Stored prompt from turn 1 — byte-identical reuse"
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(session_db=db)

        with caplog.at_level(logging.WARNING, logger="agent.conversation_loop"):
            _restore_or_build_system_prompt(agent, None, [{"role": "user", "content": "hi"}])

        assert agent._cached_system_prompt == stored
        agent._build_system_prompt.assert_not_called()
        db.update_system_prompt.assert_not_called()
        # No warnings on the happy path
        assert not [r for r in caplog.records if r.levelno >= logging.WARNING]

    def test_present_row_with_unicode_preserved(self):
        """Non-ASCII bytes in the stored prompt are not mangled."""
        stored = "Stored prompt with unicode: ☤ ⚗ ◆ — and emoji 🦊"
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(session_db=db)

        _restore_or_build_system_prompt(agent, None, [{"role": "user", "content": "hi"}])
        assert agent._cached_system_prompt == stored

    def test_present_row_with_stale_runtime_identity_rebuilds(self, caplog):
        """Stored prompts are cache gold unless their runtime identity is stale.

        A live /model switch updates the agent and DB model_config immediately.
        If the old system_prompt snapshot still says the previous model,
        blindly restoring it makes the next turn call the new model while the
        model reads old `Model:` metadata ("what model are you?" lies).
        """
        stored = (
            "You are Hermes Agent.\n\n"
            "Conversation started: Tuesday, June 16, 2026\n"
            "Session ID: test-session-id\n"
            "Model: anthropic/claude-opus-4.8-fast\n"
            "Provider: openrouter"
        )
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(
            session_db=db,
            prebuilt_prompt=(
                "You are Hermes Agent.\n\n"
                "Conversation started: Tuesday, June 16, 2026\n"
                "Session ID: test-session-id\n"
                "Model: openai/gpt-5.5\n"
                "Provider: openrouter"
            ),
        )
        agent.model = "openai/gpt-5.5"

        with caplog.at_level(logging.INFO, logger="agent.conversation_loop"):
            _restore_or_build_system_prompt(agent, None, [{"role": "user", "content": "hi"}])

        assert agent._cached_system_prompt.endswith(
            "Model: openai/gpt-5.5\nProvider: openrouter"
        )
        agent._build_system_prompt.assert_called_once_with(None)
        db.update_system_prompt.assert_called_once_with(
            agent.session_id, agent._cached_system_prompt
        )
        assert any("stale runtime identity" in r.getMessage() for r in caplog.records)


# ---------------------------------------------------------------------------
# Legitimate fresh-build paths (no history, no DB)
# ---------------------------------------------------------------------------


class TestLegitimateFreshBuild:
    def test_no_history_skips_db_and_builds_fresh(self, caplog):
        """First turn with empty history → build fresh, don't touch the DB."""
        db = MagicMock()
        agent = _make_agent(session_db=db)

        with caplog.at_level(logging.WARNING, logger="agent.conversation_loop"):
            _restore_or_build_system_prompt(agent, None, [])

        # No history → DB read skipped entirely
        db.get_session.assert_not_called()
        agent._build_system_prompt.assert_called_once_with(None)
        assert agent._cached_system_prompt == "BUILT_PROMPT"
        # Persisted to DB
        db.update_system_prompt.assert_called_once_with(agent.session_id, "BUILT_PROMPT")
        assert not [r for r in caplog.records if r.levelno >= logging.WARNING]

    def test_no_db_skips_persistence(self):
        """When session DB is None, build and skip persistence silently."""
        agent = _make_agent(session_db=None)
        _restore_or_build_system_prompt(agent, None, [])
        agent._build_system_prompt.assert_called_once()
        assert agent._cached_system_prompt == "BUILT_PROMPT"


# ---------------------------------------------------------------------------
# Silent-failure recovery — these are the new A/B logging paths
# ---------------------------------------------------------------------------


class TestSilentFailureWarnings:



    def test_db_write_failure_warns_loudly(self, caplog):
        """update_system_prompt raising → WARNING (was DEBUG before)."""
        db = MagicMock()
        # No prior row (first turn)
        db.get_session.return_value = None
        db.update_system_prompt.side_effect = RuntimeError("database is locked")
        agent = _make_agent(session_db=db)

        with caplog.at_level(logging.WARNING, logger="agent.conversation_loop"):
            _restore_or_build_system_prompt(agent, None, [])

        # Built and assigned the cache anyway
        agent._build_system_prompt.assert_called_once()
        assert agent._cached_system_prompt == "BUILT_PROMPT"
        # Warning surfaced
        warnings = [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING]
        assert any(
            "update_system_prompt failed" in m and "database is locked" in m
            for m in warnings
        ), f"Expected write-failure warning, got: {warnings}"

    def test_no_history_with_null_row_does_not_warn(self, caplog):
        """First turn (no history) hitting a null row is not surprising — no warn."""
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": None}
        agent = _make_agent(session_db=db)

        with caplog.at_level(logging.WARNING, logger="agent.conversation_loop"):
            # Empty history → DB read is skipped entirely
            _restore_or_build_system_prompt(agent, None, [])

        db.get_session.assert_not_called()
        # No "rebuilding from scratch" warning because history is empty
        warnings = [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING]
        assert not any("rebuilding" in m for m in warnings)


# ---------------------------------------------------------------------------
# Byte-stability invariant
# ---------------------------------------------------------------------------


class TestPromptStabilityInvariant:
    def test_restored_prompt_is_byte_identical_to_stored(self):
        """The restored prompt must equal the stored bytes exactly — no
        normalization, trimming, or concat that could shift the prefix.

        This is the core invariant: any byte-level change at this point
        invalidates KV cache on every prefix-cache backend.
        """
        stored = (
            "You are Hermes Agent.\n"
            "\n"
            "Conversation started: Sunday, May 17, 2026\n"
            "Session ID: 20260517_153500_abc123\n"
        )
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(session_db=db)

        _restore_or_build_system_prompt(agent, None, [{"role": "user", "content": "hi"}])

        # Identity check — must be the same object reference for maximum
        # confidence we're not slicing/copying/normalizing.
        assert agent._cached_system_prompt == stored
        # Byte-level check
        assert agent._cached_system_prompt.encode("utf-8") == stored.encode("utf-8")


# ---------------------------------------------------------------------------
# Cross-session static prefix reconstruction (issue #68191 follow-up)
# ---------------------------------------------------------------------------


class TestStaticPrefixReconstructionOnRestore:
    """The two-block cache layout must survive session restore.

    Gateway surfaces construct a fresh AIAgent per turn and restore the
    persisted prompt from the session DB; the cross-session-stable prefix
    (``_cached_system_prompt_static``) is only set on fresh builds, so
    without reconstruction the wire layout silently degrades to the legacy
    single-breakpoint layout after turn 1 (flagged on PR #68258 review).
    """

    def test_restore_reconstructs_static_prefix_when_it_matches(self):
        stable = "STATIC IDENTITY AND GUIDANCE"
        stored = stable + "\n\nper-session context\n\nvolatile tail"
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(session_db=db)
        agent._use_prompt_caching = True
        agent._cached_system_prompt_static = None

        from unittest.mock import patch as _patch

        with _patch(
            "agent.system_prompt.build_system_prompt_parts",
            return_value={"stable": stable, "context": "", "volatile": ""},
        ):
            _restore_or_build_system_prompt(
                agent, None, [{"role": "user", "content": "hi"}]
            )

        # Restored prompt bytes untouched; static prefix reconstructed.
        assert agent._cached_system_prompt == stored
        assert agent._cached_system_prompt_static == stable

    def test_restore_leaves_static_unset_on_prefix_mismatch(self):
        """Stable-tier drift (skills edited since persist) → no static prefix,
        legacy layout, restored bytes still authoritative."""
        stored = "OLD STATIC HEAD\n\nper-session context"
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(session_db=db)
        agent._use_prompt_caching = True
        agent._cached_system_prompt_static = None

        from unittest.mock import patch as _patch

        with _patch(
            "agent.system_prompt.build_system_prompt_parts",
            return_value={"stable": "NEW STATIC HEAD", "context": "", "volatile": ""},
        ):
            _restore_or_build_system_prompt(
                agent, None, [{"role": "user", "content": "hi"}]
            )

        assert agent._cached_system_prompt == stored
        assert agent._cached_system_prompt_static is None

    def test_restore_survives_parts_builder_exception(self):
        """Prefix reconstruction is fail-open: a parts-builder crash must not
        break the byte-identical restore."""
        stored = "Stored prompt — must survive"
        db = MagicMock()
        db.get_session.return_value = {"system_prompt": stored}
        agent = _make_agent(session_db=db)
        agent._use_prompt_caching = True
        agent._cached_system_prompt_static = None

        from unittest.mock import patch as _patch

        with _patch(
            "agent.system_prompt.build_system_prompt_parts",
            side_effect=RuntimeError("boom"),
        ):
            _restore_or_build_system_prompt(
                agent, None, [{"role": "user", "content": "hi"}]
            )

        assert agent._cached_system_prompt == stored
        assert agent._cached_system_prompt_static is None


class TestReconstructStaticPrefixMemoization:
    """A failed static rebuild must not re-run the parts builder every call.

    ``reconstruct_static_prefix`` sits on the retry-loop hot path via the
    failover redecoration chokepoint (#72626); ``build_system_prompt_parts``
    does real file I/O (SOUL.md, context files, memory), so a persistent
    stable-tier mismatch must be checked once per stored prompt, not on
    every attempt of every API call.
    """

    def _agent(self, stored):
        agent = _make_agent()
        agent._use_prompt_caching = True
        agent._cached_system_prompt = stored
        agent._cached_system_prompt_static = None
        return agent

    def test_failed_rebuild_is_memoized_per_stored_prompt(self):
        from unittest.mock import patch as _patch

        from agent.system_prompt import reconstruct_static_prefix

        stored = "STORED PROMPT\n\ntail"
        agent = self._agent(stored)
        with _patch(
            "agent.system_prompt.build_system_prompt_parts",
            return_value={"stable": "MISMATCH", "context": "", "volatile": ""},
        ) as build:
            reconstruct_static_prefix(agent)
            reconstruct_static_prefix(agent)
            reconstruct_static_prefix(agent)
        assert build.call_count == 1
        assert agent._cached_system_prompt_static is None

    def test_changed_stored_prompt_retries_once(self):
        from unittest.mock import patch as _patch

        from agent.system_prompt import reconstruct_static_prefix

        agent = self._agent("OLD STORED")
        with _patch(
            "agent.system_prompt.build_system_prompt_parts",
            return_value={"stable": "MISMATCH", "context": "", "volatile": ""},
        ) as build:
            reconstruct_static_prefix(agent)
            # A new stored prompt (e.g. after compression) invalidates the
            # failure memo and gets exactly one fresh attempt.
            agent._cached_system_prompt = "NEW STORED"
            reconstruct_static_prefix(agent)
            reconstruct_static_prefix(agent)
        assert build.call_count == 2

    def test_success_clears_failure_memo_and_early_returns(self):
        from unittest.mock import patch as _patch

        from agent.system_prompt import reconstruct_static_prefix

        stable = "STATIC HEAD"
        stored = stable + "\n\nvolatile"
        agent = self._agent(stored)
        with _patch(
            "agent.system_prompt.build_system_prompt_parts",
            return_value={"stable": stable, "context": "", "volatile": ""},
        ) as build:
            reconstruct_static_prefix(agent)
            reconstruct_static_prefix(agent)
        # Second call early-returns on the already-valid static prefix.
        assert build.call_count == 1
        assert agent._cached_system_prompt_static == stable
        assert getattr(agent, "_static_rebuild_failed_for", None) is None


class TestPerResponseSessionWritePath:
    """The write path under an embedding host's per-response session (#96570).

    Hermes Studio group chat pre-creates the SQLite row and pre-persists the
    user message BEFORE ``run_conversation()``, then runs one turn under a
    session id it destroys afterwards. The row therefore starts with a null
    system prompt and a non-empty history on its own genuine FIRST turn, which
    is what trips the "stored system prompt is null" warning — the warning is
    a first-turn artifact of that lifecycle, not evidence of a lost write.

    This pins the write path against that exact lifecycle: the freshly built
    prompt must land in the pre-created row within the same run.
    """

    def _agent(self, db, session_id):
        agent = _make_agent(session_db=db, prebuilt_prompt="GROUP_PROMPT")
        agent.session_id = session_id
        return agent

    def test_prepersisted_row_stores_the_freshly_built_prompt(self, tmp_path):
        from hermes_state import SessionDB

        session_id = "gc_run_room42_default_Worker_5f2c1ab9d4e34f7a8b0c6d1e2f3a4b5c"
        with SessionDB(db_path=tmp_path / "state.db") as db:
            # What the bridge does before the turn starts.
            db.create_session(session_id, source="studio")
            db.append_message(session_id=session_id, role="user", content="hi")

            _restore_or_build_system_prompt(
                self._agent(db, session_id),
                None,
                [{"role": "user", "content": "hi"}],
            )

            assert db.get_session(session_id)["system_prompt"] == "GROUP_PROMPT"

    def test_warning_is_a_first_turn_artifact_not_a_lost_write(
        self, tmp_path, caplog
    ):
        """Second turn of the SAME id restores — so nothing was dropped."""
        from hermes_state import SessionDB

        session_id = "gc_run_room42_default_Worker_9a7e3b1c05d24e6fb83a1c7d9e0f2a4b"
        history = [{"role": "user", "content": "hi"}]
        with SessionDB(db_path=tmp_path / "state.db") as db:
            db.create_session(session_id, source="studio")
            db.append_message(session_id=session_id, role="user", content="hi")

            with caplog.at_level(
                logging.WARNING, logger="agent.conversation_loop"
            ):
                _restore_or_build_system_prompt(
                    self._agent(db, session_id), None, history
                )
            assert "is null" in caplog.text

            caplog.clear()
            second = self._agent(db, session_id)
            with caplog.at_level(
                logging.WARNING, logger="agent.conversation_loop"
            ):
                _restore_or_build_system_prompt(second, None, history)

            assert second._cached_system_prompt == "GROUP_PROMPT"
            second._build_system_prompt.assert_not_called()
            assert "is null" not in caplog.text


if __name__ == "__main__":
    pytest.main([__file__, "-v"])