hub-core/hub_core/routers/messages.py

from collections.abc import Callable
from datetime import datetime, timezone
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from hub_core.models.agent_message import AgentMessage
from hub_core.models.message_identity_alias import MessageIdentityAlias
from hub_core.message_identity import resolve_message_reference
from hub_core.schemas.agent_message import MessageCreate, MessageRead, MessageReply


def create_messages_router(
    get_session: Callable[..., AsyncSession],
    *,
    message_model: type[AgentMessage] = AgentMessage,
    message_alias_model: type[MessageIdentityAlias] | None = MessageIdentityAlias,
) -> APIRouter:
    router = APIRouter(prefix="/messages", tags=["messages"])

    async def _get_message(message_reference: str, session: AsyncSession) -> Any:
        message_id = await resolve_message_reference(
            session,
            message_reference,
            alias_model=message_alias_model,
        )
        if message_id is None:
            raise HTTPException(
                status_code=404,
                detail=f"Message reference {message_reference!r} not found",
            )
        msg = await session.get(message_model, message_id)
        if msg is None:
            raise HTTPException(
                status_code=404,
                detail=f"Message reference {message_reference!r} not found",
            )
        return msg

    @router.post("/", response_model=MessageRead, status_code=status.HTTP_201_CREATED)
    async def send_message(
        body: MessageCreate,
        session: AsyncSession = Depends(get_session),
    ) -> Any:
        if body.thread_id:
            root = await session.get(message_model, body.thread_id)
            if root is None:
                raise HTTPException(status_code=404, detail=f"Thread root {body.thread_id} not found")
        msg = message_model(**body.model_dump())
        session.add(msg)
        await session.commit()
        await session.refresh(msg)
        return msg

    @router.get("/", response_model=list[MessageRead])
    async def list_messages(
        to_agent: str | None = None,
        from_agent: str | None = None,
        unread_only: bool = False,
        limit: int = 50,
        session: AsyncSession = Depends(get_session),
    ) -> list[Any]:
        q = select(message_model).where(message_model.archived_at.is_(None))
        if to_agent:
            q = q.where(
                (message_model.to_agent == to_agent) | (message_model.to_agent == "broadcast")
            )
        if from_agent:
            q = q.where(message_model.from_agent == from_agent)
        if unread_only:
            q = q.where(message_model.read_at.is_(None))
        q = q.order_by(message_model.created_at.desc()).limit(limit)
        result = await session.execute(q)
        return list(result.scalars().all())

    @router.get("/thread/{thread_id}", response_model=list[MessageRead])
    async def get_thread(
        thread_id: str,
        session: AsyncSession = Depends(get_session),
    ) -> list[Any]:
        resolved_thread_id = await resolve_message_reference(
            session,
            thread_id,
            alias_model=message_alias_model,
        )
        if resolved_thread_id is None:
            raise HTTPException(
                status_code=404,
                detail=f"Message reference {thread_id!r} not found",
            )
        q = select(message_model).where(
            (message_model.id == resolved_thread_id)
            | (message_model.thread_id == resolved_thread_id)
        ).order_by(message_model.created_at)
        result = await session.execute(q)
        return list(result.scalars().all())

    @router.patch("/{message_id}/read", response_model=MessageRead)
    async def mark_read(
        message_id: str,
        session: AsyncSession = Depends(get_session),
    ) -> Any:
        msg = await _get_message(message_id, session)
        if msg.read_at is None:
            msg.read_at = datetime.now(timezone.utc)
            await session.commit()
            await session.refresh(msg)
        return msg

    @router.patch("/{message_id}/archive", response_model=MessageRead)
    async def archive_message(
        message_id: str,
        session: AsyncSession = Depends(get_session),
    ) -> Any:
        msg = await _get_message(message_id, session)
        msg.archived_at = datetime.now(timezone.utc)
        if msg.read_at is None:
            msg.read_at = msg.archived_at
        await session.commit()
        await session.refresh(msg)
        return msg

    @router.post("/{message_id}/reply", response_model=MessageRead, status_code=status.HTTP_201_CREATED)
    async def reply_to_message(
        message_id: str,
        body: MessageReply,
        session: AsyncSession = Depends(get_session),
    ) -> Any:
        original = await _get_message(message_id, session)
        if original.read_at is None:
            original.read_at = datetime.now(timezone.utc)
        thread_root = original.thread_id or original.id
        reply = message_model(
            from_agent=body.from_agent,
            to_agent=original.from_agent,
            subject=f"Re: {original.subject}",
            body=body.body,
            thread_id=thread_root,
        )
        session.add(reply)
        await session.commit()
        await session.refresh(reply)
        return reply

    return router