Back to the login screen when the session has ended

The ComfyUI page is served with no-store and carries a tiny script, run
before ComfyUI's own code, that goes to the login page when a request is
answered 401 with X-RigShare-Login: required (now sent by the login gate).

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+### Fixed
+
+- When the session had ended (logged out elsewhere, expired, cookie cleared), a ComfyUI page already open or restored from the browser cache showed "API key missing" errors instead of going back to the login screen. The page is no longer cached, and it returns to the login screen as soon as a request says login is required
+
 ## 1.2.2 (2026-09-26)
 
 A Queue tab showing who launched what, and chat clean-up for admins.

rigshare_core/routes.py

@@ -17,6 +17,21 @@ from .workspace import Workspace, norm
 
 log = logging.getLogger("ComfyUI-RigShare")
 
+# Runs in the ComfyUI page before anything else: any request answered "login
+# required" (session ended, cookie cleared) sends the browser to the login page.
+LOGIN_WATCH = """<script id="rigshare-login-watch">(() => {
+  const fetch0 = window.fetch;
+  let leaving = false;
+  window.fetch = async (...args) => {
+    const res = await fetch0(...args);
+    if (res.status === 401 && !leaving && res.headers.get("X-RigShare-Login") === "required") {
+      leaving = true;
+      location.href = "/rigshare/login?next=" + encodeURIComponent(location.pathname + location.search);
+    }
+    return res;
+  };
+})();</script>"""
+
 # (method, path) prefixes that need a permission when api_protection is on.
 PROTECTED = [
     ("POST", "/prompt", "queue"),
@@ -617,6 +632,25 @@ def setup(server, store, hub):
         action = {"GET": "get", "HEAD": "get", "POST": "save", "DELETE": "delete"}.get(request.method)
         return (action, src, None) if action else None
 
+    def watch_login(response):
+        try:
+            file_path = getattr(response, "_path", None)
+            if file_path:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    page = f.read()
+            elif isinstance(response, web.Response) and response.body and "html" in (response.content_type or ""):
+                page = response.body.decode("utf-8")
+            else:
+                return response
+            if "<head>" not in page or "rigshare-login-watch" in page:
+                return response
+            page = page.replace("<head>", "<head>" + LOGIN_WATCH, 1)
+            return web.Response(text=page, content_type="text/html",
+                                headers={"Cache-Control": "no-store, must-revalidate"})
+        except Exception as e:
+            log.debug(f"[RigShare] could not add the login watch: {e}")
+            return response
+
     def api_path(request):
         return request.path[4:] if request.path.startswith("/api/") else request.path
 
@@ -729,8 +763,10 @@ def setup(server, store, hub):
                     path in ("/", "/index.html") or "text/html" in request.headers.get("Accept", ""))
                 if wants_page:
                     raise web.HTTPFound(f"/rigshare/login?next={path}")
+                # The header lets the ComfyUI page (see LOGIN_WATCH) go back to the login screen.
                 return web.json_response({"error": "RigShare: login required. API clients: send "
-                                                   "'Authorization: Bearer <api key>'."}, status=401)
+                                                   "'Authorization: Bearer <api key>'."}, status=401,
+                                         headers={"X-RigShare-Login": "required"})
         userdata = parse_userdata(request)
         if userdata and not trusted(request.remote or ""):
             refusal = check_userdata(request, *userdata)
@@ -793,6 +829,11 @@ def setup(server, store, hub):
                     response = web.json_response(filter_templates(index), headers={"Cache-Control": "no-store"})
             except Exception as e:
                 log.debug(f"[RigShare] could not filter templates: {e}")
+        # The ComfyUI page: never cached (an ended session must reach the redirect
+        # above), and watched, so a session that ends while it is open goes back
+        # to the login screen instead of showing API errors.
+        if request.method == "GET" and request.path in ("/", "/index.html") and store.config.get("require_login"):
+            response = watch_login(response)
         # Sliding session: renew the cookie JWT when it is getting old.
         cookie = request.cookies.get(COOKIE)
         if cookie and isinstance(response, web.Response) and not response.prepared:

tests/python/test_server.py

@@ -495,3 +495,30 @@ def test_queue_shows_who_launched_and_guards_cancelling(run):
             assert not rig.ps.prompt_queue.pending
             assert (await h.post(rig.url("/api/queue"), json={"clear": True}, headers=admin)).status == 200
     run(scenario())
+
+
+def test_ended_session_goes_back_to_login(run, tmp_path):
+    index = tmp_path / "index.html"
+    index.write_text("<!doctype html><html><head><title>ComfyUI</title></head><body></body></html>")
+
+    async def scenario():
+        async def page(request):
+            from aiohttp import web
+            return web.FileResponse(str(index))
+
+        async with rig_server(setup=lambda app: app.router.add_get("/", page)) as rig:
+            admin = await rig.setup_admin()
+            h = rig.http
+            r = await h.get(rig.url("/"), headers=admin)
+            text = await r.text()
+            assert r.status == 200 and "rigshare-login-watch" in text and text.index("rigshare-login-watch") < text.index("<title>")
+            assert "no-store" in r.headers.get("Cache-Control", ""), "never served from the browser cache"
+
+            # Session gone: API calls say so in a header the page watches for...
+            await h.post(rig.url("/rigshare/api/logout"), headers=admin)
+            r = await h.get(rig.url("/api/prompt"), headers={"Authorization": "Bearer nope"})
+            assert r.status == 401 and r.headers.get("X-RigShare-Login") == "required"
+            # ...and the page itself redirects to the login screen.
+            r = await h.get(rig.url("/"), headers={"Accept": "text/html"}, allow_redirects=False)
+            assert r.status == 302 and r.headers["Location"].startswith("/rigshare/login")
+    run(scenario())