Disable physical control through legacy HTTP/M30 producers, typed admission, queue dispatch and upgraded Windows executors. Restrict diagnostic transports to exact read-only operations, with ordinal PowerShell validation before sockets. Retire unsupported pending/claimed work with atomic, idempotent security audits. Preserve prior delivery evidence and report unknown outcomes after a prior claim, including explicitly repaired pending rows. Include containment, cancellation, concurrency and malformed-payload regressions and the old-Agent cutover runbook. Validation: independent final application-security revalidation PASS. Two unchanged baseline source-assertion failures are documented as exceptions. Production release still requires native Windows validation and verified fencing/stop/drain/upgrade; server cancellation alone cannot revoke a cached command held by an old Agent.
323 lines
14 KiB
Python
323 lines
14 KiB
Python
"""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
|
|
|
|
|
|
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:
|
|
self._retire_unsupported(cur, agent_id)
|
|
cur.execute(CLAIM_COMMANDS_SQL, (agent_id, list(READ_ONLY_COMMAND_TYPES), limit))
|
|
rows = [dict(row) for row in cur.fetchall()]
|
|
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()
|