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

# retoor <retoor@molodetz.nl>

import multiprocessing
import tempfile

import dataset

from devplacepy.services.devii.agents.store import (
    AgentStore,
    claim,
    due_rows,
    prune_terminal,
    snapshot_for,
)
from devplacepy.services.devii.tasks.store import memory_db
from devplacepy.services.devii.tasks.schedule import now_utc, to_iso
from devplacepy.utils import generate_uid


def _record(**overrides):
    now = to_iso(now_utc())
    record = {
        "uid": generate_uid(),
        "channel": "main",
        "category": "research",
        "label": None,
        "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 _store(db=None):
    return AgentStore(memory_db() if db is None else db, "user", "owner-1")


def test_create_get_list_update_scope():
    store = _store()
    first = _record()
    store.create(first)
    assert store.get(first["uid"])["task"] == "work"
    assert store.get("missing") is None
    other = AgentStore(store.db, "user", "owner-2")
    assert other.get(first["uid"]) is None
    store.create(_record(status="running"))
    assert len(store.list()) == 2
    assert len(store.list(status="running")) == 1
    store.update(first["uid"], {"status": "cancelled"})
    assert store.get(first["uid"])["status"] == "cancelled"


def test_count_active_ignores_terminal():
    store = _store()
    for status in ("queued", "running", "needs_user", "done", "cancelled"):
        store.create(_record(status=status))
    assert store.count_active() == 3


def test_due_rows_and_claim():
    db = memory_db()
    store = AgentStore(db, "user", "owner-1")
    now = to_iso(now_utc())
    store.create(_record(status="queued", check_at=now))
    store.create(_record(status="done", check_at=now))
    rows = due_rows(db, now, 10)
    assert [row["status"] for row in rows] == ["queued"]
    assert claim(db, rows[0]["uid"], now) is True
    assert claim(db, rows[0]["uid"], now) is False
    assert store.get(rows[0]["uid"])["status"] == "checking"


def test_complete_only_wins_from_live_status():
    store = _store()
    record = _record(status="running")
    store.create(record)
    won = store.complete(
        record["uid"],
        {
            "status": "done",
            "result": "ok",
            "error": None,
            "iterations": 3,
            "tokens": 100,
            "cost_usd": 0.01,
            "updated_at": to_iso(now_utc()),
        },
    )
    assert won is True
    assert store.get(record["uid"])["result"] == "ok"
    lost = store.complete(
        record["uid"],
        {
            "status": "done",
            "result": "late",
            "error": None,
            "iterations": 4,
            "tokens": 200,
            "cost_usd": 0.02,
            "updated_at": to_iso(now_utc()),
        },
    )
    assert lost is False
    assert store.get(record["uid"])["result"] == "ok"


def test_prune_terminal_keeps_live_rows():
    db = memory_db()
    store = AgentStore(db, "user", "owner-1")
    store.create(_record(status="done", updated_at="2020-01-01T00:00:00"))
    store.create(_record(status="running", updated_at="2020-01-01T00:00:00"))
    assert prune_terminal(db, "2021-01-01T00:00:00") == 1
    assert len(store.list()) == 1


def test_snapshot_for_measures_age_and_idle():
    row = _record(created_at="2026-01-01T00:00:00", last_progress_at="2026-01-01T00:10:00")
    from datetime import datetime, timezone

    snap = snapshot_for(row, datetime(2026, 1, 1, 0, 20, tzinfo=timezone.utc))
    assert snap["age_seconds"] == 1200
    assert snap["idle_seconds"] == 600
    assert snap["answered"] is False


def _claim_worker(path, uid, queue):
    import sqlalchemy

    db = dataset.connect(
        f"sqlite:///{path}", engine_kwargs={"connect_args": {"timeout": 30}}
    )
    try:
        for _ in range(50):
            try:
                queue.put(claim(db, uid, to_iso(now_utc())))
                return
            except sqlalchemy.exc.OperationalError as exc:
                if "locked" not in str(exc).lower():
                    raise
        queue.put(False)
    finally:
        queue.put(None)


def test_claim_race_exactly_one_winner():
    with tempfile.NamedTemporaryFile(suffix=".db") as handle:
        db = dataset.connect(
            f"sqlite:///{handle.name}",
            engine_kwargs={"connect_args": {"timeout": 30}},
        )
        store = AgentStore(db, "user", "owner-1")
        record = _record()
        store.create(record)
        uid = record["uid"]
        ctx = multiprocessing.get_context("fork")
        queue = ctx.Queue()
        workers = [
            ctx.Process(target=_claim_worker, args=(handle.name, uid, queue))
            for _ in range(8)
        ]
        for worker in workers:
            worker.start()
        for worker in workers:
            worker.join(60)
            assert worker.exitcode == 0
        results = []
        while len(results) < 16:
            results.append(queue.get(timeout=30))
        wins = [item for item in results if item is True]
        assert len(wins) == 1