LLeMbas/tests/test_friend.py

"""Putting a question to one of the other models, and getting its answer back.

Shares its machinery with `subagent_run` on purpose, so most of what is asserted
here is the *differences* — which model answers, with whose reasoning effort, in
what kind of chat, and what it may not do in turn. The generation loop is stubbed
exactly as `test_subagent.py` stubs it; what matters is the chat the friend is
given.
"""

from __future__ import annotations

import json

import pytest
from sqlalchemy import select

from lembas.db.models import (
    KIND_AGENT,
    KIND_CHAT,
    ROLE_USER,
    Chat,
    Connection,
    Group,
    Model,
    User,
)
from lembas.services import chat as chat_service
from lembas.services import settings_store
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 asking_allowed(db, registered):
    """The instance switch on, the permission granted, and three models to ask.

    The gates get their own test below, which asserts both directions.
    """
    settings_store.update(db, {"enabled": True}, key=settings_store.SUBAGENTS)
    settings_store.update(db, {"default_permissions": {"tools.friend": True}})
    connection = Connection(
        name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")
    )
    db.add(connection)
    db.commit()
    for index, (name, label, note) in enumerate(
        [
            ("test-model", "The asker", ""),
            ("big-model", "Big", "70B, good at maths"),
            ("small-model", "Small", ""),
        ]
    ):
        db.add(
            Model(
                connection_id=connection.id,
                model_id=name,
                display_name=label,
                notes=note,
                position=index,
                capabilities_json={"tools": True},
            )
        )
    db.commit()
    subagent_service.clear()
    yield
    subagent_service.clear()


def _user(db) -> User:
    return db.scalars(select(User).order_by(User.created_at)).first()


def _chat(db, **kwargs) -> Chat:
    chat = Chat(user_id=_user(db).id, title="t", model_id="test-model", **kwargs)
    db.add(chat)
    db.commit()
    return chat


class _Fake:
    def __init__(self, spawned: int = 0):
        self.subagents = spawned


def _spawn(monkeypatch, *, answer: str = "I disagree, and here is why.", finish: bool = True):
    from lembas.db.models import ROLE_ASSISTANT, ROLE_USER
    from lembas.db.session import session_scope

    seen: dict[str, str] = {}

    async def fake_wake(chat_id: str, content: str, *, model_id: str = "") -> str:
        seen["chat_id"] = chat_id
        seen["turn"] = content
        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, answer)
            seen["message_id"] = reply.id
        return seen["message_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 _ask(db, chat: Chat, args: dict, *, generation=None):
    """Through `resolve_tools`, never by hand — what may be run is what was
    offered, and a hand-built context falls back to the import-time registry,
    which has never held this tool."""
    from lembas.services import generation as generation_service

    user = _user(db)
    resolved = tools_service.resolve_tools(db, chat, user)
    context = tools_service.context_for(db, user, chat, tools=resolved)
    fake = generation if generation is not None else _Fake()
    original = generation_service.running_for

    def running_for(chat_id):
        return fake if chat_id == chat.id else original(chat_id)

    generation_service.running_for = running_for
    try:
        return await tools_service.run_tool(context, "ask_friend", json.dumps(args))
    finally:
        generation_service.running_for = original


# --- Whose chat it is ---------------------------------------------------------
async def test_the_friend_answers_as_itself_not_as_the_asking_model(db, monkeypatch):
    """The whole feature. `generation` resolves the endpoint from the child chat
    row, so the model on that row is the one that answers."""
    parent = _chat(db)
    seen = _spawn(monkeypatch)
    settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS)

    await _ask(db, parent, {"model": "big-model", "question": "Is this right?"})

    child = db.get(Chat, seen["chat_id"])
    assert child.model_id == "big-model"
    assert child.parent_chat_id == parent.id
    assert child.unattended is True
    assert child.temporary is True


async def test_the_friend_can_be_named_by_its_label_as_well_as_its_id(db, monkeypatch):
    """The roster prints both, so a model will sometimes type back the pretty
    one. Refusing that is a round spent on a spelling."""
    parent = _chat(db)
    seen = _spawn(monkeypatch)
    settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS)

    await _ask(db, parent, {"model": "Big", "question": "Is this right?"})

    assert db.get(Chat, seen["chat_id"]).model_id == "big-model"


async def test_the_friend_does_not_inherit_the_askers_reasoning_effort(db, monkeypatch):
    """The 1.3.0 bug with a new door: the vocabularies differ per model, and an
    effort a model does not take is rendered into its chat template and raises
    there. `high` from the asker must not follow the question to a model whose
    list says low/medium/xhigh."""
    parent = _chat(db)
    parent.params_json = {"reasoning_effort": "high"}
    friend = db.scalar(select(Model).where(Model.model_id == "big-model"))
    friend.reasoning_efforts = ["low", "medium", "xhigh"]
    friend.params_json = {"reasoning_effort": "xhigh"}
    db.commit()
    seen = _spawn(monkeypatch)
    settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS)

    await _ask(db, parent, {"model": "big-model", "question": "Is this right?"})

    child = db.get(Chat, seen["chat_id"])
    assert chat_service.resolved_effort(child) == "xhigh"


async def test_an_effort_the_friend_does_not_take_is_not_sent_at_all(db, monkeypatch):
    parent = _chat(db)
    friend = db.scalar(select(Model).where(Model.model_id == "big-model"))
    friend.reasoning_efforts = ["low", "medium"]
    friend.params_json = {"reasoning_effort": "high"}
    db.commit()
    seen = _spawn(monkeypatch)
    settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS)

    await _ask(db, parent, {"model": "big-model", "question": "?"})

    assert chat_service.resolved_effort(db.get(Chat, seen["chat_id"])) == ""


async def test_a_friend_of_an_agent_chat_is_not_given_the_machine(db, monkeypatch):
    """A peer is asked what it thinks, not put to work. An agent chat's harness
    is about the box it is working on, and handing that to somebody asked a
    question invites it to plan around a shell it has not got."""
    parent = _chat(db, kind=KIND_AGENT, project_dir="/srv/app", ssh_profile_id="nope")
    seen = _spawn(monkeypatch)
    settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS)

    await _ask(db, parent, {"model": "big-model", "question": "?"})

    child = db.get(Chat, seen["chat_id"])
    assert child.kind == KIND_CHAT
    assert not child.ssh_profile_id
    assert not child.project_dir
    # And the consequence, which is the thing that actually matters: an
    # ordinary chat resolves no agent tools, whatever the mode column says.
    offered = tools_service.resolve_tools(db, child, _user(db))
    assert not [name for name in offered.by_name if name.startswith(("shell_", "file_"))]


# --- What it may not do -------------------------------------------------------
async def test_a_friend_cannot_ask_a_friend(db, monkeypatch):
    """Otherwise one question is a fan-out with no bound anybody set. Both halves:
    the family is withdrawn from the offered set, and the runner refuses a call
    that arrived by any other route."""
    parent = _chat(db)
    seen = _spawn(monkeypatch)
    settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS)
    await _ask(db, parent, {"model": "big-model", "question": "?"})
    child = db.get(Chat, seen["chat_id"])

    offered = tools_service.resolve_tools(db, child, _user(db))
    assert "ask_friend" not in offered.by_name
    assert "subagent_run" not in offered.by_name
    assert "ask_user" not in offered.by_name

    # And the runner's own guard, reached by offering it the tool anyway --
    # which is what "a call that arrived by some other route" means. Two halves,
    # because the withdrawal is the one a prompt cannot argue with and this is
    # the one that holds if the withdrawal is ever got round.
    forced = tools_service.ToolSet(tuple(subagent_service.friend_tool_defs()))
    context = tools_service.context_for(db, _user(db), child, tools=forced)
    outcome = await tools_service.run_tool(
        context, "ask_friend", json.dumps({"model": "small-model", "question": "?"})
    )
    assert "may not pass it on" in outcome.content


async def test_a_friend_cannot_rewrite_its_own_personality(db, monkeypatch):
    """A question is written by a model that may have been reading a page, and
    the persona is carried into every conversation it will ever have."""
    settings_store.update(db, {"default_permissions": {"tools.persona": True}})
    parent = _chat(db)
    seen = _spawn(monkeypatch)
    settings_store.update(db, {"keep_transcript": True}, key=settings_store.SUBAGENTS)
    await _ask(db, parent, {"model": "big-model", "question": "?"})

    offered = tools_service.resolve_tools(db, db.get(Chat, seen["chat_id"]), _user(db))
    assert "persona_write" not in offered.by_name
    assert "impression_write" not in offered.by_name
    # And it is genuinely on for the chat somebody is present in.
    assert "persona_write" in tools_service.resolve_tools(db, parent, _user(db)).by_name


# --- Refusals that name what could have been asked ----------------------------
async def test_an_unknown_model_is_refused_with_the_list_of_real_ones(db, monkeypatch):
    """The name arrives in a tool call, so it is model-written input. A refusal
    that does not say what the valid answers are costs another round."""
    parent = _chat(db)
    _spawn(monkeypatch)

    outcome = await _ask(db, parent, {"model": "gpt-9", "question": "?"})

    assert outcome.event["status"] == "error"
    assert "big-model" in outcome.content
    assert "small-model" in outcome.content


async def test_asking_itself_is_refused_in_those_words(db, monkeypatch):
    parent = _chat(db)
    _spawn(monkeypatch)
    outcome = await _ask(db, parent, {"model": "test-model", "question": "?"})
    assert "That is you" in outcome.content


async def test_a_model_the_reader_cannot_use_is_neither_listed_nor_reachable(db, monkeypatch):
    """A roster is filtered through what this account can see, so naming a
    restricted model must fail for the same reason it is absent — not by a
    second, looser check."""
    group = Group(name="Wheel")
    db.add(group)
    restricted = db.scalar(select(Model).where(Model.model_id == "big-model"))
    restricted.public = False
    restricted.groups = [group]
    db.commit()
    user = _user(db)
    user.role = ROLE_USER
    db.commit()
    parent = _chat(db)
    _spawn(monkeypatch)

    assert "big-model" not in chat_service.roster_block(db, user, exclude="test-model")
    outcome = await _ask(db, parent, {"model": "big-model", "question": "?"})
    assert outcome.event["status"] == "error"
    assert "no model called" in outcome.content


async def test_an_empty_question_is_refused_before_anything_is_created(db, monkeypatch):
    parent = _chat(db)
    seen = _spawn(monkeypatch)
    outcome = await _ask(db, parent, {"model": "big-model", "question": "   "})
    assert outcome.event["status"] == "error"
    assert "chat_id" not in seen, "a chat was created for a call that could not work"


# --- The budget ---------------------------------------------------------------
async def test_questions_and_helpers_share_one_allowance(db, monkeypatch):
    """Two counters would let one reply spend both. `Generation.subagents` is the
    only object that knows what "this reply" means."""
    parent = _chat(db)
    _spawn(monkeypatch)
    settings_store.update(db, {"max_per_reply": 1}, key=settings_store.SUBAGENTS)

    generation = _Fake(spawned=1)
    outcome = await _ask(
        db, parent, {"model": "big-model", "question": "?"}, generation=generation
    )

    assert outcome.event["status"] == "error"
    assert "already used its 1 helpers" in outcome.content


async def test_a_successful_question_spends_one_of_the_allowance(db, monkeypatch):
    parent = _chat(db)
    _spawn(monkeypatch)
    generation = _Fake()
    await _ask(db, parent, {"model": "big-model", "question": "?"}, generation=generation)
    assert generation.subagents == 1


# --- The gates ----------------------------------------------------------------
def test_the_tool_needs_the_permission_and_the_instance_switch(db):
    parent = _chat(db)
    user = _user(db)
    assert "ask_friend" in tools_service.resolve_tools(db, parent, user).by_name

    settings_store.update(db, {"enabled": False}, key=settings_store.SUBAGENTS)
    assert "ask_friend" not in tools_service.resolve_tools(db, parent, user).by_name

    settings_store.update(db, {"enabled": True}, key=settings_store.SUBAGENTS)
    # An administrator bypasses every permission, so the permission half can
    # only be asserted on somebody who is not one.
    user.role = ROLE_USER
    settings_store.update(db, {"default_permissions": {"tools.friend": False}})
    db.commit()
    assert "ask_friend" not in tools_service.resolve_tools(db, parent, user).by_name


def test_the_model_switch_turns_it_off_for_that_model_alone(db):
    parent = _chat(db)
    asker = db.scalar(select(Model).where(Model.model_id == "test-model"))
    asker.capabilities_json = {"tools": True, "tool_friend": False}
    db.commit()
    assert "ask_friend" not in tools_service.resolve_tools(db, parent, _user(db)).by_name


# --- The answer ---------------------------------------------------------------
async def test_the_answer_comes_back_named_and_marked_as_an_opinion(db, monkeypatch):
    """A model handing on another's answer as its own is the failure worth
    wording against, so the tool result says whose it is."""
    parent = _chat(db)
    _spawn(monkeypatch, answer="No. The second premise is wrong.")

    outcome = await _ask(db, parent, {"model": "big-model", "question": "Is this right?"})

    assert outcome.event["status"] == "ok"
    assert "Big answered" in outcome.content
    assert "The second premise is wrong." in outcome.content
    assert "opinion" in outcome.content
    assert outcome.event["why"] == "Big"


async def test_the_question_says_who_is_asking_and_that_nobody_is_reading(db, monkeypatch):
    parent = _chat(db)
    seen = _spawn(monkeypatch)
    await _ask(db, parent, {"model": "big-model", "question": "Is this right?", "context": "ctx"})

    assert "test-model" in seen["turn"]
    assert "Nobody is reading" in seen["turn"]
    assert "Is this right?" in seen["turn"]
    assert "ctx" in seen["turn"]