""" Backend routes for the Load node's search dialog. The browser cannot call styleref.io directly from the ComfyUI frontend (CORS aside, an auth token must never reach page JS), so the widget asks the ComfyUI server, which asks StyleRef using the same client the nodes use. Everything here is optional. If ComfyUI's server module isn't importable — a headless test run, an API-only embed, a future frontend rewrite — importing this fails, `__init__` catches it, and the nodes carry on with the plain text field. The JS side probes /styleref/search once at startup and hides its button when these routes are absent, so both halves degrade together. """ from __future__ import annotations import asyncio import os from aiohttp import web from server import PromptServer import styleref_api as api from styleref_api import StyleRefError from styleref_auth import auth_status, clear_credentials, credentials_path, run_login_flow from styleref_nodes.load import bump_refresh, fetch_style def _card(entry: dict) -> dict: """The picker-row shape. heroImage feeds the row thumbnail.""" return { "slug": entry.get("slug"), "name": entry.get("name"), "author": entry.get("author"), "category": entry.get("category"), "tags": (entry.get("tags") or [])[:4], "heroImage": entry.get("heroImage"), "url": entry.get("url"), } def _int_param(request: web.Request, name: str, default: int) -> int: """A query int that tolerates junk — a bad value is a default, not a 400.""" try: return int(request.query.get(name, str(default))) except ValueError: return default @PromptServer.instance.routes.get("/styleref/search") async def styleref_search(request: web.Request) -> web.Response: """GET /styleref/search?query=…&category=…&sort=…&limit=…&offset=… → {styles, hasMore}""" query = request.query.get("query", "").strip() category = request.query.get("category", "").strip() or None sort = request.query.get("sort", "").strip() or None if sort not in (None, "recent", "popular"): sort = None limit = _int_param(request, "limit", 8) offset = _int_param(request, "offset", 0) try: payload = api.search_styles( query, category=category, sort=sort, limit=limit, offset=offset ) except StyleRefError as err: # 200 with an error field: the dialog renders the message inline rather # than the frontend surfacing a generic red toast with no detail. return web.json_response({"styles": [], "error": err.message}) return web.json_response( { "styles": [_card(s) for s in (payload.get("styles") or [])], "availableCategories": payload.get("availableCategories") or [], # Drives the dialog's Next button without it having to guess from a # short page — a full last page is indistinguishable otherwise. "hasMore": bool(payload.get("hasMore")), } ) @PromptServer.instance.routes.get("/styleref/saved") async def styleref_saved(request: web.Request) -> web.Response: """ The picker's "Saved" sort: gallery styles this user saved. Signed-in only, and reported as such rather than as an error — signing out is not a failure. """ signed_in, message = auth_status() if not signed_in: return web.json_response({"styles": [], "signedIn": False, "message": message}) query = request.query.get("query", "").strip() category = request.query.get("category", "").strip() or None try: payload = api.list_saved_styles( query, category=category, limit=_int_param(request, "limit", 8), offset=_int_param(request, "offset", 0), ) except StyleRefError as err: return web.json_response({"styles": [], "signedIn": True, "error": err.message}) return web.json_response( { "signedIn": True, "styles": [_card(s) for s in (payload.get("styles") or [])], "hasMore": bool(payload.get("hasMore")), } ) @PromptServer.instance.routes.get("/styleref/my-styles") async def styleref_my_styles(request: web.Request) -> web.Response: """The caller's own library, shown above gallery results when signed in.""" signed_in, message = auth_status() if not signed_in: return web.json_response({"styles": [], "signedIn": False, "message": message}) cursor = request.query.get("cursor", "").strip() or None query = request.query.get("query", "").strip() try: payload = api.list_my_styles( limit=_int_param(request, "limit", 8), cursor=cursor, query=query ) except StyleRefError as err: return web.json_response({"styles": [], "signedIn": True, "error": err.message}) return web.json_response( { "signedIn": True, "styles": [ { "slug": s.get("ref"), "name": s.get("name"), "author": None, "tags": [], "visibility": s.get("visibility"), # null → never generated on styleref.io, so loading it would # 409. The dialog badges these instead of offering them as # ready to use. "lastGeneratedAt": s.get("lastGeneratedAt"), } for s in (payload.get("styles") or []) ], # Passed back as ?cursor= for the dialog's Next button. "nextCursor": payload.get("nextCursor"), } ) @PromptServer.instance.routes.get("/styleref/resolve") async def styleref_resolve(request: web.Request) -> web.Response: """ GET /styleref/resolve?ref=… → {ok, name, url} Backs the name line under the Load node's `style_ref` field. A slug is opaque — `72e1zdae-e2d54a18d090` tells you nothing about which style is wired in, and until the graph is queued the node cannot say either. This answers the question the moment the field changes. It goes through the same client (and therefore the same session cache) the node uses, so a resolve for a style that is about to be loaded costs one request and the queue after it costs none. """ ref = request.query.get("ref", "").strip() if not ref: return web.json_response({"ok": False}) try: style = fetch_style(ref) except StyleRefError as err: # 200 with an error field, as with search: an unresolvable ref is # ordinary (a half-typed slug, a private style while signed out) and # belongs inline on the node, not in a red toast. return web.json_response({"ok": False, "error": err.message}) # No author: the structured lane carries name/sections/referenceImages only, # so the node body's attribution line is where a handle would come from. return web.json_response( {"ok": True, "name": style.get("name"), "url": style.get("url")} ) @PromptServer.instance.routes.post("/styleref/refresh") async def styleref_refresh(request: web.Request) -> web.Response: """ The ↻ Refresh style button's backend: forget the cached style and advance the generation counter IS_CHANGED reads, so the next queue re-downloads. """ try: payload = await request.json() except Exception: # noqa: BLE001 - malformed body is a client bug, not a crash payload = {} ref = str(payload.get("ref") or "").strip() if not ref: return web.json_response({"ok": False, "error": "No style_ref to refresh."}) generation = bump_refresh(ref) return web.json_response({"ok": True, "generation": generation}) # ── auth: the Login node's button backend ──────────────────────────────────── # The Login node's `action` dropdown stays as the queue-driven fallback (it is # the only path for API-only installs); these routes power the node's buttons. def _status_payload() -> dict: signed_in, message = auth_status() return {"signedIn": signed_in, "message": message} @PromptServer.instance.routes.get("/styleref/auth/status") async def styleref_auth_status(_request: web.Request) -> web.Response: return web.json_response(_status_payload()) @PromptServer.instance.routes.post("/styleref/auth/signin") async def styleref_auth_signin(_request: web.Request) -> web.Response: """ Run the OAuth loopback flow. Blocking by nature (it waits for the browser redirect), so it runs in an executor; the JS button shows a waiting state. """ loop = asyncio.get_event_loop() ok, message = await loop.run_in_executor(None, run_login_flow, api.api_base()) return web.json_response({"ok": ok, "message": message, **_status_payload()}) @PromptServer.instance.routes.post("/styleref/auth/signout") async def styleref_auth_signout(_request: web.Request) -> web.Response: removed = clear_credentials() message = ( f"Signed out — removed {credentials_path()}." if removed else "No stored credentials to remove." ) if os.environ.get("STYLEREF_TOKEN"): message += ( " Note: STYLEREF_TOKEN is still set in this environment and takes " "precedence, so you remain signed in via that token." ) return web.json_response({"ok": True, "message": message, **_status_payload()})