feat: add agent zones
This commit is contained in:
lussierj committed
1 parent
d8cd7ad058
commit
82309db79a
29 files changed
+2242
-255
No files matched your search
@@ -0,0 +1 @@
|
||||
"""Agent Zone routing domain."""
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Agent Zone dependency wiring."""
|
||||
|
||||
from app.agent_zones.repository import AgentZoneRepository
|
||||
from app.agent_zones.service import AgentZoneService
|
||||
|
||||
|
||||
def get_agent_zone_service() -> AgentZoneService:
|
||||
return AgentZoneService(AgentZoneRepository())
|
||||
@@ -0,0 +1,17 @@
|
||||
"""Agent Zone domain exceptions."""
|
||||
|
||||
|
||||
class AgentZoneError(Exception):
|
||||
"""Base Agent Zone error."""
|
||||
|
||||
|
||||
class AgentZoneNotFound(AgentZoneError):
|
||||
"""A Site or Agent Zone does not exist."""
|
||||
|
||||
|
||||
class AgentZoneValidationError(AgentZoneError):
|
||||
"""Agent Zone input or assignment is invalid."""
|
||||
|
||||
|
||||
class AgentZoneConflict(AgentZoneError):
|
||||
"""An Agent Zone uniqueness or assignment invariant conflicts."""
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Typed Agent Zone and check-in contracts."""
|
||||
|
||||
from ipaddress import ip_network
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
|
||||
class AgentZoneCreateRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
site_id: UUID
|
||||
name: str = Field(min_length=1, max_length=120)
|
||||
description: str | None = Field(default=None, max_length=1000)
|
||||
networks: list[str] = Field(default_factory=list, max_length=64)
|
||||
|
||||
@field_validator("name")
|
||||
@classmethod
|
||||
def normalize_name(cls, value: str) -> str:
|
||||
normalized = value.strip()
|
||||
if not normalized:
|
||||
raise ValueError("name cannot be blank")
|
||||
return normalized
|
||||
|
||||
@field_validator("description")
|
||||
@classmethod
|
||||
def normalize_description(cls, value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
return value.strip() or None
|
||||
|
||||
@field_validator("networks")
|
||||
@classmethod
|
||||
def normalize_networks(cls, values: list[str]) -> list[str]:
|
||||
normalized: list[str] = []
|
||||
for value in values:
|
||||
try:
|
||||
network = str(ip_network(value.strip(), strict=True))
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"invalid CIDR network: {value}") from exc
|
||||
if network in normalized:
|
||||
raise ValueError(f"duplicate CIDR network: {network}")
|
||||
normalized.append(network)
|
||||
return normalized
|
||||
|
||||
|
||||
class AgentZoneNetwork(BaseModel):
|
||||
id: UUID
|
||||
network: str
|
||||
|
||||
|
||||
class AgentZoneSummary(BaseModel):
|
||||
id: UUID
|
||||
site_id: UUID
|
||||
name: str
|
||||
description: str | None = None
|
||||
is_enabled: bool
|
||||
assigned_agent_id: str | None = None
|
||||
networks: list[AgentZoneNetwork] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentZoneListResponse(BaseModel):
|
||||
data: list[AgentZoneSummary]
|
||||
meta: dict[str, object] = Field(default_factory=dict)
|
||||
error: None = None
|
||||
|
||||
|
||||
class AgentZoneCreateResponse(BaseModel):
|
||||
data: AgentZoneSummary
|
||||
meta: dict[str, object] = Field(default_factory=dict)
|
||||
error: None = None
|
||||
|
||||
|
||||
class AgentCheckinRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
agent_id: str | None = Field(default=None, max_length=255)
|
||||
site_id: UUID | None = None
|
||||
agent_zone_id: UUID | None = None
|
||||
site_name: str | None = Field(default=None, max_length=255)
|
||||
hostname: str | None = Field(default=None, max_length=255)
|
||||
version: str | None = Field(default=None, max_length=120)
|
||||
ip_address: str | None = Field(default=None, max_length=255)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def require_site_for_explicit_zone(self) -> "AgentCheckinRequest":
|
||||
if self.agent_zone_id is not None and self.site_id is None:
|
||||
raise ValueError("site_id is required when agent_zone_id is supplied")
|
||||
return self
|
||||
|
||||
|
||||
class AgentCheckinResponse(BaseModel):
|
||||
ok: bool = True
|
||||
hub_mode: bool = True
|
||||
agent_id: str
|
||||
site_id: UUID | None = None
|
||||
agent_zone_id: UUID | None = None
|
||||
routing_available: bool
|
||||
routing_reason: str | None = None
|
||||
site_name: str
|
||||
message: str = "Agent check-in received."
|
||||
|
||||
|
||||
class AgentZoneRoutingContext(BaseModel):
|
||||
site_id: UUID | None = None
|
||||
agent_zone_id: UUID | None = None
|
||||
routing_reason: str | None = None
|
||||
|
||||
|
||||
class MinerZoneAssociationResult(BaseModel):
|
||||
agent_zone_id: UUID | None = None
|
||||
routing_reason: str | None = None
|
||||
outcomes: dict[UUID, str] = Field(default_factory=dict)
|
||||
@@ -0,0 +1,300 @@
|
||||
"""PostgreSQL persistence for Agent Zones and ownership assignments."""
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
import psycopg
|
||||
from psycopg.errors import ForeignKeyViolation, UniqueViolation
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
from app.agent_zones.exceptions import (
|
||||
AgentZoneConflict,
|
||||
AgentZoneNotFound,
|
||||
)
|
||||
from app.config import settings
|
||||
|
||||
|
||||
ConnectionFactory = Callable[[], psycopg.Connection]
|
||||
|
||||
|
||||
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 AgentZoneRepository:
|
||||
def __init__(self, connection_factory: ConnectionFactory | None = None):
|
||||
self._connection_factory = connection_factory or default_connection_factory
|
||||
|
||||
def get_agent(self, agent_id: str) -> dict | None:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT agent_id, site_id, agent_zone_id, site_name, hostname,
|
||||
version, ip_address, last_checkin_at, last_payload
|
||||
FROM minerops_site_agents
|
||||
WHERE agent_id = %s
|
||||
""",
|
||||
(agent_id,),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def get_zone(self, zone_id: UUID) -> dict | None:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT id, site_id, name, description, is_enabled,
|
||||
created_at, updated_at
|
||||
FROM agent_zones
|
||||
WHERE id = %s
|
||||
""",
|
||||
(zone_id,),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def list_usable_zones(self, site_id: UUID) -> list[dict]:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT id, site_id, name, description, is_enabled
|
||||
FROM agent_zones
|
||||
WHERE site_id = %s AND is_enabled
|
||||
ORDER BY lower(name), id
|
||||
""",
|
||||
(site_id,),
|
||||
)
|
||||
return [dict(row) for row in cur.fetchall()]
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def list_zones(self, site_id: UUID | None = None) -> list[dict]:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT zone.id, zone.site_id, zone.name, zone.description,
|
||||
zone.is_enabled, agent.agent_id AS assigned_agent_id,
|
||||
range.id AS range_id, range.network::text AS network
|
||||
FROM agent_zones AS zone
|
||||
LEFT JOIN minerops_site_agents AS agent
|
||||
ON agent.agent_zone_id = zone.id
|
||||
LEFT JOIN agent_zone_network_ranges AS range
|
||||
ON range.agent_zone_id = zone.id
|
||||
WHERE (%s::uuid IS NULL OR zone.site_id = %s::uuid)
|
||||
ORDER BY lower(zone.name), zone.id, range.network, range.id
|
||||
""",
|
||||
(site_id, site_id),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
zones: dict[UUID, dict] = {}
|
||||
for row in rows:
|
||||
zone = zones.setdefault(
|
||||
row["id"],
|
||||
{
|
||||
"id": row["id"],
|
||||
"site_id": row["site_id"],
|
||||
"name": row["name"],
|
||||
"description": row["description"],
|
||||
"is_enabled": row["is_enabled"],
|
||||
"assigned_agent_id": row["assigned_agent_id"],
|
||||
"networks": [],
|
||||
},
|
||||
)
|
||||
if row["range_id"] is not None:
|
||||
zone["networks"].append(
|
||||
{"id": row["range_id"], "network": row["network"]}
|
||||
)
|
||||
return list(zones.values())
|
||||
|
||||
def create_zone(
|
||||
self,
|
||||
*,
|
||||
site_id: UUID,
|
||||
name: str,
|
||||
description: str | None,
|
||||
networks: list[str],
|
||||
) -> UUID:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_zones (site_id, name, description)
|
||||
VALUES (%s, %s, %s)
|
||||
RETURNING id
|
||||
""",
|
||||
(site_id, name, description),
|
||||
)
|
||||
zone_id = cur.fetchone()["id"]
|
||||
for network in networks:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_zone_network_ranges (agent_zone_id, network)
|
||||
VALUES (%s, %s::cidr)
|
||||
""",
|
||||
(zone_id, network),
|
||||
)
|
||||
connection.commit()
|
||||
return zone_id
|
||||
except ForeignKeyViolation as exc:
|
||||
connection.rollback()
|
||||
raise AgentZoneNotFound("Site not found") from exc
|
||||
except UniqueViolation as exc:
|
||||
connection.rollback()
|
||||
raise AgentZoneConflict(
|
||||
"An enabled Agent Zone with this name already exists at the Site"
|
||||
) from exc
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def save_checkin(
|
||||
self,
|
||||
*,
|
||||
agent_id: str,
|
||||
site_id: UUID | None,
|
||||
agent_zone_id: UUID | None,
|
||||
site_name: str,
|
||||
hostname: str | None,
|
||||
version: str | None,
|
||||
ip_address: str | None,
|
||||
payload: dict[str, Any],
|
||||
) -> None:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO minerops_site_agents (
|
||||
agent_id, site_id, agent_zone_id, site_name, hostname,
|
||||
version, ip_address, last_checkin_at, last_payload
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, now(), %s)
|
||||
ON CONFLICT (agent_id) DO UPDATE SET
|
||||
site_id = EXCLUDED.site_id,
|
||||
agent_zone_id = EXCLUDED.agent_zone_id,
|
||||
site_name = EXCLUDED.site_name,
|
||||
hostname = EXCLUDED.hostname,
|
||||
version = EXCLUDED.version,
|
||||
ip_address = EXCLUDED.ip_address,
|
||||
last_checkin_at = now(),
|
||||
last_payload = EXCLUDED.last_payload
|
||||
""",
|
||||
(
|
||||
agent_id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
site_name,
|
||||
hostname,
|
||||
version,
|
||||
ip_address,
|
||||
Jsonb(payload),
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
except UniqueViolation as exc:
|
||||
connection.rollback()
|
||||
raise AgentZoneConflict(
|
||||
"The Agent Zone already has an assigned Agent"
|
||||
) from exc
|
||||
except (ForeignKeyViolation, psycopg.errors.CheckViolation) as exc:
|
||||
connection.rollback()
|
||||
raise AgentZoneConflict("Agent Zone and Site assignment conflict") from exc
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def get_agent_assignment(self, agent_id: str) -> dict | None:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT agent.agent_id, agent.site_id, agent.agent_zone_id,
|
||||
zone.is_enabled AS zone_enabled
|
||||
FROM minerops_site_agents AS agent
|
||||
LEFT JOIN agent_zones AS zone ON zone.id = agent.agent_zone_id
|
||||
WHERE agent.agent_id = %s
|
||||
""",
|
||||
(agent_id,),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def assign_miners_to_zone(
|
||||
self, miner_ids: Iterable[UUID], agent_zone_id: UUID
|
||||
) -> dict[UUID, str]:
|
||||
unique_ids = list(dict.fromkeys(miner_ids))
|
||||
if not unique_ids:
|
||||
return {}
|
||||
connection = self._connection_factory()
|
||||
outcomes: dict[UUID, str] = {}
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"SELECT id, site_id FROM agent_zones WHERE id = %s",
|
||||
(agent_zone_id,),
|
||||
)
|
||||
zone = cur.fetchone()
|
||||
if not zone:
|
||||
raise AgentZoneNotFound("Agent Zone not found")
|
||||
for miner_id in unique_ids:
|
||||
cur.execute(
|
||||
"SELECT id, site_id, agent_zone_id FROM miners WHERE id = %s FOR UPDATE",
|
||||
(miner_id,),
|
||||
)
|
||||
miner = cur.fetchone()
|
||||
if not miner:
|
||||
outcomes[miner_id] = "miner_not_found"
|
||||
elif miner["agent_zone_id"] == agent_zone_id:
|
||||
outcomes[miner_id] = "preserved"
|
||||
elif miner["agent_zone_id"] is not None:
|
||||
outcomes[miner_id] = "agent_zone_conflict"
|
||||
elif miner["site_id"] not in {None, zone["site_id"]}:
|
||||
outcomes[miner_id] = "site_conflict"
|
||||
else:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE miners
|
||||
SET site_id = %s, agent_zone_id = %s, updated_at = now()
|
||||
WHERE id = %s AND agent_zone_id IS NULL
|
||||
""",
|
||||
(zone["site_id"], agent_zone_id, miner_id),
|
||||
)
|
||||
outcomes[miner_id] = "assigned"
|
||||
connection.commit()
|
||||
return outcomes
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
@@ -0,0 +1,75 @@
|
||||
"""HTTP mapping for Agent Zones and Agent check-in."""
|
||||
|
||||
from typing import Annotated
|
||||
from uuid import UUID
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, status
|
||||
|
||||
from app.agent_commands.dependencies import require_agent_identity
|
||||
from app.agent_zones.dependencies import get_agent_zone_service
|
||||
from app.agent_zones.exceptions import (
|
||||
AgentZoneConflict,
|
||||
AgentZoneNotFound,
|
||||
AgentZoneValidationError,
|
||||
)
|
||||
from app.agent_zones.models import (
|
||||
AgentCheckinRequest,
|
||||
AgentCheckinResponse,
|
||||
AgentZoneCreateRequest,
|
||||
AgentZoneCreateResponse,
|
||||
AgentZoneListResponse,
|
||||
)
|
||||
from app.agent_zones.service import AgentZoneService
|
||||
from app.identity import require_permission
|
||||
|
||||
|
||||
router = APIRouter(tags=["agent-zones"])
|
||||
ZoneService = Annotated[AgentZoneService, Depends(get_agent_zone_service)]
|
||||
AgentIdentity = Annotated[str, Depends(require_agent_identity)]
|
||||
ZoneViewer = Annotated[dict, Depends(require_permission("agents.view"))]
|
||||
ZoneManager = Annotated[dict, Depends(require_permission("agents.manage"))]
|
||||
|
||||
|
||||
@router.get("/api/agent-zones", response_model=AgentZoneListResponse)
|
||||
def list_agent_zones(
|
||||
current_user: ZoneViewer,
|
||||
service: ZoneService,
|
||||
site_id: UUID | None = Query(default=None),
|
||||
) -> AgentZoneListResponse:
|
||||
del current_user
|
||||
return AgentZoneListResponse(data=service.list_zones(site_id))
|
||||
|
||||
|
||||
@router.post(
|
||||
"/api/agent-zones",
|
||||
response_model=AgentZoneCreateResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
def create_agent_zone(
|
||||
request: AgentZoneCreateRequest,
|
||||
current_user: ZoneManager,
|
||||
service: ZoneService,
|
||||
) -> AgentZoneCreateResponse:
|
||||
del current_user
|
||||
try:
|
||||
return AgentZoneCreateResponse(data=service.create_zone(request))
|
||||
except AgentZoneNotFound as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
except AgentZoneConflict as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
|
||||
|
||||
@router.post("/api/agent/checkin", response_model=AgentCheckinResponse)
|
||||
def agent_checkin(
|
||||
request: AgentCheckinRequest,
|
||||
agent_id: AgentIdentity,
|
||||
service: ZoneService,
|
||||
) -> AgentCheckinResponse:
|
||||
try:
|
||||
return service.checkin(agent_id, request)
|
||||
except AgentZoneNotFound as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
except AgentZoneConflict as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
except AgentZoneValidationError as exc:
|
||||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
@@ -0,0 +1,194 @@
|
||||
"""Business rules for Agent Zone routing and transitional check-in."""
|
||||
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from app.agent_zones.exceptions import (
|
||||
AgentZoneNotFound,
|
||||
AgentZoneValidationError,
|
||||
)
|
||||
from app.agent_zones.models import (
|
||||
AgentCheckinRequest,
|
||||
AgentCheckinResponse,
|
||||
AgentZoneRoutingContext,
|
||||
AgentZoneCreateRequest,
|
||||
AgentZoneSummary,
|
||||
MinerZoneAssociationResult,
|
||||
)
|
||||
from app.agent_zones.repository import AgentZoneRepository
|
||||
|
||||
|
||||
SENSITIVE_KEYS = {
|
||||
"agent_secret",
|
||||
"api3_account",
|
||||
"api3_password",
|
||||
"password",
|
||||
"salt",
|
||||
"token",
|
||||
}
|
||||
|
||||
|
||||
def _sanitize(value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
key: _sanitize(item)
|
||||
for key, item in value.items()
|
||||
if key.lower() not in SENSITIVE_KEYS
|
||||
and not any(
|
||||
marker in key.lower()
|
||||
for marker in ("secret", "password", "token", "salt")
|
||||
)
|
||||
}
|
||||
if isinstance(value, list):
|
||||
return [_sanitize(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
class AgentZoneService:
|
||||
@staticmethod
|
||||
def validate_agent_identity(
|
||||
authenticated_agent_id: str, claimed_agent_id: str | None
|
||||
) -> None:
|
||||
if claimed_agent_id and claimed_agent_id.strip() != authenticated_agent_id:
|
||||
raise AgentZoneValidationError(
|
||||
"Payload agent_id does not match authenticated Agent identity"
|
||||
)
|
||||
|
||||
def __init__(self, repository: AgentZoneRepository):
|
||||
self.repository = repository
|
||||
|
||||
def list_zones(self, site_id: UUID | None = None) -> list[AgentZoneSummary]:
|
||||
return [
|
||||
AgentZoneSummary.model_validate(row)
|
||||
for row in self.repository.list_zones(site_id)
|
||||
]
|
||||
|
||||
def create_zone(self, request: AgentZoneCreateRequest) -> AgentZoneSummary:
|
||||
zone_id = self.repository.create_zone(
|
||||
site_id=request.site_id,
|
||||
name=request.name,
|
||||
description=request.description,
|
||||
networks=request.networks,
|
||||
)
|
||||
for zone in self.list_zones(request.site_id):
|
||||
if zone.id == zone_id:
|
||||
return zone
|
||||
raise AgentZoneNotFound("Created Agent Zone was not found")
|
||||
|
||||
def checkin(
|
||||
self, authenticated_agent_id: str, request: AgentCheckinRequest
|
||||
) -> AgentCheckinResponse:
|
||||
self.validate_agent_identity(authenticated_agent_id, request.agent_id)
|
||||
|
||||
existing = self.repository.get_agent(authenticated_agent_id)
|
||||
requested_site = request.site_id
|
||||
requested_zone = request.agent_zone_id
|
||||
final_site = requested_site or (existing or {}).get("site_id")
|
||||
final_zone: UUID | None = None
|
||||
reason: str | None = None
|
||||
|
||||
established_zone = None
|
||||
established_zone_id = (existing or {}).get("agent_zone_id")
|
||||
if established_zone_id is not None:
|
||||
established_zone = self.repository.get_zone(established_zone_id)
|
||||
if not established_zone:
|
||||
raise AgentZoneValidationError(
|
||||
"Existing Agent Zone assignment is invalid"
|
||||
)
|
||||
if (
|
||||
requested_site is not None
|
||||
and established_zone["site_id"] != requested_site
|
||||
):
|
||||
raise AgentZoneValidationError(
|
||||
"Agent cannot change Sites while assigned to an Agent Zone"
|
||||
)
|
||||
|
||||
if requested_zone is not None:
|
||||
zone = self.repository.get_zone(requested_zone)
|
||||
if not zone:
|
||||
raise AgentZoneNotFound("Agent Zone not found")
|
||||
if zone["site_id"] != requested_site:
|
||||
raise AgentZoneValidationError(
|
||||
"Agent Zone does not belong to the supplied Site"
|
||||
)
|
||||
if not zone["is_enabled"]:
|
||||
raise AgentZoneValidationError("Agent Zone is disabled")
|
||||
if established_zone_id is not None and established_zone_id != requested_zone:
|
||||
raise AgentZoneValidationError(
|
||||
"Agent is already assigned to a different Agent Zone"
|
||||
)
|
||||
final_zone = requested_zone
|
||||
elif established_zone is not None:
|
||||
final_site = established_zone["site_id"]
|
||||
final_zone = established_zone["id"]
|
||||
if not established_zone["is_enabled"]:
|
||||
reason = "agent_zone_disabled"
|
||||
elif final_site is not None:
|
||||
usable = self.repository.list_usable_zones(final_site)
|
||||
if len(usable) == 1:
|
||||
final_zone = usable[0]["id"]
|
||||
elif not usable:
|
||||
reason = "agent_zone_unavailable"
|
||||
else:
|
||||
reason = "agent_zone_ambiguous"
|
||||
else:
|
||||
reason = "site_unassigned"
|
||||
|
||||
site_name = (request.site_name or "").strip()
|
||||
if not site_name:
|
||||
site_name = str((existing or {}).get("site_name") or "Unassigned")
|
||||
payload = _sanitize(request.model_dump(mode="json"))
|
||||
self.repository.save_checkin(
|
||||
agent_id=authenticated_agent_id,
|
||||
site_id=final_site,
|
||||
agent_zone_id=final_zone,
|
||||
site_name=site_name,
|
||||
hostname=(request.hostname or "").strip() or None,
|
||||
version=(request.version or "").strip() or None,
|
||||
ip_address=(request.ip_address or "").strip() or None,
|
||||
payload=payload,
|
||||
)
|
||||
routing_available = final_zone is not None and reason is None
|
||||
return AgentCheckinResponse(
|
||||
agent_id=authenticated_agent_id,
|
||||
site_id=final_site,
|
||||
agent_zone_id=final_zone,
|
||||
routing_available=routing_available,
|
||||
routing_reason=reason,
|
||||
site_name=site_name,
|
||||
)
|
||||
|
||||
def routing_context(
|
||||
self, agent_id: str, recorded_zone_id: UUID | None = None
|
||||
) -> AgentZoneRoutingContext:
|
||||
assignment = self.repository.get_agent_assignment(agent_id)
|
||||
if not assignment or assignment.get("agent_zone_id") is None:
|
||||
return AgentZoneRoutingContext(routing_reason="agent_zone_unassigned")
|
||||
zone_id = assignment["agent_zone_id"]
|
||||
if recorded_zone_id is not None and recorded_zone_id != zone_id:
|
||||
return AgentZoneRoutingContext(routing_reason="agent_zone_conflict")
|
||||
if assignment.get("zone_enabled") is not True:
|
||||
return AgentZoneRoutingContext(
|
||||
site_id=assignment.get("site_id"),
|
||||
agent_zone_id=zone_id,
|
||||
routing_reason="agent_zone_disabled",
|
||||
)
|
||||
return AgentZoneRoutingContext(
|
||||
site_id=assignment.get("site_id"), agent_zone_id=zone_id
|
||||
)
|
||||
|
||||
def associate_miners_from_agent(
|
||||
self, agent_id: str, miner_ids: list[UUID]
|
||||
) -> MinerZoneAssociationResult:
|
||||
routing = self.routing_context(agent_id)
|
||||
if routing.routing_reason:
|
||||
return MinerZoneAssociationResult(
|
||||
agent_zone_id=routing.agent_zone_id,
|
||||
routing_reason=routing.routing_reason,
|
||||
)
|
||||
return MinerZoneAssociationResult(
|
||||
agent_zone_id=routing.agent_zone_id,
|
||||
outcomes=self.repository.assign_miners_to_zone(
|
||||
miner_ids, routing.agent_zone_id
|
||||
),
|
||||
)
|
||||
Reference in new issue
Block a user