diff --git a/.gitignore b/.gitignore index 4ec9a7e..bf6cafe 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ whitelist.txt blacklist.txt .env secret_key.txt +history.sqlite3* diff --git a/README.md b/README.md index e65838a..99b5789 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,7 @@ Edit `config.json` before starting ComfyUI. "secret_key_env": "SECRET_KEY", "secret_key_file": "secret_key.txt", "users_db": "users_db.json", + "history_db": "history.sqlite3", "access_token_expiration_hours": 8760, "max_access_token_expiration_hours": 8760, "log": "account_manager.log", @@ -47,6 +48,10 @@ Edit `config.json` before starting ComfyUI. `separate_users` enables account-specific inputs, outputs, queue history, and asset visibility. The older misspelled `seperate_users` key is still accepted for compatibility, but new installs should use `separate_users`. +Completed generation history is persisted in `history.sqlite3` and restored when ComfyUI restarts. The history database keeps the same 10,000-item limit as ComfyUI's in-memory history. + +ComfyUI's optional Assets panel is separate from queue history and requires starting ComfyUI with `--enable-assets`. + ## API Access Authenticated API calls can include either: diff --git a/account_manager.py b/account_manager.py index 30e6a84..a61cf34 100644 --- a/account_manager.py +++ b/account_manager.py @@ -16,7 +16,7 @@ sanitizer = Sanitizer() ip_filter = IPFilter(WHITELIST, BLACKLIST) timeout = Timeout(ip_filter, BLACKLIST_AFTER_ATTEMPTS) users_db = UsersDB(USERS_FILE) -access_control = AccessControl(users_db, instance) +access_control = AccessControl(users_db, instance, HISTORY_FILE) jwt_auth = JWTAuth( users_db, access_control, logger, SECRET_KEY, TOKEN_EXPIRE_MINUTES, TOKEN_ALGORITHM ) diff --git a/config.json b/config.json index 7706f3c..a1b108e 100644 --- a/config.json +++ b/config.json @@ -2,6 +2,7 @@ "secret_key_env": "SECRET_KEY", "secret_key_file": "secret_key.txt", "users_db": "users_db.json", + "history_db": "history.sqlite3", "access_token_expiration_hours": 8760, "max_access_token_expiration_hours": 8760, "log": "account_manager.log", diff --git a/tests/test_history_store.py b/tests/test_history_store.py new file mode 100644 index 0000000..84fc4ed --- /dev/null +++ b/tests/test_history_store.py @@ -0,0 +1,65 @@ +import importlib.util +import os +import tempfile +import unittest +from pathlib import Path + + +module_path = Path(__file__).parents[1] / "utils" / "history_store.py" +spec = importlib.util.spec_from_file_location("history_store", module_path) +history_store = importlib.util.module_from_spec(spec) +spec.loader.exec_module(history_store) +HistoryStore = history_store.HistoryStore + + +class HistoryStoreTests(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.database = os.path.join(self.temp_dir.name, "history.sqlite3") + self.store = HistoryStore(self.database) + + def tearDown(self): + self.temp_dir.cleanup() + + @staticmethod + def item(owner_id, value): + return { + "prompt": [0, f"prompt-{value}", {}, {}, []], + "outputs": {"1": {"images": [{"filename": f"{value}.png"}]}}, + "status": {"status_str": "success", "completed": True, "messages": []}, + "user_id": owner_id, + } + + def test_history_survives_new_store_instance(self): + self.store.save("one", self.item("user-a", 1), 10) + + restored = HistoryStore(self.database).load(10) + + self.assertEqual(["one"], list(restored)) + self.assertEqual("user-a", restored["one"]["user_id"]) + self.assertEqual("1.png", restored["one"]["outputs"]["1"]["images"][0]["filename"]) + + def test_limit_and_owner_delete(self): + self.store.save("one", self.item("user-a", 1), 2) + self.store.save("two", self.item("user-b", 2), 2) + self.store.save("three", self.item("user-a", 3), 2) + + self.assertEqual(["two", "three"], list(self.store.load(10))) + + self.store.delete_owner("user-a") + + self.assertEqual(["two"], list(self.store.load(10))) + + def test_delete_and_clear(self): + self.store.save("one", self.item("user-a", 1), 10) + self.store.save("two", self.item("user-b", 2), 10) + + self.store.delete("one") + self.assertEqual(["two"], list(self.store.load(10))) + + self.store.clear() + self.assertEqual({}, self.store.load(10)) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/access_control.py b/utils/access_control.py index 062b9f9..7a8d220 100644 --- a/utils/access_control.py +++ b/utils/access_control.py @@ -3,6 +3,7 @@ import sys import heapq import copy import contextvars +import logging import mimetypes from datetime import datetime from aiohttp import web @@ -13,10 +14,16 @@ from server import PromptServer from execution import PromptQueue, MAXIMUM_HISTORY_SIZE from .users_db import UsersDB +from .history_store import HistoryStore + + +logger = logging.getLogger("ComfyUI-Account-Manager") class AccessControl: - def __init__(self, users_db: UsersDB, server: PromptServer): + def __init__( + self, users_db: UsersDB, server: PromptServer, history_file: str = None + ): self.users_db = users_db self.server = server @@ -29,6 +36,12 @@ class AccessControl: 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( @@ -350,7 +363,7 @@ class AccessControl: 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: + 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 @@ -367,8 +380,46 @@ class AccessControl: "user_id": user_id, } self.__prompt_queue.history[prompt[1]].update(history_result) + self._save_history_item(prompt[1]) self.server.queue_updated() + 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.""" with self.__prompt_queue.mutex: @@ -466,21 +517,25 @@ class AccessControl: 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): 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.""" + 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 diff --git a/utils/config.py b/utils/config.py index c7e4efa..6b05051 100644 --- a/utils/config.py +++ b/utils/config.py @@ -42,6 +42,7 @@ TOKEN_EXPIRE_MINUTES = 60 * config.get("access_token_expiration_hours", 8760) MAX_TOKEN_EXPIRE_MINUTES = 60 * config.get("max_access_token_expiration_hours", 8760) USERS_FILE = os.path.join(EXT_PATH, config.get("users_db", "users_db.json")) +HISTORY_FILE = os.path.join(EXT_PATH, config.get("history_db", "history.sqlite3")) LOG_FILE = os.path.join(EXT_PATH, config.get("log", "account_manager.log")) LOG_LEVELS = config.get("log_levels", ["INFO"]) diff --git a/utils/history_store.py b/utils/history_store.py new file mode 100644 index 0000000..6726f14 --- /dev/null +++ b/utils/history_store.py @@ -0,0 +1,98 @@ +import json +import os +import sqlite3 +from contextlib import closing + + +class HistoryStore: + """Persist completed ComfyUI history entries in a small SQLite database.""" + + def __init__(self, database: str): + self.database = os.fspath(database) + parent = os.path.dirname(os.path.abspath(self.database)) + os.makedirs(parent, exist_ok=True) + self._initialize() + + def _connect(self) -> sqlite3.Connection: + return sqlite3.connect(self.database, timeout=10) + + def _initialize(self) -> None: + with closing(self._connect()) as connection: + with connection: + connection.execute( + """ + CREATE TABLE IF NOT EXISTS history ( + sequence INTEGER PRIMARY KEY AUTOINCREMENT, + prompt_id TEXT NOT NULL UNIQUE, + owner_id TEXT NOT NULL, + data TEXT NOT NULL + ) + """ + ) + connection.execute( + "CREATE INDEX IF NOT EXISTS history_owner_id ON history(owner_id)" + ) + + def load(self, max_items: int) -> dict: + with closing(self._connect()) as connection: + rows = connection.execute( + """ + SELECT prompt_id, owner_id, data + FROM history + ORDER BY sequence DESC + LIMIT ? + """, + (max_items,), + ).fetchall() + + history = {} + for prompt_id, owner_id, data in reversed(rows): + item = json.loads(data) + if not isinstance(item, dict): + continue + item["user_id"] = owner_id + history[prompt_id] = item + return history + + def save(self, prompt_id: str, item: dict, max_items: int) -> None: + owner_id = item.get("user_id") or "" + data = json.dumps(item, ensure_ascii=False, separators=(",", ":")) + with closing(self._connect()) as connection: + with connection: + connection.execute( + """ + INSERT INTO history (prompt_id, owner_id, data) + VALUES (?, ?, ?) + ON CONFLICT(prompt_id) DO UPDATE SET + owner_id = excluded.owner_id, + data = excluded.data + """, + (prompt_id, owner_id, data), + ) + connection.execute( + """ + DELETE FROM history + WHERE sequence IN ( + SELECT sequence + FROM history + ORDER BY sequence DESC + LIMIT -1 OFFSET ? + ) + """, + (max_items,), + ) + + def delete(self, prompt_id: str) -> None: + with closing(self._connect()) as connection: + with connection: + connection.execute("DELETE FROM history WHERE prompt_id = ?", (prompt_id,)) + + def delete_owner(self, owner_id: str) -> None: + with closing(self._connect()) as connection: + with connection: + connection.execute("DELETE FROM history WHERE owner_id = ?", (owner_id,)) + + def clear(self) -> None: + with closing(self._connect()) as connection: + with connection: + connection.execute("DELETE FROM history")