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,
)
from hub_core.runtime.repository_navigation import NavigationProjection
from hub_core.runtime.workload_projection import WorkloadProjection


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: ...

    async def get_repository_navigation(self) -> NavigationProjection | None: ...

    async def replace_repository_navigation(
        self, projection: NavigationProjection
    ) -> None: ...

    async def mark_repository_navigation_stale(
        self, *, checked_at: datetime, diagnostic: Mapping[str, Any]
    ) -> None: ...

    async def get_workload_projection(self) -> WorkloadProjection | None: ...

    async def replace_workload_projection(self, projection: WorkloadProjection) -> None: ...

    async def mark_workload_projection_stale(
        self, *, checked_at: datetime, diagnostic: Mapping[str, Any]
    ) -> 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]] = []
        self._repository_navigation: NavigationProjection | None = None
        self._workload_projection: WorkloadProjection | None = None

    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 get_repository_navigation(self) -> NavigationProjection | None:
        async with self._lock:
            return deepcopy(self._repository_navigation)

    async def replace_repository_navigation(
        self, projection: NavigationProjection
    ) -> None:
        async with self._lock:
            self._repository_navigation = deepcopy(projection)

    async def mark_repository_navigation_stale(
        self, *, checked_at: datetime, diagnostic: Mapping[str, Any]
    ) -> None:
        async with self._lock:
            current = self._repository_navigation
            if current is None:
                return
            self._repository_navigation = NavigationProjection(
                projection_status="stale",
                source_snapshot=deepcopy(current.source_snapshot),
                source_checked_at=checked_at,
                rebuilt_at=current.rebuilt_at,
                content_hash=current.content_hash,
                repositories=deepcopy(current.repositories),
                facets=deepcopy(current.facets),
                diagnostics=(deepcopy(dict(diagnostic)),),
            )

    async def get_workload_projection(self) -> WorkloadProjection | None:
        async with self._lock:
            return deepcopy(self._workload_projection)

    async def replace_workload_projection(self, projection: WorkloadProjection) -> None:
        async with self._lock:
            self._workload_projection = deepcopy(projection)

    async def mark_workload_projection_stale(
        self, *, checked_at: datetime, diagnostic: Mapping[str, Any]
    ) -> None:
        async with self._lock:
            current = self._workload_projection
            if current is None:
                return
            self._workload_projection = WorkloadProjection(
                projection_status="stale",
                source_snapshot=deepcopy(current.source_snapshot),
                source_checked_at=checked_at,
                rebuilt_at=current.rebuilt_at,
                content_hash=current.content_hash,
                workloads=deepcopy(current.workloads),
                diagnostics=(deepcopy(dict(diagnostic)),),
            )

    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]