hub-core/hub_core/runtime/inbox_projection.py

"""Opt-in, authenticated State Hub inbox read pilot; no message writers."""
from __future__ import annotations

import hashlib
import hmac
import json
from datetime import datetime, timezone
from typing import Any

from fastapi import APIRouter, Header, HTTPException, Query, Request, Response
from sqlalchemy import select, text

from hub_core.models.agent_message import AgentMessage
from hub_core.schemas.agent_message import MessageRead

SCHEMA = "hub-core.statehub-inbox-snapshot.v1"


def normalize_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    rows = []
    for row in messages:
        model = MessageRead.model_validate(row)
        for field in ("created_at", "read_at", "archived_at"):
            value = getattr(model, field)
            if value is not None:
                # SQLite test adapters return naive UTC; PostgreSQL retains tzinfo.
                setattr(model, field, value.replace(tzinfo=timezone.utc) if value.tzinfo is None
                        else value.astimezone(timezone.utc))
        rows.append(model.model_dump(mode="json"))
    if len({row["id"] for row in rows}) != len(rows):
        raise ValueError("duplicate message IDs")
    return sorted(rows, key=lambda row: row["id"])


def snapshot_hash(messages: list[dict[str, Any]]) -> str:
    return hashlib.sha256(json.dumps(normalize_messages(messages), sort_keys=True,
                                    separators=(",", ":"), ensure_ascii=False).encode()).hexdigest()


async def import_snapshot(session, bundle: dict, *, expected_hash: str, apply: bool = False) -> dict:
    """Import into an empty receiver only; exact retries are read-only no-ops.

    Caller owns the transaction. This function never sends notifications, changes
    existing message rows, or enables a runtime route.
    """
    if bundle.get("schema") != SCHEMA:
        raise ValueError("unsupported inbox snapshot schema")
    rows = normalize_messages(bundle["messages"])
    digest = snapshot_hash(rows)
    if digest != expected_hash or digest != bundle.get("content_hash"):
        raise ValueError("snapshot hash mismatch")
    if len(rows) != bundle.get("count"):
        raise ValueError("snapshot count mismatch")
    if not bundle.get("source") or not bundle.get("captured_at"):
        raise ValueError("snapshot provenance is required")
    ids = {row["id"] for row in rows}
    if any(row["thread_id"] and row["thread_id"] not in ids for row in rows):
        raise ValueError("snapshot must include all thread roots")
    # Lock before inspecting emptiness so concurrent imports cannot both pass.
    if apply and session.bind.dialect.name == "postgresql":
        await session.execute(text("LOCK TABLE agent_messages IN EXCLUSIVE MODE"))
    existing = list((await session.execute(select(AgentMessage))).scalars())
    if existing:
        current = [MessageRead.model_validate(row).model_dump(mode="json") for row in existing]
        if snapshot_hash(current) != digest:
            raise ValueError("receiver contains different messages; refusing overwrite")
        return {"status": "noop", "count": len(rows), "content_hash": digest}
    remaining = {row["id"]: row for row in rows}
    inserted: set[str] = set()
    levels: list[list[dict]] = []
    while remaining:
        level = [row for row in remaining.values() if not row["thread_id"] or row["thread_id"] in inserted]
        if not level:
            raise ValueError("cyclic message threading")
        levels.append(level)
        for row in level:
            inserted.add(row["id"])
            del remaining[row["id"]]
    if apply:
        for level in levels:
            session.add_all([AgentMessage(**MessageRead.model_validate(row).model_dump()) for row in level])
            await session.flush()
    return {"status": "applied" if apply else "validated", "count": len(rows), "content_hash": digest}


def create_inbox_projection_router() -> APIRouter:
    router = APIRouter()

    @router.get("/ports/projections/statehub-inbox", response_model=list[MessageRead], tags=["projections"])
    async def inbox(
        request: Request,
        response: Response,
        to_agent: str = "state-hub",
        from_agent: str | None = None,
        unread_only: bool = False,
        limit: int = Query(50, ge=1, le=1000),
        authorization: str | None = Header(None),
    ):
        settings = request.app.state.settings
        token = settings.api_token
        supplied = (authorization or "").removeprefix("Bearer ")
        if not token or not (authorization or "").startswith("Bearer ") or not hmac.compare_digest(supplied, token):
            raise HTTPException(401, "inbox pilot requires operator bearer authentication",
                                headers={"WWW-Authenticate": "Bearer"})
        if to_agent != settings.statehub_inbox_agent:
            raise HTTPException(422, "inbox pilot is restricted to its configured agent")
        query = select(AgentMessage).where(
            AgentMessage.archived_at.is_(None),
            (AgentMessage.to_agent == to_agent) | (AgentMessage.to_agent == "broadcast"),
        )
        if from_agent:
            query = query.where(AgentMessage.from_agent == from_agent)
        if unread_only:
            query = query.where(AgentMessage.read_at.is_(None))
        query = query.order_by(AgentMessage.created_at.desc(), AgentMessage.id).limit(limit)
        async with request.app.state.port_store.sessions() as session:
            # The pilot may never mutate its imported source history.
            if session.bind.dialect.name == "postgresql":
                await session.execute(text("SET TRANSACTION READ ONLY"))
            rows = list((await session.execute(query)).scalars())
        response.headers["Cache-Control"] = "no-store"
        response.headers["X-Hub-Core-Read-Mode"] = "snapshot-pilot"
        return rows

    return router