"""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()