add datastar SSE rendering

repub/web.py

@@ -1,12 +1,24 @@
 from __future__ import annotations
 
+import asyncio
 import hashlib
+from collections.abc import AsyncGenerator
+from contextlib import suppress
+from typing import cast
 
 import htpy as h
+from datastar_py.quart import DatastarResponse
+from datastar_py.sse import DatastarEvent
+from htpy import Renderable
 from quart import Quart, Response, request, url_for
 
+from repub.datastar import RefreshBroker, render_stream
 from repub.pages import admin_component, shim_page
 
+REFRESH_BROKER_KEY = "repub.refresh_broker"
+ACTIVE_JOBS_KEY = "repub.demo_active_jobs"
+REFRESH_TASK_KEY = "repub.demo_refresh_task"
+
 
 def _render_shim_page(*, stylesheet_href: str, datastar_src: str) -> tuple[str, str]:
     head = (
@@ -18,8 +30,27 @@ def _render_shim_page(*, stylesheet_href: str, datastar_src: str) -> tuple[str,
     return body, etag
 
 
-def create_app() -> Quart:
+def create_app(*, enable_demo_refresh: bool = True) -> Quart:
     app = Quart(__name__)
+    app.extensions[REFRESH_BROKER_KEY] = RefreshBroker()
+    app.extensions[ACTIVE_JOBS_KEY] = 12
+
+    if enable_demo_refresh:
+
+        @app.before_serving
+        async def start_demo_refresh() -> None:
+            app.extensions[REFRESH_TASK_KEY] = asyncio.create_task(
+                _demo_refresh_loop(app)
+            )
+
+        @app.after_serving
+        async def stop_demo_refresh() -> None:
+            task = cast(asyncio.Task[None] | None, app.extensions.get(REFRESH_TASK_KEY))
+            if task is None:
+                return
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
 
     @app.get("/")
     async def index() -> Response:
@@ -37,7 +68,50 @@ def create_app() -> Quart:
         return response
 
     @app.post("/")
-    async def index_patch() -> Response:
-        return Response(str(admin_component()), mimetype="text/html")
+    async def index_patch() -> DatastarResponse:
+        queue = get_refresh_broker(app).subscribe()
+        stream = render_stream(
+            queue,
+            render=lambda: render_dashboard(app),
+            last_event_id=request.headers.get("last-event-id"),
+        )
+        return DatastarResponse(_unsubscribe_on_close(queue, stream, app))
 
     return app
+
+
+def get_refresh_broker(app: Quart) -> RefreshBroker:
+    return cast(RefreshBroker, app.extensions[REFRESH_BROKER_KEY])
+
+
+def trigger_refresh(app: Quart, event: object = "refresh-event") -> None:
+    get_refresh_broker(app).publish(event)
+
+
+async def render_dashboard(app: Quart) -> Renderable:
+    return admin_component(active_jobs=str(get_active_jobs(app)))
+
+
+async def _unsubscribe_on_close(
+    queue: object, stream: AsyncGenerator[DatastarEvent, None], app: Quart
+) -> AsyncGenerator[DatastarEvent, None]:
+    try:
+        async for event in stream:
+            yield event
+    finally:
+        get_refresh_broker(app).unsubscribe(cast(asyncio.Queue[object], queue))
+
+
+def get_active_jobs(app: Quart) -> int:
+    return cast(int, app.extensions[ACTIVE_JOBS_KEY])
+
+
+def set_active_jobs(app: Quart, value: int) -> None:
+    app.extensions[ACTIVE_JOBS_KEY] = value
+
+
+async def _demo_refresh_loop(app: Quart) -> None:
+    while True:
+        await asyncio.sleep(1)
+        set_active_jobs(app, get_active_jobs(app) + 1)
+        trigger_refresh(app)