refactor: overhaul daemon server, client, and handlers

daemon/server.py

@@ -15,12 +15,14 @@ from typing import Any
 
 from aiohttp import web
 
+from daemon.iface import PathLike
 from lib.state import state as state_store
 
 logger = logging.getLogger(__name__)
 
 PROJECT_DIR = Path(__file__).resolve().parent.parent
 SOCKET_PATH = PROJECT_DIR / "data" / "daemon.sock"
+_WS_PORT = int(os.environ.get("VACUUM_WALLD_WS_PORT", "9091"))
 
 
 class Handler:
@@ -53,20 +55,28 @@ class Registry:
         """Initialize an empty route registry."""
         self._routes: dict[tuple[str, str], Callable] = {}
 
-    def register(self, method: str, path: str):
+    def register(self, method: PathLike, path: str | None = None):
         """Decorator that registers a handler for the given method and path.
 
+        Accepts either two separate arguments (``method``, ``path``) or a
+        single :class:`daemon.iface.Endpoint` tuple.
+
         Args:
-            method: HTTP method (e.g. "GET", "POST").
-            path: URL path to register the handler under.
+            method: HTTP method string, or an :class:`Endpoint` tuple.
+            path: URL path (omit when passing an :class:`Endpoint`).
 
         Returns:
             Decorator function wrapping the handler.
         """
 
+        if isinstance(method, tuple):
+            ep_method, ep_path = method
+            path = ep_path
+            method = ep_method
+
         def decorator(fn: Callable) -> Callable:
-            self._routes[(method.upper(), path)] = fn
-            fn._handler = Handler(method, path)  # type: ignore[attr-defined]
+            self._routes[(method.upper(), path)] = fn  # type: ignore[arg-type]
+            fn._handler = Handler(method, path)  # type: ignore[attr-defined,reportArgumentType]
             return fn
 
         return decorator
@@ -111,6 +121,17 @@ def refresh_state(subsystems: list[str] | None = None) -> None:
         subsystems: List of subsystem names to refresh. If None, all subsystems are refreshed.
     """
     state_store.populate(subsystems)
+    targets = subsystems or state_store.SUBSYSTEMS
+    for name in targets:
+        state_store.bump(name)
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        pass
+    else:
+        task = asyncio.create_task(broadcast_versions())
+        task.add_done_callback(_ws_tasks.discard)
+        _ws_tasks.add(task)
 
 
 class NotFoundError(Exception):
@@ -295,10 +316,57 @@ def create_app() -> web.Application:
     app.router.add_route("GET", "/status/all", get_status_all)
     app.router.add_route("POST", "/status/refresh", refresh_status)
     app.router.add_route("POST", "/batch", _handle_batch)
+    app.router.add_route("GET", "/ws", _handle_ws)
     app.router.add_route("*", "/{tail:.*}", _catch_all)
     return app
 
 
+# WebSocket subscribers
+_ws_subscribers: set[web.WebSocketResponse] = set()
+_ws_tasks: set[asyncio.Task[None]] = set()
+
+
+async def _handle_ws(request: web.Request) -> web.Response:
+    """WebSocket endpoint for real-time state change notifications.
+
+    On connect: sends current versions. On state change: broadcasts
+    updated subsystem versions. Clients disconnect to unsubscribe.
+    """
+    ws = web.WebSocketResponse()
+    await ws.prepare(request)
+    _ws_subscribers.add(ws)
+
+    await ws.send_json({"type": "init", "versions": state_store.get_versions()})
+
+    try:
+        async for msg in ws:
+            if msg.type == web.WSMsgType.ERROR:
+                break
+            if msg.type == web.WSMsgType.CLOSE:
+                break
+    finally:
+        _ws_subscribers.discard(ws)
+
+    return ws
+
+
+async def broadcast_versions() -> None:
+    """Broadcast updated subsystem versions to all WebSocket clients."""
+    updated = state_store.get_updated_versions()
+    if not updated or not _ws_subscribers:
+        return
+    data = json.dumps({"type": "versions", "updated": updated})
+    dead: set[web.WebSocketResponse] = set()
+    for ws in _ws_subscribers:
+        try:
+            await ws.send_str(data)
+        except Exception:
+            dead.add(ws)
+    _ws_subscribers.difference_update(dead)
+    if dead:
+        logger.warning("Removed %d dead WS subscribers", len(dead))
+
+
 async def _health(_request: web.Request) -> web.Response:
     """Return the health check response.
 
@@ -399,6 +467,8 @@ def main() -> None:
     loop.run_until_complete(runner.setup())
     site = web.UnixSite(runner, socket_path)
     loop.run_until_complete(site.start())
+    tcp_site = web.TCPSite(runner, "127.0.0.1", _WS_PORT)
+    loop.run_until_complete(tcp_site.start())
 
     os.chmod(socket_path, 0o660)
 
@@ -406,6 +476,7 @@ def main() -> None:
     logger.info("Populating system state...")
     state_store.populate()
     logger.info("vacuum-walld listening on %s", socket_path)
+    logger.info("WebSocket on 127.0.0.1:%d", _WS_PORT)
 
     try:
         loop.run_forever()