iris_x_hermes/gateway-plugin/protocol.py

"""Frame schemas -- the single source of truth for the wire protocol.

Every frame the plugin sends/receives is modelled here as a dataclass +
constants. ``docs/protocol/frames.schema.json`` is generated/mirrored from
this module, and the Kotlin side mirrors these shapes (see
``docs/04-wire-protocol.md``).

Milestone M1: hello/hello.ack, message, message.send, error, ping/pong,
typing (typing is pulled forward from M2 so the app gets a live
"working…" indicator during the first milestone).
Milestone M2: message.start/update/stop, reasoning (on message /
message.stop), tool.start/progress/end, commentary.
Milestone M3: channel.*, search, sync.
Milestone M4: media.*.
Milestone M5: notification, fcm.register, read.receipt.
"""

import json
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

PROTOCOL_VERSION = 1

# ---------------------------------------------------------------------------
# Frame type constants (the ``type`` field of every frame)
# ---------------------------------------------------------------------------

# Pairing / lifecycle
TYPE_HELLO = "hello"
TYPE_HELLO_ACK = "hello.ack"
TYPE_ERROR = "error"
TYPE_PING = "ping"
TYPE_PONG = "pong"

# Chat
TYPE_MESSAGE = "message"
TYPE_MESSAGE_SEND = "message.send"
TYPE_TYPING = "typing"

# Streaming (M2)
TYPE_MESSAGE_START = "message.start"
TYPE_MESSAGE_UPDATE = "message.update"
TYPE_MESSAGE_STOP = "message.stop"

# Tool activity (M2)
TYPE_TOOL_START = "tool.start"
TYPE_TOOL_PROGRESS = "tool.progress"
TYPE_TOOL_END = "tool.end"

# Intermediate assistant beat (M2)
TYPE_COMMENTARY = "commentary"

# ---------------------------------------------------------------------------
# Error codes (``error`` frame payload.code)
# ---------------------------------------------------------------------------

ERR_AUTH = "auth"
ERR_NOT_FOUND = "not_found"
ERR_RATE_LIMITED = "rate_limited"
ERR_MEDIA_TOO_LARGE = "media_too_large"
ERR_UNSUPPORTED = "unsupported"
ERR_INTERNAL = "internal"

# ---------------------------------------------------------------------------
# Message roles (``message`` frame payload.role)
# ---------------------------------------------------------------------------

ROLE_USER = "user"
ROLE_ASSISTANT = "assistant"
ROLE_SYSTEM = "system"
ROLE_CRON = "cron"


# ---------------------------------------------------------------------------
# Envelope
# ---------------------------------------------------------------------------

@dataclass
class Frame:
    """One wire frame.

    ``v`` is always serialised; ``id``/``chat_id``/``thread_id`` are
    omitted when ``None`` (events carry no ``id``; chat-scoped frames carry
    ``chat_id``/``thread_id`` at the top level for convenience).
    """

    type: str
    payload: Dict[str, Any] = field(default_factory=dict)
    id: Optional[int] = None
    chat_id: Optional[str] = None
    thread_id: Optional[str] = None
    v: int = PROTOCOL_VERSION

    def to_dict(self) -> Dict[str, Any]:
        d: Dict[str, Any] = {"v": self.v, "type": self.type}
        if self.id is not None:
            d["id"] = self.id
        if self.chat_id is not None:
            d["chat_id"] = self.chat_id
        if self.thread_id is not None:
            d["thread_id"] = self.thread_id
        d["payload"] = self.payload
        return d

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), separators=(",", ":"), ensure_ascii=False)

    @classmethod
    def from_json(cls, raw: "str | bytes") -> Optional["Frame"]:
        """Parse a text frame. Returns ``None`` for anything not a valid
        frame (bad JSON, non-object, missing/invalid ``type``) so callers
        can ignore malformed input (forward-compat)."""
        try:
            data = json.loads(raw)
        except (json.JSONDecodeError, TypeError, UnicodeDecodeError, ValueError):
            return None
        if not isinstance(data, dict):
            return None
        ftype = data.get("type")
        if not isinstance(ftype, str) or not ftype:
            return None
        payload = data.get("payload")
        if not isinstance(payload, dict):
            payload = {}
        fid = data.get("id")
        if not isinstance(fid, int) or isinstance(fid, bool):
            fid = None
        chat_id = data.get("chat_id")
        if not isinstance(chat_id, str):
            chat_id = None
        thread_id = data.get("thread_id")
        if not isinstance(thread_id, str):
            thread_id = None
        return cls(type=ftype, payload=payload, id=fid, chat_id=chat_id, thread_id=thread_id)


# ---------------------------------------------------------------------------
# Frame constructors (server -> app)
# ---------------------------------------------------------------------------

def hello_ack(
    server_caps: Dict[str, Any],
    sync_cursor: int = 0,
    channels: Optional[list] = None,
) -> Frame:
    return Frame(
        type=TYPE_HELLO_ACK,
        payload={
            "server_caps": server_caps,
            "sync_cursor": sync_cursor,
            "channels": channels or [],
        },
    )


def message(
    chat_id: str,
    message_id: str,
    role: str,
    text: str,
    *,
    thread_id: Optional[str] = None,
    reasoning: Optional[str] = None,
    media: Optional[list] = None,
    reply_to: Optional[str] = None,
    model: Optional[str] = None,
    tokens: Optional[int] = None,
    ts: Optional[int] = None,
) -> Frame:
    payload: Dict[str, Any] = {
        "message_id": message_id,
        "role": role,
        "text": text,
    }
    if reasoning:
        payload["reasoning"] = reasoning
    if media:
        payload["media"] = media
    if reply_to:
        payload["reply_to"] = reply_to
    if model:
        payload["model"] = model
    if tokens is not None:
        payload["tokens"] = tokens
    if ts is not None:
        payload["ts"] = ts
    return Frame(type=TYPE_MESSAGE, chat_id=chat_id, thread_id=thread_id, payload=payload)


def typing(chat_id: str, on: bool = True, *, thread_id: Optional[str] = None) -> Frame:
    return Frame(
        type=TYPE_TYPING,
        chat_id=chat_id,
        thread_id=thread_id,
        payload={"on": on},
    )


# ---------------------------------------------------------------------------
# Streaming frames (M2)
# ---------------------------------------------------------------------------

def message_start(
    chat_id: str,
    message_id: str,
    role: str = ROLE_ASSISTANT,
    *,
    thread_id: Optional[str] = None,
) -> Frame:
    """Open a streaming bubble."""
    return Frame(
        type=TYPE_MESSAGE_START,
        chat_id=chat_id,
        thread_id=thread_id,
        payload={"message_id": message_id, "role": role},
    )


def message_update(
    chat_id: str,
    message_id: str,
    text: str,
    *,
    thread_id: Optional[str] = None,
) -> Frame:
    """Replace the live bubble text (full snapshot)."""
    return Frame(
        type=TYPE_MESSAGE_UPDATE,
        chat_id=chat_id,
        thread_id=thread_id,
        payload={"message_id": message_id, "text": text},
    )


def message_stop(
    chat_id: str,
    message_id: str,
    final_text: str,
    *,
    thread_id: Optional[str] = None,
    reasoning: Optional[str] = None,
    model: Optional[str] = None,
    tokens: Optional[int] = None,
    ts: Optional[int] = None,
) -> Frame:
    """Finalize a streaming bubble."""
    payload: Dict[str, Any] = {
        "message_id": message_id,
        "final_text": final_text,
    }
    if reasoning:
        payload["reasoning"] = reasoning
    if model:
        payload["model"] = model
    if tokens is not None:
        payload["tokens"] = tokens
    if ts is not None:
        payload["ts"] = ts
    return Frame(
        type=TYPE_MESSAGE_STOP,
        chat_id=chat_id,
        thread_id=thread_id,
        payload=payload,
    )


# ---------------------------------------------------------------------------
# Tool activity frames (M2)
# ---------------------------------------------------------------------------

def tool_start(
    chat_id: str,
    index: int,
    name: str,
    *,
    thread_id: Optional[str] = None,
    preview: Optional[str] = None,
    args: Optional[Dict[str, Any]] = None,
) -> Frame:
    payload: Dict[str, Any] = {"index": index, "name": name}
    if preview:
        payload["preview"] = preview
    if args:
        payload["args"] = args
    return Frame(
        type=TYPE_TOOL_START,
        chat_id=chat_id,
        thread_id=thread_id,
        payload=payload,
    )


def tool_progress(
    chat_id: str,
    index: int,
    name: str,
    *,
    thread_id: Optional[str] = None,
    note: Optional[str] = None,
) -> Frame:
    payload: Dict[str, Any] = {"index": index, "name": name}
    if note:
        payload["note"] = note
    return Frame(
        type=TYPE_TOOL_PROGRESS,
        chat_id=chat_id,
        thread_id=thread_id,
        payload=payload,
    )


def tool_end(
    chat_id: str,
    index: int,
    name: str,
    *,
    thread_id: Optional[str] = None,
    ok: bool = True,
    duration: Optional[float] = None,
    output_preview: Optional[str] = None,
) -> Frame:
    payload: Dict[str, Any] = {"index": index, "name": name, "ok": ok}
    if duration is not None:
        payload["duration"] = duration
    if output_preview:
        payload["output_preview"] = output_preview
    return Frame(
        type=TYPE_TOOL_END,
        chat_id=chat_id,
        thread_id=thread_id,
        payload=payload,
    )


# ---------------------------------------------------------------------------
# Commentary frame (M2)
# ---------------------------------------------------------------------------

def commentary(
    chat_id: str,
    message_id: str,
    text: str,
    *,
    thread_id: Optional[str] = None,
) -> Frame:
    """An intermediate assistant beat (between tool iterations)."""
    return Frame(
        type=TYPE_COMMENTARY,
        chat_id=chat_id,
        thread_id=thread_id,
        payload={"message_id": message_id, "text": text},
    )


def error(code: str, message: str, *, id: Optional[int] = None) -> Frame:
    return Frame(type=TYPE_ERROR, id=id, payload={"code": code, "message": message})


def pong(ts: Optional[int] = None) -> Frame:
    payload: Dict[str, Any] = {}
    if ts is not None:
        payload["ts"] = ts
    return Frame(type=TYPE_PONG, payload=payload)