state-hub/api/services/wait_kind.py

"""Derived wait qualifiers (CUST-WP-0074 rule 5) — read model only.

A ``wait`` task is *external* when it carries at least one dependency edge
(``workplan_dependencies.from_task_id``), *human* when ``needs_human`` is set,
*both* when both hold and *unqualified* otherwise. A ``blocked`` workplan is
*human* when any of its wait tasks is human or both, else *external* when any
is external, else *none*. Nothing here is stored; the routers call the
``annotate_*`` helpers and Pydantic picks the attributes up on serialisation.
"""
from __future__ import annotations

import uuid
from collections.abc import Iterable, Sequence
from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from api.models.task import Task, TaskStatus
from api.models.workplan import Workplan
from api.models.workplan_dependency import WorkplanDependency
from api.task_status import normalize_task_status
from api.workplan_status import normalize_workplan_status

WAIT_KIND_EXTERNAL = "external"
WAIT_KIND_HUMAN = "human"
WAIT_KIND_BOTH = "both"
WAIT_KIND_UNQUALIFIED = "unqualified"

BLOCKED_KIND_EXTERNAL = "external"
BLOCKED_KIND_HUMAN = "human"
BLOCKED_KIND_NONE = "none"


def derive_wait_kind(status: Any, *, needs_human: bool, has_dependency: bool) -> str | None:
    """Rule 1 / rule 5 for one task. ``None`` unless the task is waiting."""
    if normalize_task_status(status, default="todo") != "wait":
        return None
    if has_dependency and needs_human:
        return WAIT_KIND_BOTH
    if needs_human:
        return WAIT_KIND_HUMAN
    if has_dependency:
        return WAIT_KIND_EXTERNAL
    return WAIT_KIND_UNQUALIFIED


def derive_blocked_kind(status: Any, wait_kinds: Iterable[str | None]) -> str | None:
    """Rule 5 for one workplan. ``None`` unless the workplan is blocked."""
    if normalize_workplan_status(status) != "blocked":
        return None
    kinds = {kind for kind in wait_kinds if kind}
    if WAIT_KIND_HUMAN in kinds or WAIT_KIND_BOTH in kinds:
        return BLOCKED_KIND_HUMAN
    if WAIT_KIND_EXTERNAL in kinds:
        return BLOCKED_KIND_EXTERNAL
    return BLOCKED_KIND_NONE


async def task_ids_with_dependencies(
    session: AsyncSession, task_ids: Sequence[uuid.UUID]
) -> set[uuid.UUID]:
    if not task_ids:
        return set()
    rows = await session.execute(
        select(WorkplanDependency.from_task_id).where(
            WorkplanDependency.from_task_id.in_(list(task_ids))
        )
    )
    return {row[0] for row in rows if row[0] is not None}


async def annotate_tasks(session: AsyncSession, tasks: Sequence[Task]) -> Sequence[Task]:
    """Set ``task.wait_kind`` on each ORM task (serialised by TaskRead)."""
    waiting = [t for t in tasks if normalize_task_status(t.status, default="todo") == "wait"]
    with_deps = await task_ids_with_dependencies(session, [t.id for t in waiting])
    for task in tasks:
        task.wait_kind = derive_wait_kind(
            task.status,
            needs_human=bool(task.needs_human),
            has_dependency=task.id in with_deps,
        )
    return tasks


async def annotate_workplans(
    session: AsyncSession, workplans: Sequence[Workplan]
) -> Sequence[Workplan]:
    """Set ``workplan.blocked_kind`` on each ORM workplan (serialised by WorkplanRead)."""
    blocked = [
        wp for wp in workplans if normalize_workplan_status(wp.status) == "blocked"
    ]
    kinds_by_workplan: dict[uuid.UUID, list[str | None]] = {wp.id: [] for wp in blocked}
    if blocked:
        rows = await session.execute(
            select(Task).where(
                Task.workplan_id.in_(list(kinds_by_workplan)),
                Task.status == TaskStatus.wait,
            )
        )
        wait_tasks = list(rows.scalars().all())
        await annotate_tasks(session, wait_tasks)
        for task in wait_tasks:
            kinds_by_workplan[task.workplan_id].append(task.wait_kind)
    for wp in workplans:
        wp.blocked_kind = derive_blocked_kind(wp.status, kinds_by_workplan.get(wp.id, []))
    return workplans