hermes-agent/plugins/mascot/dashboard/plugin_api.py

"""
Mascot Dashboard Plugin — Backend API routes.

Mounted at /api/plugins/mascot/ by the dashboard plugin system.

Provides:
- GET /state — Current mascot state
- POST /state — Update mascot state
- POST /reset — Reset to idle
- WebSocket /events — Live state stream

Security note:
HTTP routes are unauthenticated (dashboard binds to localhost).
WebSocket requires session token via ?token= query param.
"""

from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import asdict
from typing import Optional

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from ..mascot_state import get_manager, MascotState, VALID_STATES

log = logging.getLogger(__name__)

router = APIRouter()


def _state_to_dict(state: MascotState) -> dict:
    """Convert state to API response."""
    return {
        "status": state.status,
        "task": state.task,
        "mood": state.mood,
        "last_update": state.last_update,
    }


# ---------------------------------------------------------------------------
# HTTP Endpoints
# ---------------------------------------------------------------------------

@router.get("/state")
def get_state():
    """Get current mascot state."""
    manager = get_manager()
    state = manager.get_state()
    return _state_to_dict(state)


class UpdateStateBody:
    """Request body for POST /state."""
    status: Optional[str] = None
    task: Optional[str] = None
    mood: Optional[str] = None


@router.post("/state")
def update_state(
    status: Optional[str] = None,
    task: Optional[str] = None,
    mood: Optional[str] = None,
):
    """Update mascot state."""
    manager = get_manager()

    # Validate status if provided
    if status is not None and status not in VALID_STATES:
        return {
            "success": False,
            "error": f"Invalid status: {status}. Must be one of {VALID_STATES}",
        }

    new_state = manager.set_state(status=status, task=task, mood=mood)
    return {
        "success": True,
        "state": _state_to_dict(new_state),
    }


@router.post("/reset")
def reset_state():
    """Reset mascot to idle state."""
    manager = get_manager()
    new_state = manager.reset()
    return {
        "success": True,
        "state": _state_to_dict(new_state),
    }


# ---------------------------------------------------------------------------
# WebSocket Endpoint
# ---------------------------------------------------------------------------

def _check_ws_token(provided: Optional[str]) -> bool:
    """Validate WebSocket session token."""
    if not provided:
        return False
    try:
        from hermes_cli import web_server as _ws
    except Exception:
        # No dashboard context (tests), accept
        return True
    expected = getattr(_ws, "_SESSION_TOKEN", None)
    if not expected:
        return True
    import hmac
    return hmac.compare_digest(str(provided), str(expected))


# Track active WebSocket connections for broadcasting
_active_ws_clients: list = []


@router.websocket("/events")
async def stream_events(ws: WebSocket):
    """
    Stream mascot state changes over WebSocket.

    Client sends ?token=<session> for auth.
    Server sends: {"type": "state", "state": {...}}
    Server also sends immediate current state on connect.

    Reconnection: client should reconnect with same logic; the state
    manager will push current state immediately.
    """
    token = ws.query_params.get("token")
    if not _check_ws_token(token):
        await ws.close(code=1008, reason="Invalid token")
        return

    await ws.accept()
    _active_ws_clients.append(ws)
    log.debug("Mascot WS client connected (%d active)", len(_active_ws_clients))

    # Send current state immediately
    manager = get_manager()
    current = manager.get_state()
    try:
        await ws.send_json({
            "type": "state",
            "state": _state_to_dict(current),
        })
    except Exception:
        pass

    # Set up state change callback
    state_changed = asyncio.Event()
    latest_state = [current]

    def on_state_change(new_state: MascotState):
        latest_state[0] = new_state
        # Signal from sync thread to async loop
        try:
            loop = asyncio.get_event_loop()
            loop.call_soon_threadsafe(state_changed.set)
        except RuntimeError:
            pass

    manager.subscribe(on_state_change)

    try:
        # Polling loop with fallback (300ms)
        # This is simpler than pushing from callback and works
        # reliably across asyncio contexts
        poll_interval = 0.3  # seconds

        while True:
            try:
                # Wait for state change or timeout
                try:
                    await asyncio.wait_for(state_changed.wait(), timeout=poll_interval)
                    state_changed.clear()
                except asyncio.TimeoutError:
                    pass

                # Send current state
                current = latest_state[0]
                await ws.send_json({
                    "type": "state",
                    "state": _state_to_dict(current),
                })

            except WebSocketDisconnect:
                break
            except Exception as e:
                log.debug("Mascot WS error: %s", e)
                break

    finally:
        manager.unsubscribe(on_state_change)
        if ws in _active_ws_clients:
            _active_ws_clients.remove(ws)
        log.debug("Mascot WS client disconnected (%d active)", len(_active_ws_clients))