cleveragents-core/tools/controller/db/models.py

"""SQLAlchemy 2.0 declarative models for the controller DB.

Schema follows plan v6/v9: unified ``workflows`` table with ``kind``
discriminator (pr/issue), workflow_attempts that doubles as the worker
queue (status + lock columns), controller_events for audit + Forgejo
replay, flake_history for CI flake learning, ci_observations for raw
CI state history.

Cross-dialect (SQLite + Postgres):
- ``JSON`` columns map to JSON on SQLite and JSONB on Postgres.
- ``Integer``/``BigInteger`` are dialect-portable; we use
  ``BigInteger`` for PKs that will grow (attempts) and ``Integer``
  for workflow ids (one-per-PR; bounded).
- Partial indexes use the SQLAlchemy ``Index(..., postgresql_where=)``
  + ``sqlite_where=`` pattern so both dialects honor them.
- ``CHECK`` constraints use SQLAlchemy ``CheckConstraint`` (both
  dialects).
- Per the v6 schema decision, the parent_workflow_id kind-validation
  is app-level (SQLite can't do subquery-CHECK; Postgres trigger
  defer to v6 plan body).

Tables NOT owned here
---------------------
``_mode_marker`` (one-row table: ``mode TEXT PRIMARY KEY, stamped_at TEXT``)
is created and maintained outside SQLAlchemy by
``tools/controller/deploy/validate_db_mode.py`` — the controller launcher
stamps it on first use and asserts it on subsequent launches so a fork-mode
db and a prod-mode db can never be confused. Ignored by every model +
query here; future migrations should leave it alone.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any

from sqlalchemy import (
    BigInteger,
    Boolean,
    CheckConstraint,
    DateTime,
    Float,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    UniqueConstraint,
    func,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    mapped_column,
    relationship,
)
from sqlalchemy.types import JSON


# Cross-dialect JSON: JSONB on Postgres, JSON on SQLite. JSONB-only
# features (GIN indexes, jsonb_path_ops) are not used; v1 stays
# portable.
JsonColumn = JSON().with_variant(JSONB(), "postgresql")

# Cross-dialect autoincrement PK: SQLite's INTEGER (4-byte rowid alias
# that auto-increments); Postgres's BIGSERIAL (8-byte). SQLAlchemy's
# BigInteger doesn't autoincrement on SQLite, so this variant gives us
# the right type on each dialect.
AutoincrementPk = Integer().with_variant(BigInteger(), "postgresql")


class Base(DeclarativeBase):
    """SQLAlchemy 2.0 declarative base for controller tables."""

    type_annotation_map = {dict[str, Any]: JsonColumn}


# ─── workflows ────────────────────────────────────────────────────────


class Workflow(Base):
    """One PR or issue under controller management.

    Plan v6 unified `workflows` table (replaces v5's separate
    pr_workflows / issue_workflows) with `kind` discriminator. An
    issue workflow can spawn a PR workflow via `parent_workflow_id`
    on the new row (CREATED_PR transition).
    """

    __tablename__ = "workflows"

    workflow_id: Mapped[int] = mapped_column(
        AutoincrementPk, primary_key=True, autoincrement=True
    )

    # Discriminator + identity
    kind: Mapped[str] = mapped_column(String(8), nullable=False)  # 'pr' | 'issue'
    owner: Mapped[str] = mapped_column(String(128), nullable=False)
    repo: Mapped[str] = mapped_column(String(128), nullable=False)
    entity_number: Mapped[int] = mapped_column(Integer, nullable=False)
    parent_workflow_id: Mapped[int | None] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="RESTRICT"),
        nullable=True,
    )

    # State machine
    current_state: Mapped[str] = mapped_column(String(32), nullable=False)
    current_tier: Mapped[int | None] = mapped_column(Integer, nullable=True)
    tier_last_succeeded: Mapped[int | None] = mapped_column(Integer, nullable=True)

    # Timestamps
    started_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    last_transition_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    entered_state_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )

    # Limits + workspace
    max_attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=6)
    workspace_dir: Mapped[str | None] = mapped_column(Text, nullable=True)

    # MERGING non-blocking retries
    merging_retry_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    merging_retry_next_attempt_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )

    # Phase 1k+: pre-pause state for label-gate PAUSE/RESUME. When the
    # operator removes CONTROLLER_OPT_IN_LABEL, the workflow's current
    # state moves to PAUSED and the prior state is captured here. On
    # label re-add, the master reads this column to determine the
    # resume target. NULL means "never paused" or "resumed".
    pre_pause_state: Mapped[str | None] = mapped_column(String(32), nullable=True)

    # Phase 0 (grooming plan): one-shot semantics + defer block.
    #
    # ``grooming_evaluated_at`` is set inside the DB transaction at the
    # start of any Gate 1 action (proceed, defer, or close). Once set,
    # Gate 1 short-circuits to PROCEED on subsequent ticks — the gate
    # is one-shot per workflow. Cleared only by explicit operator action.
    grooming_evaluated_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    # ``deferred_reason`` is the defer block: scheduler skips any workflow
    # with deferred_reason IS NOT NULL even if auto/sentinel is re-added.
    # 'duplication' is the v1 value; reason-tagged for future expansion.
    deferred_reason: Mapped[str | None] = mapped_column(String(32), nullable=True)
    deferred_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    # Canonical's workflow_id — humans + the future scope-evaluator
    # pipeline navigate "who does this duplicate?" via this pointer.
    deferred_target_workflow_id: Mapped[int | None] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="SET NULL"),
        nullable=True,
    )

    # Phase 1k++ (R3): the ``ci_flake_retries_remaining`` and
    # ``awaiting_ci_started_at`` columns added in Phase 1k were never
    # actually read or written by any producer — the round-3 review
    # caught them as dead schema weight. ``ci_poll.py`` uses
    # ``entered_state_at`` (already populated on every transition) as
    # the AWAITING_CI start timestamp; the flake-retries counter is a
    # future-phase item that should ship with its producer, not as
    # speculative schema.

    # Relationships
    attempts: Mapped[list[WorkflowAttempt]] = relationship(
        back_populates="workflow", cascade="all, delete-orphan"
    )

    __table_args__ = (
        CheckConstraint(
            "kind IN ('pr', 'issue')",
            name="ck_workflows_kind_valid",
        ),
        UniqueConstraint(
            "owner",
            "repo",
            "kind",
            "entity_number",
            name="uq_workflows_entity",
        ),
        # Filtered indexes — fast worklist queries skip terminal rows.
        Index(
            "ix_workflows_active",
            "current_state",
            postgresql_where=text(
                "current_state NOT IN ('MERGED', 'ABANDONED', 'STUCK', 'CREATED_PR')"
            ),
            sqlite_where=text(
                "current_state NOT IN ('MERGED', 'ABANDONED', 'STUCK', 'CREATED_PR')"
            ),
        ),
        Index(
            "ix_workflows_merging_retry",
            "merging_retry_next_attempt_at",
            postgresql_where=text("current_state = 'MERGING'"),
            sqlite_where=text("current_state = 'MERGING'"),
        ),
    )


# ─── workflow_attempts (IS the worker queue) ──────────────────────────


class WorkflowAttempt(Base):
    """One attempt by one role (estimator / implementer / reviewer /
    conflict_resolver / summarizer) against one workflow.

    Doubles as the worker queue: workers dequeue rows where
    ``status='pending'`` via ``dequeue.py`` (Postgres FOR UPDATE
    SKIP LOCKED, SQLite BEGIN IMMEDIATE). Lock columns (status,
    locked_by_instance, locked_at, lock_heartbeat_at, lock_ttl_seconds,
    pickup_count) implement the v9 multi-machine-safe lock protocol.
    """

    __tablename__ = "workflow_attempts"

    attempt_id: Mapped[int] = mapped_column(
        AutoincrementPk, primary_key=True, autoincrement=True
    )
    workflow_id: Mapped[int] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="CASCADE"),
        nullable=False,
    )
    attempt_number: Mapped[int] = mapped_column(Integer, nullable=False)
    role: Mapped[str] = mapped_column(String(32), nullable=False)
    tier: Mapped[int | None] = mapped_column(Integer, nullable=True)

    # v5 queue + lock columns
    status: Mapped[str] = mapped_column(String(16), nullable=False, default="pending")
    locked_by_instance: Mapped[str | None] = mapped_column(String(256), nullable=True)
    locked_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    lock_heartbeat_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    lock_ttl_seconds: Mapped[int] = mapped_column(Integer, nullable=False, default=600)
    pickup_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    started_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    finished_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    wallclock_seconds: Mapped[float | None] = mapped_column(Float, nullable=True)

    # SHA bookkeeping
    head_sha_before: Mapped[str | None] = mapped_column(String(64), nullable=True)
    head_sha_after: Mapped[str | None] = mapped_column(String(64), nullable=True)

    # Payloads (JSONB on Postgres, JSON on SQLite). Application caps
    # input_payload at 4MB (see payload_guard.py).
    input_payload: Mapped[dict[str, Any]] = mapped_column(JsonColumn, nullable=False)
    input_version: Mapped[str] = mapped_column(String(8), nullable=False)
    input_payload_truncated: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=False
    )
    output_payload: Mapped[dict[str, Any] | None] = mapped_column(
        JsonColumn, nullable=True
    )
    output_version: Mapped[str | None] = mapped_column(String(8), nullable=True)
    outcome: Mapped[str | None] = mapped_column(String(64), nullable=True)

    # Worker metadata
    session_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
    worker_subprocess_pid: Mapped[int | None] = mapped_column(Integer, nullable=True)
    strict_parse_retries: Mapped[int] = mapped_column(
        Integer, nullable=False, default=0
    )

    workflow: Mapped[Workflow] = relationship(back_populates="attempts")

    __table_args__ = (
        CheckConstraint(
            "status IN ('pending', 'in_progress', 'complete', 'failed', 'reaped')",
            name="ck_attempts_status_valid",
        ),
        CheckConstraint("pickup_count >= 0", name="ck_attempts_pickup_count_nonneg"),
        # Worker-dequeue index: hot path of every worker tick.
        Index(
            "ix_attempts_pending_by_role",
            "role",
            "created_at",
            postgresql_where=text("status = 'pending'"),
            sqlite_where=text("status = 'pending'"),
        ),
        # Reaper index: master tick scans this every 60s.
        Index(
            "ix_attempts_in_progress_by_heartbeat",
            "lock_heartbeat_at",
            postgresql_where=text("status = 'in_progress'"),
            sqlite_where=text("status = 'in_progress'"),
        ),
        # Master driver: scans complete attempts to advance state machine.
        Index(
            "ix_attempts_complete_by_workflow",
            "workflow_id",
            "attempt_number",
            postgresql_where=text("status IN ('complete', 'failed')"),
            sqlite_where=text("status IN ('complete', 'failed')"),
        ),
    )


# ─── controller_events ────────────────────────────────────────────────


class ControllerEvent(Base):
    """Audit trail + Forgejo-write-replay log.

    Every state transition is a row. Forgejo writes (status comments,
    labels, merges) live here so a crash between DB commit + Forgejo
    write can be replayed on next master tick (v9 simplified protocol:
    Forgejo-first, DB-second, reconciliation syncs DB from Forgejo).

    For v1 ``forgejo_write_pending`` is unused (the v9 simplification
    dropped the 3-step protocol); kept in the schema so a future
    upgrade can re-enable it without a migration.
    """

    __tablename__ = "controller_events"

    event_id: Mapped[int] = mapped_column(
        AutoincrementPk, primary_key=True, autoincrement=True
    )
    workflow_id: Mapped[int] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="CASCADE"),
        nullable=False,
    )
    ts: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    event_type: Mapped[str] = mapped_column(String(64), nullable=False)
    from_state: Mapped[str | None] = mapped_column(String(32), nullable=True)
    to_state: Mapped[str | None] = mapped_column(String(32), nullable=True)
    attempt_id: Mapped[int | None] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflow_attempts.attempt_id", ondelete="SET NULL"),
        nullable=True,
    )
    payload: Mapped[dict[str, Any] | None] = mapped_column(JsonColumn, nullable=True)

    # v3 write protocol fields (kept for v2-style upgrade path)
    forgejo_write_pending: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=False
    )
    forgejo_fingerprint: Mapped[str | None] = mapped_column(String(128), nullable=True)
    forgejo_result: Mapped[dict[str, Any] | None] = mapped_column(
        JsonColumn, nullable=True
    )
    replay_attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

    # Phase 0 (grooming plan): action-attribution discriminator. NULL for
    # events whose event_type already uniquely identifies the trigger;
    # populated for events like ``label-pause`` that can come from either
    # reconciliation (human pulled the label) or defer (controller-driven
    # pause). Valid values come from ``contracts.causes.Cause``.
    cause: Mapped[str | None] = mapped_column(String(32), nullable=True)

    __table_args__ = (
        Index(
            "ix_events_pending_forgejo",
            "workflow_id",
            "ts",
            postgresql_where=text("forgejo_write_pending = true"),
            sqlite_where=text("forgejo_write_pending = 1"),
        ),
        Index("ix_events_by_workflow_ts", "workflow_id", "ts"),
    )


# ─── flake_history ────────────────────────────────────────────────────


class FlakeHistory(Base):
    """CI flake learning — per-(owner, repo, context) flake-score
    accumulator. Plan v6: ship empty + learning heuristic; bootstrap
    via first-failure mandatory retry. flake_score = flake_count /
    observations; threshold 0.6.
    """

    __tablename__ = "flake_history"

    owner: Mapped[str] = mapped_column(String(128), primary_key=True)
    repo: Mapped[str] = mapped_column(String(128), primary_key=True)
    context: Mapped[str] = mapped_column(String(256), primary_key=True)

    observations: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    flake_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    flake_score: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)
    last_seen: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    operator_override: Mapped[bool | None] = mapped_column(Boolean, nullable=True)


# ─── ci_observations ──────────────────────────────────────────────────


class CIObservation(Base):
    """Raw CI status observed for a head_sha at a point in time.

    Append-only; used for parser-rot detection + flake history input
    + post-mortem debugging. Cleaned up via periodic retention sweep
    (90 days).
    """

    __tablename__ = "ci_observations"

    observation_id: Mapped[int] = mapped_column(
        AutoincrementPk, primary_key=True, autoincrement=True
    )
    head_sha: Mapped[str] = mapped_column(String(64), nullable=False)
    observed_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    state: Mapped[str] = mapped_column(String(16), nullable=False)
    contexts: Mapped[dict[str, Any]] = mapped_column(JsonColumn, nullable=False)

    __table_args__ = (Index("ix_ci_obs_by_sha", "head_sha", "observed_at"),)


# ─── grooming_decisions (Phase 0 grooming plan audit table) ──────────


class GroomingDecision(Base):
    """One audit row per grooming-gate evaluation.

    The Phase 0 ``close_issue`` / ``defer_issue`` callbacks INSERT here
    inside their crash-safe transaction; the resulting ``decision_id``
    is substituted into the audit comment posted to Forgejo so the
    comment links back to the row.

    Phase 0 callbacks populate only the NOT-NULL columns (workflow_id,
    decided_at, check_name, stage, verdict, reason_category, executed)
    plus optional ``target_workflow_id`` and ``forgejo_response``.
    Phase 1's Gate 1 (deterministic + LLM) populates the additional
    fields (action, confidence, llm_reasoning, preserved_value,
    loser_head_sha_at_decision, suspicion_score, forced_proceed_reason)
    as it produces them.
    """

    __tablename__ = "grooming_decisions"

    decision_id: Mapped[int] = mapped_column(
        AutoincrementPk, primary_key=True, autoincrement=True
    )
    workflow_id: Mapped[int] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="CASCADE"),
        nullable=False,
    )
    decided_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    # 'duplicate_open_pr' / 'linked_issue_closed' / 'base_branch_deleted' / ...
    check_name: Mapped[str] = mapped_column(String(64), nullable=False)
    # 'deterministic_conclusive' | 'stage_b_llm'
    stage: Mapped[str] = mapped_column(String(32), nullable=False)
    # 'proceed' | 'defer' | 'close'
    verdict: Mapped[str] = mapped_column(String(16), nullable=False)
    # one of the abandon-reason categories
    reason_category: Mapped[str] = mapped_column(String(64), nullable=False)
    # Stage B only — the LLM's per-duplicate verdict
    action: Mapped[str | None] = mapped_column(String(32), nullable=True)
    target_workflow_id: Mapped[int | None] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="SET NULL"),
        nullable=True,
    )
    confidence: Mapped[str | None] = mapped_column(String(16), nullable=True)
    llm_reasoning: Mapped[str | None] = mapped_column(Text, nullable=True)
    # preserved_value_summary when action='needs_evaluation'
    preserved_value: Mapped[str | None] = mapped_column(Text, nullable=True)
    # Snapshot of THIS row's workflow_id's PR head SHA at decision
    # time. Unused by v1; available to Phase 6+ scope-evaluator for
    # deep-diff comparison against canonical. Known limitation: SHA
    # is captured at decision time, before the Forgejo close PATCH —
    # if the loser's branch advances between snapshot and close (rare),
    # the recorded SHA is stale by one commit. See decision #34.
    loser_head_sha_at_decision: Mapped[str | None] = mapped_column(
        String(64), nullable=True
    )
    # Deterministic Stage A suspicion score that brought this pair to
    # Stage B; NULL for deterministic-conclusive rows. Drives
    # /api/grooming/llm_agreement_rate telemetry.
    suspicion_score: Mapped[float | None] = mapped_column(Float, nullable=True)
    # Non-NULL when verdict='proceed' was FORCED (not the LLM's
    # legitimate verdict): 'semantic_contradiction' | 'low_confidence'.
    # Separates "LLM said no" from "LLM emitted garbage" for telemetry
    # (/api/grooming/semantic_rejection_rate). See decision #33.
    forced_proceed_reason: Mapped[str | None] = mapped_column(String(32), nullable=True)
    # 0 in dry-run mode (audit-only); 1 once the Forgejo write actually
    # happened.
    executed: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    # Final Forgejo response shape: HTTP status + body summary. JSON
    # dict; populated after the Forgejo write completes (or fails).
    forgejo_response: Mapped[dict[str, Any] | None] = mapped_column(
        JsonColumn, nullable=True
    )

    __table_args__ = (
        Index("idx_grooming_decisions_target", "target_workflow_id"),
        Index("idx_grooming_decisions_verdict", "verdict"),
    )