distinct_server/models.py

"""Internal server records and immutable public views."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Optional, Tuple

from distinct_protocol import AgentCapabilities, AgentSnapshot, JobSpec, JobStatus


class ControlPlaneError(RuntimeError):
    """Base class for expected control-plane failures."""


class ValidationError(ControlPlaneError):
    """A caller supplied malformed or inconsistent input."""


class AuthenticationError(ControlPlaneError):
    """An agent credential, signature, timestamp, or nonce was invalid."""


class NotFoundError(ControlPlaneError):
    """The requested server record does not exist in the caller's scope."""


class ConflictError(ControlPlaneError):
    """The requested transition conflicts with the current state."""


class CapacityError(ControlPlaneError):
    """A configured in-memory bound has been reached."""


class NoCompatibleAgentError(ControlPlaneError):
    """No currently eligible agent can run the exact requested selection."""


class DependencyError(ControlPlaneError):
    """A requested conversation dependency is invalid or cannot complete."""


@dataclass(frozen=True)
class ControlPlaneLimits:
    """Hard cardinality and payload limits for the in-memory MVP.

    There is **one** cadence. The transport is agent-pull, so an agent's
    ``POST /agent/sync`` carries the snapshot up and brings work down in the
    same round trip: there is no separate heartbeat, and there never was one.
    ``poll_interval_seconds`` is what the server advertises at pairing and what
    liveness is measured against, so the number the UI shows is the number the
    agent actually uses. The old ``heartbeat_*`` names remain as read-only
    aliases.
    """

    max_agents: int = 256
    max_sessions: int = 2_048
    max_conversations_per_session: int = 128
    max_jobs: int = 50_000
    max_jobs_per_conversation: int = 1_024
    max_pairing_codes: int = 1_024
    #: And how many of that global pool any one identity may hold at once.
    #:
    #: The global cap alone is not a limit on anybody, it is a limit on
    #: everybody: one signed-in account minting in a loop fills all 1024
    #: slots, and since they live ten minutes it can hold them full, so
    #: nobody else on the server can pair a worker. A per-owner share turns
    #: that from an outage into one account wasting its own allowance.
    max_pairing_codes_per_owner: int = 8
    #: How many workers may sit registered-but-unclaimed at once. An unclaimed
    #: worker is the one thing here that anybody can create without signing in,
    #: so it is the one thing that needs a ceiling and a short life. Both are
    #: deliberately small: a claim is something a person does within a minute
    #: of starting a worker, not something left pending.
    max_unclaimed_agents: int = 64
    max_nonces_per_agent: int = 2_048
    max_models_per_agent: int = 64
    max_tools_per_agent: int = 64
    max_result_characters: int = 1_000_000
    pairing_ttl_seconds: float = 600.0
    #: How long an unclaimed worker waits to be claimed before the server
    #: forgets it. Long enough to walk to the browser, short enough that an
    #: abandoned worker does not hold a slot.
    claim_ttl_seconds: float = 600.0
    poll_interval_seconds: float = 5.0
    missed_polls_before_offline: int = 3
    # A floor on the offline threshold so that lowering the poll interval
    # cannot make one dropped packet look like an outage.
    min_offline_seconds: float = 15.0
    request_clock_skew_seconds: float = 60.0
    nonce_ttl_seconds: float = 180.0
    job_ttl_seconds: float = 86_400.0

    def __post_init__(self) -> None:
        integer_fields = (
            "max_agents",
            "max_sessions",
            "max_conversations_per_session",
            "max_jobs",
            "max_jobs_per_conversation",
            "max_pairing_codes",
            "max_pairing_codes_per_owner",
            "max_unclaimed_agents",
            "max_nonces_per_agent",
            "max_models_per_agent",
            "max_tools_per_agent",
            "max_result_characters",
            "missed_polls_before_offline",
        )
        duration_fields = (
            "pairing_ttl_seconds",
            "claim_ttl_seconds",
            "poll_interval_seconds",
            "min_offline_seconds",
            "request_clock_skew_seconds",
            "nonce_ttl_seconds",
            "job_ttl_seconds",
        )
        for name in integer_fields:
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        for name in duration_fields:
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
                raise ValueError(f"{name} must be positive")
        if self.nonce_ttl_seconds < self.request_clock_skew_seconds:
            raise ValueError("nonce TTL must cover the accepted clock skew")

    @property
    def offline_after_seconds(self) -> float:
        """The single threshold used for liveness, in seconds."""

        return max(
            self.min_offline_seconds,
            self.poll_interval_seconds * self.missed_polls_before_offline,
        )

    # Deprecated aliases. There is no separate heartbeat; these name the poll.
    @property
    def heartbeat_interval_seconds(self) -> float:
        return self.poll_interval_seconds

    @property
    def missed_heartbeats_before_offline(self) -> int:
        return self.missed_polls_before_offline


@dataclass(frozen=True)
class PairingCode:
    code: str
    expires_at: float


@dataclass(frozen=True)
class AgentCredential:
    agent_id: str
    secret: str
    issued_at: float


@dataclass(frozen=True)
class AgentView:
    capabilities: AgentCapabilities
    owner_id: str
    snapshot: Optional[AgentSnapshot]
    online: bool
    missed_polls: int
    last_snapshot_at: float
    outstanding_jobs: int
    revoked: bool

    @property
    def missed_heartbeats(self) -> int:
        """Deprecated alias: the poll *is* the heartbeat."""

        return self.missed_polls


@dataclass(frozen=True)
class JobView:
    spec: JobSpec
    status: JobStatus
    accepted_by: Optional[str]
    queue_position: Optional[int]
    updated_at: float
    expires_at: float
    result: Any = None
    error: Optional[str] = None
    # ``spec.prompt`` reads as a tombstone once ``prompt_erased`` is set; the
    # real text is gone from the server. Callers that still need to display a
    # transcript must hold their own copy (the browser's gr.State does).
    prompt_erased: bool = False
    result_erased: bool = False


@dataclass(frozen=True)
class ConversationView:
    session_id: str
    conversation_id: str
    created_at: float
    job_ids: Tuple[str, ...]


@dataclass(frozen=True)
class SessionView:
    session_id: str
    created_at: float
    conversation_ids: Tuple[str, ...]


@dataclass(frozen=True)
class CleanupReport:
    pairing_codes_removed: int = 0
    nonces_removed: int = 0
    jobs_expired: int = 0
    agents_offline: int = 0
    results_erased: int = 0


@dataclass(frozen=True)
class RunSelection:
    """The exact dispatch selection; no mutation or reassignment API exists."""

    model_id: str
    tool_versions: Tuple[str, ...]
    target_agent_id: str


@dataclass
class PairingRecord:
    owner_id: str
    expires_at: float
    #: Visibility is fixed when the *authenticated user* generates the code, so
    #: it is bound to their identity and an agent cannot choose to be private,
    #: or to stop being private, on its own say-so.
    private: bool = False


@dataclass
class AgentRecord:
    owner_id: str
    capabilities: AgentCapabilities
    secret: bytes
    issued_at: float
    last_snapshot_at: float
    snapshot: Optional[AgentSnapshot] = None
    nonces: dict[str, float] = field(default_factory=dict)
    revoked: bool = False
    #: The agent's Ed25519 public key, pinned at pairing, 32 raw bytes. Empty
    #: for an agent that paired before asymmetric snapshots existed. The server
    #: can verify advertisements with it and **cannot** produce one, which is
    #: the whole reason the snapshot signature is not an HMAC.
    signing_public_key: bytes = b""
    #: The digest of the access code this agent printed when it started, or
    #: empty for an agent that offers none. The code itself is never stored:
    #: what is kept is a hash of a 192-bit random string, which is worth
    #: nothing to anybody who reads it.
    access_digest: str = ""
    #: The verified identities that have redeemed that code. Provider-qualified
    #: subjects, never usernames, because a username can be changed and a
    #: subject cannot.
    granted: set = field(default_factory=set)
    #: A private agent is visible only to the authenticated user who paired it.
    #: Enforced in every listing, every selection and every dispatch path, and
    #: the errors are deliberately identical to "no such agent" so that failing
    #: to find one and not being allowed to see one cannot be told apart.
    private: bool = False
    #: Highest signed-snapshot counter accepted from this agent. A snapshot is
    #: an assertion about current state, so what has to be refused is an older
    #: one; that is an ordering question rather than a freshness question, and a
    #: counter answers it without the two sides agreeing a clock.
    snapshot_counter: int = 0
    #: Set while this worker is registered but not yet claimed by anybody.
    #:
    #: A worker started without a pairing code generates its own claim code,
    #: prints it, and registers here holding only the digest. Until somebody
    #: signed in enters that code, ``owner_id`` is empty and the record is
    #: invisible to every viewer and every dispatch path -- see
    #: ``ControlPlane._visible_locked``, which treats a non-empty claim digest
    #: as "this worker does not exist yet". Claiming clears it.
    claim_digest: str = ""
    #: When an unclaimed record stops being claimable and is swept.
    claim_expires_at: float = 0.0

    #: The capability set registered at pairing, frozen. This is the ceiling an
    #: advertisement is intersected with. It must be the *pairing* set rather
    #: than the current one: an agent that advertises nothing while its
    #: operator decides would otherwise narrow its own ceiling to nothing and
    #: could never advertise its way back up after approval.
    registered_capabilities: Optional[AgentCapabilities] = None


@dataclass
class JobRecord:
    spec: JobSpec
    selection: RunSelection
    status: JobStatus
    updated_at: float
    expires_at: float
    accepted_by: Optional[str] = None
    queue_position: Optional[int] = None
    result: Any = None
    error: Optional[str] = None
    # Erasure bookkeeping. ``prompt_erased`` is set the moment the job reaches
    # a terminal state; ``delivered_at`` records when the owning session read
    # the answer, after which the answer is dropped on the next cleanup pass.
    prompt_erased: bool = False
    result_erased: bool = False
    delivered_at: Optional[float] = None
    #: When the release-time context was built into the spec, or ``None`` for
    #: a job whose spec has carried its messages since submission. Set exactly
    #: once, by :meth:`release_with_context`.
    context_built_at: Optional[float] = None
    #: The authenticated identity that submitted this job. **Server-side
    #: only**: it exists so the per-agent pseudonymous ``user_key`` can be
    #: re-derived when the job is assigned or reassigned to a different agent,
    #: and it is never serialised into the spec that crosses the wire.
    viewer_subject: Optional[str] = None

    def release_with_context(self, messages, now: float) -> None:
        """Set a blocked child's context and release it, exactly once.

        This is the ONLY sanctioned way to change ``spec`` after creation
        (reassignment aside, which changes only the target). The old code did
        ``child.spec = replace(...)`` inline and the immutability invariant
        was held by call ordering inside one method; this method holds it
        structurally: releasing a job that is not BLOCKED, or releasing it a
        second time, raises instead of silently mutating a spec an agent may
        already have been offered.
        """

        from dataclasses import replace as _replace

        if self.status is not JobStatus.BLOCKED:
            raise ConflictError(
                f"context can only be built for a blocked job, not {self.status.value}"
            )
        if self.context_built_at is not None:
            raise ConflictError("context has already been built for this job")
        self.spec = _replace(self.spec, messages=tuple(messages))
        self.context_built_at = now
        self.status = JobStatus.OFFERED
        self.updated_at = now


@dataclass
class ConversationRecord:
    session_id: str
    conversation_id: str
    created_at: float
    job_ids: list[str] = field(default_factory=list)


@dataclass
class SessionRecord:
    session_id: str
    created_at: float
    conversation_ids: list[str] = field(default_factory=list)