LLeMbas/tests/test_helpers.py

"""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"]