import os import sys import heapq import copy import contextvars import logging import mimetypes import time from datetime import datetime from aiohttp import web from typing import Optional import folder_paths from server import PromptServer from execution import PromptQueue, MAXIMUM_HISTORY_SIZE from .users_db import UsersDB from .config import INSTANCE_LOG_FILE from .history_assets import iter_temp_references, persist_temp_assets from .history_store import HistoryStore from .scheduler_store import SchedulerStore, WorkerHeartbeat, is_sqlite_busy from .api_audit import clear_current_job, set_current_job logger = logging.getLogger("ComfyUI-Account-Manager") class AccessControl: def __init__( self, users_db: UsersDB, server: PromptServer, history_file: str = None, scheduler: SchedulerStore = None, instance_port: int = 0, heartbeat_seconds: int = 5, stale_seconds: int = 60, worker_resource_class: str = "default", cloud_archive=None, ): self.users_db = users_db self.server = server self.scheduler = scheduler self.instance_port = int(instance_port or 0) self.worker_resource_class = str(worker_resource_class or "default") self.cloud_archive = cloud_archive self._worker_heartbeat = None self._current_user = contextvars.ContextVar("user_id", default=None) self.__current_user_id = None self.__get_output_directory = folder_paths.get_output_directory self.__get_temp_directory = folder_paths.get_temp_directory self.__get_input_directory = folder_paths.get_input_directory self.__get_save_image_path = folder_paths.get_save_image_path self.__prompt_queue = self.server.prompt_queue self.__history_store = None if history_file: try: self.__history_store = HistoryStore(history_file) except Exception: logger.exception("Failed to initialize persistent history") self.__prompt_queue_put = self.__prompt_queue.put self.__prompt_queue_get_flags = self.__prompt_queue.get_flags self.__user_manager_get_request_user_id = getattr( self.server.user_manager, "get_request_user_id", None ) if self.scheduler and self.instance_port: self._worker_heartbeat = WorkerHeartbeat( self.scheduler, self.instance_port, heartbeat_seconds, stale_seconds, self.worker_resource_class, ) self._worker_heartbeat.start() @property def folder_paths(self) -> tuple: return ( self.__get_output_directory(), self.__get_temp_directory(), self.__get_input_directory(), ) def set_current_user_id(self, user_id: str, set_fallback: bool = False) -> None: """Set the current authenticated user.""" self._current_user.set(user_id) if set_fallback: self.__current_user_id = user_id def get_current_user_id(self) -> str: """Retrieve the active user, falling back to the latest prompt submitter.""" if self._current_user.get(): return self._current_user.get() return self.__current_user_id def get_context_user_id(self) -> str: """Retrieve only the current request/execution context user.""" return self._current_user.get() def is_admin_user(self, user_id: str) -> bool: if not user_id: return False _, user = self.users_db.get_user(user_id=user_id) return bool(user and user.get("admin")) def _can_access_owner(self, owner_id: str, current_user_id: str = None) -> bool: current_user_id = current_user_id or self.get_current_user_id() if self.is_admin_user(current_user_id): return True return bool(owner_id and owner_id == current_user_id) def _get_user_slug(self, user_id: str = None) -> str: user_id = user_id or self.get_current_user_id() if not user_id: return "public" _, user = self.users_db.get_user(user_id=user_id) return user.get("username") or user_id @staticmethod def _get_timeline_slug() -> str: return datetime.now().strftime("%Y-%m-%d") def _get_user_path_prefix(self, user_id: str = None) -> str: """Build the user-visible folder prefix for newly generated files.""" user_id = user_id or self.get_current_user_id() return f"{self._get_timeline_slug()}/{self._get_user_slug(user_id)}" def _is_current_user_prefix(self, path_prefix: str, user_id: str = None) -> bool: user_id = user_id or self.get_current_user_id() if not user_id or not path_prefix: return False normalized = os.path.normpath(str(path_prefix)).replace("\\", "/").strip("/") current_prefix = self._get_user_path_prefix(user_id) if normalized == current_prefix or normalized.startswith(f"{current_prefix}/"): return True # Backward compatibility for files generated before the timeline/username layout. return normalized == user_id or normalized.startswith(f"{user_id}/") def _can_access_output_subfolder( self, subfolder: str, current_user_id: str = None ) -> bool: current_user_id = current_user_id or self.get_current_user_id() if self.is_admin_user(current_user_id): return True if not current_user_id: return False normalized = os.path.normpath(str(subfolder or "")).replace("\\", "/").strip("/") if not normalized: return False if normalized == current_user_id or normalized.startswith(f"{current_user_id}/"): return True username = self._get_user_slug(current_user_id) parts = normalized.split("/") return len(parts) >= 2 and parts[1] == username def _with_user_prefix(self, filename_prefix: str) -> str: user_id = self.get_current_user_id() if not user_id or not filename_prefix: return filename_prefix normalized = os.path.normpath(str(filename_prefix)).replace("\\", "/") if self._is_current_user_prefix(normalized, user_id): return filename_prefix return f"{self._get_user_path_prefix(user_id)}/{filename_prefix}" def get_user_output_directory(self) -> str: """Get the user-specific output directory.""" return os.path.join( self.__get_output_directory(), *self._get_user_path_prefix().split("/"), ) def get_user_temp_directory(self) -> str: """Get the user-specific temp directory.""" return os.path.join( self.__get_temp_directory(), self._get_user_slug(), ) def get_user_input_directory(self) -> str: """Get the user-specific input directory.""" input_directory = os.path.join( self.__get_input_directory(), self._get_user_slug(), ) os.makedirs(input_directory, exist_ok=True) return input_directory def add_user_specific_folder_paths(self, json_data) -> None: """Add user-specific output prefixes to prompt JSON data.""" if isinstance(json_data, dict): for key, value in json_data.items(): if key == "filename_prefix" and isinstance(value, str): json_data[key] = self._with_user_prefix(value) else: self.add_user_specific_folder_paths(value) elif isinstance(json_data, list): for item in json_data: self.add_user_specific_folder_paths(item) return json_data def get_user_save_image_path( self, filename_prefix: str, output_dir: str, image_width=0, image_height=0, ) -> tuple[str, str, int, str, str]: """Ensure standard ComfyUI save helpers write under the current user.""" return self.__get_save_image_path( self._with_user_prefix(filename_prefix), output_dir, image_width, image_height, ) def patch_folder_paths(self) -> None: """Patch folder paths and save helpers with user-specific behavior.""" folder_paths.get_temp_directory = self.get_user_temp_directory folder_paths.get_input_directory = self.get_user_input_directory folder_paths.get_save_image_path = self.get_user_save_image_path self.server.add_on_prompt_handler(self.add_user_specific_folder_paths) def sync_comfy_users(self) -> None: """Mirror account-manager users into ComfyUI's in-memory user map.""" user_manager = getattr(self.server, "user_manager", None) if not user_manager or not hasattr(user_manager, "users"): return for user_id, user in self.users_db.load_users().items(): username = user.get("username") or user_id user_manager.users[user_id] = username def patch_comfy_user_manager(self) -> None: """Let ComfyUI APIs resolve authenticated account-manager users.""" user_manager = getattr(self.server, "user_manager", None) if not user_manager or self.__user_manager_get_request_user_id is None: return def get_request_user_id(request): user_id = request.get("user_id") if user_id and user_id in self.users_db.load_users(): return user_id header_user_id = request.headers.get("comfy-user") if header_user_id and header_user_id in self.users_db.load_users(): return header_user_id return self.__user_manager_get_request_user_id(request) user_manager.get_request_user_id = get_request_user_id self.sync_comfy_users() def create_view_access_control_middleware(self) -> web.middleware: """Guard /view and /api/view output access by top-level user folder.""" @web.middleware async def view_access_control_middleware( request: web.Request, handler ) -> web.Response: if request.path not in ("/view", "/api/view"): return await handler(request) if self.asset_access is not None: # The record/content adapter validates resolved paths and hashes. return await handler(request) view_type = request.rel_url.query.get("type", "output") if view_type != "output": return await handler(request) user_id = request.get("user_id") if self.is_admin_user(user_id): return await handler(request) if self._can_access_output_subfolder( request.rel_url.query.get("subfolder", ""), user_id ): return await handler(request) return web.HTTPForbidden(reason="You do not have access to this output.") return view_access_control_middleware def create_folder_access_control_middleware( self, folder_paths: tuple = () ) -> web.middleware: """Create middleware for direct folder access control.""" folder_paths = folder_paths or self.folder_paths @web.middleware async def folder_access_control_middleware( request: web.Request, handler ) -> web.Response: if not request.path.startswith(folder_paths): return await handler(request) current_user_id = request.get("user_id") try: path_parts = request.path.strip("/").split("/") folder_user_id = path_parts[1] except Exception: return web.HTTPNotFound(reason="Folder not found.") if not self._can_access_owner(folder_user_id, current_user_id): return web.HTTPForbidden( reason="You do not have access to this folder." ) return await handler(request) return folder_access_control_middleware def _wrap_queue_item(self, item): if isinstance(item, tuple) and len(item) >= 7: return item return item + (self.get_current_user_id(),) @staticmethod def _unwrap_queue_item(item): if isinstance(item, tuple) and len(item) >= 7: return item[:6] return item @staticmethod def _queue_item_user_id(item): if isinstance(item, tuple) and len(item) >= 7: return item[6] return None def _visible_queue_items(self, items): current_user_id = self.get_current_user_id() visible = [] for item in items: owner_id = self._queue_item_user_id(item) if self._can_access_owner(owner_id, current_user_id): visible.append(self._unwrap_queue_item(item)) return visible def user_queue_put(self, item): """Put an item in the user-specific queue.""" if self.scheduler: user_id = self.get_current_user_id() or "" self.scheduler.enqueue( self._unwrap_queue_item(item), user_id, self.instance_port, concurrency_exempt=self.is_admin_user(user_id), ) self.server.queue_updated() with self.__prompt_queue.not_empty: self.__prompt_queue.not_empty.notify_all() return wrapped_item = self._wrap_queue_item(item) with self.__prompt_queue.mutex: heapq.heappush(self.__prompt_queue.queue, wrapped_item) self.server.queue_updated() self.__prompt_queue.not_empty.notify() def user_queue_get(self, timeout=None): """Get an item from the user-specific queue.""" if self.scheduler: started = time.monotonic() while True: try: wrapped_item = self.scheduler.claim( self.instance_port, os.getpid(), self.worker_resource_class, INSTANCE_LOG_FILE, ) except Exception as error: if not is_sqlite_busy(error): raise logger.warning( "Scheduler database busy on port %s; retrying task claim", self.instance_port, ) wrapped_item = None if wrapped_item is not None: user_id = self._queue_item_user_id(wrapped_item) self.set_current_user_id(user_id, True) set_current_job(wrapped_item[1]) with self.__prompt_queue.mutex: item_id = self.__prompt_queue.task_counter self.__prompt_queue.currently_running[item_id] = copy.deepcopy( wrapped_item ) self.__prompt_queue.task_counter += 1 self.server.queue_updated() return (self._unwrap_queue_item(wrapped_item), item_id) if timeout is not None and time.monotonic() - started >= timeout: return None time.sleep(0.2) user_queue = self.__prompt_queue.queue with self.__prompt_queue.not_empty: while len(user_queue) == 0: self.__prompt_queue.not_empty.wait(timeout=timeout) if timeout is not None and len(user_queue) == 0: return None item = heapq.heappop(user_queue) user_id = self._queue_item_user_id(item) self.set_current_user_id(user_id, True) i = self.__prompt_queue.task_counter self.__prompt_queue.currently_running[i] = copy.deepcopy(item) self.__prompt_queue.task_counter += 1 self.server.queue_updated() return (self._unwrap_queue_item(item), i) def user_queue_task_done( self, item_id, history_result, status: Optional["PromptQueue.ExecutionStatus"], process_item=None, ): """Mark a user-specific queue task as done.""" cloud_completion = None with self.__prompt_queue.mutex: wrapped_prompt = self.__prompt_queue.currently_running.pop(item_id) prompt = self._unwrap_queue_item(wrapped_prompt) user_id = self._queue_item_user_id(wrapped_prompt) if len(self.__prompt_queue.history) >= MAXIMUM_HISTORY_SIZE: self.__prompt_queue.history.pop(next(iter(self.__prompt_queue.history))) status_dict: Optional[dict] = None if status is not None: status_dict = copy.deepcopy(status._asdict()) if process_item is not None: prompt = process_item(prompt) self.__prompt_queue.history[prompt[1]] = { "prompt": prompt, "outputs": {}, "status": status_dict, "user_id": user_id, } self.__prompt_queue.history[prompt[1]].update(history_result) self._persist_history_assets(prompt[1], user_id) self._index_history_assets(prompt[1], user_id) self._save_history_item(prompt[1]) if self.scheduler: status_text = getattr(status, "status_str", "success") if status else "success" messages = getattr(status, "messages", []) if status else [] if INSTANCE_LOG_FILE: time.sleep(0.5) for attempt in range(3): try: self.scheduler.complete( prompt[1], self.instance_port, succeeded=status_text == "success", error="\n".join(str(message) for message in messages), log_file=INSTANCE_LOG_FILE, ) break except Exception: if attempt == 2: logger.exception( "Failed to complete distributed task %s", prompt[1] ) else: time.sleep(0.2 * (attempt + 1)) if self.cloud_archive: cloud_completion = { "prompt_id": prompt[1], "owner_id": user_id, "worker_port": self.instance_port, "generation_status": ( "completed" if status_text == "success" else "failed" ), "history_item": self.__prompt_queue.history[prompt[1]], "finished_at": time.time(), } clear_current_job(prompt[1]) self.server.queue_updated() if cloud_completion: try: self.cloud_archive.enqueue_completion(**cloud_completion) except Exception: logger.exception( "Failed to enqueue cloud archive for %s", cloud_completion["prompt_id"], ) def _index_history_assets(self, prompt_id: str, user_id: str) -> None: if not self.scheduler: return def visit(value): if isinstance(value, dict): asset_id = str(value.get("id") or "") if asset_id and ( value.get("filename") or value.get("file_path") or value.get("fullpath") ): self.scheduler.upsert_asset( asset_id, user_id, prompt_id, self.instance_port, value, ) for child in value.values(): visit(child) elif isinstance(value, (list, tuple)): for child in value: visit(child) try: history_item = self.__prompt_queue.history.get(prompt_id, {}) visit(history_item.get("outputs", {})) except Exception: logger.exception("Failed to index distributed assets for %s", prompt_id) def _persist_history_assets(self, prompt_id: str, user_id: str) -> None: if not user_id: return try: destination_subfolder = ( f"{self._get_user_path_prefix(user_id)}/history_assets/{prompt_id}" ) source_paths = self._get_registered_temp_paths( self.__prompt_queue.history[prompt_id], user_id ) persisted_assets = persist_temp_assets( self.__prompt_queue.history[prompt_id], self.__get_temp_directory(), self.__get_output_directory(), self._get_user_slug(user_id), destination_subfolder, source_paths, ) if persisted_assets: self._register_history_assets(persisted_assets, prompt_id, user_id) except Exception: logger.exception("Failed to persist temporary assets for %s", prompt_id) @staticmethod def _get_registered_temp_paths(history_item: dict, user_id: str) -> dict[str, str]: try: from app.assets.database.queries.asset_reference import get_reference_by_id from app.database.db import create_session except Exception: return {} paths = {} with create_session() as session: for reference in iter_temp_references(history_item): reference_id = str(reference.get("id") or "") if not reference_id: continue asset_reference = get_reference_by_id( session, reference_id=reference_id ) if ( asset_reference and asset_reference.deleted_at is None and asset_reference.owner_id == user_id and asset_reference.file_path ): paths[reference_id] = asset_reference.file_path return paths @staticmethod def _register_history_assets( persisted_assets: list[tuple[dict, str]], prompt_id: str, user_id: str ) -> None: try: from app.assets.database.queries.asset_reference import ( get_reference_by_file_path, ) from app.assets.services.asset_management import delete_asset_reference from app.assets.services.ingest import ingest_existing_file from app.database.db import create_session except Exception: return for reference, path in persisted_assets: try: old_reference_id = str(reference.get("id") or "") ingest_existing_file(path, owner_id=user_id, job_id=prompt_id) with create_session() as session: new_reference = get_reference_by_file_path(session, path) if not new_reference or new_reference.owner_id != user_id: continue reference["id"] = new_reference.id if old_reference_id and old_reference_id != reference["id"]: delete_asset_reference( old_reference_id, owner_id=user_id, delete_content_if_orphan=False, ) except Exception: logger.exception("Failed to register persistent history asset %s", path) def _save_history_item(self, prompt_id: str) -> None: if not self.__history_store: return try: self.__history_store.save( prompt_id, self.__prompt_queue.history[prompt_id], MAXIMUM_HISTORY_SIZE, ) except Exception: logger.exception("Failed to persist history item %s", prompt_id) def _load_persisted_history(self) -> None: if not self.__history_store: return try: persisted = self.__history_store.load(MAXIMUM_HISTORY_SIZE) persisted.update(self.__prompt_queue.history) self.__prompt_queue.history = persisted except Exception: logger.exception("Failed to restore persisted history") def _delete_persisted_history( self, prompt_id: str = None, owner_id: str = None, clear: bool = False ) -> None: if not self.__history_store: return try: if clear: self.__history_store.clear() elif owner_id is not None: self.__history_store.delete_owner(owner_id) elif prompt_id is not None: self.__history_store.delete(prompt_id) except Exception: logger.exception("Failed to update persistent history") def user_queue_get_current_queue(self): """Get the current user-specific queue.""" if self.scheduler: current_user_id = self.get_current_user_id() return self.scheduler.visible_items( current_user_id, admin=self.is_admin_user(current_user_id), ) with self.__prompt_queue.mutex: running = self._visible_queue_items(self.__prompt_queue.currently_running.values()) queued = self._visible_queue_items(copy.deepcopy(self.__prompt_queue.queue)) return (running, queued) def user_queue_get_current_queue_volatile(self): """Get the current user-specific queue without deep copying items.""" if self.scheduler: return self.user_queue_get_current_queue() with self.__prompt_queue.mutex: running = self._visible_queue_items(self.__prompt_queue.currently_running.values()) queued = self._visible_queue_items(copy.copy(self.__prompt_queue.queue)) return (running, queued) def user_queue_get_tasks_remaining(self): if self.scheduler: running, queued = self.user_queue_get_current_queue() return len(running) + len(queued) with self.__prompt_queue.mutex: return ( len(self._visible_queue_items(self.__prompt_queue.queue)) + len(self._visible_queue_items(self.__prompt_queue.currently_running.values())) ) def user_queue_wipe_queue(self): """Wipe the current user's queue, or all queue items for admins.""" if self.scheduler: current_user_id = self.get_current_user_id() self.scheduler.cancel_all( current_user_id, admin=self.is_admin_user(current_user_id), ) self.server.queue_updated() return with self.__prompt_queue.mutex: current_user_id = self.get_current_user_id() if self.is_admin_user(current_user_id): self.__prompt_queue.queue = [] else: self.__prompt_queue.queue = [ item for item in self.__prompt_queue.queue if self._queue_item_user_id(item) != current_user_id ] self.server.queue_updated() def user_queue_delete_queue_item(self, function): """Delete an item from the current user's queue.""" if self.scheduler: current_user_id = self.get_current_user_id() deleted = self.scheduler.delete_matching( function, current_user_id, admin=self.is_admin_user(current_user_id), ) if deleted: self.server.queue_updated() return deleted with self.__prompt_queue.mutex: for x in range(len(self.__prompt_queue.queue)): item = self.__prompt_queue.queue[x] unwrapped_item = self._unwrap_queue_item(item) if ( function(unwrapped_item) and self._can_access_owner(self._queue_item_user_id(item)) ): self.__prompt_queue.queue.pop(x) heapq.heapify(self.__prompt_queue.queue) self.server.queue_updated() return True return False def _visible_history(self): current_user_id = self.get_current_user_id() if self.is_admin_user(current_user_id): return self.__prompt_queue.history return { k: v for k, v in self.__prompt_queue.history.items() if v.get("user_id") == current_user_id } def user_job_history(self): if self.scheduler and self.__history_store: current_user_id = self.get_current_user_id() return self.__history_store.query( owner_id=None if self.is_admin_user(current_user_id) else current_user_id, for_jobs=True, ) return self.user_queue_get_history() def user_queue_get_history(self, prompt_id=None, max_items=None, offset=-1, map_function=None): """Get the current user's queue history.""" if self.scheduler and self.__history_store: current_user_id = self.get_current_user_id() history = self.__history_store.query( max_items=max_items, offset=offset, owner_id=None if self.is_admin_user(current_user_id) else current_user_id, prompt_id=prompt_id, ) if map_function is not None: return {key: map_function(value) for key, value in history.items()} return history with self.__prompt_queue.mutex: user_history = self._visible_history() if prompt_id is None: out = {} i = 0 if offset < 0 and max_items is not None: offset = len(user_history) - max_items for k in user_history: if i >= offset: p = user_history[k] if map_function is None: p = copy.deepcopy(p) else: p = map_function(p) out[k] = p if max_items is not None and len(out) >= max_items: break i += 1 return out if prompt_id in user_history: p = user_history[prompt_id] if map_function is None: p = copy.deepcopy(p) else: p = map_function(p) return {prompt_id: p} return {} def user_queue_wipe_history(self): """Wipe the current user's history, or all history for admins.""" if self.scheduler: current_user_id = self.get_current_user_id() if self.is_admin_user(current_user_id): self._delete_persisted_history(clear=True) else: self._delete_persisted_history(owner_id=current_user_id) return with self.__prompt_queue.mutex: current_user_id = self.get_current_user_id() if self.is_admin_user(current_user_id): self.__prompt_queue.history = {} self._delete_persisted_history(clear=True) else: self.__prompt_queue.history = { k: v for k, v in self.__prompt_queue.history.items() if v.get("user_id") != current_user_id } self._delete_persisted_history(owner_id=current_user_id) def user_queue_delete_history_item(self, id_to_delete): if self.scheduler and self.__history_store: current_user_id = self.get_current_user_id() visible = self.__history_store.query( owner_id=None if self.is_admin_user(current_user_id) else current_user_id, prompt_id=id_to_delete, ) if id_to_delete in visible: self._delete_persisted_history(prompt_id=id_to_delete) return with self.__prompt_queue.mutex: history_item = self.__prompt_queue.history.get(id_to_delete) if history_item and self._can_access_owner(history_item.get("user_id")): self.__prompt_queue.history.pop(id_to_delete, None) self._delete_persisted_history(prompt_id=id_to_delete) def patch_prompt_queue(self): """Patch the prompt queue with user-specific methods.""" if not self.scheduler: self._load_persisted_history() self.__prompt_queue.put = self.user_queue_put self.__prompt_queue.get = self.user_queue_get self.__prompt_queue.task_done = self.user_queue_task_done self.__prompt_queue.get_current_queue = self.user_queue_get_current_queue self.__prompt_queue.get_current_queue_volatile = self.user_queue_get_current_queue_volatile self.__prompt_queue.get_tasks_remaining = self.user_queue_get_tasks_remaining self.__prompt_queue.wipe_queue = self.user_queue_wipe_queue self.__prompt_queue.delete_queue_item = self.user_queue_delete_queue_item self.__prompt_queue.get_history = self.user_queue_get_history self.__prompt_queue.wipe_history = self.user_queue_wipe_history self.__prompt_queue.delete_history_item = self.user_queue_delete_history_item def _default_owner_id(self) -> str: return self.get_context_user_id() or "" def patch_assets(self) -> None: """Patch ComfyUI asset ownership and visibility without editing ComfyUI.""" self.asset_access = None try: from app.assets.database.models import AssetContent except ImportError: pass # The local legacy schema still uses AssetReference. else: from .asset_access import AssetAccess model_roots = [path for name, (paths, _) in folder_paths.folder_names_and_paths.items() if name not in ('custom_nodes', 'configs') for path in paths] self.asset_access = AssetAccess( self.users_db, self.folder_paths, model_roots, folder_paths.get_previews_directory(), ) try: self.asset_access.install() except Exception: logger.exception("Asset permission adapter unavailable; asset access disabled") return try: import app.assets.database.queries as queries import app.assets.database.queries.common as common_queries import app.assets.database.queries.asset_reference as asset_reference_queries import app.assets.database.queries.tags as tag_queries import app.assets.services as asset_services import app.assets.services.ingest as ingest import app.assets.services.asset_management as asset_management import app.assets.services.tagging as tagging from app.assets.database.models import AssetReference from app.database.db import create_session except Exception: logger.exception("Legacy asset permission adapter unavailable") # Never leave asset routes unprotected when the schema changes again. from .asset_access import AssetAccess self.asset_access = AssetAccess(self.users_db, self.folder_paths, []) return def build_visible_owner_clause(owner_id: str): owner_id = (owner_id or "").strip() if self.is_admin_user(owner_id): return asset_reference_queries.sa.true() if owner_id == "": return AssetReference.owner_id == "" return AssetReference.owner_id == owner_id def get_reference_with_owner_check(session, reference_id: str, owner_id: str): ref = asset_reference_queries.get_reference_by_id( session, reference_id=reference_id ) if not ref or ref.deleted_at is not None: raise ValueError(f"AssetReference {reference_id} not found") if self.is_admin_user(owner_id): return ref if not ref.owner_id or ref.owner_id != owner_id: raise PermissionError("not owner") return ref common_queries.build_visible_owner_clause = build_visible_owner_clause asset_reference_queries.build_visible_owner_clause = build_visible_owner_clause tag_queries.build_visible_owner_clause = build_visible_owner_clause queries.build_visible_owner_clause = build_visible_owner_clause asset_reference_queries.get_reference_with_owner_check = get_reference_with_owner_check queries.get_reference_with_owner_check = get_reference_with_owner_check asset_management.get_reference_with_owner_check = get_reference_with_owner_check tagging.get_reference_with_owner_check = get_reference_with_owner_check original_ingest_existing_file = ingest.ingest_existing_file original_upload_from_temp_path = ingest.upload_from_temp_path original_register_file_in_place = ingest.register_file_in_place original_create_from_hash = ingest.create_from_hash def ingest_existing_file(*args, owner_id="", **kwargs): owner_id = owner_id or self._default_owner_id() return original_ingest_existing_file(*args, owner_id=owner_id, **kwargs) def upload_from_temp_path(*args, owner_id="", **kwargs): owner_id = owner_id or self._default_owner_id() return original_upload_from_temp_path(*args, owner_id=owner_id, **kwargs) def register_file_in_place(*args, owner_id="", **kwargs): owner_id = owner_id or self._default_owner_id() return original_register_file_in_place(*args, owner_id=owner_id, **kwargs) def create_from_hash(*args, owner_id="", **kwargs): owner_id = owner_id or self._default_owner_id() return original_create_from_hash(*args, owner_id=owner_id, **kwargs) ingest.ingest_existing_file = ingest_existing_file ingest.upload_from_temp_path = upload_from_temp_path ingest.register_file_in_place = register_file_in_place ingest.create_from_hash = create_from_hash asset_services.upload_from_temp_path = upload_from_temp_path asset_services.register_file_in_place = register_file_in_place asset_services.create_from_hash = create_from_hash def resolve_hash_to_path(asset_hash: str, owner_id: str = ""): with create_session() as session: asset = asset_management.queries_get_asset_by_hash(session, asset_hash) if not asset: return None refs = asset_management.list_references_by_asset_id( session, asset_id=asset.id ) if self.is_admin_user(owner_id): visible = refs else: visible = [r for r in refs if r.owner_id == owner_id] abs_path = asset_management.select_best_live_path(visible) if not abs_path: return None display_name = os.path.basename(abs_path) for ref in visible: if ref.file_path == abs_path and ref.name: display_name = ref.name break content_type = ( asset.mime_type or mimetypes.guess_type(display_name)[0] or "application/octet-stream" ) return asset_management.DownloadResolutionResult( abs_path=abs_path, content_type=content_type, download_name=display_name, ) asset_management.resolve_hash_to_path = resolve_hash_to_path asset_services.resolve_hash_to_path = resolve_hash_to_path route_module = sys.modules.get("app.assets.api.routes") if route_module: route_module.upload_from_temp_path = upload_from_temp_path route_module.create_from_hash = create_from_hash for module_name in ("server", "__main__", "main"): module = sys.modules.get(module_name) if not module: continue if hasattr(module, "register_file_in_place"): module.register_file_in_place = register_file_in_place if hasattr(module, "resolve_hash_to_path"): module.resolve_hash_to_path = resolve_hash_to_path def create_manager_access_control_middleware( self, manager_directory: str = "/extensions/comfyui-manager", manager_routes: tuple = () ) -> web.middleware: """Create middleware for manager access control.""" @web.middleware async def manager_access_control_middleware( request: web.Request, handler ) -> web.Response: user_id = request.get("user_id") if self.users_db.get_admin_user()[0] == user_id or ( not request.path.startswith(manager_routes) and not request.path.lower().startswith(manager_directory) ): return await handler(request) return web.HTTPForbidden( reason="You do not have access to comfyui manager." ) return manager_access_control_middleware