nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(lora): report integration save failures (#220)

Browse Source 
POST /prompt_manager/lora/enable commits integrations.lora_manager
through the config store and publishes the runtime values only after
config.json holds them. Before, it set the runtime values, rewrote the
whole runtime config (dropping sections it did not know) and reported
success even when the write had failed.

A write failure is 500 config-write-failed and changes nothing at
runtime; an unparseable config.json is 409 config-unreadable. The
body goes through the capped JSON reader and the flags must be
booleans. A trigger word cache that fails to load is reported as
persisted:true, cache_loaded:false. The CivitAI key is never echoed.
The LoRA tests now write to a temporary config path.
This commit is contained in:
Vito Sansevero
2026-10-22 19:07:20 -07:00
parent fba5903049
commit 9193373eac
2 changed files with 308 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/lora_integration.py
+65 -10
View File
@@ -5,12 +5,37 @@ from pathlib import Path
from aiohttp import web
from ..config_store import ConfigUnreadable, RevisionConflict, get_config_store
# Flags of POST /lora/enable; each must be a JSON boolean when present.
_LORA_FLAGS = ("enabled", "trigger_words_enabled", "clear_civitai_api_key")
_LORA_PATH_ERROR = (
"LoraManager path must be a directory under ComfyUI's custom_nodes "
"whose name contains 'lora'"
)
def _lora_settings_mutation(values):
"""The config commit for a LoRA save: only ``integrations.lora_manager``.
The submitted keys replace their values; other LoRA keys, other
integrations and every other section stay as they are on disk.
"""
def mutate(config):
integrations = config.get("integrations")
integrations = dict(integrations) if isinstance(integrations, dict) else {}
lora = integrations.get("lora_manager")
lora = dict(lora) if isinstance(lora, dict) else {}
lora.update(values)
integrations["lora_manager"] = lora
config["integrations"] = integrations
return config
return mutate
class LoraIntegrationMixin:
"""Mixin providing LoraManager detection, scanning, and trigger word endpoints."""
@@ -131,14 +156,25 @@ class LoraIntegrationMixin:
return allowed, None
async def lora_enable(self, request):
"""Enable or disable LoraManager integration and save to config.json."""
"""Enable or disable LoraManager integration and save it to config.json.
The values are committed through the config store first (#220) and
published at runtime only once config.json holds them; a save that
cannot be written reports ``persisted: false`` and changes nothing.
The CivitAI key is never echoed back.
"""
try:
data = await request.json()
data, error = await self._read_json_body(request)
if error is not None:
return error
for flag in _LORA_FLAGS:
if flag in data and not isinstance(data[flag], bool):
return self._bad_lora_request(f"{flag} must be true or false")
enabled = data.get("enabled", False)
path = data.get("path", "")
trigger_words = data.get("trigger_words_enabled", False)
from ..config import IntegrationConfig, PromptManagerConfig
from ..config import IntegrationConfig
from ..lora_utils import get_trigger_cache
# The key is never echoed back to the UI, so a blank or absent
@@ -161,20 +197,35 @@ class LoraIntegrationMixin:
if error is not None:
return error
# Update in-memory config
values = {
"enabled": enabled,
"path": path,
"trigger_words_enabled": trigger_words,
"civitai_api_key": civitai_key,
}
try:
_saved, revision = await self._run_in_executor(
get_config_store().commit,
_lora_settings_mutation(values),
expected_revision=None,
)
except (ConfigUnreadable, RevisionConflict, OSError) as exc:
return self._config_commit_failure(exc)
IntegrationConfig.LORA_MANAGER_ENABLED = enabled
IntegrationConfig.LORA_MANAGER_PATH = path
IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = trigger_words
IntegrationConfig.CIVITAI_API_KEY = civitai_key
# Persist to the configured config.json
PromptManagerConfig.save_to_file()
# Load trigger word cache if enabling
cache = get_trigger_cache()
if enabled and trigger_words and path:
count = await self._run_in_executor(cache.load, path)
self.logger.info(f"Trigger word cache loaded: {count} LoRAs")
try:
count = await self._run_in_executor(cache.load, path)
self.logger.info(f"Trigger word cache loaded: {count} LoRAs")
except Exception as e:
# Saved and enabled; only the trigger words are missing.
self.logger.warning(f"Trigger word cache could not be loaded: {e}")
cache.clear()
elif not enabled:
cache.clear()
@@ -184,6 +235,10 @@ class LoraIntegrationMixin:
"enabled": enabled,
"path": self._public_path(path),
"trigger_words_enabled": trigger_words,
"persisted": True,
"revision": revision,
"runtime_applied": True,
"cache_loaded": cache.is_loaded,
}
)
except Exception as e:
tests/test_api_lora.py
+243 -11
View File
@@ -4,11 +4,13 @@ Filesystem detection and network downloads are patched out so the tests run
anywhere without a LoraManager install or a CivitAI connection.
"""
import asyncio
import json
import os
import shutil
import sys
import tempfile
import threading
import types
import unittest
from pathlib import Path
@@ -36,9 +38,15 @@ from database.operations import PromptDatabase
from py.api import PromptManagerAPI
from py import lora_utils
from py.config import IntegrationConfig, PromptManagerConfig
from py.api import admin as admin_module
from py.lora_utils import get_trigger_cache
from utils.hashing import generate_prompt_hash
try:
from tests.open_handles import assert_closed
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
SECRET_KEY = "sk-civitai-super-secret-0123456789"
@@ -56,7 +64,13 @@ class LoraAPITestCase(AioHTTPTestCase):
return_value=[self.custom_nodes],
)
self._custom_nodes_patch.start()
self._real_save = PromptManagerConfig.save_to_file
# Never the repository's config.json (live user data): every save in
# these tests goes to a file in this test's temp directory.
self.config_file = self.tmp_root / "config.json"
self._env_patch = patch.dict(
os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file)}
)
self._env_patch.start()
self._config_patch = patch.multiple(
IntegrationConfig,
@@ -70,8 +84,6 @@ class LoraAPITestCase(AioHTTPTestCase):
"py.lora_utils.detect_lora_manager", return_value=None
)
self.detect_mock = self._detect_patch.start()
self._save_patch = patch.object(PromptManagerConfig, "save_to_file")
self.save_mock = self._save_patch.start()
get_trigger_cache().clear()
app = web.Application()
@@ -83,7 +95,7 @@ class LoraAPITestCase(AioHTTPTestCase):
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
self._save_patch.stop()
self._env_patch.stop()
self._detect_patch.stop()
self._config_patch.stop()
self._custom_nodes_patch.stop()
@@ -96,8 +108,19 @@ class LoraAPITestCase(AioHTTPTestCase):
):
if os.path.exists(path):
os.unlink(path)
assert_closed(self, self.tmp_root)
shutil.rmtree(self.tmp_root, ignore_errors=True)
def _saved(self):
return json.loads(self.config_file.read_text())
def _make_unwritable(self):
"""Point the config path under a regular file, so no write can succeed."""
blocker = self.tmp_root / "not-a-dir"
blocker.write_text("x")
self.config_file = blocker / "config.json"
os.environ["PROMPT_MANAGER_CONFIG_PATH"] = str(self.config_file)
async def _post_json(self, url, body):
return await self.client.request(
"POST",
@@ -250,7 +273,7 @@ class TestLoraEnableRoute(LoraAPITestCase):
)
self.assertEqual(resp.status, 400)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.save_mock.assert_not_called()
self.assertFalse(self.config_file.exists())
async def test_enable_persists_and_loads_trigger_cache(self):
lm = _make_lora_manager(self.custom_nodes)
@@ -285,7 +308,15 @@ class TestLoraEnableRoute(LoraAPITestCase):
os.path.normcase(os.path.realpath(str(lm))),
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
self.save_mock.assert_called_once_with()
saved = self._saved()["integrations"]["lora_manager"]
self.assertEqual(saved["civitai_api_key"], SECRET_KEY)
self.assertEqual(saved["path"], IntegrationConfig.LORA_MANAGER_PATH)
self.assertIs(saved["enabled"], True)
self.assertIs(saved["trigger_words_enabled"], True)
self.assertEqual(
(data["persisted"], data["runtime_applied"], data["cache_loaded"]),
(True, True, True),
)
self.assertTrue(get_trigger_cache().is_loaded)
self.assertEqual(get_trigger_cache().get_trigger_words("neon"), ["glow"])
@@ -319,7 +350,7 @@ class TestLoraEnableRoute(LoraAPITestCase):
self.assertFalse(json.loads(body)["success"])
self.assertNotIn(str(self.tmp_root), body)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.save_mock.assert_not_called()
self.assertFalse(self.config_file.exists())
async def test_enable_with_custom_node_lacking_lora_in_its_name_is_400(self):
lm = _make_lora_manager(self.custom_nodes).parent / "ComfyUI-Other"
@@ -344,7 +375,6 @@ class TestLoraEnableRoute(LoraAPITestCase):
async def test_enable_writes_the_configured_config_file(self):
config_file = self.tmp_root / "cfg" / "config.json"
self.save_mock.side_effect = self._real_save
with patch.dict(os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(config_file)}):
resp = await self._post_json(
@@ -353,7 +383,6 @@ class TestLoraEnableRoute(LoraAPITestCase):
)
self.assertEqual(resp.status, 200)
self.save_mock.assert_called_once_with()
self.assertTrue(config_file.is_file())
saved = json.loads(config_file.read_text())["integrations"]["lora_manager"]
self.assertFalse(saved["enabled"])
@@ -414,9 +443,212 @@ class TestLoraEnableRoute(LoraAPITestCase):
self.assertEqual(resp.status, 400)
async def test_persist_failure_is_500(self):
self.save_mock.side_effect = OSError("read-only")
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
self._make_unwritable()
with self.assertLogs("prompt_manager.api", level="ERROR"):
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False}
)
self.assertEqual(resp.status, 500)
self.assertEqual((await resp.json())["code"], "config-write-failed")
async def test_enable_write_failure_reports_and_keeps_runtime(self):
# #220: a save that cannot be written publishes nothing and says so.
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
self._make_unwritable()
with self.assertLogs("prompt_manager.api", level="ERROR"):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
"enabled": True,
"path": str(lm),
"trigger_words_enabled": True,
"clear_civitai_api_key": True,
},
)
text = await resp.text()
body = json.loads(text)
self.assertEqual(resp.status, 500)
self.assertEqual(body["code"], "config-write-failed")
self.assertEqual((body["persisted"], body["runtime_applied"]), (False, False))
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertEqual(IntegrationConfig.LORA_MANAGER_PATH, "")
self.assertFalse(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
self.assertFalse(get_trigger_cache().is_loaded)
self.assertNotIn(SECRET_KEY, text)
self.assertNotIn(str(self.tmp_root), text)
async def test_enable_preserves_unknown_sections(self):
self.config_file.write_text(
json.dumps(
{
"custom_section": {"keep": True},
"web_ui": {"result_timeout": 11},
"integrations": {
"lora_manager": {"civitai_api_key": SECRET_KEY, "extra": 1},
"other_tool": {"on": True},
},
}
)
)
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
self.assertEqual(resp.status, 200)
saved = self._saved()
self.assertEqual(saved["custom_section"], {"keep": True})
self.assertEqual(saved["web_ui"], {"result_timeout": 11})
self.assertEqual(saved["integrations"]["other_tool"], {"on": True})
lora = saved["integrations"]["lora_manager"]
self.assertEqual(lora["extra"], 1)
self.assertEqual(lora["civitai_api_key"], SECRET_KEY)
self.assertIs(lora["trigger_words_enabled"], True)
self.assertNotIn("database", saved) # no runtime defaults written
async def test_enable_refuses_a_corrupt_config_file(self):
self.config_file.write_bytes(b"{not json")
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
body = await resp.json()
self.assertEqual((resp.status, body["code"]), (409, "config-unreadable"))
self.assertEqual(self.config_file.read_bytes(), b"{not json")
self.assertFalse(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
async def test_enable_requires_boolean_flags_and_body_cap(self):
for body in (
{"enabled": "yes"},
{"enabled": False, "trigger_words_enabled": 1},
{"enabled": False, "clear_civitai_api_key": "true"},
):
with self.subTest(body=body):
resp = await self._post_json("/prompt_manager/lora/enable", body)
self.assertEqual(resp.status, 400)
resp = await self._post_json("/prompt_manager/lora/enable", [True])
self.assertEqual(resp.status, 400)
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "civitai_api_key": "k" * (2 * 1024 * 1024)},
)
self.assertEqual(resp.status, 413)
self.assertFalse(self.config_file.exists())
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_blank_key_unchanged_and_explicit_clear(self):
self.config_file.write_text(
json.dumps({"integrations": {"lora_manager": {"civitai_api_key": "old"}}})
)
IntegrationConfig.CIVITAI_API_KEY = "old"
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False, "civitai_api_key": " "}
)
self.assertEqual(resp.status, 200)
self.assertEqual(
self._saved()["integrations"]["lora_manager"]["civitai_api_key"], "old"
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "old")
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "clear_civitai_api_key": True},
)
self.assertEqual(resp.status, 200)
self.assertEqual(
self._saved()["integrations"]["lora_manager"]["civitai_api_key"], ""
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "")
async def test_cache_load_failure_reports_persisted_but_unloaded(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
with (
patch.object(
type(get_trigger_cache()), "load", side_effect=OSError("unreadable")
),
self.assertLogs("prompt_manager.api", level="WARNING"),
):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": True, "path": str(lm), "trigger_words_enabled": True},
)
text = await resp.text()
body = json.loads(text)
self.assertEqual(resp.status, 200)
self.assertIs(body["success"], True)
self.assertEqual((body["persisted"], body["cache_loaded"]), (True, False))
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertIs(self._saved()["integrations"]["lora_manager"]["enabled"], True)
self.assertNotIn("unreadable", text)
class TestConfigWritersShareOneStore(LoraAPITestCase):
async def asyncSetUp(self):
await super().asyncSetUp()
self._orig_timeout = PromptManagerConfig.RESULT_TIMEOUT
async def asyncTearDown(self):
PromptManagerConfig.RESULT_TIMEOUT = self._orig_timeout
await super().asyncTearDown()
async def test_settings_and_lora_saves_do_not_overwrite_each_other(self):
# #220: the LoRA save is sent while the settings save is inside its
# commit (lock held, file already read). The LoRA commit re-reads the
# file under the same lock, so both edits survive and nothing else
# is lost (before, the LoRA save rewrote the whole runtime config).
self.config_file.write_text(json.dumps({"custom_section": {"keep": True}}))
inside = threading.Event()
release = threading.Event()
real_mutation = admin_module._settings_mutation
def gated_mutation(updates, new_roots):
mutate = real_mutation(updates, new_roots)
def gated(config):
inside.set()
release.wait(10)
return mutate(config)
return gated
loop = asyncio.get_running_loop()
with patch.object(admin_module, "_settings_mutation", gated_mutation):
settings = asyncio.ensure_future(
self._post_json("/prompt_manager/settings", {"result_timeout": 13})
)
self.assertTrue(await loop.run_in_executor(None, inside.wait, 10))
lora = asyncio.ensure_future(
self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
)
await asyncio.sleep(0) # let the LoRA request reach the lock
release.set()
settings_resp, lora_resp = await asyncio.gather(settings, lora)
self.assertEqual((settings_resp.status, lora_resp.status), (200, 200))
saved = self._saved()
self.assertEqual(saved["custom_section"], {"keep": True})
self.assertEqual(saved["web_ui"]["result_timeout"], 13)
self.assertIs(
saved["integrations"]["lora_manager"]["trigger_words_enabled"], True
)
settings_body = await settings_resp.json()
lora_body = await lora_resp.json()
self.assertNotEqual(settings_body["revision"], lora_body["revision"])
class TestLoraScanRoute(LoraAPITestCase):