minerops/backend/app/agent_commands/repository.py

"""PostgreSQL persistence for Agent command transport."""

from collections.abc import Callable
from typing import Any

import psycopg
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from app.agent_commands.exceptions import AgentCommandConflict, AgentCommandNotFound
from app.agent_commands.policy import READ_ONLY_COMMAND_TYPES, rejection_execution_outcome, rejection_reason
from app.config import settings
from app.discovery.repository import DiscoveryRepository
from fastapi import HTTPException


ConnectionFactory = Callable[[], psycopg.Connection]

CLAIM_COMMANDS_SQL = """
WITH candidates AS (
    SELECT id
    FROM minerops_agent_command_queue
    WHERE agent_id = %s
      AND command_status = 'pending'
      AND command_type = ANY(%s)
    ORDER BY created_at, id
    FOR UPDATE SKIP LOCKED
    LIMIT %s
)
UPDATE minerops_agent_command_queue AS queue
SET command_status = 'claimed', claimed_at = now()
FROM candidates
WHERE queue.id = candidates.id
  AND queue.command_status = 'pending'
RETURNING queue.id AS command_id, queue.command_type,
          queue.request_payload, queue.created_at
"""

# Inspect existing durable history before changing status/completed_at/error.
# A linked non-creation audit is conservatively ambiguous, even if timestamps
# were lost. Legacy results need not agree on Agent ID to prevent a false claim
# of non-delivery. Missing tables/permissions fail the transaction closed.
NO_DELIVERY_EVIDENCE_SQL = """
    q.created_at IS NOT NULL
    AND q.claimed_at IS NULL AND q.completed_at IS NULL
    AND q.result_payload IS NULL AND COALESCE(q.error_message, '') = ''
    AND NOT EXISTS (
        SELECT 1 FROM minerops_agent_command_results r
        WHERE r.command_id = q.id
           OR (r.agent_id = q.agent_id AND NULLIF(btrim(r.command_id), '') IS NULL)
    )
    AND NOT EXISTS (
        SELECT 1 FROM audit_logs a
        WHERE a.action <> 'miner.control.queued'
          AND (a.old_value_json ->> 'command_id' = q.id
               OR a.new_value_json ->> 'command_id' = q.id
               OR a.old_value_json ->> 'operation_id' = q.id
               OR a.new_value_json ->> 'operation_id' = q.id
               OR (a.entity_type = 'agent_command' AND a.entity_id::text = q.id))
    )
"""

RETIRE_COMMANDS_SQL = f"""
WITH candidates AS (
    SELECT q.id, q.command_status, ({NO_DELIVERY_EVIDENCE_SQL}) AS no_delivery_evidence
    FROM minerops_agent_command_queue q
    WHERE agent_id = %s AND command_status IN ('pending', 'claimed')
      AND NOT (command_type = ANY(%s))
    ORDER BY created_at, id
    FOR UPDATE SKIP LOCKED
    LIMIT 1000
)
UPDATE minerops_agent_command_queue AS queue
SET command_status = 'cancelled', completed_at = COALESCE(queue.completed_at, now()),
    error_message = concat_ws(chr(10), NULLIF(queue.error_message, ''), CASE WHEN queue.command_type = 'miner_control'
        THEN 'physical_control_disabled_stage_a'
        ELSE 'unsupported_command_schema_stage_a' END)
FROM candidates WHERE queue.id = candidates.id
RETURNING queue.id, queue.agent_id, queue.command_type,
          candidates.command_status AS previous_status, candidates.no_delivery_evidence, queue.error_message
"""

RETIRE_SCHEDULE_COMMANDS_SQL = f"""
WITH candidates AS (
    SELECT q.id, q.command_status, ({NO_DELIVERY_EVIDENCE_SQL}) AS no_delivery_evidence
    FROM minerops_agent_command_queue q
    WHERE agent_id = %s AND request_payload ->> 'scheduled_action_id' = %s
      AND command_status IN ('pending', 'claimed')
      AND NOT (command_type = ANY(%s))
    ORDER BY id
    FOR UPDATE
)
UPDATE minerops_agent_command_queue AS queue
SET command_status = 'cancelled', completed_at = COALESCE(queue.completed_at, now()),
    error_message = concat_ws(chr(10), NULLIF(queue.error_message, ''), CASE WHEN queue.command_type = 'miner_control'
        THEN 'physical_control_disabled_stage_a'
        ELSE 'unsupported_command_schema_stage_a' END)
FROM candidates WHERE queue.id = candidates.id
RETURNING queue.id, queue.agent_id, queue.command_type,
          candidates.command_status AS previous_status, candidates.no_delivery_evidence
"""


def default_connection_factory() -> psycopg.Connection:
    return psycopg.connect(
        host=settings.postgres_host,
        port=settings.postgres_port,
        dbname=settings.postgres_db,
        user=settings.postgres_user,
        password=settings.postgres_password,
        connect_timeout=5,
    )


class AgentCommandRepository:
    """Own queue locking and immutable terminal transitions."""

    def __init__(self, connection_factory: ConnectionFactory | None = None):
        self._connection_factory = connection_factory or default_connection_factory

    def claim_pending(self, agent_id: str, limit: int = 5) -> list[dict]:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                DiscoveryRepository.lock(cur)
                self._retire_unsupported(cur, agent_id)
                discovery = DiscoveryRepository(self._connection_factory)
                cur.execute(f"""SELECT q.*, ({NO_DELIVERY_EVIDENCE_SQL}) AS no_delivery_evidence
                    FROM minerops_agent_command_queue q WHERE agent_id=%s
                    AND command_status='pending' AND command_type=ANY(%s)
                    ORDER BY created_at,id FOR UPDATE SKIP LOCKED LIMIT 1000""",
                    (agent_id,list(READ_ONLY_COMMAND_TYPES)))
                pending=list(cur.fetchall())
                for command in pending:
                    try:
                        discovery.validate_command(cur,agent_id,command)
                    except HTTPException as exc:
                        cur.execute("""UPDATE minerops_agent_command_queue SET command_status='cancelled',
                            completed_at=COALESCE(completed_at,now()), error_message=concat_ws(chr(10),NULLIF(error_message,''),%s::text)
                            WHERE id=%s""",(exc.detail,command['id']))
                        DiscoveryRepository.audit(cur,'discovery.command_rejected',None,
                            {'command_id':command['id'],'agent_id':agent_id,'reason':exc.detail,
                             'execution_outcome':rejection_execution_outcome({**command,'previous_status':'pending'})})
                cur.execute(CLAIM_COMMANDS_SQL, (agent_id, list(READ_ONLY_COMMAND_TYPES), limit))
                rows = [dict(row) for row in cur.fetchall()]
                for row in rows:
                    discovery.validate_command(cur,agent_id,row)
            connection.commit()
            return rows
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()

    @staticmethod
    def _audit_rejection(
        cur, row: dict, *, actor_user_id: str | None = None,
        source: str = "agent_poll",
    ) -> None:
        # Queue IDs can be legacy text, not UUIDs. Preserve them in JSON rather
        # than coercing them into audit_logs.entity_id. Never log raw payloads.
        cur.execute(
            """INSERT INTO audit_logs
               (actor_user_id, action, entity_type, old_value_json, new_value_json)
               VALUES (%s, 'agent.command.rejected', 'agent_command', %s, %s)""",
            (actor_user_id, Jsonb({"command_id": row["id"], "agent_id": row["agent_id"],
                    "status": row["previous_status"]}),
             Jsonb({"status": "cancelled", "reason": rejection_reason(row["command_type"]),
                    "source": source,
                    "execution_outcome": rejection_execution_outcome(row)})),
        )

    @classmethod
    def _retire_unsupported(
        cls, cur, agent_id: str, *, source: str = "agent_poll",
        actor_user_id: str | None = None,
    ) -> int:
        cur.execute(RETIRE_COMMANDS_SQL, (agent_id, list(READ_ONLY_COMMAND_TYPES)))
        rows = list(cur.fetchall())
        for row in rows:
            cls._audit_rejection(cur, row, source=source, actor_user_id=actor_user_id)
        return len(rows)

    def retire_unsupported(self, agent_id: str, *, actor_user_id: str) -> int:
        """Audited cutover batch, including offline Agents; never claim any work.

        A zero batch can mean locked rows were skipped. Operators must separately
        verify no unsupported live rows remain before declaring a drain complete.
        """
        if not actor_user_id:
            raise ValueError("Cutover retirement requires an accountable operator actor")
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                count = self._retire_unsupported(cur, agent_id, source="stage_a_cutover",
                                                 actor_user_id=actor_user_id)
            connection.commit()
            return count
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()

    def retire_schedule(self, schedule_id: str, *, actor_user_id: str, operation: str) -> int:
        """Cancel live unsupported work and retain schedule history atomically.

        The stored schedule determines Agent scope. Delete retires the schedule;
        neither operation erases evidence or changes terminal command history.
        """
        if not actor_user_id or operation not in {"cancel", "delete"}:
            raise ValueError("An authenticated actor and cancel/delete operation are required")
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                # Serialize retries and cancel/delete workers for this schedule.
                cur.execute(
                    """SELECT id, agent_id, status FROM minerops_scheduled_bulk_actions
                       WHERE id = %s FOR UPDATE""", (schedule_id,),
                )
                schedule = cur.fetchone()
                if not schedule:
                    raise AgentCommandNotFound("Scheduled action not found")
                cur.execute(RETIRE_SCHEDULE_COMMANDS_SQL,
                            (schedule["agent_id"], schedule_id, list(READ_ONLY_COMMAND_TYPES)))
                rows = list(cur.fetchall())
                for row in rows:
                    self._audit_rejection(cur, row, actor_user_id=actor_user_id,
                                          source=f"schedule_{operation}")
                if schedule["status"] in {"pending", "queued", "running"}:
                    cur.execute(
                        """UPDATE minerops_scheduled_bulk_actions
                           SET status = 'cancelled', updated_at = now(),
                               last_error = concat_ws(chr(10), NULLIF(last_error, ''),
                                                      'physical_control_disabled_stage_a')
                           WHERE id = %s""", (schedule_id,),
                    )
                    cur.execute(
                        """INSERT INTO audit_logs
                           (actor_user_id, action, entity_type, old_value_json, new_value_json)
                           VALUES (%s, 'scheduled_action.retired', 'scheduled_action', %s, %s)""",
                        (actor_user_id, Jsonb(dict(schedule)),
                         Jsonb({"status": "cancelled", "operation": operation,
                                "reason": "physical_control_disabled_stage_a"})),
                    )
            connection.commit()
            return len(rows)
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()

    def get_command(self, command_id: str) -> dict | None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT id, agent_id, command_type, command_status, request_payload
                    FROM minerops_agent_command_queue WHERE id = %s
                    """,
                    (command_id,),
                )
                row = cur.fetchone()
                return dict(row) if row else None
        finally:
            connection.close()

    def finish_command(
        self,
        command_id: str,
        agent_id: str,
        status: str,
        result_payload: dict[str, Any],
        error_message: str | None,
    ) -> None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                # Enforce the same boundary on results: an old claimed entry
                # cannot be promoted by an old Agent, retry, or direct caller.
                cur.execute(
                    f"""SELECT q.id, q.agent_id, q.command_type, q.command_status AS previous_status,
                              ({NO_DELIVERY_EVIDENCE_SQL}) AS no_delivery_evidence
                       FROM minerops_agent_command_queue q
                       WHERE q.id = %s AND q.agent_id = %s FOR UPDATE""",
                    (command_id, agent_id),
                )
                existing = cur.fetchone()
                if existing and rejection_reason(existing["command_type"]):
                    reason = rejection_reason(existing["command_type"])
                    if existing["previous_status"] in {"pending", "claimed"}:
                        cur.execute(
                            """UPDATE minerops_agent_command_queue
                               SET command_status = 'cancelled',
                                   error_message = concat_ws(chr(10), NULLIF(error_message, ''), %s::text),
                                   completed_at = COALESCE(completed_at, now()) WHERE id = %s""",
                            (reason, command_id),
                        )
                        # A result submission itself reports possible delivery, even
                        # when earlier durable claim metadata is absent.
                        existing["delivery_reported"] = True
                        self._audit_rejection(cur, existing, source="agent_result")
                    connection.commit()
                    raise AgentCommandConflict(reason)
                cur.execute(
                    """
                    UPDATE minerops_agent_command_queue
                    SET command_status = %s, result_payload = %s,
                        error_message = %s, completed_at = now()
                    WHERE id = %s AND agent_id = %s AND command_status = 'claimed'
                    RETURNING id
                    """,
                    (status, Jsonb(result_payload), error_message, command_id, agent_id),
                )
                if cur.fetchone():
                    connection.commit()
                    return

                cur.execute(
                    """
                    SELECT agent_id, command_status
                    FROM minerops_agent_command_queue WHERE id = %s
                    """,
                    (command_id,),
                )
                row = cur.fetchone()
                connection.rollback()
                if not row:
                    raise AgentCommandNotFound("Agent command not found")
                if row["agent_id"] != agent_id:
                    raise AgentCommandConflict("Agent command belongs to another Agent")
                raise AgentCommandConflict(
                    f"Agent command is {row['command_status']}, not claimed"
                )
        except (AgentCommandNotFound, AgentCommandConflict):
            raise
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()