hub-core/hub_core/message_identity.py

from __future__ import annotations

import uuid
from dataclasses import dataclass
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from hub_core.models.agent_message import AgentMessage
from hub_core.models.message_identity_alias import MessageIdentityAlias


class MessageIdentityError(ValueError):
    """Base error for governed message identity operations."""


class MessageIdentityConflict(MessageIdentityError):
    """An alias is already bound to a different canonical message."""


class CanonicalMessageNotFound(MessageIdentityError):
    """The proposed canonical message does not exist."""


@dataclass(frozen=True)
class MessageAliasRegistration:
    alias: str
    message_id: uuid.UUID
    created: bool
    source: str
    reason: str
    created_by: str


def canonical_message_id(reference: str) -> uuid.UUID | None:
    """Return a UUID only when ``reference`` is already a complete UUID."""

    try:
        return uuid.UUID(reference)
    except (AttributeError, TypeError, ValueError):
        return None


async def resolve_message_reference(
    session: AsyncSession,
    reference: str,
    *,
    alias_model: type[MessageIdentityAlias] | None = MessageIdentityAlias,
) -> uuid.UUID | None:
    """Resolve a canonical UUID or one exact, explicitly registered alias."""

    canonical = canonical_message_id(reference)
    if canonical is not None:
        return canonical
    if alias_model is None:
        return None
    alias = await session.get(alias_model, reference)
    return alias.message_id if alias is not None else None


async def register_message_alias(
    session: AsyncSession,
    *,
    alias: str,
    message_id: uuid.UUID,
    source: str,
    reason: str,
    created_by: str,
    message_model: type[AgentMessage] = AgentMessage,
    alias_model: type[MessageIdentityAlias] = MessageIdentityAlias,
) -> MessageAliasRegistration:
    """Register one attributable alias without changing the message row."""

    _validate_alias(alias)
    source = _required_text("source", source, 255)
    reason = _required_text("reason", reason, None)
    created_by = _required_text("created_by", created_by, 100)

    message = await session.get(message_model, message_id)
    if message is None:
        raise CanonicalMessageNotFound(f"Canonical message {message_id} not found")

    existing: Any = await session.get(alias_model, alias)
    if existing is not None:
        if existing.message_id != message_id:
            raise MessageIdentityConflict(
                f"Alias {alias!r} is already mapped to {existing.message_id}"
            )
        return MessageAliasRegistration(
            alias=existing.alias,
            message_id=existing.message_id,
            created=False,
            source=existing.source,
            reason=existing.reason,
            created_by=existing.created_by,
        )

    record = alias_model(
        alias=alias,
        message_id=message_id,
        source=source,
        reason=reason,
        created_by=created_by,
    )
    session.add(record)
    await session.commit()
    return MessageAliasRegistration(
        alias=alias,
        message_id=message_id,
        created=True,
        source=source,
        reason=reason,
        created_by=created_by,
    )


def _validate_alias(alias: str) -> None:
    if not isinstance(alias, str) or not alias:
        raise MessageIdentityError("alias must be a non-empty string")
    if alias != alias.strip():
        raise MessageIdentityError("alias must not contain leading or trailing whitespace")
    if len(alias) > 255:
        raise MessageIdentityError("alias must be at most 255 characters")
    if "/" in alias:
        raise MessageIdentityError("alias must not contain a path separator")
    if canonical_message_id(alias) is not None:
        raise MessageIdentityError("canonical UUIDs cannot be registered as aliases")


def _required_text(name: str, value: str, maximum: int | None) -> str:
    if not isinstance(value, str) or not value.strip():
        raise MessageIdentityError(f"{name} must be a non-empty string")
    value = value.strip()
    if maximum is not None and len(value) > maximum:
        raise MessageIdentityError(f"{name} must be at most {maximum} characters")
    return value