devplacepy/tests/unit/services/devii/agents/controller.py

# retoor <retoor@molodetz.nl>

import json

import pytest

from devplacepy.database import get_setting, set_setting
from devplacepy.services.devii.agents.controller import AgentController
from devplacepy.services.devii.agents.policy import CATEGORIES
from devplacepy.services.devii.agents.store import AgentStore
from devplacepy.services.devii.config import FIELD_AGENTS_ENABLED
from devplacepy.services.devii.errors import ToolInputError
from devplacepy.utils import generate_uid
from tests.conftest import run_async


class _Quota:
    def __init__(self, reached=False):
        self.reached = reached


class _StubSession:
    def __init__(self, db, owner_id):
        self.store = AgentStore(db, "user", owner_id)
        self.owner_kind = "user"
        self.owner_id = owner_id
        self.channel = "main"
        self.emitted = []
        self.quota = _Quota()

    def _quota_snapshot(self):
        return {"limit_reached": self.quota.reached}

    async def _emit(self, payload, buffer=True):
        self.emitted.append(payload)
        return True

    async def emit_agent(self, payload):
        self.emitted.append(payload)
        return True


class _StubSupervisor:
    def __init__(self, store=None):
        self._store = store
        self.resumed = []
        self.cancelled = []

    async def resume(self, uid, answer):
        self.resumed.append((uid, answer))
        if self._store is not None:
            self._store.update(uid, {"status": "running"})
        return True

    async def cancel(self, uid, reason):
        self.cancelled.append((uid, reason))
        if self._store is not None:
            self._store.update(uid, {"status": "cancelled"})
        return True


@pytest.fixture
def flag_on(local_db):
    previous = get_setting(FIELD_AGENTS_ENABLED, "")
    set_setting(FIELD_AGENTS_ENABLED, "1")
    try:
        yield
    finally:
        set_setting(FIELD_AGENTS_ENABLED, previous or "")


@pytest.fixture
def session(local_db):
    return _StubSession(local_db, f"ctl-{generate_uid()}")


def _start_args(**overrides):
    args = {
        "category": "research",
        "task": "Compare static site generators.",
        "output_format": "A markdown table.",
    }
    args.update(overrides)
    return args


def test_start_validates_contract(local_db, session, flag_on):
    controller = AgentController(session, supervisor=_StubSupervisor(session.store))
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_start", _start_args(category="nope")))
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_start", _start_args(task="  ")))
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_start", _start_args(output_format="")))
    with pytest.raises(ToolInputError):
        run_async(
            controller.dispatch(
                "agent_start", _start_args(tools_extra=["missing_tool"])
            )
        )
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_start", _start_args(tools_extra=["run_js"])))
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_start", _start_args(budget_hint="huge")))


def test_start_queues_and_emits(local_db, session, flag_on):
    controller = AgentController(session, supervisor=_StubSupervisor(session.store))
    raw = run_async(controller.dispatch(
        "agent_start", _start_args(label="sg", budget_hint="quick"))
    )
    payload = json.loads(raw)
    assert payload["status"] == "queued"
    assert payload["run"]["category"] == "research"
    assert session.emitted[0]["type"] == "agent"
    assert session.emitted[0]["kind"] == "started"


def test_start_refuses_when_flag_off_or_quota_spent(local_db, session):
    set_setting(FIELD_AGENTS_ENABLED, "")
    controller = AgentController(session, supervisor=_StubSupervisor(session.store))
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_start", _start_args()))
    set_setting(FIELD_AGENTS_ENABLED, "1")
    try:
        session.quota.reached = True
        with pytest.raises(ToolInputError):
            run_async(controller.dispatch("agent_start", _start_args()))
    finally:
        set_setting(FIELD_AGENTS_ENABLED, "")


def test_start_enforces_active_cap(local_db, session, flag_on):
    controller = AgentController(session, supervisor=_StubSupervisor(session.store))
    run_async(controller.dispatch("agent_start", _start_args()))
    run_async(controller.dispatch("agent_start", _start_args()))
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_start", _start_args()))


def test_status_result_cancel_list_round_trip(local_db, session, flag_on):
    supervisor = _StubSupervisor(session.store)
    controller = AgentController(session, supervisor=supervisor)
    uid = json.loads(run_async(controller.dispatch("agent_start", _start_args())))["run"][
        "uid"
    ]
    status = json.loads(run_async(controller.dispatch("agent_status", {"uid": uid})))
    assert status["run"]["status"] == "queued"
    listed = json.loads(run_async(controller.dispatch("agent_list", {})))
    assert listed["count"] == 1
    cancelled = json.loads(run_async(controller.dispatch("agent_cancel", {"uid": uid})))
    assert cancelled["status"] == "cancelled"
    assert supervisor.cancelled == [(uid, "cancelled by orchestrator")]
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("agent_status", {"uid": "missing"}))


def test_result_resumes_needs_user(local_db, session, flag_on):
    supervisor = _StubSupervisor(session.store)
    controller = AgentController(session, supervisor=supervisor)
    uid = json.loads(run_async(controller.dispatch("agent_start", _start_args())))["run"][
        "uid"
    ]
    controller._store.update(uid, {"status": "needs_user", "needs_question": "Budget?"})
    payload = json.loads(
        run_async(controller.dispatch("agent_result", {"uid": uid, "answer": "zero"}))
    )
    assert supervisor.resumed == [(uid, "zero")]
    assert payload["run"]["uid"] == uid


def test_ask_user_returns_pause_marker(session):
    controller = AgentController(session, supervisor=_StubSupervisor(session.store))
    payload = json.loads(
        run_async(controller.dispatch("ask_user", {"question": "Which budget?"}))
    )
    assert payload["status"] == "paused_for_user"
    with pytest.raises(ToolInputError):
        run_async(controller.dispatch("ask_user", {"question": "  "}))


def test_categories_match_policy():
    assert set(CATEGORIES) == {
        "research",
        "files",
        "planning-report",
        "email-triage",
        "containers",
        "long-mutation",
    }