studio/server.py

import os
import sys
import asyncio
import traceback
import time

import nodes
import folder_paths
import execution
from studio_execution.jobs import JobStatus, get_job, get_all_jobs
import uuid
import urllib
import json
import glob
import struct
import ssl
import socket
import ipaddress
from PIL import Image, ImageOps
from PIL.PngImagePlugin import PngInfo
from io import BytesIO

import aiohttp
from aiohttp import web
import logging

import mimetypes
from studio.cli_args import args
import studio.utils
import studio.model_management
from studio_api import feature_flags
import node_helpers
from studio_version import __version__
from app.frontend_management import FrontendManager, parse_version
from studio_api.internal import _StudioNodeInternal
from app.assets.scanner import seed_assets
from app.assets.api.routes import register_assets_system

from app.user_manager import UserManager
from app.model_manager import ModelFileManager
from app.custom_node_manager import CustomNodeManager
from app.subgraph_manager import SubgraphManager
from app.node_replace_manager import NodeReplaceManager
from typing import Optional, Union
from api_server.routes.internal.internal_routes import InternalRoutes
from protocol import BinaryEventTypes

# Import cache control middleware
from middleware.cache_middleware import cache_control
from middleware.iam_auth_middleware import create_iam_auth_middleware
from middleware.rate_limit_middleware import create_rate_limit_middleware
from middleware import metrics_middleware
from middleware import billing_middleware
from middleware import compute_config
from middleware import prompt_router
from middleware import visor_client

if args.enable_manager:
    import comfyui_manager


def _remove_sensitive_from_queue(queue: list) -> list:
    """Remove sensitive data (index 5) from queue item tuples."""
    return [item[:5] for item in queue]


# ── Multi-tenant scoping for the inherited core endpoints (HIP-0506 security) ──
# /history and /queue are process-global in upstream ComfyUI. In a multi-tenant
# deployment that is cross-tenant disclosure (one org reads another's prompts,
# graphs, seeds) and a cross-tenant DoS (one org's "clear queue" wipes another's).
# Every item carries its owner in extra_data["org_id"] (bound at enqueue). These
# filter each list/dict to the caller's org. `org is None` means auth is OFF
# (single-tenant/local dev) → return everything unchanged (no behavior change).
def _item_org(item) -> str | None:
    """org_id from a queue-item tuple (index 3 = extra_data)."""
    try:
        return item[3].get("org_id") if len(item) > 3 and isinstance(item[3], dict) else None
    except Exception:
        return None


def _scope_queue_to_org(items: list, org: str | None) -> list:
    if not org:
        return items
    return [it for it in items if _item_org(it) == org]


def _scope_history_to_org(hist: dict, org: str | None) -> dict:
    if not org:
        return hist
    out = {}
    for k, entry in (hist or {}).items():
        try:
            ed = entry.get("prompt", [None, None, None, {}])[3]
            if isinstance(ed, dict) and ed.get("org_id") == org:
                out[k] = entry
        except Exception:
            continue  # malformed/legacy entry with no org → excluded (fail-closed)
    return out


async def send_socket_catch_exception(function, message):
    try:
        await function(message)
    except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError, BrokenPipeError, ConnectionError) as err:
        logging.warning("send error: {}".format(err))

# Track deprecated paths that have been warned about to only warn once per file
_deprecated_paths_warned = set()

@web.middleware
async def deprecation_warning(request: web.Request, handler):
    """Middleware to warn about deprecated frontend API paths"""
    path = request.path

    if path.startswith("/scripts/ui") or path.startswith("/extensions/core/"):
        # Only warn once per unique file path
        if path not in _deprecated_paths_warned:
            _deprecated_paths_warned.add(path)
            logging.warning(
                f"[DEPRECATION WARNING] Detected import of deprecated legacy API: {path}. "
                f"This is likely caused by a custom node extension using outdated APIs. "
                f"Please update your extensions or contact the extension author for an updated version."
            )

    response: web.Response = await handler(request)
    return response


@web.middleware
async def compress_body(request: web.Request, handler):
    accept_encoding = request.headers.get("Accept-Encoding", "")
    response: web.Response = await handler(request)
    if not isinstance(response, web.Response):
        return response
    if response.content_type not in ["application/json", "text/plain"]:
        return response
    if response.body and "gzip" in accept_encoding:
        response.enable_compression()
    return response


def create_cors_middleware(allowed_origin: str):
    @web.middleware
    async def cors_middleware(request: web.Request, handler):
        if request.method == "OPTIONS":
            # Pre-flight request. Reply successfully:
            response = web.Response()
        else:
            response = await handler(request)

        response.headers['Access-Control-Allow-Origin'] = allowed_origin
        response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, PUT, OPTIONS, PATCH'
        response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'
        response.headers['Access-Control-Allow-Credentials'] = 'true'
        return response

    return cors_middleware

def is_loopback(host):
    if host is None:
        return False
    try:
        if ipaddress.ip_address(host).is_loopback:
            return True
        else:
            return False
    except:
        pass

    loopback = False
    for family in (socket.AF_INET, socket.AF_INET6):
        try:
            r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM)
            for family, _, _, _, sockaddr in r:
                if not ipaddress.ip_address(sockaddr[0]).is_loopback:
                    return loopback
                else:
                    loopback = True
        except socket.gaierror:
            pass

    return loopback


def create_origin_only_middleware():
    @web.middleware
    async def origin_only_middleware(request: web.Request, handler):
        #this code is used to prevent the case where a random website can queue studio workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason.
        #in that case the Host and Origin hostnames won't match
        #I know the proper fix would be to add a cookie but this should take care of the problem in the meantime
        if 'Host' in request.headers and 'Origin' in request.headers:
            host = request.headers['Host']
            origin = request.headers['Origin']
            host_domain = host.lower()
            parsed = urllib.parse.urlparse(origin)
            origin_domain = parsed.netloc.lower()
            host_domain_parsed = urllib.parse.urlsplit('//' + host_domain)

            #limit the check to when the host domain is localhost, this makes it slightly less safe but should still prevent the exploit
            loopback = is_loopback(host_domain_parsed.hostname)

            if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host
                host_domain = host_domain_parsed.hostname
            if host_domain_parsed.port is None:
                origin_domain = parsed.hostname

            if loopback and host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0:
                if host_domain != origin_domain:
                    logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain))
                    return web.Response(status=403)

        if request.method == "OPTIONS":
            response = web.Response()
        else:
            response = await handler(request)

        return response

    return origin_only_middleware


def create_block_external_middleware():
    @web.middleware
    async def block_external_middleware(request: web.Request, handler):
        if request.method == "OPTIONS":
            # Pre-flight request. Reply successfully:
            response = web.Response()
        else:
            response = await handler(request)

        response.headers['Content-Security-Policy'] = "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval' blob:; style-src 'self' 'unsafe-inline'; img-src 'self' data: blob:; font-src 'self'; connect-src 'self' data:; frame-src 'self'; object-src 'self';"
        return response

    return block_external_middleware


class PromptServer():
    def __init__(self, loop):
        PromptServer.instance = self

        mimetypes.init()
        mimetypes.add_type('application/javascript; charset=utf-8', '.js')
        mimetypes.add_type('image/webp', '.webp')

        self.user_manager = UserManager()
        self.model_file_manager = ModelFileManager()
        self.custom_node_manager = CustomNodeManager()
        self.subgraph_manager = SubgraphManager()
        self.node_replace_manager = NodeReplaceManager()
        self.internal_routes = InternalRoutes(self)
        self.supports = ["custom_nodes_from_web"]
        # Queue backend precedence: SQLite (durable, zero external process) >
        # PERSIST (JSON snapshot) > memory. STUDIO_QUEUE_DB opts into the
        # crash-durable SQLite queue; without it Studio behaves exactly as
        # before (in-memory + optional JSON snapshot).
        if os.environ.get("STUDIO_QUEUE_DB"):
            from middleware.tasks_queue import SqlitePromptQueue
            self.prompt_queue = SqlitePromptQueue(self)
        else:
            self.prompt_queue = execution.PromptQueue(self)
        self.loop = loop
        self.messages = asyncio.Queue()
        self.client_session:Optional[aiohttp.ClientSession] = None
        self.number = 0

        from middleware.studio_home import home_redirect
        middlewares = [cache_control, deprecation_warning, home_redirect]

        self._iam_auth = None
        if args.enable_iam_auth:
            localhost_bypass = not args.no_localhost_bypass
            self._iam_auth = create_iam_auth_middleware(
                iam_url=args.iam_url,
                localhost_bypass=localhost_bypass,
            )
            middlewares.append(self._iam_auth)

        # Rate limiting (after auth so we have org context)
        middlewares.append(create_rate_limit_middleware(rpm=args.rate_limit_rpm))

        if args.enable_compress_response_body:
            middlewares.append(compress_body)

        if args.enable_cors_header:
            middlewares.append(create_cors_middleware(args.enable_cors_header))
        else:
            middlewares.append(create_origin_only_middleware())

        if args.disable_api_nodes:
            middlewares.append(create_block_external_middleware())

        if args.enable_manager:
            middlewares.append(comfyui_manager.create_middleware())

        max_upload_size = round(args.max_upload_size * 1024 * 1024)
        self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)
        # OIDC Authorization Code flow (standard code flow, standalone app):
        # /login begins it (the login page's sole entry point), /callback ends it.
        if self._iam_auth is not None:
            self.app.router.add_get("/login", self._iam_auth.handle_login)
            self.app.router.add_get("/callback", self._iam_auth.handle_callback)
        self.sockets = dict()
        self.sockets_metadata = dict()
        self.web_root = (
            FrontendManager.init_frontend(args.front_end_version)
            if args.front_end_root is None
            else args.front_end_root
        )
        logging.info(f"[Prompt Server] web root: {self.web_root}")
        register_assets_system(self.app, self.user_manager)
        routes = web.RouteTableDef()
        self.routes = routes
        self.last_node_id = None
        self.client_id = None

        self.on_prompt_handlers = []

        @routes.get('/ws')
        async def websocket_handler(request):
            ws = web.WebSocketResponse()
            await ws.prepare(request)
            sid = request.rel_url.query.get('clientId', '')
            if sid:
                # Reusing existing session, remove old
                self.sockets.pop(sid, None)
            else:
                sid = uuid.uuid4().hex

            # Store WebSocket for backward compatibility
            self.sockets[sid] = ws
            # Store metadata separately
            self.sockets_metadata[sid] = {"feature_flags": {}}

            try:
                # Send initial state to the new client
                await self.send("status", {"status": self.get_queue_info(), "sid": sid}, sid)
                # On reconnect if we are the currently executing client send the current node
                if self.client_id == sid and self.last_node_id is not None:
                    await self.send("executing", { "node": self.last_node_id }, sid)

                # Flag to track if we've received the first message
                first_message = True

                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.ERROR:
                        logging.warning('ws connection closed with exception %s' % ws.exception())
                    elif msg.type == aiohttp.WSMsgType.TEXT:
                        try:
                            data = json.loads(msg.data)
                            # Check if first message is feature flags
                            if first_message and data.get("type") == "feature_flags":
                                # Store client feature flags
                                client_flags = data.get("data", {})
                                self.sockets_metadata[sid]["feature_flags"] = client_flags

                                # Send server feature flags in response
                                await self.send(
                                    "feature_flags",
                                    feature_flags.get_server_features(),
                                    sid,
                                )

                                logging.debug(
                                    f"Feature flags negotiated for client {sid}: {client_flags}"
                                )
                            first_message = False
                        except json.JSONDecodeError:
                            logging.warning(
                                f"Invalid JSON received from client {sid}: {msg.data}"
                            )
                        except Exception as e:
                            logging.error(f"Error processing WebSocket message: {e}")
            finally:
                self.sockets.pop(sid, None)
                self.sockets_metadata.pop(sid, None)
            return ws

        @routes.get("/health")
        async def health_check(request):
            return web.json_response({"status": "ok"})

        @routes.get("/ready")
        async def readiness_check(request):
            # The render worker must actually be ALIVE — not merely that the queue
            # object exists (prompt_queue is set once at boot and never cleared, so
            # it is constant-true and says nothing about the executor). A dead worker
            # daemon thread accepts /prompt but runs nothing; readiness must fail so
            # the pod leaves the Service and a roll never promotes a wedged pod.
            wt = getattr(self, "prompt_worker_thread", None)
            worker_alive = wt.is_alive() if wt is not None else (self.prompt_queue is not None)
            ready = os.path.isdir(folder_paths.get_input_directory()) and worker_alive
            if ready:
                return web.json_response({"status": "ready"})
            return web.json_response({"status": "not_ready"}, status=503)

        @routes.get("/metrics")
        async def prometheus_metrics(request):
            if not args.enable_metrics:
                return web.Response(status=404)
            # Update live gauges before rendering
            metrics_middleware.set_queue_depth(
                self.prompt_queue.get_tasks_remaining()
                if self.prompt_queue else 0
            )
            metrics_middleware.set_websocket_connections(len(self.sockets))
            body = metrics_middleware.render_metrics()
            return web.Response(
                text=body,
                content_type="text/plain",
                charset="utf-8",
                headers={"X-Content-Type-Options": "nosniff"},
            )

        @routes.get("/")
        async def get_root(request):
            # If IAM auth is enabled, show login page to unauthenticated users
            if args.enable_iam_auth:
                has_token = bool(
                    request.headers.get("Authorization", "").startswith("Bearer ")
                    or request.cookies.get("hanzo_token")
                    or request.cookies.get("access_token")
                )
                if not has_token:
                    login_page = os.path.join(
                        os.path.dirname(os.path.realpath(__file__)), "web", "login.html"
                    )
                    if os.path.isfile(login_page):
                        with open(login_page, "r") as f:
                            html = f.read()
                        return web.Response(
                            text=html,
                            content_type="text/html",
                            headers={
                                "Cache-Control": "no-cache",
                                "Pragma": "no-cache",
                                "Expires": "0",
                            },
                        )

            # Serve the editor SPA with the shared Studio shell injected (the ONE
            # chrome — Studio/Editor toggle, wallet, GPUs, chat sidebar — that the
            # Studio home also includes). No-cache, same as before.
            from middleware.studio_home import editor_index
            return editor_index(self.web_root)

        @routes.get("/embeddings")
        def get_embeddings(request):
            embeddings = folder_paths.get_filename_list("embeddings")
            return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings)))

        @routes.get("/models")
        def list_model_types(request):
            model_types = list(folder_paths.folder_names_and_paths.keys())

            return web.json_response(model_types)

        @routes.get("/models/{folder}")
        async def get_models(request):
            folder = request.match_info.get("folder", None)
            if folder not in folder_paths.folder_names_and_paths:
                return web.Response(status=404)
            files = folder_paths.get_filename_list(folder)
            return web.json_response(files)

        @routes.get("/extensions")
        async def get_extensions(request):
            files = glob.glob(os.path.join(
                glob.escape(self.web_root), 'extensions/**/*.js'), recursive=True)

            extensions = list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files))

            for name, dir in nodes.EXTENSION_WEB_DIRS.items():
                files = glob.glob(os.path.join(glob.escape(dir), '**/*.js'), recursive=True)
                extensions.extend(list(map(lambda f: "/extensions/" + urllib.parse.quote(
                    name) + "/" + os.path.relpath(f, dir).replace("\\", "/"), files)))

            return web.json_response(extensions)

        def get_dir_by_type(dir_type, request=None):
            if dir_type is None:
                dir_type = "input"

            # Use org-scoped directories if multi-tenant
            org_id = None
            if request and request.get("iam_user"):
                org_id = request["iam_user"].get("org_id")

            if dir_type == "input":
                type_dir = folder_paths.get_org_input_directory(org_id)
            elif dir_type == "temp":
                type_dir = folder_paths.get_org_temp_directory(org_id)
            elif dir_type == "output":
                type_dir = folder_paths.get_org_output_directory(org_id)
            else:
                type_dir = folder_paths.get_input_directory()

            return type_dir, dir_type

        def compare_image_hash(filepath, image):
            hasher = node_helpers.hasher()

            # function to compare hashes of two images to see if it already exists, fix to #3465
            if os.path.exists(filepath):
                a = hasher()
                b = hasher()
                with open(filepath, "rb") as f:
                    a.update(f.read())
                    b.update(image.file.read())
                    image.file.seek(0)
                return a.hexdigest() == b.hexdigest()
            return False

        def image_upload(post, image_save_function=None, request=None):
            image = post.get("image")
            overwrite = post.get("overwrite")
            image_is_duplicate = False

            image_upload_type = post.get("type")
            # Forward the request so uploads land in the caller's org-scoped
            # directory (orgs/{org}/output) in multi-tenant mode.
            upload_dir, image_upload_type = get_dir_by_type(image_upload_type, request)

            if image and image.file:
                filename = image.filename
                if not filename:
                    return web.Response(status=400)

                subfolder = post.get("subfolder", "")
                full_output_folder = os.path.join(upload_dir, os.path.normpath(subfolder))
                filepath = os.path.abspath(os.path.join(full_output_folder, filename))

                if os.path.commonpath((upload_dir, filepath)) != upload_dir:
                    return web.Response(status=400)

                if not os.path.exists(full_output_folder):
                    os.makedirs(full_output_folder)

                split = os.path.splitext(filename)

                if overwrite is not None and (overwrite == "true" or overwrite == "1"):
                    pass
                else:
                    i = 1
                    while os.path.exists(filepath):
                        if compare_image_hash(filepath, image): #compare hash to prevent saving of duplicates with same name, fix for #3465
                            image_is_duplicate = True
                            break
                        filename = f"{split[0]} ({i}){split[1]}"
                        filepath = os.path.join(full_output_folder, filename)
                        i += 1

                if not image_is_duplicate:
                    if image_save_function is not None:
                        image_save_function(image, post, filepath)
                    else:
                        with open(filepath, "wb") as f:
                            f.write(image.file.read())

                return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type})
            else:
                return web.Response(status=400)

        @routes.post("/upload/image")
        async def upload_image(request):
            post = await request.post()
            return image_upload(post, request=request)

        @routes.post("/upload/output")
        async def upload_output(request):
            # Token-authorized gallery ingest for BYO-GPU workers. A remote worker
            # (hanzod's gpu-fleet loop) renders a studio.render job on its own GPU
            # and POSTs each finished image here with the user's IAM bearer. The IAM
            # middleware validates the token and sets request["iam_user"]; org is
            # derived from that (owner/groups), so the file lands in this org's
            # orgs/{org}/output — which the PVC + S3 mirror persist into the gallery.
            # No S3/rclone credentials ever touch the worker box: the user's session
            # token is the only credential. Replaces the old rclone/S3 sidecar sync.
            post = await request.post()
            # Force type=output regardless of what the client sent.
            forced = dict(post)
            forced["type"] = "output"
            resp = image_upload(forced, request=request)
            # Fix outputs: restore the render's global colour to its source, undoing
            # the full-regen (denoise 1.0) white-balance/exposure drift and its
            # fix-of-a-fix compounding. Pod-side + fail-safe, so it can never affect a
            # render's success; a render that didn't drift matches to itself untouched.
            try:
                saved = json.loads(resp.body) if getattr(resp, "body", None) else {}
                if saved.get("name"):
                    from middleware.studio_home import colormatch_fix_ingest
                    colormatch_fix_ingest(request, saved.get("subfolder", ""), saved["name"])
            except Exception:
                pass
            return resp

        @routes.post("/upload/mask")
        async def upload_mask(request):
            post = await request.post()

            def image_save_function(image, post, filepath):
                original_ref = json.loads(post.get("original_ref"))
                filename, output_dir = folder_paths.annotated_filepath(original_ref['filename'])

                if not filename:
                    return web.Response(status=400)

                # validation for security: prevent accessing arbitrary path
                if filename[0] == '/' or '..' in filename:
                    return web.Response(status=400)

                if output_dir is None:
                    type = original_ref.get("type", "output")
                    output_dir = folder_paths.get_directory_by_type(type)

                if output_dir is None:
                    return web.Response(status=400)

                if original_ref.get("subfolder", "") != "":
                    full_output_dir = os.path.join(output_dir, original_ref["subfolder"])
                    if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:
                        return web.Response(status=403)
                    output_dir = full_output_dir

                file = os.path.join(output_dir, filename)

                if os.path.isfile(file):
                    with Image.open(file) as original_pil:
                        metadata = PngInfo()
                        if hasattr(original_pil,'text'):
                            for key in original_pil.text:
                                metadata.add_text(key, original_pil.text[key])
                        original_pil = original_pil.convert('RGBA')
                        mask_pil = Image.open(image.file).convert('RGBA')

                        # alpha copy
                        new_alpha = mask_pil.getchannel('A')
                        original_pil.putalpha(new_alpha)
                        original_pil.save(filepath, compress_level=4, pnginfo=metadata)

            return image_upload(post, image_save_function, request=request)

        @routes.get("/view")
        async def view_image(request):
            if "filename" in request.rel_url.query:
                filename = request.rel_url.query["filename"]
                filename, output_dir = folder_paths.annotated_filepath(filename)

                if not filename:
                    return web.Response(status=400)

                # validation for security: prevent accessing arbitrary path
                if filename[0] == '/' or '..' in filename:
                    return web.Response(status=400)

                if output_dir is None:
                    type = request.rel_url.query.get("type", "output")
                    output_dir = folder_paths.get_directory_by_type(type)

                if output_dir is None:
                    return web.Response(status=400)

                if "subfolder" in request.rel_url.query:
                    full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"])
                    if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:
                        return web.Response(status=403)
                    output_dir = full_output_dir

                filename = os.path.basename(filename)
                file = os.path.join(output_dir, filename)

                if os.path.isfile(file):
                    if 'preview' in request.rel_url.query:
                        with Image.open(file) as img:
                            preview_info = request.rel_url.query['preview'].split(';')
                            image_format = preview_info[0]
                            if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''):
                                image_format = 'webp'

                            quality = 90
                            if preview_info[-1].isdigit():
                                quality = int(preview_info[-1])

                            buffer = BytesIO()
                            if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb':
                                img = img.convert("RGB")
                            img.save(buffer, format=image_format, quality=quality)
                            buffer.seek(0)

                            return web.Response(body=buffer.read(), content_type=f'image/{image_format}',
                                                headers={"Content-Disposition": f"filename=\"{filename}\""})

                    if 'channel' not in request.rel_url.query:
                        channel = 'rgba'
                    else:
                        channel = request.rel_url.query["channel"]

                    if channel == 'rgb':
                        with Image.open(file) as img:
                            if img.mode == "RGBA":
                                r, g, b, a = img.split()
                                new_img = Image.merge('RGB', (r, g, b))
                            else:
                                new_img = img.convert("RGB")

                            buffer = BytesIO()
                            new_img.save(buffer, format='PNG')
                            buffer.seek(0)

                            return web.Response(body=buffer.read(), content_type='image/png',
                                                headers={"Content-Disposition": f"filename=\"{filename}\""})

                    elif channel == 'a':
                        with Image.open(file) as img:
                            if img.mode == "RGBA":
                                _, _, _, a = img.split()
                            else:
                                a = Image.new('L', img.size, 255)

                            # alpha img
                            alpha_img = Image.new('RGBA', img.size)
                            alpha_img.putalpha(a)
                            alpha_buffer = BytesIO()
                            alpha_img.save(alpha_buffer, format='PNG')
                            alpha_buffer.seek(0)

                            return web.Response(body=alpha_buffer.read(), content_type='image/png',
                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
                    else:
                        # Get content type from mimetype, defaulting to 'application/octet-stream'
                        content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream'

                        # For security, force certain mimetypes to download instead of display
                        if content_type in {'text/html', 'text/html-sandboxed', 'application/xhtml+xml', 'text/javascript', 'text/css'}:
                            content_type = 'application/octet-stream'  # Forces download

                        return web.FileResponse(
                            file,
                            headers={
                                "Content-Disposition": f"filename=\"{filename}\"",
                                "Content-Type": content_type
                            }
                        )

            return web.Response(status=404)

        @routes.get("/view_metadata/{folder_name}")
        async def view_metadata(request):
            folder_name = request.match_info.get("folder_name", None)
            if folder_name is None:
                return web.Response(status=404)
            if "filename" not in request.rel_url.query:
                return web.Response(status=404)

            filename = request.rel_url.query["filename"]
            if not filename.endswith(".safetensors"):
                return web.Response(status=404)

            safetensors_path = folder_paths.get_full_path(folder_name, filename)
            if safetensors_path is None:
                return web.Response(status=404)
            out = studio.utils.safetensors_header(safetensors_path, max_size=1024*1024)
            if out is None:
                return web.Response(status=404)
            dt = json.loads(out)
            if "__metadata__" not in dt:
                return web.Response(status=404)
            return web.json_response(dt["__metadata__"])

        @routes.get("/system_stats")
        async def system_stats(request):
            device = studio.model_management.get_torch_device()
            device_name = studio.model_management.get_torch_device_name(device)
            cpu_device = studio.model_management.torch.device("cpu")
            ram_total = studio.model_management.get_total_memory(cpu_device)
            ram_free = studio.model_management.get_free_memory(cpu_device)
            vram_total, torch_vram_total = studio.model_management.get_total_memory(device, torch_total_too=True)
            vram_free, torch_vram_free = studio.model_management.get_free_memory(device, torch_free_too=True)
            required_frontend_version = FrontendManager.get_required_frontend_version()
            installed_templates_version = FrontendManager.get_installed_templates_version()
            required_templates_version = FrontendManager.get_required_templates_version()

            system_stats = {
                "system": {
                    "os": sys.platform,
                    "ram_total": ram_total,
                    "ram_free": ram_free,
                    "studio_version": __version__,
                    "required_frontend_version": required_frontend_version,
                    "installed_templates_version": installed_templates_version,
                    "required_templates_version": required_templates_version,
                    "python_version": sys.version,
                    "pytorch_version": studio.model_management.torch_version,
                    "embedded_python": os.path.split(os.path.split(sys.executable)[0])[1] == "python_embeded",
                    "argv": sys.argv
                },
                "devices": [
                    {
                        "name": device_name,
                        "type": device.type,
                        "index": device.index,
                        "vram_total": vram_total,
                        "vram_free": vram_free,
                        "torch_vram_total": torch_vram_total,
                        "torch_vram_free": torch_vram_free,
                    }
                ]
            }
            return web.json_response(system_stats)

        @routes.get("/features")
        async def get_features(request):
            return web.json_response(feature_flags.get_server_features())

        @routes.get("/prompt")
        async def get_prompt(request):
            return web.json_response(self.get_queue_info())

        def node_info(node_class):
            obj_class = nodes.NODE_CLASS_MAPPINGS[node_class]
            if issubclass(obj_class, _StudioNodeInternal):
                return obj_class.GET_NODE_INFO_V1()
            info = {}
            info['input'] = obj_class.INPUT_TYPES()
            info['input_order'] = {key: list(value.keys()) for (key, value) in obj_class.INPUT_TYPES().items()}
            info['is_input_list'] = getattr(obj_class, "INPUT_IS_LIST", False)
            info['output'] = obj_class.RETURN_TYPES
            info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)
            info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']
            info['name'] = node_class
            info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class
            info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else ''
            info['python_module'] = getattr(obj_class, "RELATIVE_PYTHON_MODULE", "nodes")
            info['category'] = 'sd'
            if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True:
                info['output_node'] = True
            else:
                info['output_node'] = False

            if hasattr(obj_class, 'CATEGORY'):
                info['category'] = obj_class.CATEGORY

            if hasattr(obj_class, 'OUTPUT_TOOLTIPS'):
                info['output_tooltips'] = obj_class.OUTPUT_TOOLTIPS

            if getattr(obj_class, "DEPRECATED", False):
                info['deprecated'] = True
            if getattr(obj_class, "EXPERIMENTAL", False):
                info['experimental'] = True
            if getattr(obj_class, "DEV_ONLY", False):
                info['dev_only'] = True

            if hasattr(obj_class, 'API_NODE'):
                info['api_node'] = obj_class.API_NODE

            info['search_aliases'] = getattr(obj_class, 'SEARCH_ALIASES', [])

            if hasattr(obj_class, 'ESSENTIALS_CATEGORY'):
                info['essentials_category'] = obj_class.ESSENTIALS_CATEGORY

            return info

        @routes.get("/object_info")
        async def get_object_info(request):
            try:
                seed_assets(["models"])
            except Exception as e:
                logging.error(f"Failed to seed assets: {e}")
            with folder_paths.cache_helper:
                out = {}
                for x in nodes.NODE_CLASS_MAPPINGS:
                    try:
                        out[x] = node_info(x)
                    except Exception:
                        logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
                        logging.error(traceback.format_exc())
                return web.json_response(out)

        @routes.get("/object_info/{node_class}")
        async def get_object_info_node(request):
            node_class = request.match_info.get("node_class", None)
            out = {}
            if (node_class is not None) and (node_class in nodes.NODE_CLASS_MAPPINGS):
                out[node_class] = node_info(node_class)
            return web.json_response(out)

        @routes.get("/api/jobs")
        async def get_jobs(request):
            """List all jobs with filtering, sorting, and pagination.

            Query parameters:
                status: Filter by status (comma-separated): pending, in_progress, completed, failed
                workflow_id: Filter by workflow ID
                sort_by: Sort field: created_at (default), execution_duration
                sort_order: Sort direction: asc, desc (default)
                limit: Max items to return (positive integer)
                offset: Items to skip (non-negative integer, default 0)
            """
            query = request.rel_url.query

            status_param = query.get('status')
            workflow_id = query.get('workflow_id')
            sort_by = query.get('sort_by', 'created_at').lower()
            sort_order = query.get('sort_order', 'desc').lower()

            status_filter = None
            if status_param:
                status_filter = [s.strip().lower() for s in status_param.split(',') if s.strip()]
                invalid_statuses = [s for s in status_filter if s not in JobStatus.ALL]
                if invalid_statuses:
                    return web.json_response(
                        {"error": f"Invalid status value(s): {', '.join(invalid_statuses)}. Valid values: {', '.join(JobStatus.ALL)}"},
                        status=400
                    )

            if sort_by not in {'created_at', 'execution_duration'}:
                return web.json_response(
                    {"error": "sort_by must be 'created_at' or 'execution_duration'"},
                    status=400
                )

            if sort_order not in {'asc', 'desc'}:
                return web.json_response(
                    {"error": "sort_order must be 'asc' or 'desc'"},
                    status=400
                )

            limit = None

            # If limit is provided, validate that it is a positive integer, else continue without a limit
            if 'limit' in query:
                try:
                    limit = int(query.get('limit'))
                    if limit <= 0:
                        return web.json_response(
                            {"error": "limit must be a positive integer"},
                            status=400
                        )
                except (ValueError, TypeError):
                    return web.json_response(
                        {"error": "limit must be an integer"},
                        status=400
                    )

            offset = 0
            if 'offset' in query:
                try:
                    offset = int(query.get('offset'))
                    if offset < 0:
                        offset = 0
                except (ValueError, TypeError):
                    return web.json_response(
                        {"error": "offset must be an integer"},
                        status=400
                    )

            running, queued = self.prompt_queue.get_current_queue_volatile()
            history = self.prompt_queue.get_history()

            running = _remove_sensitive_from_queue(running)
            queued = _remove_sensitive_from_queue(queued)

            jobs, total = get_all_jobs(
                running, queued, history,
                status_filter=status_filter,
                workflow_id=workflow_id,
                sort_by=sort_by,
                sort_order=sort_order,
                limit=limit,
                offset=offset
            )

            has_more = (offset + len(jobs)) < total

            return web.json_response({
                'jobs': jobs,
                'pagination': {
                    'offset': offset,
                    'limit': limit,
                    'total': total,
                    'has_more': has_more
                }
            })

        @routes.get("/api/jobs/{job_id}")
        async def get_job_by_id(request):
            """Get a single job by ID."""
            job_id = request.match_info.get("job_id", None)
            if not job_id:
                return web.json_response(
                    {"error": "job_id is required"},
                    status=400
                )

            running, queued = self.prompt_queue.get_current_queue_volatile()
            history = self.prompt_queue.get_history(prompt_id=job_id)

            running = _remove_sensitive_from_queue(running)
            queued = _remove_sensitive_from_queue(queued)

            job = get_job(job_id, running, queued, history)
            if job is None:
                return web.json_response(
                    {"error": "Job not found"},
                    status=404
                )

            return web.json_response(job)

        @routes.get("/history")
        async def get_history(request):
            max_items = request.rel_url.query.get("max_items", None)
            if max_items is not None:
                max_items = int(max_items)

            offset = request.rel_url.query.get("offset", None)
            if offset is not None:
                offset = int(offset)
            else:
                offset = -1

            hist = self.prompt_queue.get_history(max_items=max_items, offset=offset)
            hist = _scope_history_to_org(hist, self._caller_org(request))
            return web.json_response(hist)

        @routes.get("/history/{prompt_id}")
        async def get_history_prompt_id(request):
            prompt_id = request.match_info.get("prompt_id", None)
            hist = self.prompt_queue.get_history(prompt_id=prompt_id)
            hist = _scope_history_to_org(hist, self._caller_org(request))
            return web.json_response(hist)

        @routes.get("/queue")
        async def get_queue(request):
            org = self._caller_org(request)
            queue_info = {}
            current_queue = self.prompt_queue.get_current_queue_volatile()
            queue_info['queue_running'] = _scope_queue_to_org(_remove_sensitive_from_queue(current_queue[0]), org)
            queue_info['queue_pending'] = _scope_queue_to_org(_remove_sensitive_from_queue(current_queue[1]), org)
            return web.json_response(queue_info)

        @routes.post("/prompt")
        async def post_prompt(request):
            logging.info("got prompt")

            # Execution-engine (worker) mode: this process IS a BYO-GPU render backend,
            # not a front. A direct /prompt POST without the coordinator token is refused
            # so no render reaches a GPU without appearing in the org's visible queue.
            # One way onto a GPU: enqueue -> claim -> execute. (Policy: worker_client.)
            from middleware.worker_client import reject_untrusted_worker_submit
            _gate = reject_untrusted_worker_submit(request, args.worker_mode)
            if _gate is not None:
                return _gate

            # Billing: check balance before accepting prompt
            if args.enable_billing and args.billing_check_balance:
                iam_user = request.get("iam_user")
                if iam_user:
                    user_id = iam_user.get("sub", "anonymous")
                    org_id = iam_user.get("org_id", "default")
                    has_balance, available = await billing_middleware.check_balance(
                        user_id, org_id
                    )
                    if not has_balance:
                        return web.json_response(
                            {"error": "Insufficient balance", "available": available},
                            status=402,
                        )

            json_data =  await request.json()
            json_data = self.trigger_on_prompt(json_data)

            # --- Prompt routing: check if this should go to a GPU worker ---
            # Legacy in-cluster PUSH federation (prompt_router -> compute_config
            # workers) is OFF by default: the gpu-jobs enqueue below (dispatch_if_worker)
            # is the ONE submit path. Set STUDIO_LEGACY_PUSH_ROUTER=1 only to resurrect
            # the old push seam. Keeps exactly one way onto a GPU in prod.
            if not args.worker_mode and os.environ.get("STUDIO_LEGACY_PUSH_ROUTER") == "1":
                org_id = self._get_org_id(request)
                try:
                    route_result = await prompt_router.route_prompt(org_id, json_data)
                except Exception as e:
                    logging.warning("Prompt routing error: %s", e)
                    route_result = None

                if route_result is not None:
                    action = route_result.get("action")
                    if action == "forward":
                        # Successfully forwarded to a GPU worker
                        return web.json_response(route_result["response"])
                    elif action == "provisioning":
                        return web.json_response(
                            {"status": "provisioning", "message": "GPU worker is being provisioned. Retry in 30-60 seconds."},
                            status=202,
                        )
                    elif action == "unavailable":
                        return web.json_response(
                            {"error": "GPU requested but no GPU worker available. Use POST /api/compute/provision to launch one."},
                            status=503,
                        )

            if "number" in json_data:
                number = float(json_data['number'])
            else:
                number = self.number
                if "front" in json_data:
                    if json_data['front']:
                        number = -number

                self.number += 1

            if "prompt" in json_data:
                prompt = json_data["prompt"]
                prompt_id = str(json_data.get("prompt_id", uuid.uuid4()))

                partial_execution_targets = None
                if "partial_execution_targets" in json_data:
                    partial_execution_targets = json_data["partial_execution_targets"]

                self.node_replace_manager.apply_replacements(prompt)

                org_id = self._get_org_id(request)  # IAM org — scopes outputs + billing
                # BYO-GPU dispatch: cloud pods only. A worker-mode node IS the
                # render backend — it must queue locally, never re-dispatch.
                # Runs BEFORE local validation: the job executes on the org's GPU
                # box, which validates against ITS node classes and model files. A
                # GPU-less pod has an empty model list and would wrongly reject
                # every checkpoint-referencing graph ("not in []").
                if not args.worker_mode:
                    try:
                        from middleware.gpu_dispatch import dispatch_if_worker, has_models_for
                    except ImportError:  # package-style checkout
                        from .middleware.gpu_dispatch import dispatch_if_worker, has_models_for
                    if dispatch_if_worker(request, org_id, prompt_id, prompt):
                        return web.json_response({"prompt_id": prompt_id, "number": number, "node_errors": {}})
                    # Dispatch didn't happen. If THIS box lacks the models THIS graph
                    # needs (a coordinator pod, even one with a stray SD1.5), local
                    # validation is GUARANTEED to fail with a cryptic "not in []" — the
                    # render belongs on the org's GPU worker, momentarily unreachable.
                    # Return a clear, retryable error instead of the doomed validation.
                    # A box that CAN load the graph falls through and renders it.
                    if not has_models_for(prompt):
                        return web.json_response({"error": {
                            "type": "gpu_worker_unavailable",
                            "message": "GPU render worker is momentarily unavailable — please retry.",
                            "details": "", "extra_info": {}}, "node_errors": {}}, status=503)

                valid = await execution.validate_prompt(prompt_id, prompt, partial_execution_targets)
                extra_data = {}
                if "extra_data" in json_data:
                    extra_data = json_data["extra_data"]

                if "client_id" in json_data:
                    extra_data["client_id"] = json_data["client_id"]
                if valid[0]:
                    outputs_to_execute = valid[2]
                    sensitive = {}
                    for sensitive_val in execution.SENSITIVE_EXTRA_DATA_KEYS:
                        if sensitive_val in extra_data:
                            sensitive[sensitive_val] = extra_data.pop(sensitive_val)
                    extra_data["create_time"] = int(time.time() * 1000)  # timestamp in milliseconds
                    extra_data["org_id"] = org_id
                    # Carry the requester's verified IAM token to the worker so the
                    # finished render is recorded in the content lane as this user,
                    # into this user's org. `sensitive` is stripped from history (see
                    # SENSITIVE_EXTRA_DATA_KEYS) — the token is never persisted.
                    iam_token = request.get("iam_token")
                    if iam_token:
                        sensitive["iam_token"] = iam_token
                    self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute, sensitive))
                    response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
                    return web.json_response(response)
                else:
                    logging.warning("invalid prompt: {}".format(valid[1]))
                    return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400)
            else:
                error = {
                    "type": "no_prompt",
                    "message": "No prompt provided",
                    "details": "No prompt provided",
                    "extra_info": {}
                }
                return web.json_response({"error": error, "node_errors": {}}, status=400)

        @routes.post("/queue")
        async def post_queue(request):
            json_data =  await request.json()
            org = self._caller_org(request)  # None = auth off → upstream behavior
            if "clear" in json_data:
                if json_data["clear"]:
                    if org is None:
                        self.prompt_queue.wipe_queue()
                    else:
                        # SCOPED clear: only THIS org's pending items, never others'.
                        self.prompt_queue.delete_queue_item(lambda a: _item_org(a) == org)
            if "delete" in json_data:
                to_delete = json_data['delete']
                for id_to_delete in to_delete:
                    # Only delete an item the caller owns (org must match when auth on).
                    delete_func = (lambda a, i=id_to_delete: a[1] == i) if org is None \
                        else (lambda a, i=id_to_delete: a[1] == i and _item_org(a) == org)
                    self.prompt_queue.delete_queue_item(delete_func)

            return web.Response(status=200)

        @routes.post("/interrupt")
        async def post_interrupt(request):
            try:
                json_data = await request.json()
            except json.JSONDecodeError:
                json_data = {}

            org = self._caller_org(request)  # None = auth off → upstream behavior
            # Check if a specific prompt_id was provided for targeted interruption
            prompt_id = json_data.get('prompt_id')
            if prompt_id:
                currently_running, _ = self.prompt_queue.get_current_queue()

                # Check if the prompt_id matches any currently running prompt — and
                # (multi-tenant) that it belongs to the caller's org, so one tenant
                # cannot interrupt another tenant's render.
                should_interrupt = False
                for item in currently_running:
                    # item structure: (number, prompt_id, prompt, extra_data, outputs_to_execute)
                    if item[1] == prompt_id and (org is None or _item_org(item) == org):
                        logging.info(f"Interrupting prompt {prompt_id}")
                        should_interrupt = True
                        break

                if should_interrupt:
                    nodes.interrupt_processing()
                else:
                    logging.info(f"Prompt {prompt_id} is not currently running / not owned by caller, skipping interrupt")
            elif org is None:
                # No prompt_id + auth off → upstream global interrupt.
                logging.info("Global interrupt (no prompt_id specified)")
                nodes.interrupt_processing()
            else:
                # Multi-tenant: a global interrupt would kill another tenant's render.
                # Only interrupt if the running job belongs to the caller's org.
                currently_running, _ = self.prompt_queue.get_current_queue()
                if any(_item_org(item) == org for item in currently_running):
                    logging.info(f"Interrupting current render for org {org}")
                    nodes.interrupt_processing()

            return web.Response(status=200)

        @routes.post("/free")
        async def post_free(request):
            json_data = await request.json()
            unload_models = json_data.get("unload_models", False)
            free_memory = json_data.get("free_memory", False)
            if unload_models:
                self.prompt_queue.set_flag("unload_models", unload_models)
            if free_memory:
                self.prompt_queue.set_flag("free_memory", free_memory)
            return web.Response(status=200)

        @routes.post("/history")
        async def post_history(request):
            json_data =  await request.json()
            if "clear" in json_data:
                if json_data["clear"]:
                    self.prompt_queue.wipe_history()
            if "delete" in json_data:
                to_delete = json_data['delete']
                for id_to_delete in to_delete:
                    self.prompt_queue.delete_history_item(id_to_delete)

            return web.Response(status=200)

        # --- Compute Config API ---

        @routes.get("/compute/config")
        async def get_compute_config(request):
            """Get the org's compute profile."""
            org_id = self._get_org_id(request)
            config = compute_config.load_config(org_id)
            return web.json_response(config.to_dict())

        @routes.put("/compute/config")
        async def put_compute_config(request):
            """Update the org's compute profile."""
            org_id = self._get_org_id(request)
            try:
                updates = await request.json()
            except Exception:
                return web.json_response({"error": "Invalid JSON"}, status=400)

            try:
                config = compute_config.update_config(org_id, updates)
            except ValueError as e:
                return web.json_response({"error": str(e)}, status=400)

            return web.json_response(config.to_dict())

        @routes.get("/v1/workers")
        async def get_compute_workers(request):
            """List workers and their status for the org."""
            org_id = self._get_org_id(request)
            config = compute_config.load_config(org_id)
            workers = []
            for w in config.workers:
                wd = {
                    "worker_id": w.worker_id,
                    "url": w.url,
                    "device": w.device,
                    "status": w.status if w.is_alive() else "offline",
                    "gpu_model": w.gpu_model,
                    "vram_gb": w.vram_gb,
                    "last_heartbeat": w.last_heartbeat,
                }
                workers.append(wd)
            return web.json_response({"workers": workers})

        @routes.post("/v1/workers/register")
        async def post_worker_register(request):
            """Worker self-registration / heartbeat endpoint (coordinator trust)."""
            from middleware.worker_client import verify_worker_token
            if not verify_worker_token(request):
                return web.json_response({"error": "Unauthorized worker"}, status=401)
            try:
                data = await request.json()
            except Exception:
                return web.json_response({"error": "Invalid JSON"}, status=400)

            org_id = data.get("org_id") or self._get_org_id(request)
            worker_data = {k: v for k, v in data.items() if k != "org_id"}
            try:
                worker = compute_config.register_worker(org_id, worker_data)
            except ValueError as e:
                return web.json_response({"error": str(e)}, status=400)

            return web.json_response({
                "status": "registered",
                "worker_id": worker.worker_id,
            })

        @routes.post("/compute/provision")
        async def post_provision_gpu(request):
            """Request GPU provisioning via Visor."""
            org_id = self._get_org_id(request)
            config = compute_config.load_config(org_id)

            if not config.gpu_enabled:
                return web.json_response(
                    {"error": "GPU not enabled in compute profile. Update config first."},
                    status=400,
                )

            try:
                body = await request.json()
            except Exception:
                body = {}

            gpu_type = body.get("gpu_type", config.gpu_type or "t4")
            worker_id = body.get("worker_id", f"gpu-{org_id}-{int(time.time())}")
            default_coordinator = os.environ.get(
                "STUDIO_COORDINATOR_URL",
                f"http://{args.listen}:{args.port}",
            )
            coordinator_url = body.get("coordinator_url", default_coordinator)

            result = await visor_client.launch_gpu_worker(
                org_id=org_id,
                gpu_type=gpu_type,
                coordinator_url=coordinator_url,
                worker_id=worker_id,
            )

            if result:
                return web.json_response({
                    "status": "provisioning",
                    "machine_id": result.get("machine_id"),
                    "worker_id": worker_id,
                    "gpu_type": gpu_type,
                    "message": "GPU worker is being provisioned. It will register automatically.",
                }, status=202)
            else:
                return web.json_response(
                    {"error": "Failed to provision GPU worker. Check Visor connectivity."},
                    status=503,
                )

        @routes.delete("/compute/provision")
        async def delete_provision_gpu(request):
            """Tear down a provisioned GPU worker."""
            try:
                body = await request.json()
            except Exception:
                return web.json_response({"error": "Invalid JSON"}, status=400)

            machine_id = body.get("machine_id")
            if not machine_id:
                return web.json_response({"error": "machine_id required"}, status=400)

            worker_id = body.get("worker_id")
            org_id = self._get_org_id(request)

            success = await visor_client.terminate_gpu_worker(machine_id)

            # Also remove worker from registry
            if worker_id:
                compute_config.remove_worker(org_id, worker_id)

            if success:
                return web.json_response({"status": "terminated", "machine_id": machine_id})
            else:
                return web.json_response(
                    {"error": "Failed to terminate GPU worker"},
                    status=503,
                )

    def _get_org_id(self, request) -> str:
        """Extract org_id from IAM user context or fall back to default."""
        iam_user = request.get("iam_user")
        if iam_user:
            return iam_user.get("org_id", "default")
        return os.environ.get("STUDIO_ORG_ID", "default")

    def _caller_org(self, request) -> str | None:
        """The verified caller org to SCOPE cross-tenant reads by — or None when
        multi-tenant auth is OFF (single-tenant/local dev) so /history + /queue
        behave exactly as upstream. Only returns a value when there is a real IAM
        user, so scoping never hides another tenant's data by accident."""
        iam_user = request.get("iam_user")
        if iam_user and iam_user.get("org_id"):
            return iam_user.get("org_id")
        return None

    async def setup(self):
        timeout = aiohttp.ClientTimeout(total=None) # no timeout
        self.client_session = aiohttp.ClientSession(timeout=timeout)

        # Register shutdown hooks for compute subsystem sessions
        async def _on_shutdown(app):
            await prompt_router.close_session()
            await visor_client.close_session()
            from middleware import copilot
            await copilot.close_session()
            wc = getattr(self, "_worker_client", None)
            if wc:
                await wc.stop()
        self.app.on_shutdown.append(_on_shutdown)

    def add_routes(self):
        self.user_manager.add_routes(self.routes)
        self.model_file_manager.add_routes(self.routes)
        self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items())
        self.subgraph_manager.add_routes(self.routes, nodes.LOADED_MODULE_DIRS.items())
        self.node_replace_manager.add_routes(self.routes)
        from middleware.engine_selector import add_engine_routes
        add_engine_routes(self.routes, self)
        from middleware.copilot import add_copilot_routes
        add_copilot_routes(self.routes, self)
        from middleware.session import add_session_routes
        add_session_routes(self.routes, self)
        from middleware.studio_home import add_studio_home_routes
        add_studio_home_routes(self.routes, self)
        from middleware.mcp import add_mcp_routes
        add_mcp_routes(self.routes, self)
        self.app.add_subapp('/internal', self.internal_routes.get_app())

        # Prefix every route with /api for easier matching for delegation.
        # This is very useful for frontend dev server, which need to forward
        # everything except serving of static files.
        # Currently both the old endpoints without prefix and new endpoints with
        # prefix are supported.
        api_routes = web.RouteTableDef()
        for route in self.routes:
            # Custom nodes might add extra static routes. Only process non-static
            # routes to add /api prefix.
            if isinstance(route, web.RouteDef):
                api_routes.route(route.method, "/api" + route.path)(route.handler, **route.kwargs)
        self.app.add_routes(api_routes)
        self.app.add_routes(self.routes)

        # Add routes from web extensions.
        for name, dir in nodes.EXTENSION_WEB_DIRS.items():
            self.app.add_routes([web.static('/extensions/' + name, dir)])

        installed_templates_version = FrontendManager.get_installed_templates_version()
        use_legacy_templates = True
        if installed_templates_version:
            try:
                use_legacy_templates = (
                    parse_version(installed_templates_version)
                    < parse_version("0.3.0")
                )
            except Exception as exc:
                logging.warning(
                    "Unable to parse templates version '%s': %s",
                    installed_templates_version,
                    exc,
                )

        if use_legacy_templates:
            workflow_templates_path = FrontendManager.legacy_templates_path()
            if workflow_templates_path:
                self.app.add_routes([
                    web.static('/templates', workflow_templates_path)
                ])
        else:
            handler = FrontendManager.template_asset_handler()
            if handler:
                self.app.router.add_get("/templates/{path:.*}", handler)

        # Serve embedded documentation from the package
        embedded_docs_path = FrontendManager.embedded_docs_path()
        if embedded_docs_path:
            self.app.add_routes([
                web.static('/docs', embedded_docs_path)
            ])

        # Serve the unified Hanzo marketing shell (header + footer bundle) used by
        # the logged-out login page. Built from web/shell into web/marketing and
        # committed; assets are public (.js/.css) so they load pre-auth.
        marketing_path = os.path.join(
            os.path.dirname(os.path.realpath(__file__)), "web", "marketing"
        )
        if os.path.isdir(marketing_path):
            self.app.add_routes([
                web.static('/marketing', marketing_path)
            ])

        self.app.add_routes([
            web.static('/', self.web_root),
        ])

    def get_queue_info(self):
        prompt_info = {}
        exec_info = {}
        exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining()
        prompt_info['exec_info'] = exec_info
        return prompt_info

    async def send(self, event, data, sid=None):
        if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE:
            await self.send_image(data, sid=sid)
        elif event == BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA:
            # data is (preview_image, metadata)
            preview_image, metadata = data
            await self.send_image_with_metadata(preview_image, metadata, sid=sid)
        elif isinstance(data, (bytes, bytearray)):
            await self.send_bytes(event, data, sid)
        else:
            await self.send_json(event, data, sid)

    def encode_bytes(self, event, data):
        if not isinstance(event, int):
            raise RuntimeError(f"Binary event types must be integers, got {event}")

        packed = struct.pack(">I", event)
        message = bytearray(packed)
        message.extend(data)
        return message

    async def send_image(self, image_data, sid=None):
        image_type = image_data[0]
        image = image_data[1]
        max_size = image_data[2]
        if max_size is not None:
            if hasattr(Image, 'Resampling'):
                resampling = Image.Resampling.BILINEAR
            else:
                resampling = Image.Resampling.LANCZOS

            image = ImageOps.contain(image, (max_size, max_size), resampling)
        type_num = 1
        if image_type == "JPEG":
            type_num = 1
        elif image_type == "PNG":
            type_num = 2

        bytesIO = BytesIO()
        header = struct.pack(">I", type_num)
        bytesIO.write(header)
        image.save(bytesIO, format=image_type, quality=95, compress_level=1)
        preview_bytes = bytesIO.getvalue()
        await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid)

    async def send_image_with_metadata(self, image_data, metadata=None, sid=None):
        image_type = image_data[0]
        image = image_data[1]
        max_size = image_data[2]
        if max_size is not None:
            if hasattr(Image, 'Resampling'):
                resampling = Image.Resampling.BILINEAR
            else:
                resampling = Image.Resampling.LANCZOS

            image = ImageOps.contain(image, (max_size, max_size), resampling)

        mimetype = "image/png" if image_type == "PNG" else "image/jpeg"

        # Prepare metadata
        if metadata is None:
            metadata = {}
        metadata["image_type"] = mimetype

        # Serialize metadata as JSON
        import json
        metadata_json = json.dumps(metadata).encode('utf-8')
        metadata_length = len(metadata_json)

        # Prepare image data
        bytesIO = BytesIO()
        image.save(bytesIO, format=image_type, quality=95, compress_level=1)
        image_bytes = bytesIO.getvalue()

        # Combine metadata and image
        combined_data = bytearray()
        combined_data.extend(struct.pack(">I", metadata_length))
        combined_data.extend(metadata_json)
        combined_data.extend(image_bytes)

        await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA, combined_data, sid=sid)

    async def send_bytes(self, event, data, sid=None):
        message = self.encode_bytes(event, data)

        if sid is None:
            sockets = list(self.sockets.values())
            for ws in sockets:
                await send_socket_catch_exception(ws.send_bytes, message)
        elif sid in self.sockets:
            await send_socket_catch_exception(self.sockets[sid].send_bytes, message)

    async def send_json(self, event, data, sid=None):
        message = {"type": event, "data": data}

        if sid is None:
            sockets = list(self.sockets.values())
            for ws in sockets:
                await send_socket_catch_exception(ws.send_json, message)
        elif sid in self.sockets:
            await send_socket_catch_exception(self.sockets[sid].send_json, message)

    def send_sync(self, event, data, sid=None):
        self.loop.call_soon_threadsafe(
            self.messages.put_nowait, (event, data, sid))

    def queue_updated(self):
        self.send_sync("status", { "status": self.get_queue_info() })

    async def publish_loop(self):
        while True:
            msg = await self.messages.get()
            await self.send(*msg)

    async def start(self, address, port, verbose=True, call_on_start=None):
        await self.start_multi_address([(address, port)], call_on_start=call_on_start)

    async def start_multi_address(self, addresses, call_on_start=None, verbose=True):
        runner = web.AppRunner(self.app, access_log=None)
        await runner.setup()
        ssl_ctx = None
        scheme = "http"
        if args.tls_keyfile and args.tls_certfile:
            ssl_ctx = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_SERVER, verify_mode=ssl.CERT_NONE)
            ssl_ctx.load_cert_chain(certfile=args.tls_certfile,
                                keyfile=args.tls_keyfile)
            scheme = "https"

        if verbose:
            logging.info("Starting server\n")
        for addr in addresses:
            address = addr[0]
            port = addr[1]
            site = web.TCPSite(runner, address, port, ssl_context=ssl_ctx)
            await site.start()

            if not hasattr(self, 'address'):
                self.address = address #TODO: remove this
                self.port = port

            if ':' in address:
                address_print = "[{}]".format(address)
            else:
                address_print = address

            if verbose:
                logging.info("To see the GUI go to: {}://{}:{}".format(scheme, address_print, port))

        if call_on_start is not None:
            call_on_start(scheme, self.address, self.port)

    def add_on_prompt_handler(self, handler):
        self.on_prompt_handlers.append(handler)

    def trigger_on_prompt(self, json_data):
        for handler in self.on_prompt_handlers:
            try:
                json_data = handler(json_data)
            except Exception:
                logging.warning("[ERROR] An error occurred during the on_prompt_handler processing")
                logging.warning(traceback.format_exc())

        return json_data

    def send_progress_text(
        self, text: Union[bytes, bytearray, str], node_id: str, sid=None
    ):
        if isinstance(text, str):
            text = text.encode("utf-8")
        node_id_bytes = str(node_id).encode("utf-8")

        # Pack the node_id length as a 4-byte unsigned integer, followed by the node_id bytes
        message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text

        self.send_sync(BinaryEventTypes.TEXT, message, sid)