ai-rpg/app/engine/sse.py

"""SSE event emitter — wraps sse-starlette to emit typed events."""

from __future__ import annotations

import asyncio
import json
import uuid
from collections.abc import AsyncIterator
from typing import Any

from app.core.logging import get_logger

_logger = get_logger(__name__)


class SseEmitter:
    """Async queue-based SSE emitter.

    Usage:
        emitter = SseEmitter()
        async with emitter.stream() as stream:
            async for event in stream:
                yield event

    In a producer task:
        await emitter.emit("tool_call", {...})
        await emitter.done({"result": "ok"})

    Production mode: when `debug=False`, raw `tool_call` and `llm_call_*`
    events are filtered or transformed into friendly `status` events.
    """

    def __init__(self, debug: bool = True) -> None:
        self._queue: asyncio.Queue[tuple[str, str, str] | None] = asyncio.Queue()
        self._event_counter = 0
        self._closed = False
        self._debug = debug

    async def emit(self, event_type: str, data: Any) -> None:
        if self._closed:
            return
        # In production mode, transform/fiter debug-only events
        if not self._debug:
            event_type, data = self._transform_for_prod(event_type, data)
            if event_type is None:
                return  # event filtered out
        self._event_counter += 1
        event_id = f"evt_{self._event_counter}"
        try:
            data_str = json.dumps(data, ensure_ascii=False, default=str)
        except (TypeError, ValueError):
            data_str = json.dumps({"error": "serialization_failed"})
        await self._queue.put((event_type, data_str, event_id))

    def _transform_for_prod(self, event_type: str, data: Any) -> tuple[str | None, Any]:
        """Transform debug events into user-friendly status events for production."""
        if event_type == "tool_call":
            # Transform tool_call into a friendly status message
            tool = data.get("tool", "") if isinstance(data, dict) else ""
            result = data.get("result", {}) if isinstance(data, dict) else {}
            is_success = data.get("is_success", True) if isinstance(data, dict) else True
            # Friendly message based on tool type
            friendly = self._friendly_tool_message(tool, result, is_success)
            if friendly:
                return ("status", {"message": friendly, "type": "tool"})
            return (None, None)  # filter out
        elif event_type in ("llm_call_start", "llm_call_end"):
            # Filter out raw LLM call events in production
            return (None, None)
        elif event_type == "phase_start":
            # Keep but with friendly name
            phase = data.get("phase") if isinstance(data, dict) else None
            friendly_names = {
                1: "planning",
                2: "writing",
                3: "sending",
            }
            name = friendly_names.get(phase, data.get("name", ""))
            return ("phase_start", {"phase": phase, "name": name, "status": name})
        elif event_type == "phase_end":
            return (event_type, data)
        elif event_type == "warning":
            return (event_type, data)
        else:
            return (event_type, data)

    def _friendly_tool_message(self, tool: str, result: dict, is_success: bool) -> str | None:
        """Generate a user-friendly message for a tool call."""
        if not is_success:
            return None  # hide failed tool calls in production
        data = result.get("data", {}) if isinstance(result, dict) else {}
        msg = result.get("message", "") if isinstance(result, dict) else ""
        if tool == "entity_create":
            name = data.get("entity_id", "")
            return f"Added new entity" + (f": {name}" if name else "")
        elif tool == "entity_update":
            return "Updated entity"
        elif tool == "entity_delete":
            return "Removed entity"
        elif tool == "env_update":
            return "Updated game state"
        elif tool == "update_plot_rails":
            return "Updated story progress"
        elif tool == "advance_time":
            new_time = data.get("new_time", "")
            return f"Time advanced" + (f" to {new_time}" if new_time else "")
        elif tool == "schedule_trigger":
            return "Scheduled future event"
        elif tool == "rag_query":
            return None  # hide RAG queries in production
        elif tool == "rag_add":
            return "Recorded a new fact"
        elif tool == "calc":
            return None  # hide calculations
        elif tool == "random_choice":
            return None
        elif tool == "submit_plan":
            return None
        elif tool == "submit_step":
            return None
        elif tool == "suggest_actions":
            return None
        return msg if msg else None

    async def ping(self) -> None:
        await self.emit("ping", {"ts": _now_iso()})

    async def done(self, result: Any = None) -> None:
        await self.emit("done", result if result is not None else {})
        await self._queue.put(None)  # sentinel
        self._closed = True

    async def error(self, code: str, message: str, details: Any = None) -> None:
        payload: dict[str, Any] = {"code": code, "message": message}
        if details is not None:
            payload["details"] = details
        await self.emit("error", payload)
        await self._queue.put(None)
        self._closed = True

    async def stream(self) -> AsyncIterator[dict[str, str]]:
        """Yield SSE-formatted dicts until the emitter is closed."""
        try:
            while True:
                item = await self._queue.get()
                if item is None:
                    break
                event_type, data_str, event_id = item
                yield {
                    "event": event_type,
                    "data": data_str,
                    "id": event_id,
                }
        except asyncio.CancelledError:
            _logger.info("sse_stream_cancelled")
            raise


def _now_iso() -> str:
    from datetime import datetime, timezone

    return datetime.now(timezone.utc).isoformat()