hub-core/hub_core/runtime/store.py

from __future__ import annotations

import asyncio
import hashlib
import json
from collections.abc import Mapping
from copy import deepcopy
from datetime import datetime, timezone
from typing import Any, Protocol
from uuid import UUID, uuid4

from hub_core.contracts import CONTRACT_VERSION
from hub_core.runtime.models import (
    EventCommand,
    MessageCommand,
    PortAccepted,
    PortCollection,
    PortRecord,
    Provenance,
    RegistryRegistration,
)


class PortStore(Protocol):
    """Persistence boundary for the initial hub-core runtime ports."""

    backend_name: str

    async def readiness_checks(self) -> dict[str, str]: ...

    async def register_extension(
        self,
        registration: RegistryRegistration,
        correlation_id: UUID,
    ) -> PortAccepted: ...

    async def send_message(self, command: MessageCommand) -> PortAccepted: ...

    async def list_messages(self, address: str, conversation_id: UUID | None) -> PortCollection: ...

    async def append_progress(self, command: EventCommand) -> PortAccepted: ...

    async def append_interaction(self, command: EventCommand) -> PortAccepted: ...

    async def query_projection(self, projection_id: str) -> PortRecord | None: ...


class InMemoryPortStore:
    """Deterministic ephemeral backend for local runtime and conformance tests.

    Production readiness rejects this backend unless explicitly allowed. The
    store deliberately keeps progress and interaction event families separate.
    """

    backend_name = "memory"

    def __init__(self) -> None:
        self._lock = asyncio.Lock()
        self._registrations: dict[str, dict[str, Any]] = {}
        self._messages: list[dict[str, Any]] = []
        self._progress_events: list[dict[str, Any]] = []
        self._interaction_events: list[dict[str, Any]] = []

    async def readiness_checks(self) -> dict[str, str]:
        return {"database": "not_applicable"}

    async def register_extension(
        self,
        registration: RegistryRegistration,
        correlation_id: UUID,
    ) -> PortAccepted:
        hub_slug = str(registration.descriptor["hub_slug"])
        value = registration.model_dump(mode="json")
        async with self._lock:
            duplicate = self._registrations.get(hub_slug) == value
            self._registrations[hub_slug] = deepcopy(value)
        return PortAccepted(
            id=hub_slug,
            status="duplicate" if duplicate else "accepted",
            correlation_id=correlation_id,
        )

    async def send_message(self, command: MessageCommand) -> PortAccepted:
        message_id = uuid4()
        value = {
            "id": str(message_id),
            "created_at": _now().isoformat(),
            **command.model_dump(mode="json"),
        }
        async with self._lock:
            self._messages.append(value)
        return PortAccepted(
            id=str(message_id),
            status="accepted",
            correlation_id=command.correlation_id,
        )

    async def list_messages(self, address: str, conversation_id: UUID | None) -> PortCollection:
        async with self._lock:
            values = [
                deepcopy(message)
                for message in self._messages
                if address in message["to_addresses"]
                and (
                    conversation_id is None
                    or message.get("conversation_id") == str(conversation_id)
                )
            ]
        return PortCollection(items=[self._record("message", value) for value in values])

    async def append_progress(self, command: EventCommand) -> PortAccepted:
        return await self._append_event(command, self._progress_events, "progress")

    async def append_interaction(self, command: EventCommand) -> PortAccepted:
        return await self._append_event(command, self._interaction_events, "interaction")

    async def query_projection(self, projection_id: str) -> PortRecord | None:
        async with self._lock:
            sources: Mapping[str, Any] = {
                "hub_registry": list(self._registrations.values()),
                "messages": self._messages,
                "progress_events": self._progress_events,
                "interaction_events": self._interaction_events,
            }
            if projection_id not in sources:
                return None
            items = deepcopy(sources[projection_id])
        return self._record(
            projection_id,
            {
                "projection_id": projection_id,
                "items": items,
                "rebuild_from": _rebuild_sources(projection_id),
            },
        )

    async def _append_event(
        self,
        command: EventCommand,
        target: list[dict[str, Any]],
        family: str,
    ) -> PortAccepted:
        event_id = uuid4()
        value = {
            "id": str(event_id),
            "family": family,
            "recorded_at": _now().isoformat(),
            **command.model_dump(mode="json"),
        }
        async with self._lock:
            target.append(value)
        return PortAccepted(
            id=str(event_id),
            status="accepted",
            correlation_id=command.correlation_id,
        )

    def _record(self, kind: str, value: dict[str, Any]) -> PortRecord:
        encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
        record_id = str(value.get("id") or kind)
        return PortRecord(
            id=record_id,
            data=deepcopy(value),
            provenance=Provenance(
                source_system="hub-core-memory",
                source_ref=f"memory://{kind}/{record_id}",
                schema_version=CONTRACT_VERSION,
                content_hash=hashlib.sha256(encoded).hexdigest(),
                indexed_at=_now(),
            ),
        )


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _rebuild_sources(projection_id: str) -> list[str]:
    return {
        "hub_registry": ["hub_descriptors", "hub_manifests"],
        "messages": ["messages"],
        "progress_events": ["progress_events"],
        "interaction_events": ["interaction_events"],
    }[projection_id]