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``), an owner-return marker
(``blocked_on: message-from:<agent>``, rule 6) or a ``wait_until`` date (a
time window clears by itself); *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 re
import uuid
from collections.abc import Iterable, Sequence
from datetime import date
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"

# Rule 6 owner-return marker; mirrors _BLOCKED_ON_RE in scripts/consistency_check.py.
MESSAGE_FROM_RE = re.compile(r"^message-from:.+$")


def is_external_wait(
    *,
    has_dependency: bool,
    blocked_on: str | None = None,
    wait_until: date | None = None,
) -> bool:
    """Case A holds for a dependency edge, an owner-return marker or a wait window."""
    if has_dependency or wait_until is not None:
        return True
    return bool(blocked_on and MESSAGE_FROM_RE.match(blocked_on.strip()))


def derive_wait_kind(
    status: Any,
    *,
    needs_human: bool,
    has_dependency: bool,
    blocked_on: str | None = None,
    wait_until: date | None = None,
) -> 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
    external = is_external_wait(
        has_dependency=has_dependency, blocked_on=blocked_on, wait_until=wait_until
    )
    if external and needs_human:
        return WAIT_KIND_BOTH
    if needs_human:
        return WAIT_KIND_HUMAN
    if external:
        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,
            blocked_on=task.blocked_on,
            wait_until=task.wait_until,
        )
    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