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

# retoor <retoor@molodetz.nl>

import json

from devplacepy.database import get_table
from devplacepy.services.devii.agents.runner import AgentRunner
from devplacepy.services.devii.agents.store import AgentStore
from devplacepy.services.devii.agents.supervisor import AgentSupervisor
from devplacepy.services.devii.cost.tracker import CostTracker, Pricing
from devplacepy.services.devii.store import UsageLedger
from devplacepy.services.devii.tasks.schedule import now_utc, to_iso
from tests.conftest import run_async
from devplacepy.utils import generate_uid


class _Settings:
    daily_limit_usd = 1.0
    plan_required = True
    verify_required = True
    ai_model = "test-model"


class _StubSession:
    def __init__(self, owner_id):
        self.cost = CostTracker(pricing=Pricing())
        self.settings = _Settings()
        self._ledger = UsageLedger()
        self._llm = None
        self.dispatcher = None
        self.owner_kind = "user"
        self.owner_id = owner_id
        self.channel = "main"
        self._agents = {}
        self.emitted = []

    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._ledger.record(
            self.owner_kind, self.owner_id, uid, usage, cost, self.settings.ai_model
        )


def _record(owner, **overrides):
    now = to_iso(now_utc())
    record = {
        "uid": generate_uid(),
        "channel": "main",
        "category": "research",
        "label": "sg",
        "task": "work",
        "output_format": "text",
        "tools_extra": "[]",
        "effort": "standard",
        "task_scope": 0,
        "status": "running",
        "progress": "started",
        "result": None,
        "error": None,
        "checks": 1,
        "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):
    owner = f"drv-{generate_uid()}"
    store = AgentStore(local_db, "user", owner)
    session = _StubSession(owner)
    supervisor = AgentSupervisor(local_db, lambda row: session)
    return owner, store, session, supervisor


def _runner(session, uid):
    return AgentRunner(
        session, uid, "sg", "research", "standard", "work", "text", [], False
    )


def test_drive_done_finishes_and_ledgers(local_db, monkeypatch):
    async def fake_loop(**kwargs):
        kwargs["state"].iteration = 2
        return "final report"

    monkeypatch.setattr(
        "devplacepy.services.devii.agents.runner.react_loop", fake_loop
    )
    owner, store, session, supervisor = _setup(local_db)
    record = _record(owner)
    store.create(record)
    runner = _runner(session, record["uid"])
    runner.cost.prompt_tokens = 50
    runner.cost.completion_tokens = 10
    supervisor._runners[record["uid"]] = runner
    supervisor._running[record["uid"]] = {
        "task": None,
        "owner_id": owner,
        "ledgered": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "cache_hit_tokens": 0,
            "cache_miss_tokens": 0,
            "cost": 0.0,
        },
    }
    run_async(supervisor._drive(record["uid"], session))
    row = store.get(record["uid"])
    assert row["status"] == "done"
    assert row["result"] == "final report"
    kinds = [frame["kind"] for frame in session.emitted]
    assert "done" in kinds
    ledgered = get_table("devii_usage_ledger").find_one(turn_id=record["uid"])
    assert ledgered is not None
    assert int(ledgered["prompt_tokens"]) == 50
    assert int(ledgered["completion_tokens"]) == 10


def test_drive_needs_user_parks_with_question(local_db, monkeypatch):
    async def fake_loop(**kwargs):
        messages = kwargs["messages"]
        messages.append(
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [
                    {
                        "id": "c9",
                        "function": {
                            "name": "ask_user",
                            "arguments": json.dumps({"question": "Budget?"}),
                        },
                    }
                ],
            }
        )
        messages.append(
            {
                "role": "tool",
                "tool_call_id": "c9",
                "name": "ask_user",
                "content": json.dumps({"status": "paused_for_user"}),
            }
        )
        return "awaiting user"

    monkeypatch.setattr(
        "devplacepy.services.devii.agents.runner.react_loop", fake_loop
    )
    owner, store, session, supervisor = _setup(local_db)
    record = _record(owner)
    store.create(record)
    supervisor._runners[record["uid"]] = _runner(session, record["uid"])
    run_async(supervisor._drive(record["uid"], session))
    row = store.get(record["uid"])
    assert row["status"] == "needs_user"
    assert row["needs_question"] == "Budget?"
    assert row["pending_ask_id"] == "c9"
    assert session.emitted[-1]["kind"] == "needs_user"


def test_drive_error_records_and_stays_runnable(local_db, monkeypatch):
    async def fake_loop(**kwargs):
        raise RuntimeError("boom")

    monkeypatch.setattr(
        "devplacepy.services.devii.agents.runner.react_loop", fake_loop
    )
    owner, store, session, supervisor = _setup(local_db)
    record = _record(owner)
    store.create(record)
    supervisor._runners[record["uid"]] = _runner(session, record["uid"])
    run_async(supervisor._drive(record["uid"], session))
    row = store.get(record["uid"])
    assert row["status"] == "running"
    assert row["error"] == "boom"