EvoScientist/tests/test_auxiliary_model.py

"""Tests for the auxiliary-model resolver and its middleware scoping.

Covers ``EvoScientist.EvoScientist._ensure_auxiliary_chat_model`` — now
resolved through the model registry role mapping (design doc 6.1:
``defaults.auxiliary ?? defaults.primary``) instead of the legacy
``cfg.auxiliary_model``/``auxiliary_provider`` free strings — and the wiring
in ``_get_default_middleware`` that routes the main agent's tool selector to
the auxiliary model while keeping context editing — and async sub-agents —
on the main model.
"""

from unittest.mock import MagicMock, patch

import pytest

import EvoScientist.EvoScientist as E
from EvoScientist.model_registry.errors import (
    MODEL_REGISTRY_NOT_READY,
    ModelRegistryError,
)
from tests.registry_fixtures import activate_store


@pytest.fixture(autouse=True)
def _reset_model_caches(monkeypatch):
    """Isolate the module-level model caches per test."""
    monkeypatch.setattr(E, "_chat_model", None, raising=False)
    monkeypatch.setattr(E, "_chat_model_key", None, raising=False)
    monkeypatch.setattr(E, "_auxiliary_chat_model", None, raising=False)
    monkeypatch.setattr(E, "_auxiliary_chat_model_key", None, raising=False)


class TestAuxiliaryResolver:
    def test_bootstrap_registry_fails_closed(self, isolated_snapshot_runtime):
        """Bootstrap registry: auxiliary resolution propagates NOT_READY."""
        with pytest.raises(ModelRegistryError) as excinfo:
            E._ensure_auxiliary_chat_model()
        assert excinfo.value.code == MODEL_REGISTRY_NOT_READY

    def test_no_auxiliary_default_returns_main_instance(
        self, monkeypatch, isolated_snapshot_runtime
    ):
        main = object()
        activate_store(isolated_snapshot_runtime.store, auxiliary_default=False)
        monkeypatch.setattr(E, "_ensure_chat_model", lambda: main)
        assert E._ensure_auxiliary_chat_model() is main

    def test_auxiliary_equal_to_primary_reuses_main_instance(
        self, monkeypatch, isolated_snapshot_runtime
    ):
        main = object()
        store = activate_store(isolated_snapshot_runtime.store)
        registry = store.load_registry()
        registry.defaults.auxiliary = registry.defaults.primary
        store.save_registry(expected_revision=registry.revision, registry=registry)
        monkeypatch.setattr(E, "_ensure_chat_model", lambda: main)
        assert E._ensure_auxiliary_chat_model() is main

    def test_auxiliary_default_builds_via_role_mapping(
        self, monkeypatch, isolated_snapshot_runtime
    ):
        fake = object()
        build = MagicMock(return_value=fake)
        activate_store(isolated_snapshot_runtime.store)
        monkeypatch.setattr(
            isolated_snapshot_runtime, "build_default_role_model", build
        )
        monkeypatch.setattr(
            E, "_ensure_chat_model", lambda: pytest.fail("must not build main")
        )
        assert E._ensure_auxiliary_chat_model() is fake
        build.assert_called_once_with("auxiliary")

    def test_auxiliary_cache_reused_within_same_registry_revision(
        self, monkeypatch, isolated_snapshot_runtime
    ):
        build = MagicMock(side_effect=[object(), object()])
        activate_store(isolated_snapshot_runtime.store)
        monkeypatch.setattr(
            isolated_snapshot_runtime, "build_default_role_model", build
        )
        first = E._ensure_auxiliary_chat_model()
        assert E._ensure_auxiliary_chat_model() is first
        assert build.call_count == 1

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


class TestAuxiliaryMiddlewareScope:
    """``_get_default_middleware`` routes only the right components to aux."""

    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_aux_context_editing_main(self):
        cap, fake_ts, fake_ce = self._capture()
        main_model, aux_model = object(), object()
        with (
            patch.object(E, "_ensure_config", return_value=_mock_cfg()),
            patch.object(E, "_ensure_chat_model", return_value=main_model),
            patch.object(E, "_ensure_auxiliary_chat_model", return_value=aux_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()

        assert cap["tool_selector"] is aux_model
        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, aux_model = object(), object()
        with (
            patch.object(E, "_ensure_config", return_value=_mock_cfg()),
            patch.object(E, "_ensure_chat_model", return_value=main_model),
            patch.object(E, "_ensure_auxiliary_chat_model", return_value=aux_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):
        """The pure path never resolves auxiliary free strings (Task 7 wires
        the local snapshot entry); the threaded model stands in."""
        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.object(E, "_ensure_auxiliary_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.object(E, "_ensure_auxiliary_chat_model", return_value=main_model),
            patch(
                "EvoScientist.middleware.create_tool_selector_middleware",
                return_value=[MagicMock()],
            ),
        ):
            mw = E._get_default_middleware(snapshot_role="auxiliary")

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


def test_memory_agent_factory_uses_auxiliary_role(monkeypatch):
    """Memory workers bind the auxiliary role model at graph build."""
    sentinel = object()
    monkeypatch.setattr(E, "_ensure_auxiliary_chat_model", lambda: 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