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

# retoor <retoor@molodetz.nl>

from devplacepy.services.devii.agents.store import AgentStore
from devplacepy.services.devii.agents.supervisor import AgentSupervisor
from devplacepy.services.devii.tasks.schedule import now_utc, to_iso
from devplacepy.utils import generate_uid
from tests.conftest import run_async


class _StubSession:
    def __init__(self):
        self.emitted = []
        self._agents = {}
        self.limit_reached = False

    def _quota_snapshot(self):
        return {"limit_reached": self.limit_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

    def record_agent_spend(self, uid, usage, cost):
        self.spent = (uid, usage, cost)


def _record(owner, **overrides):
    now = to_iso(now_utc())
    record = {
        "uid": generate_uid(),
        "owner_kind": "user",
        "owner_id": owner,
        "channel": "main",
        "category": "research",
        "label": "run",
        "task": "work",
        "output_format": "text",
        "tools_extra": "[]",
        "effort": "standard",
        "task_scope": 0,
        "status": "queued",
        "progress": "queued",
        "result": None,
        "error": None,
        "checks": 0,
        "nudges_used": 0,
        "retries_used": 0,
        "error_streak": 0,
        "iterations": 0,
        "tokens": 0,
        "cost_usd": 0.0,
        "needs_question": None,
        "pending_ask_id": None,
        "check_at": now,
        "last_progress_at": now,
        "created_at": now,
        "updated_at": now,
    }
    record.update(overrides)
    return record


def _setup(local_db, resolve):
    owner = f"sup-{generate_uid()}"
    store = AgentStore(local_db, "user", owner)
    supervisor = AgentSupervisor(local_db, resolve)
    return owner, store, supervisor


def test_reap_stuck_on_start(local_db):
    owner, store, supervisor = _setup(local_db, lambda row: None)
    store.create(_record(owner, status="running"))
    store.create(_record(owner, status="checking"))
    store.create(_record(owner, status="needs_user"))
    supervisor._reap_stuck()
    statuses = sorted(row["status"] for row in store.list())
    assert statuses == ["interrupted", "interrupted", "needs_user"]


def test_cancel_terminal_and_missing(local_db):
    session = _StubSession()
    owner, store, supervisor = _setup(local_db, lambda row: session)
    record = _record(owner)
    store.create(record)
    assert run_async(supervisor.cancel(record["uid"], "stop")) is True
    assert store.get(record["uid"])["status"] == "cancelled"
    assert session.emitted[0]["kind"] == "cancelled"
    assert run_async(supervisor.cancel(record["uid"], "stop")) is False
    assert run_async(supervisor.cancel("missing", "stop")) is False


def test_check_interrupts_ownerless_rows(local_db):
    owner, store, supervisor = _setup(local_db, lambda row: None)
    record = _record(owner)
    store.create(record)
    run_async(supervisor._tick())
    assert store.get(record["uid"])["status"] == "interrupted"


def test_queued_with_spent_quota_cancels(local_db):
    session = _StubSession()
    session.limit_reached = True
    owner, store, supervisor = _setup(local_db, lambda row: session)
    record = _record(owner)
    store.create(record)
    run_async(supervisor._tick())
    assert store.get(record["uid"])["status"] == "cancelled"


def test_needs_user_waits_then_times_out(local_db):
    session = _StubSession()
    owner, store, supervisor = _setup(local_db, lambda row: session)
    record = _record(owner, status="needs_user", needs_question="Budget?")
    store.create(record)
    run_async(supervisor._tick())
    row = store.get(record["uid"])
    assert row["status"] == "needs_user"
    assert row["checks"] == 1
    store.update(
        record["uid"], {"created_at": "2020-01-01T00:00:00", "check_at": to_iso(now_utc())}
    )
    run_async(supervisor._tick())
    assert store.get(record["uid"])["status"] == "cancelled"


def test_cancel_ledgers_accrued_spend(local_db):
    from types import SimpleNamespace

    from devplacepy.services.devii.agents.runner import AgentRunner
    from devplacepy.services.devii.cost.tracker import CostTracker, Pricing

    session = _StubSession()
    session.cost = CostTracker(pricing=Pricing())
    owner, store, supervisor = _setup(local_db, lambda row: session)
    record = _record(owner)
    store.create(record)
    runner = AgentRunner(
        session, record["uid"], "run", "research", "standard", "work", "text", [], False
    )
    runner.cost.prompt_tokens = 40
    supervisor._runners[record["uid"]] = runner
    supervisor._running[record["uid"]] = {
        "task": SimpleNamespace(cancel=lambda: None),
        "owner_id": owner,
        "ledgered": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "cache_hit_tokens": 0,
            "cache_miss_tokens": 0,
            "cost": 0.0,
        },
    }
    assert run_async(supervisor.cancel(record["uid"], "stop")) is True
    assert session.spent[0] == record["uid"]
    assert session.spent[1]["prompt_tokens"] == 40


def test_counts_reports_statuses(local_db):
    owner, store, supervisor = _setup(local_db, lambda row: None)
    store.create(_record(owner, status="running"))
    store.create(_record(owner, status="done"))
    counts = supervisor.counts()
    assert counts.get("running", 0) >= 1
    assert counts.get("done", 0) >= 1