letsbe-orchestrator/app/schemas/agent.py

"""Agent schemas for API validation."""

import uuid
from datetime import datetime
from typing import Any

from pydantic import BaseModel, ConfigDict, Field


class AgentRegisterRequest(BaseModel):
    """Schema for agent registration request (new secure flow)."""

    hostname: str = Field(..., min_length=1, max_length=255)
    version: str = Field(..., min_length=1, max_length=50)
    metadata: dict[str, Any] | None = None
    registration_token: str = Field(
        ...,
        min_length=1,
        description="Registration token issued by the orchestrator",
    )


class AgentRegisterRequestLegacy(BaseModel):
    """Schema for legacy agent registration request (deprecated).

    This schema is kept for backward compatibility during migration.
    New agents should use AgentRegisterRequest with registration_token.
    """

    hostname: str = Field(..., min_length=1, max_length=255)
    version: str = Field(..., min_length=1, max_length=50)
    metadata: dict[str, Any] | None = None
    tenant_id: uuid.UUID | None = Field(
        default=None,
        description="Tenant UUID to associate the agent with (DEPRECATED)",
    )


class AgentRegisterResponse(BaseModel):
    """Schema for agent registration response."""

    agent_id: uuid.UUID
    agent_secret: str = Field(
        ...,
        description="Agent secret for authentication. Store securely - shown only once.",
    )
    tenant_id: uuid.UUID = Field(
        ...,
        description="Tenant this agent is associated with",
    )


class AgentRegisterResponseLegacy(BaseModel):
    """Schema for legacy agent registration response (deprecated)."""

    agent_id: uuid.UUID
    token: str


class AgentHeartbeatResponse(BaseModel):
    """Schema for agent heartbeat response."""

    status: str = "ok"


class AgentResponse(BaseModel):
    """Schema for agent response."""

    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    tenant_id: uuid.UUID | None
    name: str
    version: str
    status: str
    last_heartbeat: datetime | None
    created_at: datetime
    updated_at: datetime