smart-org/apps/api/app/infrastructure/db/models_engine.py

"""ORM tables owned by Agent 2 (AI workforce engine).

Use ``from app.core.database import Base`` and the helpers in
``app.infrastructure.db.columns`` (``pk()``, ``created()``). Do NOT import from
``app.infrastructure.db.models`` here: that module imports this one (import cycle).
Foreign keys to existing tables are written as plain strings, e.g.
``ForeignKey("organizations.id", ondelete="CASCADE")``. See docs/plan/PLAN-v2.md §6.

Four tables, deliberately split so ``tasks`` and ``employees`` stay untouched
(they are locked):

* ``task_pipeline_state`` — the engine's own view of a task: pipeline, stage,
  cycle counters, branch and the optimistic ``version``. One row per task.
* ``employee_profiles`` — which granted model an employee runs on, and how
  busy it is. Model identity never leaves the server.
* ``engine_turns`` — one row per agent turn: the audit trail of who thought
  what, with a bounded reply excerpt.
* ``gate_results`` — one row per gate verdict (tests, diff guard, no-progress,
  green), with the detail that justifies it.

Both side tables cascade from their parent, so deleting a task or an employee
never leaves an orphaned engine row behind.
"""

from __future__ import annotations

from datetime import datetime

from sqlalchemy import (
    JSON,
    Boolean,
    DateTime,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
)
from sqlalchemy.orm import Mapped, mapped_column

from app.core.database import Base
from app.domain.base import utcnow
from app.infrastructure.db.columns import created, pk

#: The engine's default limits, mirroring the settings block in ``config.py``.
#: They are column defaults (not settings) so that a database written by an
#: older deployment still carries a sane value.
DEFAULT_MAX_CYCLES = 6


class TaskPipelineStateRow(Base):
    __tablename__ = "task_pipeline_state"

    task_id: Mapped[str] = mapped_column(
        ForeignKey("tasks.id", ondelete="CASCADE"), primary_key=True
    )
    #: Denormalised from the task so the tenant filter stays a single-table
    #: WHERE: the scheduler's hot query never joins back to ``tasks``.
    organization_id: Mapped[str] = mapped_column(
        ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False, index=True
    )
    project_id: Mapped[str | None] = mapped_column(
        ForeignKey("projects.id", ondelete="CASCADE"), nullable=True
    )
    pipeline: Mapped[str] = mapped_column(
        String(32), default="plan_code_test", nullable=False
    )
    kind: Mapped[str] = mapped_column(String(16), default="feature", nullable=False)
    risk: Mapped[str] = mapped_column(String(8), default="medium", nullable=False)
    #: Indexed: every scheduler tick asks "which tasks are at a turn stage".
    stage: Mapped[str] = mapped_column(
        String(32), default="pending_manager_review", nullable=False, index=True
    )
    cycle: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    max_cycles: Mapped[int] = mapped_column(
        Integer, default=DEFAULT_MAX_CYCLES, nullable=False
    )
    branch: Mapped[str] = mapped_column(String(120), default="", nullable=False)
    last_diff_hash: Mapped[str] = mapped_column(String(64), default="", nullable=False)
    same_diff_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    retry_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    #: A waived test pass must always be explainable: the reason is stored,
    #: never inferred (PLAN-v2 §3, green gate).
    waiver_reason: Mapped[str] = mapped_column(Text, default="", nullable=False)
    #: The latest accepted diff (the writer caps it). The worktree is gone after the
    #: merge, so the task page reads the diff from here.
    last_diff: Mapped[str] = mapped_column(Text, default="", nullable=False)
    #: The sha the merge landed on, once the state is COMPLETED.
    merged_sha: Mapped[str] = mapped_column(String(64), default="", nullable=False)
    parent_task_id: Mapped[str | None] = mapped_column(
        ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True
    )
    depends_on: Mapped[list] = mapped_column(JSON, default=list, nullable=False)
    #: Optimistic lock. Every stage move bumps it; a stale write is refused.
    version: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
    )


class EmployeeProfileRow(Base):
    __tablename__ = "employee_profiles"
    __table_args__ = (Index("ix_employee_profile_busy", "busy_task_id"),)

    employee_id: Mapped[str] = mapped_column(
        ForeignKey("employees.id", ondelete="CASCADE"), primary_key=True
    )
    #: A granted model id (``org_model_grants.ai_model_id``). Nullable until the
    #: owner grants something: an employee exists but is inactive, which is a
    #: state the UI shows instead of hiding the employee.
    ai_model_id: Mapped[str | None] = mapped_column(String(32), nullable=True)
    effort: Mapped[str] = mapped_column(String(16), default="medium", nullable=False)
    roles: Mapped[str] = mapped_column(String(255), default="", nullable=False)
    hired_by: Mapped[str] = mapped_column(String(8), default="user", nullable=False)
    busy_task_id: Mapped[str | None] = mapped_column(
        ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
    )


class EngineTurnRow(Base):
    __tablename__ = "engine_turns"
    __table_args__ = (
        Index("ix_engine_turn_task_stage", "task_id", "stage"),
        Index("ix_engine_turn_run", "run_id"),
    )

    id: Mapped[str] = pk()
    task_id: Mapped[str] = mapped_column(
        ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False
    )
    run_id: Mapped[str] = mapped_column(
        ForeignKey("runs.id", ondelete="CASCADE"), nullable=False
    )
    employee_id: Mapped[str | None] = mapped_column(
        ForeignKey("employees.id", ondelete="SET NULL"), nullable=True
    )
    role: Mapped[str] = mapped_column(String(24), default="developer", nullable=False)
    mode: Mapped[str] = mapped_column(String(16), default="work", nullable=False)
    stage: Mapped[str] = mapped_column(String(32), default="", nullable=False)
    prompt_chars: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    #: Bounded on write: a reply is unbounded, a row is not.
    reply_excerpt: Mapped[str] = mapped_column(Text, default="", nullable=False)
    tags: Mapped[dict] = mapped_column(JSON, default=dict, nullable=False)
    tool_calls: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    ok: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
    error_kind: Mapped[str] = mapped_column(String(16), default="", nullable=False)
    started_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=utcnow, nullable=False
    )
    finished_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    created_at: Mapped[datetime] = created()


class GateResultRow(Base):
    __tablename__ = "gate_results"
    __table_args__ = (Index("ix_gate_result_task_gate", "task_id", "gate"),)

    id: Mapped[str] = pk()
    task_id: Mapped[str] = mapped_column(
        ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False
    )
    gate: Mapped[str] = mapped_column(String(16), default="tests", nullable=False)
    ok: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
    detail: Mapped[str] = mapped_column(Text, default="", nullable=False)
    created_at: Mapped[datetime] = created()


__all__ = [
    "DEFAULT_MAX_CYCLES",
    "EmployeeProfileRow",
    "EngineTurnRow",
    "GateResultRow",
    "TaskPipelineStateRow",
]