bifroest/app/tasks/engine.py

"""Runs anlegen, beanspruchen, verlängern und abschließen (→ Konzept 9.3).

Alle Zeitpunkte der Warteschlange (available_at, locked_until) kommen von der Uhr
der Datenbank, mit der auch verglichen wird; die Uhren von Web und Worker können
abweichen. Ein Run ist zugleich der Job. Worker beanspruchen fällige Runs mit
`FOR UPDATE SKIP LOCKED` und halten sie über eine Lease, die sie per Heartbeat
verlängern. Läuft eine Lease ab, ist der Worker vermutlich abgestürzt; der Run
wird erneut eingereiht oder nach zu vielen Versuchen als fehlgeschlagen beendet.
"""

import uuid
from dataclasses import dataclass
from datetime import timedelta
from typing import Any

from sqlalchemy import func, select, update
from sqlalchemy.orm import Session, aliased

from app import audit
from app.approvals.models import CANCELLED, EXECUTING, PENDING, UNKNOWN, ApprovalRequest
from app.core.config import Settings
from app.core.database import Database, now
from app.core.errors import Forbidden, Invalid, NotFound
from app.identity import Principal
from app.tasks import events, states
from app.tasks.models import TaskRun
from app.workspaces import service as workspaces
from app.workspaces.models import Workspace


def create(session: Session, workspace: Workspace, created_by: int, kind: str, input_: dict[str, Any],
           chat_session_id: int | None = None, parent_run_id: uuid.UUID | None = None) -> TaskRun:
    run = TaskRun(workspace_id=workspace.id, team_id=workspace.team_id, kind=kind, input=input_,
                  chat_session_id=chat_session_id, parent_run_id=parent_run_id, created_by=created_by)
    session.add(run)
    session.flush()
    events.append(session, run.id, "status", {"status": states.QUEUED})
    audit.record(session, "run.created", user_id=created_by, workspace_id=workspace.id, target_kind="run",
                 target_id=run.id, kind=kind)
    return run


def get(session: Session, principal: Principal, run_id: uuid.UUID) -> TaskRun:
    run = session.get(TaskRun, run_id)
    if run is None:
        raise NotFound
    workspaces.get(session, principal, run.workspace_id)
    return run


def set_status(session: Session, run: TaskRun, status: str, **data: Any) -> None:
    states.check(run.status, status)
    run.status = status
    if status in states.RESTING:
        run.finished = now()
    if status != states.RUNNING:
        run.locked_by = run.locked_until = None
    events.append(session, run.id, "status", {"status": status, **data})


# Namensraum der Advisory Locks für das Teamlimit (beliebige feste Zahl)
TEAM_LOCK = 7101

def claim(session: Session, worker_id: str, settings: Settings) -> TaskRun | None:
    """Nächsten fälligen Run beanspruchen, unter Beachtung des Limits gleichzeitiger Runs je Team.

    Das Teamlimit ist hart: Nach der Auswahl wird je Team ein Advisory Lock bis zum Ende der
    Transaktion gehalten und erneut gezählt. So sieht der nächste Worker desselben Teams den
    Run bereits als laufend, statt dass mehrere Worker gleichzeitig unter dem Limit zählen.
    """
    other = aliased(TaskRun)
    running = (select(func.count()).select_from(other)
               .where(other.team_id == TaskRun.team_id, other.status == states.RUNNING)
               .scalar_subquery())
    run = session.scalar(select(TaskRun)
                         .where(TaskRun.status == states.QUEUED, TaskRun.available_at <= func.now(),
                                running < settings.max_runs_per_team)
                         .order_by(TaskRun.available_at, TaskRun.created)
                         .limit(1)
                         .with_for_update(skip_locked=True))
    if run is None:
        return None
    session.execute(select(func.pg_advisory_xact_lock(TEAM_LOCK, run.team_id)))
    running_now = session.scalar(select(func.count()).select_from(TaskRun).where(
        TaskRun.team_id == run.team_id, TaskRun.status == states.RUNNING))
    if running_now >= settings.max_runs_per_team:
        return None  # ein anderer Worker war schneller; der Run bleibt eingereiht
    set_status(session, run, states.RUNNING)
    run.attempts += 1
    run.locked_by = worker_id
    run.locked_until = func.now() + timedelta(seconds=settings.run_lease_seconds)
    run.started = run.started or now()
    return run


@dataclass(frozen=True)
class Beat:
    alive: bool  # Lease noch bei diesem Worker
    cancel: bool  # Abbruch angefordert


def heartbeat(db: Database, run_id: uuid.UUID, worker_id: str, settings: Settings) -> Beat:
    with db.session() as session:
        row = session.execute(
            update(TaskRun)
            .where(TaskRun.id == run_id, TaskRun.locked_by == worker_id, TaskRun.status == states.RUNNING)
            .values(locked_until=func.now() + timedelta(seconds=settings.run_lease_seconds))
            .returning(TaskRun.cancel_requested)).first()
    return Beat(alive=row is not None, cancel=bool(row and row[0]))


def owned(session: Session, run_id: uuid.UUID, worker_id: str) -> TaskRun | None:
    """Den Run nur zurückgeben, wenn dieser Worker die Lease noch hält (gesperrt bis Commit)."""
    run = session.get(TaskRun, run_id, with_for_update=True)
    if run is None or run.status != states.RUNNING or run.locked_by != worker_id:
        return None
    return run


def recover_expired(db: Database, settings: Settings) -> int:
    """Runs mit abgelaufener Lease wieder einreihen oder nach zu vielen Versuchen beenden.

    Steckte der Run in einer Schreibaktion, ist ihr Ausgang unklar: Der Run wechselt
    nach needs_reconciliation, statt die Aktion blind zu wiederholen (→ Konzept 9.2).
    """
    count = 0
    with db.session() as session:
        expired = session.scalars(select(TaskRun)
                                  .where(TaskRun.status == states.RUNNING, TaskRun.locked_until < func.now())
                                  .with_for_update(skip_locked=True)).all()
        for run in expired:
            worker = run.locked_by
            writing = session.scalars(select(ApprovalRequest).where(ApprovalRequest.run_id == run.id,
                                                                    ApprovalRequest.status == EXECUTING)).all()
            if writing:
                for approval in writing:
                    approval.status = UNKNOWN
                state = dict(run.state or {})
                waiting = dict(state.get("waiting", {}))
                waiting.update({a.call_id: a.id for a in writing if not a.call_id.startswith("rerequest-")})
                run.state = {**state, "waiting": waiting}
                run.result = {**run.result, "message": "Der Worker fiel während einer Schreibaktion aus. "
                              "Ein Workspace-Admin muss den Abgleich mit dem Zielsystem starten."}
                set_status(session, run, states.NEEDS_RECONCILIATION, reason="lease_expired")
            elif run.attempts >= settings.run_max_attempts:
                run.error = f"Abgebrochen nach {run.attempts} Versuchen (Worker nicht mehr erreichbar)."
                set_status(session, run, states.FAILED, reason="lease_expired")
            else:
                set_status(session, run, states.QUEUED, reason="lease_expired")
                run.available_at = func.now()
            audit.record(session, "run.lease_expired", workspace_id=run.workspace_id, target_kind="run",
                         target_id=run.id, attempts=run.attempts, worker=worker)
            count += 1
    return count


def request_cancel(session: Session, principal: Principal, run_id: uuid.UUID) -> TaskRun:
    run = get(session, principal, run_id)
    if run.status in (states.QUEUED, states.WAITING_APPROVAL, states.WAITING_INPUT):
        set_status(session, run, states.CANCELLED, by=principal.username)
        session.execute(update(ApprovalRequest).where(ApprovalRequest.run_id == run.id,
                                                      ApprovalRequest.status == PENDING)
                        .values(status=CANCELLED, decided_at=now()))
    elif run.status == states.RUNNING:
        run.cancel_requested = True
    else:
        raise Invalid("Dieser Run läuft nicht mehr.")
    audit.record(session, "run.cancel_requested", user_id=principal.user_id, workspace_id=run.workspace_id,
                 target_kind="run", target_id=run.id)
    return run


def reconcile(session: Session, principal: Principal, run_id: uuid.UUID, give_up: bool = False) -> TaskRun:
    """Nach unklarem Ausgang einer Schreibaktion (→ Konzept 9.2): Workspace-Admin startet den
    Abgleich (die Tools prüfen dabei den Zielzustand, bevor sie etwas tun) oder beendet den Run."""
    run = get(session, principal, run_id)
    if not workspaces.get(session, principal, run.workspace_id).is_admin:
        raise Forbidden
    if run.status != states.NEEDS_RECONCILIATION:
        raise Invalid("Für diesen Run ist kein Abgleich nötig.")
    if give_up:
        run.result = {**run.result, "message": "Nach unklarer Schreibaktion beendet. Bitte das Zielsystem "
                      "(Branch ai/…, Pull Request) von Hand prüfen."}
        set_status(session, run, states.FAILED, by=principal.username, message=run.result["message"])
    else:
        set_status(session, run, states.QUEUED, resumed=True)
        run.available_at = func.now()
    audit.record(session, "run.reconcile", user_id=principal.user_id, workspace_id=run.workspace_id,
                 target_kind="run", target_id=run.id, give_up=give_up)
    return run