hermes-agent/tests/agent/test_context_engine_select_context.py

"""Tests for the per-turn ``ContextEngine.select_context()`` hook.

``select_context()`` is the *selection / routing* verb — distinct from
compression — that lets an external context engine replace which context
enters the prompt for a single request, every turn, independent of
``should_compress()``. It is additive and no-op by default, and the host
call site (``_apply_context_engine_selection``) is fail-open: a missing hook,
an exception, or an invalid return value must leave the assembled request
untouched and must never mutate persisted history.

This pins the contract that engines such as retrieval-augmented, topic-routed,
and role-switching engines rely on (RFC #36765), consolidating the per-turn
request-assembly surface proposed across #41918, #24949, #47109, and #50053.
"""

from __future__ import annotations

from typing import Any, Dict, List
from unittest.mock import MagicMock

from agent.context_engine import ContextEngine
from agent.conversation_loop import (
    _apply_context_engine_selection,
    _notify_context_engine_turn_complete,
)


class _MinimalEngine(ContextEngine):
    """Concrete engine implementing only the abstract methods."""

    @property
    def name(self) -> str:
        return "minimal"

    def update_from_response(self, usage: Dict[str, Any]) -> None:
        pass

    def should_compress(self, prompt_tokens: int = None) -> bool:
        return False

    def compress(
        self,
        messages: List[Dict[str, Any]],
        current_tokens: int = None,
        focus_topic: str = None,
    ) -> List[Dict[str, Any]]:
        return messages


def _agent_with(engine) -> Any:
    agent = MagicMock()
    agent.session_id = "test-session"
    agent.context_compressor = engine
    return agent


REQUEST = [
    {"role": "system", "content": "sys"},
    {"role": "user", "content": "hello"},
]
HISTORY = [{"role": "user", "content": "hello"}]


# -- ABC default -----------------------------------------------------------

def test_default_select_context_is_noop():
    """The base implementation returns None (no replacement)."""
    engine = _MinimalEngine()
    assert (
        engine.select_context(
            REQUEST,
            conversation_messages=HISTORY,
            incoming_message=HISTORY[-1],
            budget_tokens=0,
        )
        is None
    )


# -- Host call site: _apply_context_engine_selection -----------------------

def test_none_return_leaves_request_unchanged():
    """An engine returning None falls through to the assembled request."""
    engine = _MinimalEngine()  # default select_context -> None
    agent = _agent_with(engine)
    out = _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
    )
    assert out is REQUEST


def test_missing_hook_leaves_request_unchanged():
    """An engine without select_context (older/stub base) is a no-op."""
    engine = object()  # no select_context attribute
    agent = _agent_with(engine)
    out = _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
    )
    assert out is REQUEST


def test_no_engine_leaves_request_unchanged():
    agent = MagicMock()
    agent.session_id = "test-session"
    agent.context_compressor = None
    out = _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
    )
    assert out is REQUEST


def test_valid_list_replaces_request():
    """A valid list of dicts replaces the request messages for this call."""
    replacement = [
        {"role": "system", "content": "sys"},
        {"role": "user", "content": "routed-context"},
    ]

    class _Engine(_MinimalEngine):
        def select_context(self, request_messages, **kwargs):
            return replacement

    agent = _agent_with(_Engine())
    out = _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
    )
    assert out is replacement


def test_exception_fails_open():
    """A raising hook is swallowed; the unmodified request is used."""

    class _Engine(_MinimalEngine):
        def select_context(self, request_messages, **kwargs):
            raise RuntimeError("backend offline")

    logger = MagicMock()
    agent = _agent_with(_Engine())
    out = _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=logger
    )
    assert out is REQUEST
    assert logger.warning.called


def test_non_list_return_is_ignored():
    """A non-list return value is rejected and logged, request unchanged."""

    class _Engine(_MinimalEngine):
        def select_context(self, request_messages, **kwargs):
            return {"role": "user", "content": "oops not a list"}

    logger = MagicMock()
    agent = _agent_with(_Engine())
    out = _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=logger
    )
    assert out is REQUEST
    assert logger.warning.called


def test_list_of_non_dicts_is_ignored():
    """A list that isn't all dicts is rejected, request unchanged."""

    class _Engine(_MinimalEngine):
        def select_context(self, request_messages, **kwargs):
            return ["not", "dicts"]

    agent = _agent_with(_Engine())
    out = _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
    )
    assert out is REQUEST


def test_persisted_history_not_mutated():
    """The hook must not mutate the persisted conversation history."""

    class _Engine(_MinimalEngine):
        def select_context(self, request_messages, *, conversation_messages=None, **kwargs):
            # Even a misbehaving engine touching its inputs must not affect
            # what the host persists — the host passes the live list, so we
            # assert the host contract by checking the engine received it and
            # the canonical copy is unchanged after the call.
            return list(request_messages)

    history_snapshot = [dict(m) for m in HISTORY]
    agent = _agent_with(_Engine())
    _apply_context_engine_selection(
        agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
    )
    assert HISTORY == history_snapshot


# -- on_turn_complete (post-turn observation) ------------------------------

def test_default_on_turn_complete_is_noop():
    """The base on_turn_complete returns None and does nothing."""
    assert _MinimalEngine().on_turn_complete(HISTORY, usage=None) is None


def test_on_turn_complete_called_with_snapshot_and_meta():
    """Host forwards a transcript copy + metadata; base no-op is skipped."""
    captured = {}

    class _Engine(_MinimalEngine):
        def on_turn_complete(self, messages, usage=None, **kwargs):
            captured["messages"] = messages
            captured["usage"] = usage
            captured["kwargs"] = kwargs

    agent = _agent_with(_Engine())
    _notify_context_engine_turn_complete(
        agent, HISTORY, usage={"total_tokens": 12}, logger=MagicMock(),
        turn_id="t1", api_call_count=1,
    )
    assert captured["messages"] == HISTORY
    assert captured["messages"] is not HISTORY  # shallow copy
    assert captured["usage"] == {"total_tokens": 12}
    assert captured["kwargs"]["turn_id"] == "t1"
    assert captured["kwargs"]["api_call_count"] == 1


def test_on_turn_complete_base_noop_is_skipped():
    """An engine that only inherits the base no-op is handled safely.

    The helper short-circuits the base implementation (so non-implementing
    engines pay nothing), and in any case must not raise.
    """
    agent = _agent_with(_MinimalEngine())  # inherits base on_turn_complete
    _notify_context_engine_turn_complete(agent, HISTORY, logger=MagicMock())


def test_on_turn_complete_fails_open():
    """A raising observation hook is swallowed and logged."""

    class _Engine(_MinimalEngine):
        def on_turn_complete(self, messages, usage=None, **kwargs):
            raise RuntimeError("indexing backend down")

    logger = MagicMock()
    agent = _agent_with(_Engine())
    _notify_context_engine_turn_complete(agent, HISTORY, logger=logger)
    assert logger.warning.called


def test_on_turn_complete_missing_engine_is_safe():
    agent = MagicMock()
    agent.session_id = "s"
    agent.context_compressor = None
    # No engine -> silent return, no raise.
    _notify_context_engine_turn_complete(agent, HISTORY, logger=MagicMock())