"""Helpers on another model: capacity, designations, the choice, and the screens. `test_subagent.py` covers what a helper is *not* given; this covers which model it runs on. The generation loop is stubbed the way that file stubs it, and the tool is always reached through `resolve_tools`, because what may be run is what was offered. """ from __future__ import annotations import json import pytest from sqlalchemy import select from lembas.db.models import ( Chat, ChatHelper, Connection, DataGroup, HelperDesignation, Model, User, ) from lembas.services import helpers, settings_store, talk from lembas.services import subagent as subagent_service from lembas.services import tools as tools_service from lembas.services.crypto import encrypt @pytest.fixture(autouse=True) def setup(db, registered): """Helpers on, three models on one local connection and one hosted model.""" settings_store.update(db, {"enabled": True}, key=settings_store.SUBAGENTS) settings_store.update(db, {"default_permissions": {"tools.subagent": True}}) local = Connection(name="Local", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")) cloud = Connection(name="Cloud", base_url="http://127.0.0.1:2", api_key_encrypted=encrypt("")) db.add_all([local, cloud]) db.flush() tools = {"tools": True} db.add_all( [ Model(connection_id=local.id, model_id="gpt-oss", capabilities_json=tools, position=0, params_json={"reasoning_effort": "high"}), Model(connection_id=local.id, model_id="qwen35", capabilities_json=tools, position=1), Model(connection_id=local.id, model_id="bonsai", capabilities_json=tools, position=2, params_json={"reasoning_effort": "xhigh"}, reasoning_efforts=["low", "medium", "xhigh"]), Model(connection_id=cloud.id, model_id="deepseek", capabilities_json=tools, position=3), ] ) db.commit() subagent_service.clear() yield {"local": local, "cloud": cloud} subagent_service.clear() def _user(db) -> User: return db.scalars(select(User).order_by(User.created_at)).first() def _chat(db, model_id="gpt-oss", connection=None) -> Chat: connection_id = connection.id if connection else db.scalar( select(Model.connection_id).where(Model.model_id == model_id) ) chat = Chat(user_id=_user(db).id, title="t", model_id=model_id, connection_id=connection_id) db.add(chat) db.commit() return chat def _row(db, model_id) -> Model: return db.scalar(select(Model).where(Model.model_id == model_id)) def _ids(found) -> list[str]: return [c.model.model_id for c in found] def _spawn(monkeypatch): from lembas.db.models import ROLE_ASSISTANT, ROLE_USER from lembas.db.session import session_scope from lembas.services import chat as chat_service seen: dict[str, str] = {} async def fake_wake(chat_id: str, content: str, *, model_id: str = "") -> str: seen["chat_id"] = chat_id with session_scope() as db: child = db.get(Chat, chat_id) chat_service.create_message(db, child, ROLE_USER, content) reply = chat_service.create_message(db, child, ROLE_ASSISTANT, "Done.") return reply.id monkeypatch.setattr("lembas.services.wake.wake_chat", fake_wake) monkeypatch.setattr("lembas.services.generation.running_for", lambda chat_id: None) return seen async def _run(db, chat: Chat, args: dict): from lembas.services import generation as generation_service class _Fake: subagents = 0 user = _user(db) resolved = tools_service.resolve_tools(db, chat, user) context = tools_service.context_for(db, user, chat, tools=resolved) original = generation_service.running_for generation_service.running_for = ( lambda chat_id: _Fake() if chat_id == chat.id else original(chat_id) ) try: return await tools_service.run_tool(context, "subagent_run", json.dumps(args)) finally: generation_service.running_for = original def _schema(db, chat): resolved = tools_service.resolve_tools(db, chat, _user(db)) tool = resolved.by_name.get("subagent_run") return tool.parameters["properties"] if tool else None # --- Capacity --------------------------------------------------------------------------- def test_by_default_a_model_is_its_own_helper_and_the_tool_is_unchanged(db): """Both switches off, nothing designated: exactly the 1.11 behaviour.""" chat = _chat(db) assert _ids(helpers.candidates(db, chat, _user(db))) == ["gpt-oss"] assert "model" not in _schema(db, chat) def test_a_single_session_model_cannot_help_itself(db): _row(db, "gpt-oss").single_session = True db.commit() chat = _chat(db) assert helpers.candidates(db, chat, _user(db)) == [] # Nothing to send a helper to, so the tool is not offered at all. assert _schema(db, chat) is None def test_one_model_at_a_time_forbids_another_model_on_that_connection(db, setup): main, other = _row(db, "gpt-oss"), _row(db, "qwen35") assert helpers.capacity_refusal(main, other) == "" setup["local"].one_model_at_a_time = True db.commit() assert "holds one model at a time" in helpers.capacity_refusal(main, other) # ...but itself, and a model on another connection, are fine. assert helpers.capacity_refusal(main, main) == "" assert helpers.capacity_refusal(main, _row(db, "deepseek")) == "" # --- Designations ------------------------------------------------------------------------ def test_an_offered_designation_joins_the_candidates_and_the_enum(db): helpers.set_designation(db, None, "gpt-oss", "qwen35", offer=True) chat = _chat(db) assert _ids(helpers.candidates(db, chat, _user(db))) == ["gpt-oss", "qwen35"] assert _schema(db, chat)["model"]["enum"] == ["gpt-oss", "qwen35"] def test_a_by_hand_designation_is_used_only_once_added_to_the_chat(db): helpers.set_designation(db, None, "gpt-oss", "qwen35", offer=False) chat = _chat(db) assert _ids(helpers.candidates(db, chat, _user(db))) == ["gpt-oss"] helpers.apply_chat_helpers(db, chat, _user(db), ["qwen35"]) db.commit() assert _ids(helpers.candidates(db, chat, _user(db))) == ["gpt-oss", "qwen35"] def test_designations_are_per_main_model(db): helpers.set_designation(db, None, "bonsai", "qwen35", offer=True) assert _ids(helpers.candidates(db, _chat(db), _user(db))) == ["gpt-oss"] def test_a_persons_designations_need_the_permission_and_replace_the_instances(db): user = _user(db) user.role = "user" db.commit() helpers.set_designation(db, None, "gpt-oss", "qwen35", offer=True) helpers.set_designation(db, user, "gpt-oss", "qwen35", offer=False) assert [d.owner_id for d in helpers.designations(db, user, "gpt-oss")] == [None] settings_store.update( db, {"default_permissions": {"tools.subagent": True, helpers.PERMISSION: True}} ) rows = helpers.designations(db, user, "gpt-oss") assert [(d.owner_id, d.offer) for d in rows] == [(user.id, False)] def test_the_talk_rules_filter_designations(db): user = _user(db) user.role = "user" db.commit() helpers.set_designation(db, None, "gpt-oss", "qwen35", offer=True) talk.set_rule(db, None, "gpt-oss", "qwen35", "deny") assert _ids(helpers.candidates(db, _chat(db), user)) == ["gpt-oss"] def test_another_data_group_needs_a_rule(db, setup): user = _user(db) user.role = "user" db.add(DataGroup(id="hosted", name="Hosted")) setup["cloud"].data_group_id = "hosted" db.commit() helpers.set_designation(db, None, "gpt-oss", "deepseek", offer=True) assert "deepseek" not in _ids(helpers.candidates(db, _chat(db), user)) talk.set_rule(db, None, "gpt-oss", "deepseek", "allow") assert "deepseek" in _ids(helpers.candidates(db, _chat(db), user)) # --- The choice, at call time ------------------------------------------------------------- async def test_a_named_helper_runs_on_that_model_with_its_own_effort(db, monkeypatch): helpers.set_designation(db, None, "gpt-oss", "bonsai", offer=True) chat = _chat(db) chat.params_json = {"reasoning_effort": "high"} db.commit() seen = _spawn(monkeypatch) settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS) outcome = await _run(db, chat, {"task": "Check it.", "model": "bonsai"}) assert outcome.event["status"] == "ok" child = db.get(Chat, seen["chat_id"]) assert child.model_id == "bonsai" # Bonsai's own default, never the parent's `high`, which it would refuse. assert child.params_json["reasoning_effort"] == "xhigh" assert child.parent_chat_id == chat.id async def test_no_model_named_is_the_main_model_itself(db, monkeypatch): helpers.set_designation(db, None, "gpt-oss", "bonsai", offer=True) seen = _spawn(monkeypatch) settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS) await _run(db, _chat(db), {"task": "Check it."}) assert db.get(Chat, seen["chat_id"]).model_id == "gpt-oss" async def test_a_model_that_is_not_a_candidate_is_refused_with_the_list(db, monkeypatch): helpers.set_designation(db, None, "gpt-oss", "bonsai", offer=True) _spawn(monkeypatch) outcome = await _run(db, _chat(db), {"task": "Check it.", "model": "deepseek"}) assert outcome.event["status"] == "error" assert "bonsai" in outcome.content def test_without_itself_the_default_is_the_first_hand_added_helper(db): _row(db, "gpt-oss").single_session = True helpers.set_designation(db, None, "gpt-oss", "qwen35", offer=False) db.commit() chat = _chat(db) helpers.apply_chat_helpers(db, chat, _user(db), ["qwen35"]) db.commit() model, refusal = helpers.choose(db, chat, _user(db), "") assert model.model_id == "qwen35" and refusal == "" def test_the_harness_lists_where_a_helper_can_run(db): from lembas.services import harness helpers.set_designation(db, None, "gpt-oss", "qwen35", offer=True) chat = _chat(db) resolved = tools_service.resolve_tools(db, chat, _user(db)) text = harness.compose(db, _user(db), resolved.schemas, chat=chat) assert "Where a helper can run" in text assert "qwen35 (qwen35)" in text # --- The screens ----------------------------------------------------------------------------- def test_the_model_page_designates_a_helper(client, db): model = _row(db, "gpt-oss") client.post(f"/admin/models/{model.id}/helpers", data={"helper_model": "qwen35", "offer": "true"}) row = db.scalar(select(HelperDesignation)) assert (row.main_model, row.helper_model, row.offer) == ("gpt-oss", "qwen35", True) assert "offered to the model" in client.get(f"/admin/models/{model.id}/edit").text def test_the_model_page_saves_single_session(client, db): model = _row(db, "gpt-oss") page = client.get(f"/admin/models/{model.id}/edit").text assert 'name="single_session"' in page client.post(f"/admin/models/{model.id}", data={"display_name": "", "enabled": "true", "public": "true", "single_session": "true"}) db.expire_all() assert _row(db, "gpt-oss").single_session is True def test_the_connection_form_saves_one_model_at_a_time(client, db, setup): local = setup["local"] client.post(f"/admin/connections/{local.id}", data={"name": "Local", "base_url": local.base_url, "one_model_at_a_time": "true"}) db.expire_all() assert db.get(Connection, local.id).one_model_at_a_time is True def test_the_composer_picker_adds_a_helper_by_hand(client, db): helpers.set_designation(db, None, "gpt-oss", "qwen35", offer=False) chat = _chat(db) page = client.get(f"/chat/{chat.id}").text assert 'name="helper_model_ids"' in page and 'id="helpers-form"' in page client.patch(f"/api/chats/{chat.id}", data={"helper_model_ids": ["qwen35"]}) assert [h.model_id for h in db.scalars(select(ChatHelper))] == ["qwen35"] client.patch(f"/api/chats/{chat.id}", data={"helper_model_ids": [""]}) db.expire_all() assert list(db.scalars(select(ChatHelper))) == [] def test_a_person_adds_their_own_designation_with_the_permission(client, db): user = _user(db) user.role = "user" db.commit() data = {"main_model": "gpt-oss", "helper_model": "qwen35", "offer": "true"} client.post("/api/preferences/helpers", data=data) assert list(db.scalars(select(HelperDesignation))) == [] settings_store.update( db, {"default_permissions": {"tools.subagent": True, helpers.PERMISSION: True}} ) client.post("/api/preferences/helpers", data=data) assert db.scalar(select(HelperDesignation)).owner_id == user.id # --- Upgrading --------------------------------------------------------------------------------- def test_an_upgraded_database_reads_both_switches_as_off(db, setup): """Both are NOT NULL booleans, which `sync_schema` backfills with False -- which is why they name the restrictive state: False has to mean "as before".""" from sqlalchemy import text from lembas.db.migrations import sync_schema from lembas.db.session import get_engine engine = get_engine() db.close() with engine.begin() as connection: connection.execute(text("DROP TABLE IF EXISTS chat_helpers")) connection.execute(text("DROP TABLE IF EXISTS helper_designations")) connection.execute(text("ALTER TABLE models DROP COLUMN single_session")) connection.execute(text("ALTER TABLE connections DROP COLUMN one_model_at_a_time")) sync_schema(engine) assert _row(db, "gpt-oss").single_session is False assert db.get(Connection, setup["local"].id).one_model_at_a_time is False assert _ids(helpers.candidates(db, _chat(db), _user(db))) == ["gpt-oss"]