bifroest/tests/test_runs.py

"""Runs im Worker: Lebenszyklus, Queue, Lease, Abbruch, Wiederholung, SSE (→ Konzept 4.4, 9)."""

import re
import threading
import time
import uuid
from datetime import timedelta

import pytest
from sqlalchemy import select, update

from app.chat import runner
from app.chat.models import Message
from app.core.database import now
from app.tasks import engine, states
from app.tasks.models import RunEvent, TaskRun
from app.tasks.worker import HANDLERS, Worker
from helpers import connect_repo, create_workspace


@pytest.fixture
def chat(browser, people, git_repo):
    """anna mit Workspace, getesteter Quelle und einem Chat."""
    anna = browser("anna")
    workspace_id = create_workspace(anna, people["team_demo"])
    connect_repo(anna, workspace_id, git_repo)
    location = anna.post(f"/workspaces/{workspace_id}/chats").headers["location"]
    return anna, int(location.rsplit("/", 1)[1])


def send(b, chat_id: int, text: str = "Hallo") -> uuid.UUID:
    response = b.post(f"/chats/{chat_id}/messages", {"content": text}, headers={"HX-Request": "true"})
    assert response.status_code == 200, response.text
    return uuid.UUID(re.search(r'data-stream="/runs/([0-9a-f-]+)/events"', response.text).group(1))


def run_of(db, run_id) -> TaskRun:
    with db.session() as session:
        run = session.get(TaskRun, run_id)
        session.expunge(run)
    return run


def event_types(db, run_id) -> list[str]:
    with db.session() as session:
        rows = session.scalars(select(RunEvent).where(RunEvent.run_id == run_id).order_by(RunEvent.seq)).all()
        return [r.type if r.type != "status" else f"status:{r.data['status']}" for r in rows]


def test_lebenszyklus(chat, db, settings):
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    assert run_of(db, run_id).status == states.QUEUED
    assert Worker(db, settings).run_once()
    assert run_of(db, run_id).status == states.COMPLETED
    types = event_types(db, run_id)
    assert types[:2] == ["status:queued", "status:running"]
    assert "token" in types and types[-2:] == ["message", "status:completed"]
    answer = anna.get(f"/runs/{run_id}/answer")
    assert "Antwort vom Modell." in answer.text
    assert "Antwort vom Modell." in anna.get(f"/chats/{chat_id}").text


def test_nichts_zu_tun(db, settings, people):
    assert not Worker(db, settings).run_once()


def test_nur_ein_aktiver_run_je_chat(chat):
    anna, chat_id = chat
    send(anna, chat_id)
    response = anna.post(f"/chats/{chat_id}/messages", {"content": "Noch was"}, headers={"HX-Request": "true"})
    assert response.status_code == 422


def test_fremde_sehen_run_nicht(chat, browser, db, settings):
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    dora = browser("dora")
    assert dora.get(f"/chats/{chat_id}").status_code == 404
    assert dora.get(f"/runs/{run_id}/events").status_code == 404
    assert dora.get(f"/runs/{run_id}/answer").status_code == 404
    assert dora.post(f"/runs/{run_id}/cancel").status_code == 404


def test_skip_locked(chat, db, settings):
    anna, chat_id = chat
    first = send(anna, chat_id)
    with db.session() as a:
        claimed = engine.claim(a, "worker-a", settings)
        assert claimed.id == first
        # Zweiter Worker sieht den gesperrten Run nicht, solange A die Transaktion hält
        with db.session() as b:
            assert engine.claim(b, "worker-b", settings) is None


def test_teamlimit(chat, db, settings):
    anna, chat_id = chat
    send(anna, chat_id)
    with db.session() as session:
        workspace_id = session.scalar(select(TaskRun.workspace_id))
    other_chat = int(anna.post(f"/workspaces/{workspace_id}/chats").headers["location"].rsplit("/", 1)[1])
    send(anna, other_chat)
    limited = settings.model_copy(update={"max_runs_per_team": 1})
    with db.session() as session:
        assert engine.claim(session, "worker-a", limited) is not None
    with db.session() as session:
        assert engine.claim(session, "worker-b", limited) is None


def test_lease_abgelaufen_wird_neu_eingereiht(chat, db, settings):
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    with db.session() as session:
        engine.claim(session, "worker-a", settings)
    with db.session() as session:
        session.execute(update(TaskRun).values(locked_until=now() - timedelta(seconds=1)))
    assert engine.recover_expired(db, settings) == 1
    run = run_of(db, run_id)
    assert run.status == states.QUEUED and run.locked_by is None and run.attempts == 1
    # Der alte Worker darf danach nichts mehr speichern
    with db.session() as session:
        assert engine.owned(session, run_id, "worker-a") is None
    # Ein anderer Worker beantwortet den Run vollständig
    assert Worker(db, settings, "worker-b").run_once()
    assert run_of(db, run_id).status == states.COMPLETED


def test_lease_zu_oft_abgelaufen(chat, db, settings):
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    for _ in range(settings.run_max_attempts):
        with db.session() as session:
            assert engine.claim(session, "worker-a", settings) is not None
        with db.session() as session:
            session.execute(update(TaskRun).values(locked_until=now() - timedelta(seconds=1)))
        engine.recover_expired(db, settings)
    run = run_of(db, run_id)
    assert run.status == states.FAILED and "Versuchen" in run.error


def test_alter_worker_speichert_nichts_nach_uebernahme(chat, db, settings, monkeypatch):
    anna, chat_id = chat
    run_id = send(anna, chat_id)

    def taken_over(context):
        # Während der Arbeit übernimmt ein anderer Worker den Run, dann will dieser speichern
        with db.session() as session:
            session.execute(update(TaskRun).values(locked_by="worker-b"))
        runner.answer(context)

    monkeypatch.setitem(HANDLERS, "chat", taken_over)
    Worker(db, settings, "worker-a").run_once()
    with db.session() as session:
        assert session.scalar(select(Message).where(Message.role == "assistant")) is None
    run = run_of(db, run_id)
    assert run.status == states.RUNNING and run.locked_by == "worker-b"


def test_abbruch_wartender_run(chat, db):
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    assert anna.post(f"/runs/{run_id}/cancel").status_code == 204
    assert run_of(db, run_id).status == states.CANCELLED
    # Ein beendeter Run lässt sich nicht noch einmal abbrechen
    assert anna.post(f"/runs/{run_id}/cancel", headers={"HX-Request": "true"}).status_code == 422


def test_abbruch_laufender_run(chat, db, settings, monkeypatch):
    anna, chat_id = chat
    run_id = send(anna, chat_id)

    def slow(context):
        with db.session() as session:
            session.execute(update(TaskRun).values(cancel_requested=True))
        context.cancel.set()  # sonst erst beim nächsten Heartbeat
        context.check()

    monkeypatch.setitem(HANDLERS, "chat", slow)
    Worker(db, settings).run_once()
    assert run_of(db, run_id).status == states.CANCELLED


def test_fehler_wird_wiederholt_und_endet_als_failed(chat, db, settings, monkeypatch):
    anna, chat_id = chat
    run_id = send(anna, chat_id)

    def broken(context):
        raise RuntimeError("intern: geheimes Detail")

    monkeypatch.setitem(HANDLERS, "chat", broken)
    worker = Worker(db, settings)
    for attempt in range(1, settings.run_max_attempts + 1):
        with db.session() as session:
            session.execute(update(TaskRun).values(available_at=now() - timedelta(seconds=1)))
        assert worker.run_once()
        run = run_of(db, run_id)
        expected = states.FAILED if attempt == settings.run_max_attempts else states.QUEUED
        assert run.status == expected and run.attempts == attempt
    # Interne Details nur im Run, nicht in den Ereignissen für den Browser
    with db.session() as session:
        data = [r.data for r in session.scalars(select(RunEvent).where(RunEvent.run_id == run_id))]
    assert "geheimes Detail" not in str(data)
    assert "queued" in str(data) and "retry" in str(data)


def parse_sse(body: str) -> list[dict[str, str]]:
    events = []
    for block in body.strip().split("\n\n"):
        fields = dict(line.split(": ", 1) for line in block.splitlines() if ": " in line and not line.startswith(":"))
        if fields:
            events.append(fields)
    return events


def test_sse_liefert_alles_und_endet(chat, db, settings):
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    Worker(db, settings).run_once()
    response = anna.get(f"/runs/{run_id}/events")
    assert response.headers["content-type"].startswith("text/event-stream")
    events = parse_sse(response.text)
    assert events[0]["event"] == "status" and events[0]["id"] == "1"
    assert events[-1]["event"] == "end" and "completed" in events[-1]["data"]
    tokens = "".join(re.search(r'"text": ?"([^"]*)"', e["data"]).group(1) for e in events if e["event"] == "token")
    assert tokens.startswith("Antwort vom Modell.")


def test_sse_fortsetzen_mit_last_event_id(chat, db, settings):
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    Worker(db, settings).run_once()
    total = len(event_types(db, run_id))
    events = parse_sse(anna.get(f"/runs/{run_id}/events", headers={"Last-Event-ID": "3"}).text)
    ids = [int(e["id"]) for e in events if "id" in e]
    assert ids == list(range(4, total + 1))


def test_sse_live_waehrend_der_worker_arbeitet(chat, db, settings, monkeypatch):
    """Stream ist offen, bevor der Worker beginnt: NOTIFY weckt ihn auf (kein Warten auf das Polling)."""
    anna, chat_id = chat
    run_id = send(anna, chat_id)
    worker = threading.Timer(0.3, Worker(db, settings).run_once)
    started = time.monotonic()
    worker.start()
    with anna.client.stream("GET", f"/runs/{run_id}/events") as response:
        body = "".join(response.iter_text())
    worker.join()
    events = parse_sse(body)
    assert [e["event"] for e in events][-3:] == ["message", "status", "end"], body[-600:]
    assert time.monotonic() - started < 5