EvoScientist/tests/test_auxiliary_model.py

"""Tests for the helper-model wiring in ``_get_default_middleware``.

After the auxiliary-role removal (design doc 6.1: every role resolves to the
snapshot's frozen primary), the main agent's tool selector no longer binds a
compile-time auxiliary model — it resolves its helper model from the run
snapshot on every call (``model=None``). Async sub-agents and the pure path
keep their threaded model. Context editing always keeps the main model.
"""

from unittest.mock import MagicMock, patch

import EvoScientist.EvoScientist as E


def _mock_cfg():
    cfg = MagicMock()
    cfg.enable_ask_user = False
    cfg.auto_mode = False
    cfg.auto_approve = False
    return cfg


class TestHelperModelScope:
    """``_get_default_middleware`` routes only the right components."""

    def _capture(self):
        cap: dict[str, object] = {}

        def fake_tool_selector(*args, model=None, **kwargs):
            cap["tool_selector"] = model
            return [MagicMock()]

        def fake_context_editing(model=None, *args, **kwargs):
            cap["context_editing"] = model
            return MagicMock()

        return cap, fake_tool_selector, fake_context_editing

    def test_main_agent_tool_selector_uses_snapshot_resolution(self):
        cap, fake_ts, fake_ce = self._capture()
        main_model = object()
        with (
            patch.object(E, "_ensure_config", return_value=_mock_cfg()),
            patch.object(
                E, "_compile_time_role_model", return_value=main_model
            ),
            patch(
                "EvoScientist.middleware.create_tool_selector_middleware",
                side_effect=fake_ts,
            ),
            patch(
                "EvoScientist.middleware.create_context_editing_middleware",
                side_effect=fake_ce,
            ),
        ):
            E._get_default_middleware()

        # None = per-call resolution from the run snapshot.
        assert cap["tool_selector"] is None
        assert cap["context_editing"] is main_model

    def test_async_subagent_tool_selector_stays_main(self):
        cap, fake_ts, fake_ce = self._capture()
        main_model = object()
        with (
            patch.object(E, "_ensure_config", return_value=_mock_cfg()),
            patch.object(
                E, "_compile_time_role_model", return_value=main_model
            ),
            patch(
                "EvoScientist.middleware.create_tool_selector_middleware",
                side_effect=fake_ts,
            ),
            patch(
                "EvoScientist.middleware.create_context_editing_middleware",
                side_effect=fake_ce,
            ),
        ):
            E._get_default_middleware(for_async_subagent=True)

        assert cap["tool_selector"] is main_model
        assert cap["context_editing"] is main_model

    def test_pure_path_tool_selector_uses_threaded_model(self):
        cap, fake_ts, fake_ce = self._capture()
        cfg = _mock_cfg()
        main_model = object()

        with (
            patch.object(E, "_ensure_config", side_effect=AssertionError),
            patch.object(E, "_ensure_chat_model", side_effect=AssertionError),
            patch(
                "EvoScientist.middleware.create_tool_selector_middleware",
                side_effect=fake_ts,
            ),
            patch(
                "EvoScientist.middleware.create_context_editing_middleware",
                side_effect=fake_ce,
            ),
        ):
            E._get_default_middleware(cfg=cfg, chat_model=main_model)

        assert cap["tool_selector"] is main_model
        assert cap["context_editing"] is main_model

    def test_snapshot_role_forwarded_to_configurable_model_middleware(self):
        cfg = _mock_cfg()
        main_model = object()
        with (
            patch.object(E, "_ensure_config", return_value=cfg),
            patch.object(E, "_ensure_chat_model", return_value=main_model),
            patch(
                "EvoScientist.middleware.create_tool_selector_middleware",
                return_value=[MagicMock()],
            ),
        ):
            mw = E._get_default_middleware(snapshot_role="primary")

        configurable = next(
            m for m in mw if type(m).__name__ == "ConfigurableModelMiddleware"
        )
        assert configurable._role == "primary"


def test_memory_agent_factory_binds_primary_placeholder(monkeypatch):
    """Memory workers bind the primary placeholder plus snapshot middleware."""
    sentinel = object()
    monkeypatch.setattr(E, "_compile_time_role_model", lambda role="primary": sentinel)
    captured: dict[str, object] = {}

    def fake_create_deep_agent(**kwargs):
        captured.update(kwargs)
        return MagicMock()

    monkeypatch.setattr("deepagents.create_deep_agent", fake_create_deep_agent)

    from EvoScientist.memory.agents._factory import build_memory_agent_graph

    build_memory_agent_graph(
        name="worker",
        system_prompt="",
        memory_dir="/tmp/m",
        workspace_dir="/tmp/w",
        tools=[],
        middleware=[],
        backend=MagicMock(),
    )
    assert captured["model"] is sentinel
    middleware = captured["middleware"]
    assert type(middleware[0]).__name__ == "ConfigurableModelMiddleware"