smart-org/apps/api/app/infrastructure/repositories/engine.py

"""Repositories for the AI-workforce engine's own tables.

Four repositories, one per engine table, all following the house rule: the
ORM row never leaves this module — the domain dataclasses in
``app.domain.workforce.pipeline`` / ``engine`` are what the application layer
holds.

The interesting one is :class:`PipelineStateRepository`. Its ``advance`` is an
optimistic compare-and-set: the ``UPDATE`` carries the caller's expected stage
and version in its ``WHERE`` clause, so two workers moving the same task
cannot both win — the loser sees zero updated rows and gets a
:class:`StaleStageError` instead of silently overwriting the winner.
"""

from __future__ import annotations

from datetime import datetime

from sqlalchemy import func, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.domain.base import utcnow
from app.domain.workforce.engine import (
    EmployeeProfile,
    EngineTurn,
    GateName,
    GateResult,
)
from app.domain.workforce.models import RoleKey
from app.domain.workforce.pipeline import (
    TRANSITIONS,
    Mode,
    PipelineKey,
    Risk,
    Stage,
    TaskKind,
    TaskPipelineState,
)
from app.infrastructure.db.models import RunRow, TaskRow, UsageEventRow
from app.infrastructure.db.models_engine import (
    EmployeeProfileRow,
    EngineTurnRow,
    GateResultRow,
    TaskPipelineStateRow,
)


class StaleStageError(Exception):
    """Another worker moved the task first; the caller must re-read it."""


class UnknownTaskError(LookupError):
    """The task has no pipeline state row (it was never registered)."""


class InvalidTransitionError(ValueError):
    """The requested move is not in the stage transition table."""


# ---------------------------------------------------------------------------
# task_pipeline_state
# ---------------------------------------------------------------------------

#: Columns ``patch`` may touch. An allow-list, so a typo can never become SQL.
_PATCHABLE = frozenset(
    {
        "pipeline",
        "kind",
        "risk",
        "stage",
        "cycle",
        "max_cycles",
        "branch",
        "last_diff_hash",
        "same_diff_count",
        "retry_count",
        "waiver_reason",
        "last_diff",
        "merged_sha",
        "parent_task_id",
        "depends_on",
    }
)


def _state(row: TaskPipelineStateRow) -> TaskPipelineState:
    return TaskPipelineState(
        task_id=row.task_id,
        organization_id=row.organization_id,
        project_id=row.project_id,
        pipeline=PipelineKey(row.pipeline),
        kind=TaskKind(row.kind),
        risk=Risk(row.risk),
        stage=Stage(row.stage),
        cycle=row.cycle,
        max_cycles=row.max_cycles,
        branch=row.branch,
        last_diff_hash=row.last_diff_hash,
        same_diff_count=row.same_diff_count,
        retry_count=row.retry_count,
        waiver_reason=row.waiver_reason,
        last_diff=row.last_diff,
        merged_sha=row.merged_sha,
        parent_task_id=row.parent_task_id,
        depends_on=list(row.depends_on or []),
        version=row.version,
        updated_at=row.updated_at,
    )


def _legal_targets(stage: Stage) -> frozenset[Stage]:
    """The stage moves the domain allows from here."""
    return TRANSITIONS.get(stage, frozenset())


class PipelineStateRepository:
    """1:1 with a task: pipeline, stage, cycle guards and the version lock."""

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get(self, task_id: str) -> TaskPipelineState | None:
        row = await self.session.get(TaskPipelineStateRow, task_id)
        return _state(row) if row else None

    async def get_or_raise(self, task_id: str) -> TaskPipelineState:
        state = await self.get(task_id)
        if state is None:
            raise UnknownTaskError(task_id)
        return state

    async def list_for_org(
        self, organization_id: str, *, limit: int = 500
    ) -> list[TaskPipelineState]:
        stmt = (
            select(TaskPipelineStateRow)
            .where(TaskPipelineStateRow.organization_id == organization_id)
            .order_by(TaskPipelineStateRow.updated_at.desc())
            .limit(min(limit, 1000))
        )
        return [_state(r) for r in (await self.session.execute(stmt)).scalars()]

    async def list_at_stages(
        self, stages: tuple[Stage, ...], *, limit: int = 50
    ) -> list[TaskPipelineState]:
        """The scheduler's hot query: tasks currently at a turn stage."""
        if not stages:
            return []
        stmt = (
            select(TaskPipelineStateRow)
            .where(TaskPipelineStateRow.stage.in_([str(s) for s in stages]))
            .order_by(TaskPipelineStateRow.updated_at.asc())
            .limit(min(limit, 500))
        )
        return [_state(r) for r in (await self.session.execute(stmt)).scalars()]

    async def list_dependents(self, task_id: str) -> list[TaskPipelineState]:
        """Tasks that declare ``task_id`` in ``depends_on`` and are not closed.

        JSON columns have no portable containment operator (SQLite differs from
        PostgreSQL), so the match happens in Python. The set is small — it is
        the tasks waiting on one specific task — so this is not a scan of the
        table.
        """
        stmt = select(TaskPipelineStateRow).where(
            TaskPipelineStateRow.stage.notin_(
                [str(Stage.COMPLETED), str(Stage.CANCELLED), str(Stage.FAILED)]
            )
        )
        dependents: list[TaskPipelineState] = []
        for row in (await self.session.execute(stmt)).scalars():
            if task_id in list(row.depends_on or []):
                dependents.append(_state(row))
        return dependents

    async def create(self, state: TaskPipelineState) -> TaskPipelineState:
        row = TaskPipelineStateRow(
            task_id=state.task_id,
            organization_id=state.organization_id,
            project_id=state.project_id,
            pipeline=str(state.pipeline),
            kind=str(state.kind),
            risk=str(state.risk),
            stage=str(state.stage),
            cycle=state.cycle,
            max_cycles=state.max_cycles,
            branch=state.branch,
            last_diff_hash=state.last_diff_hash,
            same_diff_count=state.same_diff_count,
            retry_count=state.retry_count,
            waiver_reason=state.waiver_reason,
            last_diff=state.last_diff,
            merged_sha=state.merged_sha,
            parent_task_id=state.parent_task_id,
            depends_on=list(state.depends_on),
            version=state.version,
        )
        self.session.add(row)
        await self.session.flush()
        return _state(row)

    async def advance(
        self, *, task_id: str, expected_from: Stage, to: Stage
    ) -> TaskPipelineState:
        """Move a task one stage forward, but only if nobody moved it yet.

        The row is read for its current version, then written with both the
        expected stage and that version in the ``WHERE``. If another worker
        got there first, zero rows match and :class:`StaleStageError` is
        raised — the caller re-reads and retries (or gives up cleanly).
        """
        row = await self.session.get(TaskPipelineStateRow, task_id)
        if row is None:
            raise UnknownTaskError(task_id)
        if Stage(row.stage) is not Stage(expected_from):
            raise StaleStageError(
                f"task {task_id} is at {row.stage}, not {expected_from.value}"
            )
        if Stage(to) not in _legal_targets(Stage(row.stage)):
            raise InvalidTransitionError(f"{row.stage} -> {to.value} is not allowed")
        expected_version = row.version
        # synchronize_session="fetch" keeps the identity-mapped object truthful
        # (SQLite has no RETURNING for this), so the check below reads what the
        # database actually did rather than what we hoped it would do.
        await self.session.execute(
            update(TaskPipelineStateRow)
            .where(
                TaskPipelineStateRow.task_id == task_id,
                TaskPipelineStateRow.stage == str(expected_from),
                TaskPipelineStateRow.version == expected_version,
            )
            .values(stage=str(to), version=expected_version + 1, updated_at=utcnow())
            .execution_options(synchronize_session="fetch")
        )
        await self.session.flush()
        # A version that did not move means the WHERE matched nothing: another
        # worker wrote this row first and this move never happened.
        if row.version == expected_version or Stage(row.stage) is not Stage(to):
            raise StaleStageError(f"task {task_id} was moved by another worker")
        return _state(row)

    async def patch(
        self, *, task_id: str, expected_version: int, **fields: object
    ) -> TaskPipelineState:
        """Change counters without moving the stage, under the version lock."""
        unknown = set(fields) - _PATCHABLE
        if unknown:
            raise ValueError(f"non-patchable columns: {sorted(unknown)}")
        row = await self.session.get(TaskPipelineStateRow, task_id)
        if row is None:
            raise UnknownTaskError(task_id)
        if not fields:
            return _state(row)
        before = row.version
        await self.session.execute(
            update(TaskPipelineStateRow)
            .where(
                TaskPipelineStateRow.task_id == task_id,
                TaskPipelineStateRow.version == expected_version,
            )
            .values(**fields, version=expected_version + 1, updated_at=utcnow())
            .execution_options(synchronize_session="fetch")
        )
        await self.session.flush()
        if row.version == before:
            raise StaleStageError(f"task {task_id} was moved by another worker")
        return _state(row)

    async def count_at(self, stage: Stage) -> int:
        stmt = (
            select(func.count())
            .select_from(TaskPipelineStateRow)
            .where(TaskPipelineStateRow.stage == str(stage))
        )
        return int((await self.session.execute(stmt)).scalar_one())

    async def count_at_by_org(self, stage: Stage) -> dict[str, int]:
        """How many tasks sit at one stage, per organization.

        Cross-organization on purpose: only the platform owner's views may use it.
        """
        stmt = (
            select(TaskPipelineStateRow.organization_id, func.count())
            .where(TaskPipelineStateRow.stage == str(stage))
            .group_by(TaskPipelineStateRow.organization_id)
        )
        return {org: int(n) for org, n in (await self.session.execute(stmt)).all()}

    async def stage_counts(self, organization_id: str) -> dict[str, int]:
        """How many of one organization's tasks sit at each stage."""
        stmt = (
            select(TaskPipelineStateRow.stage, func.count())
            .where(TaskPipelineStateRow.organization_id == organization_id)
            .group_by(TaskPipelineStateRow.stage)
        )
        return {stage: int(n) for stage, n in (await self.session.execute(stmt)).all()}


# ---------------------------------------------------------------------------
# employee_profiles
# ---------------------------------------------------------------------------


def _profile(row: EmployeeProfileRow) -> EmployeeProfile:
    return EmployeeProfile(
        employee_id=row.employee_id,
        ai_model_id=row.ai_model_id,
        effort=row.effort,
        roles=[r.strip() for r in (row.roles or "").split(",") if r.strip()],
        hired_by=row.hired_by,
        busy_task_id=row.busy_task_id,
        updated_at=row.updated_at,
    )


class EmployeeProfileRepository:
    """The engine's view of an employee: granted model, roles, current task."""

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get(self, employee_id: str) -> EmployeeProfile | None:
        row = await self.session.get(EmployeeProfileRow, employee_id)
        return _profile(row) if row else None

    async def list_free_with_role(self, role: str, *, limit: int = 10) -> list[str]:
        """Employee ids that hold ``role`` and are not busy, least busy first."""
        stmt = select(EmployeeProfileRow).where(
            EmployeeProfileRow.busy_task_id.is_(None),
            EmployeeProfileRow.roles.like(f"%{role}%"),
        )
        return [r.employee_id for r in (await self.session.execute(stmt)).scalars()][:limit]

    async def upsert(self, profile: EmployeeProfile) -> EmployeeProfile:
        row = await self.session.get(EmployeeProfileRow, profile.employee_id)
        if row is None:
            row = EmployeeProfileRow(
                employee_id=profile.employee_id,
                ai_model_id=profile.ai_model_id,
                effort=profile.effort,
                roles=",".join(profile.role_list()),
                hired_by=profile.hired_by,
                busy_task_id=profile.busy_task_id,
            )
            self.session.add(row)
        else:
            row.ai_model_id = profile.ai_model_id
            row.effort = profile.effort
            row.roles = ",".join(profile.role_list())
            row.hired_by = profile.hired_by
            row.busy_task_id = profile.busy_task_id
        await self.session.flush()
        return _profile(row)

    async def set_busy(self, *, employee_id: str, task_id: str | None) -> None:
        row = await self.session.get(EmployeeProfileRow, employee_id)
        if row is None:
            return
        row.busy_task_id = task_id
        await self.session.flush()

    async def set_model(self, *, employee_id: str, ai_model_id: str | None) -> None:
        row = await self.session.get(EmployeeProfileRow, employee_id)
        if row is None:
            return
        row.ai_model_id = ai_model_id
        await self.session.flush()


# ---------------------------------------------------------------------------
# engine_turns
# ---------------------------------------------------------------------------


def _turn(row: EngineTurnRow) -> EngineTurn:
    return EngineTurn(
        id=row.id,
        task_id=row.task_id,
        run_id=row.run_id,
        employee_id=row.employee_id,
        role=RoleKey(row.role),
        mode=Mode(row.mode),
        stage=Stage(row.stage) if row.stage else Stage.PENDING_IMPLEMENTATION,
        prompt_chars=row.prompt_chars,
        reply_excerpt=row.reply_excerpt,
        tags=dict(row.tags or {}),
        tool_calls=row.tool_calls,
        ok=row.ok,
        error_kind=row.error_kind,
        started_at=row.started_at,
        finished_at=row.finished_at,
    )


class EngineTurnRepository:
    """One row per agent turn: the audit trail of who thought what."""

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def create(self, turn: EngineTurn) -> EngineTurn:
        row = EngineTurnRow(
            id=turn.id,
            task_id=turn.task_id,
            run_id=turn.run_id,
            employee_id=turn.employee_id,
            role=str(turn.role),
            mode=str(turn.mode),
            stage=str(turn.stage),
            prompt_chars=turn.prompt_chars,
            reply_excerpt=turn.reply_excerpt,
            tags=dict(turn.tags),
            tool_calls=turn.tool_calls,
            ok=turn.ok,
            error_kind=turn.error_kind,
            started_at=turn.started_at,
            finished_at=turn.finished_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _turn(row)

    async def list_for_task(self, task_id: str, *, limit: int = 200) -> list[EngineTurn]:
        stmt = (
            select(EngineTurnRow)
            .where(EngineTurnRow.task_id == task_id)
            .order_by(EngineTurnRow.started_at.asc())
            .limit(min(limit, 1000))
        )
        return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]

    async def list_for_run(self, run_id: str) -> list[EngineTurn]:
        stmt = (
            select(EngineTurnRow)
            .where(EngineTurnRow.run_id == run_id)
            .order_by(EngineTurnRow.started_at.asc())
        )
        return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]

    async def list_for_org(self, organization_id: str, *, limit: int = 50) -> list[EngineTurn]:
        """The newest turns across one organization's tasks.

        Turns carry no organization column, so the tenant filter joins through
        the task: another organization's turn can never match.
        """
        stmt = (
            select(EngineTurnRow)
            .join(TaskRow, EngineTurnRow.task_id == TaskRow.id)
            .where(TaskRow.organization_id == organization_id)
            .order_by(EngineTurnRow.started_at.desc())
            .limit(min(limit, 200))
        )
        return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]

    async def finish(
        self,
        turn_id: str,
        *,
        ok: bool,
        tool_calls: int,
        error_kind: str = "",
        reply_excerpt: str | None = None,
        tags: dict | None = None,
    ) -> None:
        """Close a turn: outcome, tool count and the bounded reply excerpt."""
        row = await self.session.get(EngineTurnRow, turn_id)
        if row is None:
            return
        row.ok = ok
        row.tool_calls = tool_calls
        row.error_kind = error_kind
        if reply_excerpt is not None:
            row.reply_excerpt = reply_excerpt
        if tags is not None:
            row.tags = dict(tags)
        row.finished_at = utcnow()
        await self.session.flush()


# ---------------------------------------------------------------------------
# gate_results
# ---------------------------------------------------------------------------


def _gate(row: GateResultRow) -> GateResult:
    return GateResult(
        id=row.id,
        task_id=row.task_id,
        gate=GateName(row.gate),
        ok=row.ok,
        detail=row.detail,
        created_at=row.created_at,
    )


class GateResultRepository:
    """One row per gate verdict, with the detail that justifies it."""

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def create(self, result: GateResult) -> GateResult:
        row = GateResultRow(
            id=result.id,
            task_id=result.task_id,
            gate=str(result.gate),
            ok=result.ok,
            detail=result.detail,
            created_at=result.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _gate(row)

    async def list_for_task(self, task_id: str, *, limit: int = 200) -> list[GateResult]:
        stmt = (
            select(GateResultRow)
            .where(GateResultRow.task_id == task_id)
            .order_by(GateResultRow.created_at.asc())
            .limit(min(limit, 1000))
        )
        return [_gate(r) for r in (await self.session.execute(stmt)).scalars()]

    async def latest_for(self, task_id: str, gate: str) -> GateResult | None:
        stmt = (
            select(GateResultRow)
            .where(GateResultRow.task_id == task_id, GateResultRow.gate == gate)
            .order_by(GateResultRow.created_at.desc())
            .limit(1)
        )
        row = (await self.session.execute(stmt)).scalar_one_or_none()
        return _gate(row) if row else None


class TaskSpendRepository:
    """Spend, read from the usage ledger (never written here).

    The scheduler's budget guards need "how much has this task cost so far"
    (joining the ledger's events to the task's runs) and "how much has this
    organization spent since a date" (the owner's monthly credit budget). The
    ledger stays the source of truth; this only reads it.
    """

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def spend_for_task(self, task_id: str) -> float:
        run_ids = select(RunRow.id).where(RunRow.task_id == task_id)
        stmt = select(func.coalesce(func.sum(UsageEventRow.total_price), 0.0)).where(
            UsageEventRow.run_id.in_(run_ids)
        )
        return float((await self.session.execute(stmt)).scalar_one())

    async def spend_for_org_since(self, organization_id: str, since: datetime) -> float:
        """What the organization's usage has cost since ``since``, over all its tasks."""
        stmt = select(func.coalesce(func.sum(UsageEventRow.total_price), 0.0)).where(
            UsageEventRow.organization_id == organization_id,
            UsageEventRow.created_at >= since,
        )
        return float((await self.session.execute(stmt)).scalar_one())