diff --git a/.gitignore b/.gitignore index 8e87659..a90ed8d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ *.py[cod] node_modules/ +config.json diff --git a/README.md b/README.md index b57a263..ec7c5ac 100644 --- a/README.md +++ b/README.md @@ -40,15 +40,13 @@ git clone https://github.com/ethanfel/ComfyUI-LM-Remote.git ## Configuration -Edit `config.json` in the package directory: +Open the configuration panel from any of these places: -```json -{ - "remote_url": "http://192.168.1.3:8188", - "timeout": 30, - "path_mappings": {} -} -``` +- Click the gear in the **LoRA Info** sidebar +- Open **Settings**, select **LM Remote**, then click **Configure LM Remote** +- Run **Configure LM Remote** from the command palette + +Enter the remote URL, adjust the timeout if needed, and click **Test connection**. **Save** validates the settings and applies them to new requests immediately; the ComfyUI process does not need to restart. If the remote URL changes, the panel offers a page reload so Manager assets and live progress reconnect to the new server. Save workflow edits before using it. | Field | Type | Default | Description | |-------|------|---------|-------------| @@ -56,14 +54,19 @@ Edit `config.json` in the package directory: | `timeout` | int | `30` | HTTP request timeout in seconds | | `path_mappings` | object | `{}` | Remote-to-local path prefix mapping (see below) | +Configuration is stored under `/ComfyUI-LM-Remote/config.json`, outside the custom-node checkout. An existing package-level `config.json` is loaded as a legacy fallback and is left untouched. The first save through the panel migrates its values into ComfyUI user data. + +The panel is the normal way to manage this file; manual editing is not required. For a headless or automated deployment, start from [`config.example.json`](config.example.json), place the copy outside the custom-node checkout, and point `LM_REMOTE_CONFIG` to it. + ### Environment Variable Overrides -Environment variables take priority over `config.json`: +Environment variables take priority over the stored configuration without rewriting it. Overridden fields are shown as managed in the panel. -| Variable | Overrides | -|----------|-----------| -| `LM_REMOTE_URL` | `remote_url` | -| `LM_REMOTE_TIMEOUT` | `timeout` | +| Variable | Purpose | +|----------|---------| +| `LM_REMOTE_URL` | Overrides `remote_url`; the field is shown as managed in the panel | +| `LM_REMOTE_TIMEOUT` | Overrides `timeout`; the field is shown as managed in the panel | +| `LM_REMOTE_CONFIG` | Uses an explicit configuration file instead of the ComfyUI user-data path | ### Path Mappings @@ -108,6 +111,8 @@ ComfyUI does not currently expose an extension API for adding custom tabs to the Auto-open is enabled by default. Disable it under **Settings > LM Remote > LoRA Info > Auto-open** if you prefer to open **LoRA Info** manually from the sidebar or command palette. +The gear beside the refresh button opens LM Remote connection settings without leaving the sidebar. + ## How It Works ### Reverse Proxy @@ -145,12 +150,14 @@ After fetching the relative path from the remote metadata, LoRA files are loaded After installation and configuration: -1. Restart ComfyUI -2. Check logs for: `[LM-Remote] Proxy routes registered -> http://192.168.1.3:8188` -3. Open ComfyUI -- the LoRA Manager web UI should load (proxied from remote) -4. Add a stock or remote LoRA loader and click the node -- **LoRA Info** should open -5. Select a LoRA -- its Manager card (or external search links) should appear and remote trigger words should populate where supported -6. Run the workflow -- the LoRA loads from local shared storage +1. Restart ComfyUI once after installing the custom node +2. Open **Configure LM Remote**, enter the URL, and run **Test connection** +3. Save the configuration; no ComfyUI process restart is required +4. If prompted after a remote URL change, save workflow edits and reload the browser page +5. Open the LoRA Manager web UI -- it should load through the remote proxy +6. Add a stock or remote LoRA loader and click the node -- **LoRA Info** should open +7. Select a LoRA -- its Manager card (or external search links) should appear and remote trigger words should populate where supported +8. Run the workflow -- the LoRA loads from local shared storage ## License diff --git a/config.example.json b/config.example.json new file mode 100644 index 0000000..57cab36 --- /dev/null +++ b/config.example.json @@ -0,0 +1,5 @@ +{ + "remote_url": "", + "timeout": 30, + "path_mappings": {} +} diff --git a/config.json b/config.json deleted file mode 100644 index e4f936e..0000000 --- a/config.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "remote_url": "http://192.168.1.3:8188", - "timeout": 30, - "path_mappings": {} -} diff --git a/config.py b/config.py index a559d05..0cc3b33 100644 --- a/config.py +++ b/config.py @@ -1,62 +1,535 @@ -"""Configuration for ComfyUI-LM-Remote.""" +"""Validated, reloadable configuration for ComfyUI-LM-Remote.""" + from __future__ import annotations +import hashlib import json import logging import os +import re +import stat +import tempfile +import threading +from dataclasses import dataclass from pathlib import Path +from typing import Mapping +from urllib.parse import urlsplit, urlunsplit logger = logging.getLogger(__name__) _PACKAGE_DIR = Path(__file__).resolve().parent -_CONFIG_FILE = _PACKAGE_DIR / "config.json" +_LEGACY_CONFIG_FILE = _PACKAGE_DIR / "config.json" +_CONFIG_DIRECTORY_NAME = "ComfyUI-LM-Remote" +_CONFIG_FILE_NAME = "config.json" +_KNOWN_FIELDS = frozenset({"remote_url", "timeout", "path_mappings"}) +_MAX_TIMEOUT = 300 +_MAX_MAPPINGS = 100 +_MAX_VALUE_LENGTH = 4096 +_CONTROL_CHARACTERS = re.compile(r"[\x00-\x1f\x7f]") -class RemoteConfig: - """Holds remote LoRA Manager connection settings.""" +class ConfigValidationError(ValueError): + """Raised when a proposed configuration value is invalid.""" - def __init__(self): - self.remote_url: str = "" - self.timeout: int = 30 - self.path_mappings: dict[str, str] = {} - self._load() + def __init__(self, field: str, message: str): + super().__init__(message) + self.field = field - # ------------------------------------------------------------------ - def _load(self): - # Environment variable takes priority - env_url = os.environ.get("LM_REMOTE_URL", "") - env_timeout = os.environ.get("LM_REMOTE_TIMEOUT", "") - # Load config.json defaults - if _CONFIG_FILE.exists(): - try: - with open(_CONFIG_FILE, "r", encoding="utf-8") as f: - data = json.load(f) - self.remote_url = data.get("remote_url", "") - self.timeout = int(data.get("timeout", 30)) - self.path_mappings = data.get("path_mappings", {}) - except Exception as exc: - logger.warning("[LM-Remote] Failed to read config.json: %s", exc) +class ConfigConflictError(RuntimeError): + """Raised when a browser attempts to replace a stale configuration.""" - # Env overrides - if env_url: - self.remote_url = env_url - if env_timeout: - self.timeout = int(env_timeout) - # Strip trailing slash - self.remote_url = self.remote_url.rstrip("/") +@dataclass(frozen=True) +class ConfigSnapshot: + """One coherent set of effective runtime values.""" - @property - def is_configured(self) -> bool: - return bool(self.remote_url) + generation: int + remote_url: str + timeout: int + path_mappings: tuple[tuple[str, str], ...] + + def mappings_dict(self) -> dict[str, str]: + return dict(self.path_mappings) def map_path(self, remote_path: str) -> str: - """Apply remote->local path prefix mappings.""" - for remote_prefix, local_prefix in self.path_mappings.items(): - if remote_path.startswith(remote_prefix): - return local_prefix + remote_path[len(remote_prefix):] + """Map a path using only the values captured by this snapshot.""" + if not isinstance(remote_path, str): + return remote_path + normalized_path = remote_path.replace("\\", "/") + for remote_prefix, local_prefix in self.path_mappings: + is_root = remote_prefix == "/" + if normalized_path == remote_prefix: + remainder = "" + elif is_root and normalized_path.startswith("/"): + remainder = normalized_path[1:] + elif normalized_path.startswith(f"{remote_prefix}/"): + remainder = normalized_path[len(remote_prefix) + 1 :] + else: + continue + + if not remainder: + return local_prefix + separator = ( + "\\" if "\\" in local_prefix and "/" not in local_prefix else os.sep + ) + local_base = local_prefix.rstrip("/\\") + return f"{local_base}{separator}{remainder.replace('/', separator)}" return remote_path +def _default_user_config_file(environ: Mapping[str, str]) -> Path: + explicit_path = environ.get("LM_REMOTE_CONFIG", "").strip() + if explicit_path: + return Path(explicit_path).expanduser() + + try: + import folder_paths # type: ignore + + user_directory = Path(folder_paths.get_user_directory()) + return user_directory / _CONFIG_DIRECTORY_NAME / _CONFIG_FILE_NAME + except Exception: + # Outside ComfyUI (for example, documentation tools), preserve the + # historical package-level behaviour. + return _LEGACY_CONFIG_FILE + + +def _normalize_url(value: object, *, allow_empty: bool = True) -> str: + if not isinstance(value, str): + raise ConfigValidationError("remote_url", "Remote URL must be text.") + value = value.strip() + if not value: + if allow_empty: + return "" + raise ConfigValidationError("remote_url", "Enter a remote LoRA Manager URL.") + if _CONTROL_CHARACTERS.search(value): + raise ConfigValidationError( + "remote_url", "Remote URL contains invalid characters." + ) + + try: + parsed = urlsplit(value) + # Accessing .port performs its own range and syntax validation. + parsed.port + except ValueError as exc: + raise ConfigValidationError( + "remote_url", "Remote URL has an invalid port." + ) from exc + + if parsed.scheme.lower() not in {"http", "https"}: + raise ConfigValidationError( + "remote_url", "Remote URL must use http:// or https://." + ) + if not parsed.hostname: + raise ConfigValidationError( + "remote_url", "Remote URL must include a host name." + ) + if parsed.username is not None or parsed.password is not None: + raise ConfigValidationError( + "remote_url", "Credentials are not allowed in the remote URL." + ) + if parsed.query or parsed.fragment: + raise ConfigValidationError( + "remote_url", "Remote URL cannot contain a query or fragment." + ) + + path = parsed.path.rstrip("/") + return urlunsplit((parsed.scheme.lower(), parsed.netloc, path, "", "")) + + +def _normalize_timeout(value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ConfigValidationError("timeout", "Timeout must be a whole number.") + if value < 1 or value > _MAX_TIMEOUT: + raise ConfigValidationError( + "timeout", f"Timeout must be between 1 and {_MAX_TIMEOUT} seconds." + ) + return value + + +def _normalize_remote_prefix(value: str) -> str: + normalized = value.replace("\\", "/") + if normalized != "/": + normalized = normalized.rstrip("/") + return normalized + + +def _normalize_mapping_value(field: str, value: object) -> str: + if not isinstance(value, str): + raise ConfigValidationError("path_mappings", f"{field} path must be text.") + value = value.strip() + if not value: + raise ConfigValidationError("path_mappings", f"{field} path cannot be empty.") + if len(value) > _MAX_VALUE_LENGTH: + raise ConfigValidationError("path_mappings", f"{field} path is too long.") + if _CONTROL_CHARACTERS.search(value): + raise ConfigValidationError( + "path_mappings", f"{field} path contains invalid characters." + ) + return value + + +def _normalize_mappings(value: object) -> tuple[tuple[str, str], ...]: + if not isinstance(value, dict): + raise ConfigValidationError("path_mappings", "Path mappings must be an object.") + if len(value) > _MAX_MAPPINGS: + raise ConfigValidationError( + "path_mappings", f"No more than {_MAX_MAPPINGS} path mappings are allowed." + ) + + normalized: dict[str, str] = {} + for remote_value, local_value in value.items(): + remote_prefix = _normalize_remote_prefix( + _normalize_mapping_value("Remote", remote_value) + ) + if not remote_prefix: + raise ConfigValidationError( + "path_mappings", + "Remote path cannot consist only of path separators.", + ) + local_prefix = _normalize_mapping_value("Local", local_value) + if remote_prefix in normalized: + raise ConfigValidationError( + "path_mappings", f"Duplicate remote path prefix: {remote_prefix}" + ) + normalized[remote_prefix] = local_prefix + + # Specific mappings must win over broader parent mappings. + return tuple( + sorted(normalized.items(), key=lambda pair: len(pair[0]), reverse=True) + ) + + +def validate_config(data: object, *, allow_empty_url: bool = True) -> dict[str, object]: + """Validate and normalize a complete stored configuration.""" + if not isinstance(data, dict): + raise ConfigValidationError("config", "Configuration must be a JSON object.") + unknown = set(data) - _KNOWN_FIELDS + if unknown: + names = ", ".join(sorted(str(name) for name in unknown)) + raise ConfigValidationError( + "config", f"Unknown configuration field(s): {names}" + ) + + return { + "remote_url": _normalize_url( + data.get("remote_url", ""), allow_empty=allow_empty_url + ), + "timeout": _normalize_timeout(data.get("timeout", 30)), + "path_mappings": dict(_normalize_mappings(data.get("path_mappings", {}))), + } + + +def _revision_for(data: Mapping[str, object]) -> str: + encoded = json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _revision_for_file(path: Path) -> str: + return f"sha256:{hashlib.sha256(path.read_bytes()).hexdigest()}" + + +class RemoteConfig: + """Thread-safe configuration with legacy fallback and atomic persistence.""" + + def __init__( + self, + config_file: str | Path | None = None, + legacy_config_file: str | Path | None = None, + environ: Mapping[str, str] | None = None, + ) -> None: + self._environ = environ if environ is not None else os.environ + self._explicit_target = config_file is None and bool( + self._environ.get("LM_REMOTE_CONFIG", "").strip() + ) + self._config_file = ( + Path(config_file) + if config_file + else _default_user_config_file(self._environ) + ) + self._legacy_config_file = ( + Path(legacy_config_file) if legacy_config_file else _LEGACY_CONFIG_FILE + ) + self._lock = threading.RLock() + self._snapshot = ConfigSnapshot(0, "", 30, ()) + self._configured: dict[str, object] = { + "remote_url": "", + "timeout": 30, + "path_mappings": {}, + } + self._overrides: dict[str, str | None] = { + "remote_url": None, + "timeout": None, + } + self._revision = _revision_for({}) + self._source = "defaults" + self._warnings: list[str] = [] + self.reload() + + def _read_source(self) -> tuple[dict[str, object], str]: + if self._config_file.exists(): + path = self._config_file + source = "explicit" if self._explicit_target else "user" + elif self._explicit_target: + return {}, "explicit" + elif self._legacy_config_file.exists(): + path = self._legacy_config_file + source = "legacy" + else: + return {}, "defaults" + + with path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + if not isinstance(data, dict): + raise ConfigValidationError( + "config", "Configuration file must contain an object." + ) + return data, source + + def _source_location(self) -> tuple[Path | None, str]: + if self._config_file.exists(): + source = "explicit" if self._explicit_target else "user" + return self._config_file, source + if self._explicit_target: + return None, "explicit" + if self._legacy_config_file.exists(): + return self._legacy_config_file, "legacy" + return None, "defaults" + + def _effective_values( + self, configured: dict[str, object] + ) -> tuple[dict[str, object], dict[str, str | None], list[str]]: + effective = { + "remote_url": configured["remote_url"], + "timeout": configured["timeout"], + "path_mappings": dict(configured["path_mappings"]), + } + overrides: dict[str, str | None] = {"remote_url": None, "timeout": None} + warnings: list[str] = [] + + env_url = self._environ.get("LM_REMOTE_URL", "") + if env_url: + try: + effective["remote_url"] = _normalize_url(env_url) + overrides["remote_url"] = "LM_REMOTE_URL" + except ConfigValidationError as exc: + warnings.append(f"Ignoring invalid LM_REMOTE_URL: {exc}") + + env_timeout = self._environ.get("LM_REMOTE_TIMEOUT", "") + if env_timeout: + try: + if isinstance(env_timeout, str) and env_timeout.strip().isdigit(): + parsed_timeout: object = int(env_timeout.strip()) + else: + parsed_timeout = env_timeout + effective["timeout"] = _normalize_timeout(parsed_timeout) + overrides["timeout"] = "LM_REMOTE_TIMEOUT" + except ConfigValidationError as exc: + warnings.append(f"Ignoring invalid LM_REMOTE_TIMEOUT: {exc}") + + return effective, overrides, warnings + + def reload(self) -> ConfigSnapshot: + """Reload persisted and environment-managed values without partial mutation.""" + with self._lock: + warnings: list[str] = [] + try: + raw, source = self._read_source() + known_values = {key: raw[key] for key in _KNOWN_FIELDS if key in raw} + configured = validate_config(known_values) + revision = _revision_for(raw) + except (OSError, json.JSONDecodeError, ConfigValidationError) as exc: + logger.warning("[LM-Remote] Failed to read configuration: %s", exc) + source_path, source = self._source_location() + raw = {} + configured = validate_config({}) + try: + revision = ( + _revision_for_file(source_path) + if source_path is not None + else _revision_for(raw) + ) + except OSError: + revision = _revision_for(raw) + warnings.append(f"Stored configuration could not be loaded: {exc}") + + effective, overrides, env_warnings = self._effective_values(configured) + warnings.extend(env_warnings) + generation = self._snapshot.generation + 1 + snapshot = ConfigSnapshot( + generation=generation, + remote_url=str(effective["remote_url"]), + timeout=int(effective["timeout"]), + path_mappings=_normalize_mappings(effective["path_mappings"]), + ) + self._configured = configured + self._overrides = overrides + self._revision = revision + self._source = source + self._warnings = warnings + self._snapshot = snapshot + return snapshot + + @property + def snapshot(self) -> ConfigSnapshot: + with self._lock: + return self._snapshot + + @property + def generation(self) -> int: + return self.snapshot.generation + + @property + def remote_url(self) -> str: + return self.snapshot.remote_url + + @property + def timeout(self) -> int: + return self.snapshot.timeout + + @property + def path_mappings(self) -> dict[str, str]: + return self.snapshot.mappings_dict() + + @property + def is_configured(self) -> bool: + return bool(self.snapshot.remote_url) + + def as_dict(self) -> dict[str, object]: + """Return browser-safe configured/effective values and source metadata.""" + with self._lock: + snapshot = self._snapshot + return { + "configured": { + "remote_url": self._configured["remote_url"], + "timeout": self._configured["timeout"], + "path_mappings": dict(self._configured["path_mappings"]), + }, + "effective": { + "remote_url": snapshot.remote_url, + "timeout": snapshot.timeout, + "path_mappings": snapshot.mappings_dict(), + }, + "overrides": dict(self._overrides), + "revision": self._revision, + "generation": snapshot.generation, + "storage": { + "source": self._source, + "writable": self._storage_writable(), + }, + "warnings": list(self._warnings), + "restart_required": False, + } + + def _storage_writable(self) -> bool: + """Return whether an atomic write can be created beside the target.""" + target = self._config_file + try: + if target.exists() and target.is_dir(): + return False + + ancestor = target.parent + while not ancestor.exists(): + parent = ancestor.parent + if parent == ancestor: + return False + ancestor = parent + if not ancestor.is_dir(): + return False + + mode = ancestor.stat().st_mode + write_bits = stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH + execute_bits = stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + if not mode & write_bits or not mode & execute_bits: + return False + return os.access(ancestor, os.W_OK | os.X_OK) + except OSError: + return False + + def save( + self, data: object, *, expected_revision: str | None = None + ) -> ConfigSnapshot: + """Atomically persist a complete configuration and activate it.""" + normalized = validate_config(data) + serializable = { + "remote_url": normalized["remote_url"], + "timeout": normalized["timeout"], + "path_mappings": normalized["path_mappings"], + } + + with self._lock: + try: + current_raw, _ = self._read_source() + except (OSError, json.JSONDecodeError, ConfigValidationError): + source_path, _ = self._source_location() + try: + current_revision = ( + _revision_for_file(source_path) + if source_path is not None + else _revision_for({}) + ) + except OSError as revision_exc: + raise ConfigConflictError( + "The stored configuration changed and can no longer be read." + ) from revision_exc + current_raw = {} + else: + current_revision = _revision_for(current_raw) + if expected_revision is not None and expected_revision != current_revision: + raise ConfigConflictError( + "The configuration changed in another window. Reload it before saving." + ) + + # Preserve future/third-party keys from an existing file while replacing + # only fields owned by LM Remote. + output = dict(current_raw) + output.update(serializable) + self._write_atomic(output) + return self.reload() + + def _write_atomic(self, data: Mapping[str, object]) -> None: + target = self._config_file + target.parent.mkdir(parents=True, exist_ok=True) + existing_mode = ( + stat.S_IMODE(target.stat().st_mode) if target.exists() else 0o600 + ) + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=target.parent, + prefix=f".{target.name}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary_path = Path(handle.name) + json.dump(data, handle, indent=2, ensure_ascii=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.chmod(temporary_path, existing_mode) + os.replace(temporary_path, target) + temporary_path = None + try: + directory_fd = os.open(target.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + except OSError: + # Directory fsync is not supported by every filesystem. + pass + finally: + if temporary_path is not None: + try: + temporary_path.unlink() + except FileNotFoundError: + pass + + def map_path(self, remote_path: str) -> str: + """Apply the longest boundary-aware remote-to-local path mapping.""" + return self.snapshot.map_path(remote_path) + + remote_config = RemoteConfig() diff --git a/proxy.py b/proxy.py index 150786c..2e8afda 100644 --- a/proxy.py +++ b/proxy.py @@ -8,19 +8,35 @@ Docker instance. Non-matching requests fall through to the regular ComfyUI rout Routes that use ``send_sync`` are handled locally so that events are broadcast to the local ComfyUI frontend (the remote instance has no connected browsers). """ + from __future__ import annotations import asyncio +import json import logging +import time +from contextlib import asynccontextmanager +from dataclasses import dataclass import aiohttp from aiohttp import web, WSMsgType -from .config import remote_config +from .config import ( + ConfigConflictError, + ConfigSnapshot, + ConfigValidationError, + remote_config, +) from .remote_client import RemoteLoraClient logger = logging.getLogger(__name__) +_CONFIG_ROUTE = "/api/lm-remote/config" +_TEST_CONNECTION_ROUTE = "/api/lm-remote/test-connection" +_PROXY_HOP_HEADER = "X-LM-Remote-Proxy" +_MAX_CONFIG_BODY = 64 * 1024 +_MAX_TEST_RESPONSE = 64 * 1024 + # --------------------------------------------------------------------------- # URL prefixes that should be forwarded to the remote LoRA Manager # --------------------------------------------------------------------------- @@ -59,6 +75,7 @@ _WS_ROUTES = { def _get_prompt_server(): """Lazily import PromptServer to avoid circular imports at module level.""" from server import PromptServer # type: ignore + return PromptServer.instance @@ -110,9 +127,7 @@ async def _handle_get_trigger_words(request: web.Request) -> web.Response: return web.json_response({"success": True}) except Exception as exc: logger.error("[LM-Remote] Error getting trigger words: %s", exc) - return web.json_response( - {"success": False, "error": str(exc)}, status=500 - ) + return web.json_response({"success": False, "error": str(exc)}, status=500) async def _handle_update_lora_code(request: web.Request) -> web.Response: @@ -179,17 +194,413 @@ _SEND_SYNC_HANDLERS = { "/api/lm/register-nodes": _handle_register_nodes, } -# Shared HTTP session for proxied requests (connection pooling) -_proxy_session: aiohttp.ClientSession | None = None + +def _config_response() -> web.Response: + return web.json_response({"success": True, **remote_config.as_dict()}) -async def _get_proxy_session() -> aiohttp.ClientSession: - """Return a shared aiohttp session for HTTP proxy requests.""" - global _proxy_session - if _proxy_session is None or _proxy_session.closed: - timeout = aiohttp.ClientTimeout(total=remote_config.timeout) - _proxy_session = aiohttp.ClientSession(timeout=timeout) - return _proxy_session +class _RequestBodyTooLarge(ValueError): + pass + + +async def _read_limited_json(request: web.Request) -> object: + """Read JSON without ever buffering more than the configuration limit.""" + if request.content_length is not None and request.content_length > _MAX_CONFIG_BODY: + raise _RequestBodyTooLarge + + body = bytearray() + while True: + remaining = _MAX_CONFIG_BODY + 1 - len(body) + chunk = await request.content.read(min(8192, remaining)) + if not chunk: + break + body.extend(chunk) + if len(body) > _MAX_CONFIG_BODY: + raise _RequestBodyTooLarge + return json.loads(body) + + +async def _activate_runtime_generation() -> None: + """Invalidate and rotate all resources tied to the previous snapshot.""" + RemoteLoraClient.get_instance().invalidate_caches() + await _retire_proxy_sessions(remote_config.generation) + await _rotate_active_websockets(remote_config.generation) + + +async def _handle_config(request: web.Request) -> web.Response: + """Read or atomically replace LM Remote's server-side configuration.""" + if request.method == "GET": + return _config_response() + if request.method != "PUT": + return web.json_response( + {"success": False, "error": "Method not allowed."}, + status=405, + headers={"Allow": "GET, PUT"}, + ) + if request.content_type != "application/json": + return web.json_response( + {"success": False, "error": "Content-Type must be application/json."}, + status=415, + ) + + try: + payload = await _read_limited_json(request) + except (_RequestBodyTooLarge, web.HTTPRequestEntityTooLarge): + return web.json_response( + {"success": False, "error": "Configuration request is too large."}, + status=413, + ) + except (json.JSONDecodeError, UnicodeDecodeError): + return web.json_response( + {"success": False, "error": "Request body is not valid JSON."}, status=400 + ) + if not isinstance(payload, dict): + return web.json_response( + {"success": False, "error": "Request body must be an object."}, status=400 + ) + unknown = set(payload) - {"revision", "config"} + if unknown: + return web.json_response( + { + "success": False, + "error": f"Unknown request field(s): {', '.join(sorted(unknown))}", + }, + status=400, + ) + + revision = payload.get("revision") + proposed = payload.get("config") + if not isinstance(revision, str) or not revision: + return web.json_response( + { + "success": False, + "field": "revision", + "error": "Reload configuration before saving.", + }, + status=400, + ) + if not isinstance(proposed, dict): + return web.json_response( + { + "success": False, + "field": "config", + "error": "Configuration must be an object.", + }, + status=400, + ) + + current = remote_config.as_dict() + configured = current["configured"] + overrides = current["overrides"] + for field in ("remote_url", "timeout"): + if overrides.get(field) and proposed.get(field) != configured.get(field): + variable = overrides[field] + return web.json_response( + { + "success": False, + "field": field, + "error": f"{field.replace('_', ' ').title()} is managed by {variable}.", + }, + status=409, + ) + + try: + remote_config.save(proposed, expected_revision=revision) + await _activate_runtime_generation() + except ConfigValidationError as exc: + return web.json_response( + {"success": False, "field": exc.field, "error": str(exc)}, status=400 + ) + except ConfigConflictError as exc: + remote_config.reload() + await _activate_runtime_generation() + return web.json_response( + {"success": False, "error": str(exc), "latest": remote_config.as_dict()}, + status=409, + ) + except OSError: + logger.exception("[LM-Remote] Failed to persist configuration") + return web.json_response( + {"success": False, "error": "Could not save LM Remote configuration."}, + status=500, + ) + + logger.info( + "[LM-Remote] Configuration applied without restart (generation %s)", + remote_config.generation, + ) + return _config_response() + + +class _ConnectionTestError(RuntimeError): + pass + + +async def _read_small_json(response: aiohttp.ClientResponse) -> object: + body = bytearray() + while True: + remaining = _MAX_TEST_RESPONSE + 1 - len(body) + chunk = await response.content.read(min(8192, remaining)) + if not chunk: + break + body.extend(chunk) + if len(body) > _MAX_TEST_RESPONSE: + raise _ConnectionTestError("Remote response was unexpectedly large.") + try: + return json.loads(body.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise _ConnectionTestError("Remote response was not valid JSON.") from exc + + +async def _perform_connection_test(remote_url: str, timeout_seconds: int) -> int: + """Probe a fixed Manager endpoint without changing the active configuration.""" + started = time.monotonic() + bounded_timeout = min(timeout_seconds, 30) + timeout = aiohttp.ClientTimeout( + total=bounded_timeout, + connect=min(bounded_timeout, 10), + ) + headers = {"Accept": "application/json", _PROXY_HOP_HEADER: "probe"} + async with aiohttp.ClientSession(timeout=timeout) as session: + health_url = f"{remote_url}/api/lm/health-check" + async with session.get( + health_url, headers=headers, allow_redirects=False + ) as response: + if response.status == 404: + payload = None + elif not 200 <= response.status < 300: + raise _ConnectionTestError( + f"Remote LoRA Manager returned HTTP {response.status}." + ) + else: + payload = await _read_small_json(response) + if not isinstance(payload, dict) or payload.get("status") != "ok": + raise _ConnectionTestError( + "The server answered, but it is not a compatible LoRA Manager." + ) + + # Older Manager releases may not expose health-check. + if payload is None: + list_url = f"{remote_url}/api/lm/loras/list?page=1&page_size=1" + async with session.get( + list_url, headers=headers, allow_redirects=False + ) as response: + if not 200 <= response.status < 300: + raise _ConnectionTestError( + f"Remote LoRA Manager returned HTTP {response.status}." + ) + list_payload = await _read_small_json(response) + if not isinstance(list_payload, dict) or "items" not in list_payload: + raise _ConnectionTestError( + "The server answered, but it is not a compatible LoRA Manager." + ) + + return max(1, round((time.monotonic() - started) * 1000)) + + +async def _handle_test_connection(request: web.Request) -> web.Response: + if request.method != "POST": + return web.json_response( + {"success": False, "error": "Method not allowed."}, + status=405, + headers={"Allow": "POST"}, + ) + if request.content_type != "application/json": + return web.json_response( + {"success": False, "error": "Content-Type must be application/json."}, + status=415, + ) + try: + payload = await _read_limited_json(request) + if not isinstance(payload, dict): + raise ConfigValidationError("config", "Request body must be an object.") + unknown = set(payload) - {"remote_url", "timeout"} + if unknown: + raise ConfigValidationError( + "config", f"Unknown request field(s): {', '.join(sorted(unknown))}" + ) + from .config import _normalize_timeout, _normalize_url + + remote_url = _normalize_url(payload.get("remote_url", ""), allow_empty=False) + timeout_seconds = _normalize_timeout(payload.get("timeout", 30)) + latency_ms = await _perform_connection_test(remote_url, timeout_seconds) + except (_RequestBodyTooLarge, web.HTTPRequestEntityTooLarge): + return web.json_response( + {"success": False, "error": "Connection test request is too large."}, + status=413, + ) + except ConfigValidationError as exc: + return web.json_response( + {"success": False, "field": exc.field, "error": str(exc)}, status=400 + ) + except (json.JSONDecodeError, UnicodeDecodeError): + return web.json_response( + {"success": False, "error": "Request body is not valid JSON."}, status=400 + ) + except asyncio.TimeoutError: + return web.json_response( + {"success": False, "error": "Connection test timed out."}, status=504 + ) + except aiohttp.ClientError: + return web.json_response( + {"success": False, "error": "Could not reach the remote LoRA Manager."}, + status=502, + ) + except _ConnectionTestError as exc: + return web.json_response({"success": False, "error": str(exc)}, status=502) + + return web.json_response( + { + "success": True, + "message": "Connected to LoRA Manager.", + "latency_ms": latency_ms, + } + ) + + +# Generation-scoped HTTP sessions allow an old request to finish while a newly +# saved configuration starts using its own connection pool immediately. +@dataclass(eq=False) +class _ProxySessionState: + session: aiohttp.ClientSession + active_requests: int = 0 + retired: bool = False + + +_proxy_sessions: dict[int, _ProxySessionState] = {} +_proxy_session_lock: asyncio.Lock | None = None + + +def _get_proxy_session_lock() -> asyncio.Lock: + global _proxy_session_lock + if _proxy_session_lock is None: + _proxy_session_lock = asyncio.Lock() + return _proxy_session_lock + + +@asynccontextmanager +async def _proxy_session_lease(snapshot: ConfigSnapshot): + """Lease the pool for one generation without closing active requests.""" + async with _get_proxy_session_lock(): + state = _proxy_sessions.get(snapshot.generation) + if state is None or state.session.closed: + if state is not None: + state.retired = True + state = _ProxySessionState( + session=aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=snapshot.timeout), + cookie_jar=aiohttp.DummyCookieJar(), + ), + retired=snapshot.generation != remote_config.generation, + ) + _proxy_sessions[snapshot.generation] = state + state.active_requests += 1 + + try: + yield state.session + finally: + close_session: aiohttp.ClientSession | None = None + async with _get_proxy_session_lock(): + state.active_requests -= 1 + if state.retired and state.active_requests == 0: + if _proxy_sessions.get(snapshot.generation) is state: + _proxy_sessions.pop(snapshot.generation, None) + close_session = state.session + if close_session is not None and not close_session.closed: + await close_session.close() + + +async def _retire_proxy_sessions(active_generation: int) -> None: + """Retire old pools and close only those with no requests in flight.""" + close_sessions: list[aiohttp.ClientSession] = [] + async with _get_proxy_session_lock(): + for generation, state in list(_proxy_sessions.items()): + if generation == active_generation: + continue + state.retired = True + if state.active_requests == 0: + _proxy_sessions.pop(generation, None) + close_sessions.append(state.session) + for session in close_sessions: + if not session.closed: + await session.close() + + +async def _close_all_proxy_sessions() -> None: + """Force-close every pool during application shutdown.""" + async with _get_proxy_session_lock(): + sessions = [state.session for state in _proxy_sessions.values()] + _proxy_sessions.clear() + for session in sessions: + if not session.closed: + await session.close() + + +@dataclass(eq=False) +class _ActiveWebSocket: + generation: int + local_ws: web.WebSocketResponse + session: aiohttp.ClientSession + remote_ws: aiohttp.ClientWebSocketResponse | None = None + retired: bool = False + + +_active_proxy_websockets: set[_ActiveWebSocket] = set() +_active_proxy_websockets_lock: asyncio.Lock | None = None + + +def _get_active_websockets_lock() -> asyncio.Lock: + global _active_proxy_websockets_lock + if _active_proxy_websockets_lock is None: + _active_proxy_websockets_lock = asyncio.Lock() + return _active_proxy_websockets_lock + + +async def _register_active_websocket(bridge: _ActiveWebSocket) -> bool: + async with _get_active_websockets_lock(): + if bridge.generation != remote_config.generation: + bridge.retired = True + return False + _active_proxy_websockets.add(bridge) + return True + + +async def _unregister_active_websocket(bridge: _ActiveWebSocket) -> None: + async with _get_active_websockets_lock(): + _active_proxy_websockets.discard(bridge) + + +async def _close_websocket_bridge(bridge: _ActiveWebSocket) -> None: + bridge.retired = True + try: + if bridge.remote_ws is not None and not bridge.remote_ws.closed: + await bridge.remote_ws.close( + code=1012, message=b"LM Remote configuration changed" + ) + if not bridge.local_ws.closed: + await bridge.local_ws.close( + code=1012, message=b"LM Remote configuration changed" + ) + finally: + if not bridge.session.closed: + await bridge.session.close() + + +async def _rotate_active_websockets(active_generation: int | None) -> None: + """Close bridges for retired generations so browsers reconnect.""" + async with _get_active_websockets_lock(): + bridges = [ + bridge + for bridge in _active_proxy_websockets + if active_generation is None or bridge.generation != active_generation + ] + for bridge in bridges: + _active_proxy_websockets.discard(bridge) + for bridge in bridges: + try: + await _close_websocket_bridge(bridge) + except Exception: + logger.exception("[LM-Remote] Failed to rotate a proxied WebSocket") def _should_proxy(path: str) -> bool: @@ -205,9 +616,13 @@ def _is_ws_route(path: str) -> bool: return path in _WS_ROUTES -async def _proxy_ws(request: web.Request) -> web.WebSocketResponse: +async def _proxy_ws( + request: web.Request, snapshot: ConfigSnapshot +) -> web.WebSocketResponse: """Proxy a WebSocket connection to the remote LoRA Manager.""" - remote_url = remote_config.remote_url.replace("http://", "ws://").replace("https://", "wss://") + remote_url = snapshot.remote_url.replace("http://", "ws://", 1).replace( + "https://", "wss://", 1 + ) remote_ws_url = f"{remote_url}{request.path}" if request.query_string: remote_ws_url += f"?{request.query_string}" @@ -215,10 +630,21 @@ async def _proxy_ws(request: web.Request) -> web.WebSocketResponse: local_ws = web.WebSocketResponse() await local_ws.prepare(request) - timeout = aiohttp.ClientTimeout(total=None) + timeout = aiohttp.ClientTimeout( + total=None, + sock_connect=min(snapshot.timeout, 30), + ) session = aiohttp.ClientSession(timeout=timeout) + bridge = _ActiveWebSocket(snapshot.generation, local_ws, session) + if not await _register_active_websocket(bridge): + await _close_websocket_bridge(bridge) + return local_ws + try: - async with session.ws_connect(remote_ws_url) as remote_ws: + async with session.ws_connect( + remote_ws_url, headers={_PROXY_HOP_HEADER: "1"} + ) as remote_ws: + bridge.remote_ws = remote_ws async def forward_local_to_remote(): async for msg in local_ws: @@ -226,7 +652,11 @@ async def _proxy_ws(request: web.Request) -> web.WebSocketResponse: await remote_ws.send_str(msg.data) elif msg.type == WSMsgType.BINARY: await remote_ws.send_bytes(msg.data) - elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): + elif msg.type in ( + WSMsgType.CLOSE, + WSMsgType.CLOSING, + WSMsgType.CLOSED, + ): return async def forward_remote_to_local(): @@ -235,7 +665,11 @@ async def _proxy_ws(request: web.Request) -> web.WebSocketResponse: await local_ws.send_str(msg.data) elif msg.type == WSMsgType.BINARY: await local_ws.send_bytes(msg.data) - elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): + elif msg.type in ( + WSMsgType.CLOSE, + WSMsgType.CLOSING, + WSMsgType.CLOSED, + ): return # Run both directions concurrently. When either side closes, @@ -260,16 +694,22 @@ async def _proxy_ws(request: web.Request) -> web.WebSocketResponse: await local_ws.close() except Exception as exc: - logger.warning("[LM-Remote] WebSocket proxy error for %s: %s", request.path, exc) + if not bridge.retired: + logger.warning( + "[LM-Remote] WebSocket proxy error for %s: %s", request.path, exc + ) finally: + await _unregister_active_websocket(bridge) + if not local_ws.closed: + await local_ws.close() await session.close() return local_ws -async def _proxy_http(request: web.Request) -> web.Response: +async def _proxy_http(request: web.Request, snapshot: ConfigSnapshot) -> web.Response: """Forward an HTTP request to the remote LoRA Manager and return its response.""" - remote_url = f"{remote_config.remote_url}{request.path}" + remote_url = f"{snapshot.remote_url}{request.path}" if request.query_string: remote_url += f"?{request.query_string}" @@ -278,33 +718,52 @@ async def _proxy_http(request: web.Request) -> web.Response: # Filter hop-by-hop headers headers = {} - skip = {"host", "transfer-encoding", "connection", "keep-alive", "upgrade"} + skip = { + "host", + "transfer-encoding", + "connection", + "keep-alive", + "upgrade", + "authorization", + "proxy-authorization", + "cookie", + "origin", + "referer", + } for k, v in request.headers.items(): if k.lower() not in skip: headers[k] = v + headers[_PROXY_HOP_HEADER] = "1" - session = await _get_proxy_session() try: - async with session.request( - method=request.method, - url=remote_url, - headers=headers, - data=body, - ) as resp: - resp_body = await resp.read() - resp_headers = {} - for k, v in resp.headers.items(): - if k.lower() not in ("transfer-encoding", "content-encoding", "content-length"): - resp_headers[k] = v - return web.Response( - status=resp.status, - body=resp_body, - headers=resp_headers, - ) + async with _proxy_session_lease(snapshot) as session: + async with session.request( + method=request.method, + url=remote_url, + headers=headers, + data=body, + ) as resp: + resp_body = await resp.read() + resp_headers = {} + for k, v in resp.headers.items(): + if k.lower() not in ( + "transfer-encoding", + "content-encoding", + "content-length", + "set-cookie", + ): + resp_headers[k] = v + return web.Response( + status=resp.status, + body=resp_body, + headers=resp_headers, + ) except Exception as exc: - logger.error("[LM-Remote] Proxy error for %s %s: %s", request.method, request.path, exc) + logger.error( + "[LM-Remote] Proxy error for %s %s: %s", request.method, request.path, exc + ) return web.json_response( - {"error": f"Remote LoRA Manager unavailable: {exc}"}, + {"error": "Remote LoRA Manager unavailable."}, status=502, ) @@ -313,13 +772,28 @@ async def _proxy_http(request: web.Request) -> web.Response: # Middleware factory # --------------------------------------------------------------------------- + @web.middleware async def lm_remote_proxy_middleware(request: web.Request, handler): """aiohttp middleware that intercepts LoRA Manager requests.""" - if not remote_config.is_configured: + path = request.path + + # Configuration remains local and available even before a remote is set. + if path == _CONFIG_ROUTE: + return await _handle_config(request) + if path == _TEST_CONNECTION_ROUTE: + return await _handle_test_connection(request) + + snapshot = remote_config.snapshot + if not snapshot.remote_url: return await handler(request) - path = request.path + if request.headers.get(_PROXY_HOP_HEADER) and ( + _should_proxy(path) or _is_ws_route(path) + ): + return web.json_response( + {"error": "LM Remote proxy loop detected."}, status=508 + ) # Routes that need send_sync are handled locally so events reach # the local browser (the remote instance has no connected browsers). @@ -329,31 +803,32 @@ async def lm_remote_proxy_middleware(request: web.Request, handler): # WebSocket routes if _is_ws_route(path): - return await _proxy_ws(request) + return await _proxy_ws(request, snapshot) # Regular proxy routes if _should_proxy(path): - return await _proxy_http(request) + return await _proxy_http(request, snapshot) # Not a LoRA Manager route — fall through return await handler(request) async def _cleanup_proxy_session(app) -> None: - """Shutdown hook to close the shared proxy session.""" - global _proxy_session - if _proxy_session and not _proxy_session.closed: - await _proxy_session.close() - _proxy_session = None + """Shutdown hook to close every HTTP pool and WebSocket bridge.""" + await _rotate_active_websockets(None) + await _close_all_proxy_sessions() + await RemoteLoraClient.get_instance().close() def register_proxy(app) -> None: - """Insert the proxy middleware into the aiohttp app.""" - if not remote_config.is_configured: - logger.warning("[LM-Remote] No remote_url configured — proxy disabled") - return - - # Insert at position 0 so we run before the original LoRA Manager routes - app.middlewares.insert(0, lm_remote_proxy_middleware) - app.on_shutdown.append(_cleanup_proxy_session) - logger.info("[LM-Remote] Proxy routes registered -> %s", remote_config.remote_url) + """Append the proxy after ComfyUI's origin and security guards.""" + if lm_remote_proxy_middleware not in app.middlewares: + app.middlewares.append(lm_remote_proxy_middleware) + if _cleanup_proxy_session not in app.on_shutdown: + app.on_shutdown.append(_cleanup_proxy_session) + if remote_config.is_configured: + logger.info( + "[LM-Remote] Proxy routes registered -> %s", remote_config.remote_url + ) + else: + logger.info("[LM-Remote] Configuration API ready; remote URL is not set yet") diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..698ac51 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +addopts = --import-mode=importlib --confcutdir=tests +testpaths = tests diff --git a/remote_client.py b/remote_client.py index 35f4d83..8fc28f5 100644 --- a/remote_client.py +++ b/remote_client.py @@ -1,26 +1,35 @@ """HTTP client for the remote LoRA Manager instance.""" + from __future__ import annotations -import asyncio import logging +import os +import posixpath +import threading import time from typing import Any import aiohttp -from .config import remote_config +from .config import ConfigSnapshot, remote_config logger = logging.getLogger(__name__) # Cache TTL in seconds — how long before we re-fetch the full LoRA list _CACHE_TTL = 60 +_LIST_PAGE_SIZE = 100 +_MAX_LIST_PAGES = 1000 + + +class _ConfigurationChanged(RuntimeError): + """Raised when a multi-page read outlives its configuration snapshot.""" class RemoteLoraClient: """Singleton HTTP client that talks to the remote LoRA Manager. Uses the actual LoRA Manager REST API endpoints: - - ``GET /api/lm/loras/list?page_size=9999`` — paginated LoRA list + - ``GET /api/lm/loras/list?page=N&page_size=100`` — paginated LoRA list - ``GET /api/lm/loras/get-trigger-words?name=X`` — trigger words - ``POST /api/lm/loras/random-sample`` — random LoRA selection - ``POST /api/lm/loras/cycler-list`` — sorted LoRA list for cycler @@ -30,13 +39,15 @@ class RemoteLoraClient: """ _instance: RemoteLoraClient | None = None - _session: aiohttp.ClientSession | None = None def __init__(self): self._lora_cache: list[dict] = [] self._lora_cache_ts: float = 0 + self._lora_cache_generation: int = -1 self._checkpoint_cache: list[dict] = [] self._checkpoint_cache_ts: float = 0 + self._checkpoint_cache_generation: int = -1 + self._cache_lock = threading.RLock() @classmethod def get_instance(cls) -> RemoteLoraClient: @@ -44,71 +55,176 @@ class RemoteLoraClient: cls._instance = cls() return cls._instance - async def _get_session(self) -> aiohttp.ClientSession: - if self._session is None or self._session.closed: - timeout = aiohttp.ClientTimeout(total=remote_config.timeout) - self._session = aiohttp.ClientSession(timeout=timeout) - return self._session - async def close(self): - if self._session and not self._session.closed: - await self._session.close() - self._session = None + """Compatibility hook; requests use loop-safe, short-lived sessions.""" + + def invalidate_caches(self) -> None: + """Forget results associated with a previous remote configuration.""" + with self._cache_lock: + self._lora_cache = [] + self._lora_cache_ts = 0 + self._lora_cache_generation = -1 + self._checkpoint_cache = [] + self._checkpoint_cache_ts = 0 + self._checkpoint_cache_generation = -1 # ------------------------------------------------------------------ # Core HTTP helpers # ------------------------------------------------------------------ - async def _get_json(self, path: str, params: dict | None = None) -> Any: - url = f"{remote_config.remote_url}{path}" - session = await self._get_session() - async with session.get(url, params=params) as resp: - resp.raise_for_status() - return await resp.json() + async def _get_json( + self, + path: str, + params: dict | None = None, + *, + snapshot: ConfigSnapshot | None = None, + ) -> Any: + snapshot = snapshot or remote_config.snapshot + url = f"{snapshot.remote_url}{path}" + timeout = aiohttp.ClientTimeout(total=snapshot.timeout) + # Node execution can invoke this singleton from several short-lived event + # loops. A request-scoped session avoids retaining a loop-bound session. + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.get(url, params=params) as resp: + resp.raise_for_status() + return await resp.json() - async def _post_json(self, path: str, json_body: dict | None = None) -> Any: - url = f"{remote_config.remote_url}{path}" - session = await self._get_session() - async with session.post(url, json=json_body) as resp: - resp.raise_for_status() - return await resp.json() + async def _post_json( + self, + path: str, + json_body: dict | None = None, + *, + snapshot: ConfigSnapshot | None = None, + ) -> Any: + snapshot = snapshot or remote_config.snapshot + url = f"{snapshot.remote_url}{path}" + timeout = aiohttp.ClientTimeout(total=snapshot.timeout) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.post(url, json=json_body) as resp: + resp.raise_for_status() + return await resp.json() + + async def _get_all_pages( + self, path: str, *, snapshot: ConfigSnapshot + ) -> list[dict]: + """Fetch a complete bounded listing from an API capped at 100 rows.""" + items: list[dict] = [] + page = 1 + + while page <= _MAX_LIST_PAGES: + if remote_config.generation != snapshot.generation: + raise _ConfigurationChanged + data = await self._get_json( + path, + params={"page": str(page), "page_size": str(_LIST_PAGE_SIZE)}, + snapshot=snapshot, + ) + if not isinstance(data, dict): + raise ValueError("Remote listing response must be an object.") + page_items = data.get("items", []) + if not isinstance(page_items, list): + raise ValueError("Remote listing items must be an array.") + items.extend(item for item in page_items if isinstance(item, dict)) + + raw_total_pages = data.get("total_pages") + total_pages: int | None = None + if raw_total_pages is not None: + try: + total_pages = int(raw_total_pages) + except (TypeError, ValueError) as exc: + raise ValueError( + "Remote listing has an invalid total_pages value." + ) from exc + if total_pages < 0: + raise ValueError("Remote listing has an invalid total_pages value.") + if total_pages > _MAX_LIST_PAGES: + raise ValueError( + f"Remote listing exceeds the {_MAX_LIST_PAGES}-page safety limit." + ) + + if not page_items: + break + if total_pages is not None: + if page >= total_pages: + break + elif len(page_items) < _LIST_PAGE_SIZE: + break + page += 1 + else: + raise ValueError( + f"Remote listing exceeds the {_MAX_LIST_PAGES}-page safety limit." + ) + + if remote_config.generation != snapshot.generation: + raise _ConfigurationChanged + return items # ------------------------------------------------------------------ # Cached list helpers # ------------------------------------------------------------------ - async def _get_lora_list_cached(self) -> list[dict]: + async def _get_lora_list_cached( + self, *, snapshot: ConfigSnapshot | None = None + ) -> list[dict]: """Return the full LoRA list, using a short-lived cache.""" now = time.monotonic() - if self._lora_cache and (now - self._lora_cache_ts) < _CACHE_TTL: - return self._lora_cache + snapshot = snapshot or remote_config.snapshot + with self._cache_lock: + if ( + self._lora_cache_generation == snapshot.generation + and (now - self._lora_cache_ts) < _CACHE_TTL + ): + return list(self._lora_cache) try: - data = await self._get_json( - "/api/lm/loras/list", params={"page_size": "9999"} - ) - self._lora_cache = data.get("items", []) - self._lora_cache_ts = now + items = await self._get_all_pages("/api/lm/loras/list", snapshot=snapshot) + if remote_config.generation == snapshot.generation: + with self._cache_lock: + self._lora_cache = list(items) + self._lora_cache_ts = now + self._lora_cache_generation = snapshot.generation + return list(items) + except _ConfigurationChanged: + pass except Exception as exc: logger.warning("[LM-Remote] Failed to fetch LoRA list: %s", exc) # Return stale cache on error, or empty list - return self._lora_cache + with self._cache_lock: + if self._lora_cache_generation == snapshot.generation: + return list(self._lora_cache) + return [] - async def _get_checkpoint_list_cached(self) -> list[dict]: + async def _get_checkpoint_list_cached( + self, *, snapshot: ConfigSnapshot | None = None + ) -> list[dict]: """Return the full checkpoint list, using a short-lived cache.""" now = time.monotonic() - if self._checkpoint_cache and (now - self._checkpoint_cache_ts) < _CACHE_TTL: - return self._checkpoint_cache + snapshot = snapshot or remote_config.snapshot + with self._cache_lock: + if ( + self._checkpoint_cache_generation == snapshot.generation + and (now - self._checkpoint_cache_ts) < _CACHE_TTL + ): + return list(self._checkpoint_cache) try: - data = await self._get_json( - "/api/lm/checkpoints/list", params={"page_size": "9999"} + items = await self._get_all_pages( + "/api/lm/checkpoints/list", snapshot=snapshot ) - self._checkpoint_cache = data.get("items", []) - self._checkpoint_cache_ts = now + if remote_config.generation == snapshot.generation: + with self._cache_lock: + self._checkpoint_cache = list(items) + self._checkpoint_cache_ts = now + self._checkpoint_cache_generation = snapshot.generation + return list(items) + except _ConfigurationChanged: + pass except Exception as exc: logger.warning("[LM-Remote] Failed to fetch checkpoint list: %s", exc) - return self._checkpoint_cache + with self._cache_lock: + if self._checkpoint_cache_generation == snapshot.generation: + return list(self._checkpoint_cache) + return [] def _find_item_by_name(self, items: list[dict], name: str) -> dict | None: """Find an item in a list by file_name.""" @@ -117,6 +233,31 @@ class RemoteLoraClient: return item return None + @staticmethod + def _relative_lora_path(mapped_file_path: str, folder: str) -> str: + """Convert a mapped absolute path to a local ComfyUI LoRA name.""" + try: + import folder_paths # type: ignore + + candidate = os.path.normpath(mapped_file_path) + for root in folder_paths.get_folder_paths("loras"): + try: + relative = os.path.relpath(candidate, os.path.normpath(str(root))) + except ValueError: + # Windows paths on different drives cannot be relativized. + continue + if relative == os.pardir or relative.startswith(f"{os.pardir}{os.sep}"): + continue + return relative.replace(os.sep, "/") + except Exception: + # ComfyUI's folder registry is not present in lightweight tooling. + pass + + normalized_path = mapped_file_path.replace("\\", "/") + basename = posixpath.basename(normalized_path) + normalized_folder = str(folder or "").replace("\\", "/").strip("/") + return f"{normalized_folder}/{basename}" if normalized_folder else basename + # ------------------------------------------------------------------ # LoRA metadata # ------------------------------------------------------------------ @@ -127,35 +268,17 @@ class RemoteLoraClient: Uses the cached ``/api/lm/loras/list`` data. Falls back to the per-LoRA ``get-trigger-words`` endpoint if the list lookup fails. """ - import posixpath - try: - items = await self._get_lora_list_cached() + snapshot = remote_config.snapshot + items = await self._get_lora_list_cached(snapshot=snapshot) item = self._find_item_by_name(items, lora_name) if item: file_path = item.get("file_path", "") - file_path = remote_config.map_path(file_path) + file_path = snapshot.map_path(file_path) - # file_path is the absolute path (forward-slashed) from - # the remote. We need a relative path that the local - # folder_paths.get_full_path("loras", ...) can resolve. - # - # The ``folder`` field gives the subfolder within the - # model root (e.g. "anime" or "anime/characters"). - # The basename of file_path has the extension. - # - # Example: file_path="/mnt/loras/anime/test.safetensors" - # folder="anime" - # -> basename="test.safetensors" - # -> relative="anime/test.safetensors" folder = item.get("folder", "") - basename = posixpath.basename(file_path) # "test.safetensors" - - if folder: - relative = f"{folder}/{basename}" - else: - relative = basename + relative = self._relative_lora_path(file_path, folder) civitai = item.get("civitai") or {} trigger_words = civitai.get("trainedWords", []) if civitai else [] @@ -165,6 +288,7 @@ class RemoteLoraClient: tw_data = await self._get_json( "/api/lm/loras/get-trigger-words", params={"name": lora_name}, + snapshot=snapshot, ) trigger_words = tw_data.get("trigger_words", []) return lora_name, trigger_words @@ -176,7 +300,8 @@ class RemoteLoraClient: async def get_lora_hash(self, lora_name: str) -> str | None: """Return the SHA-256 hash for a LoRA by display name.""" try: - items = await self._get_lora_list_cached() + snapshot = remote_config.snapshot + items = await self._get_lora_list_cached(snapshot=snapshot) item = self._find_item_by_name(items, lora_name) if item: return item.get("sha256") or item.get("hash") @@ -187,18 +312,23 @@ class RemoteLoraClient: async def get_checkpoint_hash(self, checkpoint_name: str) -> str | None: """Return the SHA-256 hash for a checkpoint by display name.""" try: - items = await self._get_checkpoint_list_cached() + snapshot = remote_config.snapshot + items = await self._get_checkpoint_list_cached(snapshot=snapshot) item = self._find_item_by_name(items, checkpoint_name) if item: return item.get("sha256") or item.get("hash") except Exception as exc: - logger.warning("[LM-Remote] get_checkpoint_hash(%s) failed: %s", checkpoint_name, exc) + logger.warning( + "[LM-Remote] get_checkpoint_hash(%s) failed: %s", checkpoint_name, exc + ) return None async def get_random_loras(self, **kwargs) -> list[dict]: """Ask the remote to generate random LoRAs (for Randomizer node).""" try: - result = await self._post_json("/api/lm/loras/random-sample", json_body=kwargs) + result = await self._post_json( + "/api/lm/loras/random-sample", json_body=kwargs + ) return result if isinstance(result, list) else result.get("loras", []) except Exception as exc: logger.warning("[LM-Remote] get_random_loras failed: %s", exc) @@ -207,7 +337,9 @@ class RemoteLoraClient: async def get_cycler_list(self, **kwargs) -> list[dict]: """Ask the remote for a sorted LoRA list (for Cycler node).""" try: - result = await self._post_json("/api/lm/loras/cycler-list", json_body=kwargs) + result = await self._post_json( + "/api/lm/loras/cycler-list", json_body=kwargs + ) return result if isinstance(result, list) else result.get("loras", []) except Exception as exc: logger.warning("[LM-Remote] get_cycler_list failed: %s", exc) diff --git a/tests/frontend/remote_config_utils.test.js b/tests/frontend/remote_config_utils.test.js new file mode 100644 index 0000000..dfc71f4 --- /dev/null +++ b/tests/frontend/remote_config_utils.test.js @@ -0,0 +1,149 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { + RemoteConfigFormError, + buildConfigDraft, + buildConnectionDraft, + didEffectiveRemoteUrlChange, + isConfigWritable, + mappingsToRows, +} from "../../web/comfyui/remote_config_utils.js"; + +test("buildConfigDraft normalizes form values and skips blank mappings", () => { + assert.deepEqual( + buildConfigDraft({ + remoteUrl: " http://manager.local:8188/ ", + timeout: "45", + mappingRows: [ + { remote: "/data/loras/", local: "/mnt/loras" }, + { remote: "", local: "" }, + ], + }), + { + remote_url: "http://manager.local:8188", + timeout: 45, + path_mappings: { "/data/loras": "/mnt/loras" }, + } + ); +}); + +test("buildConfigDraft allows an empty URL to disable remote mode", () => { + const draft = buildConfigDraft({ + remoteUrl: "", + timeout: 30, + mappingRows: [], + }); + assert.equal(draft.remote_url, ""); +}); + +test("buildConfigDraft rejects unsafe URLs and invalid mappings", () => { + assert.throws( + () => + buildConfigDraft({ + remoteUrl: "ftp://manager.local", + timeout: 30, + mappingRows: [], + }), + (error) => + error instanceof RemoteConfigFormError && error.field === "remote_url" + ); + assert.throws( + () => + buildConfigDraft({ + remoteUrl: "http://manager.local", + timeout: 30, + mappingRows: [{ remote: "/remote", local: "" }], + }), + (error) => + error instanceof RemoteConfigFormError && error.field === "path_mappings" + ); + assert.throws( + () => + buildConfigDraft({ + remoteUrl: "http://manager.local", + timeout: 301, + mappingRows: [], + }), + (error) => + error instanceof RemoteConfigFormError && error.field === "timeout" + ); +}); + +test("buildConfigDraft rejects normalized duplicate path prefixes", () => { + assert.throws( + () => + buildConfigDraft({ + remoteUrl: "http://manager.local", + timeout: 30, + mappingRows: [ + { remote: "/data/loras", local: "/mnt/a" }, + { remote: "\\data\\loras\\", local: "/mnt/b" }, + ], + }), + /Duplicate remote path prefix/ + ); +}); + +test("buildConfigDraft preserves __proto__ as a path mapping key", () => { + const draft = buildConfigDraft({ + remoteUrl: "http://manager.local", + timeout: 30, + mappingRows: [{ remote: "__proto__", local: "/mnt/prototype" }], + }); + + assert.equal(Object.hasOwn(draft.path_mappings, "__proto__"), true); + assert.equal(draft.path_mappings.__proto__, "/mnt/prototype"); + assert.equal( + JSON.parse(JSON.stringify(draft.path_mappings)).__proto__, + "/mnt/prototype" + ); +}); + +test("connection tests use effective environment-managed values", () => { + assert.deepEqual( + buildConnectionDraft({ + remoteUrl: "http://stored.local:8188", + timeout: "30", + effective: { + remote_url: "https://managed.local:443", + timeout: 12, + }, + overrides: { + remote_url: "LM_REMOTE_URL", + timeout: "LM_REMOTE_TIMEOUT", + }, + }), + { remote_url: "https://managed.local:443", timeout: 12 } + ); +}); + +test("mappingsToRows creates editable row values", () => { + assert.deepEqual(mappingsToRows({ "/remote": "/local" }), [ + { remote: "/remote", local: "/local" }, + ]); +}); + +test("didEffectiveRemoteUrlChange only tracks the effective URL", () => { + assert.equal( + didEffectiveRemoteUrlChange( + { effective: { remote_url: "http://old:8188", timeout: 30 } }, + { effective: { remote_url: "http://new:8188", timeout: 30 } } + ), + true + ); + assert.equal( + didEffectiveRemoteUrlChange( + { effective: { remote_url: "http://same:8188", timeout: 30 } }, + { effective: { remote_url: "http://same:8188", timeout: 60 } } + ), + false + ); +}); + +test("isConfigWritable only disables saving for an explicit false value", () => { + assert.equal(isConfigWritable({ storage: { writable: false } }), false); + assert.equal(isConfigWritable({ storage: { writable: true } }), true); + assert.equal(isConfigWritable({ storage: {} }), true); + assert.equal(isConfigWritable(null), true); +}); diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..ec05a5f --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + + +@pytest.fixture(scope="module") +def config_module(): + module_name = "lm_remote_config_test_module" + module_path = Path(__file__).resolve().parents[1] / "config.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + yield module + sys.modules.pop(module_name, None) + + +def write_json(path: Path, data: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data), encoding="utf-8") + + +def test_legacy_config_loads_and_first_save_migrates(config_module, tmp_path): + legacy = tmp_path / "package" / "config.json" + user = tmp_path / "user" / "config.json" + original = { + "remote_url": "http://legacy.local:8188", + "timeout": 30, + "path_mappings": {}, + "future_setting": {"preserve": True}, + } + write_json(legacy, original) + + config = config_module.RemoteConfig(user, legacy, environ={}) + state = config.as_dict() + assert state["storage"]["source"] == "legacy" + assert state["effective"]["remote_url"] == "http://legacy.local:8188" + + config.save( + { + "remote_url": "http://new.local:8188/", + "timeout": 45, + "path_mappings": {"/remote": "/local"}, + }, + expected_revision=state["revision"], + ) + + assert json.loads(legacy.read_text(encoding="utf-8")) == original + saved = json.loads(user.read_text(encoding="utf-8")) + assert saved["remote_url"] == "http://new.local:8188" + assert saved["future_setting"] == {"preserve": True} + assert config.as_dict()["storage"]["source"] == "user" + + +def test_environment_values_override_without_rewriting_stored_values( + config_module, tmp_path +): + user = tmp_path / "config.json" + write_json( + user, + { + "remote_url": "http://stored.local:8188", + "timeout": 30, + "path_mappings": {}, + }, + ) + config = config_module.RemoteConfig( + user, + tmp_path / "missing.json", + environ={ + "LM_REMOTE_URL": "https://managed.local:443/", + "LM_REMOTE_TIMEOUT": "12", + }, + ) + + state = config.as_dict() + assert state["configured"]["remote_url"] == "http://stored.local:8188" + assert state["effective"]["remote_url"] == "https://managed.local:443" + assert state["effective"]["timeout"] == 12 + assert state["overrides"] == { + "remote_url": "LM_REMOTE_URL", + "timeout": "LM_REMOTE_TIMEOUT", + } + + +def test_explicit_config_environment_reports_explicit_source(config_module, tmp_path): + explicit = tmp_path / "managed" / "remote.json" + config = config_module.RemoteConfig( + legacy_config_file=tmp_path / "missing-legacy.json", + environ={"LM_REMOTE_CONFIG": str(explicit)}, + ) + initial = config.as_dict() + config.save( + {"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}}, + expected_revision=initial["revision"], + ) + + assert config.as_dict()["storage"]["source"] == "explicit" + + +def test_explicit_missing_target_does_not_fall_back_to_legacy(config_module, tmp_path): + explicit = tmp_path / "managed" / "remote.json" + legacy = tmp_path / "package" / "config.json" + legacy_data = { + "remote_url": "http://legacy.local", + "timeout": 30, + "path_mappings": {}, + } + write_json(legacy, legacy_data) + + config = config_module.RemoteConfig( + legacy_config_file=legacy, + environ={"LM_REMOTE_CONFIG": str(explicit)}, + ) + state = config.as_dict() + + assert state["storage"]["source"] == "explicit" + assert state["configured"]["remote_url"] == "" + config.save( + {"remote_url": "http://explicit.local", "timeout": 30, "path_mappings": {}}, + expected_revision=state["revision"], + ) + assert json.loads(legacy.read_text(encoding="utf-8")) == legacy_data + assert ( + json.loads(explicit.read_text(encoding="utf-8"))["remote_url"] + == "http://explicit.local" + ) + + +def test_invalid_environment_timeout_is_ignored(config_module, tmp_path): + config = config_module.RemoteConfig( + tmp_path / "missing-user.json", + tmp_path / "missing-legacy.json", + environ={"LM_REMOTE_TIMEOUT": "not-a-number"}, + ) + state = config.as_dict() + assert state["effective"]["timeout"] == 30 + assert state["overrides"]["timeout"] is None + assert any("LM_REMOTE_TIMEOUT" in warning for warning in state["warnings"]) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("remote_url", "ftp://manager.local"), + ("remote_url", "http://user:secret@manager.local"), + ("remote_url", "http://manager.local?query=yes"), + ("timeout", True), + ("timeout", 0), + ("timeout", 301), + ("path_mappings", []), + ], +) +def test_validation_rejects_invalid_values(config_module, field, value): + candidate = { + "remote_url": "http://manager.local", + "timeout": 30, + "path_mappings": {}, + } + candidate[field] = value + with pytest.raises(config_module.ConfigValidationError) as caught: + config_module.validate_config(candidate) + assert caught.value.field == field + + +def test_mapping_is_longest_first_and_path_boundary_aware(config_module, tmp_path): + user = tmp_path / "config.json" + write_json( + user, + { + "remote_url": "http://manager.local", + "timeout": 30, + "path_mappings": { + "/models": "/mnt/general", + "/models/special": "/mnt/special", + }, + }, + ) + config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={}) + assert ( + config.map_path("/models/special/a.safetensors") == "/mnt/special/a.safetensors" + ) + assert ( + config.map_path("/models/base.safetensors") == "/mnt/general/base.safetensors" + ) + assert config.map_path("/models-old/a.safetensors") == "/models-old/a.safetensors" + + +def test_mapping_rejects_separator_only_prefix_but_allows_root(config_module): + candidate = { + "remote_url": "http://manager.local", + "timeout": 30, + "path_mappings": {"////": "/mnt/invalid"}, + } + with pytest.raises(config_module.ConfigValidationError) as caught: + config_module.validate_config(candidate) + assert caught.value.field == "path_mappings" + + root_mapping = config_module.validate_config( + {**candidate, "path_mappings": {"/": "/mnt/root"}} + ) + snapshot = config_module.ConfigSnapshot( + 1, + root_mapping["remote_url"], + root_mapping["timeout"], + config_module._normalize_mappings(root_mapping["path_mappings"]), + ) + assert ( + snapshot.map_path("/models/a.safetensors") == "/mnt/root/models/a.safetensors" + ) + + +def test_snapshot_mapping_does_not_change_after_reload(config_module, tmp_path): + user = tmp_path / "config.json" + write_json( + user, + { + "remote_url": "http://one.local", + "timeout": 30, + "path_mappings": {"/remote": "/local-one"}, + }, + ) + config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={}) + original = config.snapshot + revision = config.as_dict()["revision"] + + config.save( + { + "remote_url": "http://two.local", + "timeout": 30, + "path_mappings": {"/remote": "/local-two"}, + }, + expected_revision=revision, + ) + + assert original.map_path("/remote/a.safetensors") == "/local-one/a.safetensors" + assert config.map_path("/remote/a.safetensors") == "/local-two/a.safetensors" + + +def test_storage_reports_unwritable_when_parent_is_not_a_directory( + config_module, tmp_path +): + blocking_file = tmp_path / "not-a-directory" + blocking_file.write_text("blocked", encoding="utf-8") + config = config_module.RemoteConfig( + blocking_file / "config.json", + tmp_path / "missing-legacy.json", + environ={}, + ) + + assert config.as_dict()["storage"]["writable"] is False + + +def test_stale_revision_does_not_overwrite_external_change(config_module, tmp_path): + user = tmp_path / "config.json" + initial = {"remote_url": "http://one.local", "timeout": 30, "path_mappings": {}} + write_json(user, initial) + config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={}) + revision = config.as_dict()["revision"] + write_json( + user, + {"remote_url": "http://two.local", "timeout": 30, "path_mappings": {}}, + ) + + with pytest.raises(config_module.ConfigConflictError): + config.save(initial, expected_revision=revision) + assert ( + json.loads(user.read_text(encoding="utf-8"))["remote_url"] == "http://two.local" + ) + + +def test_atomic_write_failure_preserves_file_and_live_snapshot( + config_module, tmp_path, monkeypatch +): + user = tmp_path / "config.json" + initial = {"remote_url": "http://one.local", "timeout": 30, "path_mappings": {}} + write_json(user, initial) + config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={}) + before = config.snapshot + revision = config.as_dict()["revision"] + + def fail_replace(source, target): + raise OSError("simulated replace failure") + + monkeypatch.setattr(config_module.os, "replace", fail_replace) + with pytest.raises(OSError, match="simulated"): + config.save( + {"remote_url": "http://two.local", "timeout": 50, "path_mappings": {}}, + expected_revision=revision, + ) + + assert json.loads(user.read_text(encoding="utf-8")) == initial + assert config.snapshot == before + assert not list(tmp_path.glob(".config.json.*.tmp")) + + +def test_invalid_json_can_be_repaired_from_loaded_revision(config_module, tmp_path): + user = tmp_path / "config.json" + user.write_text("{not valid", encoding="utf-8") + config = config_module.RemoteConfig(user, tmp_path / "missing.json", environ={}) + state = config.as_dict() + assert state["effective"]["remote_url"] == "" + assert state["warnings"] + + config.save( + {"remote_url": "http://fixed.local", "timeout": 30, "path_mappings": {}}, + expected_revision=state["revision"], + ) + assert ( + json.loads(user.read_text(encoding="utf-8"))["remote_url"] + == "http://fixed.local" + ) diff --git a/tests/test_proxy_config.py b/tests/test_proxy_config.py new file mode 100644 index 0000000..4e33ba6 --- /dev/null +++ b/tests/test_proxy_config.py @@ -0,0 +1,617 @@ +from __future__ import annotations + +import importlib.util +import json +import sys +import types +from pathlib import Path + +import pytest +from aiohttp import web + + +@pytest.fixture(scope="module") +def modules(): + root = Path(__file__).resolve().parents[1] + package_name = "lm_remote_proxy_test_package" + package = types.ModuleType(package_name) + package.__path__ = [str(root)] + sys.modules[package_name] = package + + loaded = {} + for name in ("config", "remote_client", "proxy"): + full_name = f"{package_name}.{name}" + spec = importlib.util.spec_from_file_location(full_name, root / f"{name}.py") + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[full_name] = module + spec.loader.exec_module(module) + loaded[name] = module + + yield types.SimpleNamespace(**loaded) + for name in ("proxy", "remote_client", "config"): + sys.modules.pop(f"{package_name}.{name}", None) + sys.modules.pop(package_name, None) + + +class DummyContent: + def __init__(self, body: bytes): + self.body = body + self.offset = 0 + + async def read(self, size: int) -> bytes: + chunk = self.body[self.offset : self.offset + size] + self.offset += len(chunk) + return chunk + + +class FragmentedContent: + def __init__(self, chunks: list[bytes]): + self.chunks = list(chunks) + + async def read(self, size: int) -> bytes: + if not self.chunks: + return b"" + chunk = self.chunks.pop(0) + if len(chunk) <= size: + return chunk + self.chunks.insert(0, chunk[size:]) + return chunk[:size] + + +class DummyRequest: + def __init__( + self, + method="GET", + path="/", + payload=None, + headers=None, + *, + raw_body: bytes | None = None, + content_length: int | None | object = ..., + ): + self.method = method + self.path = path + self._payload = payload + self.headers = headers or {} + self.query_string = "" + self.can_read_body = payload is not None + self.content_type = ( + "application/json" if payload is not None else "application/octet-stream" + ) + self._body = ( + raw_body + if raw_body is not None + else json.dumps(payload).encode("utf-8") + if payload is not None + else b"" + ) + self.content = DummyContent(self._body) + if content_length is ...: + self.content_length = len(self._body) if self._body else None + else: + self.content_length = content_length + if raw_body is not None: + self.content_type = "application/json" + + async def json(self): + return self._payload + + async def read(self): + return self._body + + +def response_json(response: web.Response) -> dict: + return json.loads(response.body.decode("utf-8")) + + +@pytest.fixture +def isolated_proxy(modules, tmp_path, monkeypatch): + config = modules.config.RemoteConfig( + tmp_path / "user" / "config.json", + tmp_path / "missing-legacy.json", + environ={}, + ) + monkeypatch.setattr(modules.proxy, "remote_config", config) + monkeypatch.setattr(modules.remote_client, "remote_config", config) + modules.proxy.RemoteLoraClient._instance = None + modules.proxy._proxy_sessions = {} + modules.proxy._proxy_session_lock = None + modules.proxy._active_proxy_websockets = set() + modules.proxy._active_proxy_websockets_lock = None + return modules.proxy, config + + +@pytest.mark.asyncio +async def test_config_endpoint_saves_and_hot_enables(isolated_proxy): + proxy, config = isolated_proxy + get_response = await proxy._handle_config(DummyRequest()) + initial = response_json(get_response) + assert initial["configured"]["remote_url"] == "" + + put_response = await proxy._handle_config( + DummyRequest( + "PUT", + proxy._CONFIG_ROUTE, + { + "revision": initial["revision"], + "config": { + "remote_url": "http://manager.local:8188/", + "timeout": 40, + "path_mappings": {"/remote": "/local"}, + }, + }, + ) + ) + payload = response_json(put_response) + assert put_response.status == 200 + assert payload["effective"]["remote_url"] == "http://manager.local:8188" + assert payload["restart_required"] is False + assert config.is_configured + + +@pytest.mark.asyncio +async def test_middleware_falls_through_when_disabled_then_uses_new_url( + isolated_proxy, monkeypatch +): + proxy, config = isolated_proxy + + async def local_handler(request): + return web.Response(text="local") + + disabled_response = await proxy.lm_remote_proxy_middleware( + DummyRequest(path="/loras"), local_handler + ) + assert disabled_response.text == "local" + + state = config.as_dict() + config.save( + {"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}}, + expected_revision=state["revision"], + ) + captured = {} + + async def fake_proxy_http(request, snapshot): + captured["url"] = snapshot.remote_url + return web.Response(text="remote") + + monkeypatch.setattr(proxy, "_proxy_http", fake_proxy_http) + enabled_response = await proxy.lm_remote_proxy_middleware( + DummyRequest(path="/loras"), local_handler + ) + assert enabled_response.text == "remote" + assert captured["url"] == "http://manager.local" + + +@pytest.mark.asyncio +async def test_config_endpoint_rejects_changes_to_environment_managed_field( + modules, tmp_path, monkeypatch +): + config = modules.config.RemoteConfig( + tmp_path / "user.json", + tmp_path / "legacy.json", + environ={"LM_REMOTE_URL": "http://managed.local"}, + ) + monkeypatch.setattr(modules.proxy, "remote_config", config) + initial = config.as_dict() + response = await modules.proxy._handle_config( + DummyRequest( + "PUT", + modules.proxy._CONFIG_ROUTE, + { + "revision": initial["revision"], + "config": { + "remote_url": "http://changed.local", + "timeout": 30, + "path_mappings": {}, + }, + }, + ) + ) + payload = response_json(response) + assert response.status == 409 + assert payload["field"] == "remote_url" + assert "LM_REMOTE_URL" in payload["error"] + + +@pytest.mark.asyncio +async def test_conflict_reload_rotates_runtime_generation(isolated_proxy, monkeypatch): + proxy, config = isolated_proxy + initial = config.as_dict() + config._config_file.parent.mkdir(parents=True, exist_ok=True) + config._config_file.write_text( + json.dumps( + { + "remote_url": "http://external.local", + "timeout": 20, + "path_mappings": {}, + } + ), + encoding="utf-8", + ) + retired = [] + rotated = [] + + async def fake_retire(generation): + retired.append(generation) + + async def fake_rotate(generation): + rotated.append(generation) + + monkeypatch.setattr(proxy, "_retire_proxy_sessions", fake_retire) + monkeypatch.setattr(proxy, "_rotate_active_websockets", fake_rotate) + response = await proxy._handle_config( + DummyRequest( + "PUT", + proxy._CONFIG_ROUTE, + { + "revision": initial["revision"], + "config": { + "remote_url": "http://browser.local", + "timeout": 30, + "path_mappings": {}, + }, + }, + ) + ) + payload = response_json(response) + + assert response.status == 409 + assert payload["latest"]["effective"]["remote_url"] == "http://external.local" + assert retired == [config.generation] + assert rotated == [config.generation] + + +@pytest.mark.asyncio +async def test_connection_test_uses_unsaved_draft(isolated_proxy, monkeypatch): + proxy, _ = isolated_proxy + captured = {} + + async def fake_test(remote_url, timeout): + captured.update(remote_url=remote_url, timeout=timeout) + return 17 + + monkeypatch.setattr(proxy, "_perform_connection_test", fake_test) + response = await proxy._handle_test_connection( + DummyRequest( + "POST", + proxy._TEST_CONNECTION_ROUTE, + {"remote_url": "http://draft.local:8188/", "timeout": 8}, + ) + ) + payload = response_json(response) + assert response.status == 200 + assert payload["latency_ms"] == 17 + assert captured == {"remote_url": "http://draft.local:8188", "timeout": 8} + + +@pytest.mark.asyncio +async def test_connection_test_rejects_empty_url(isolated_proxy): + proxy, _ = isolated_proxy + response = await proxy._handle_test_connection( + DummyRequest( + "POST", + proxy._TEST_CONNECTION_ROUTE, + {"remote_url": "", "timeout": 30}, + ) + ) + assert response.status == 400 + assert response_json(response)["field"] == "remote_url" + + +@pytest.mark.asyncio +async def test_connection_response_reader_accepts_fragmented_json(isolated_proxy): + proxy, _ = isolated_proxy + response = types.SimpleNamespace( + content=FragmentedContent([b'{"sta', b'tus":"', b'ok"}']) + ) + + assert await proxy._read_small_json(response) == {"status": "ok"} + + +@pytest.mark.asyncio +async def test_connection_response_reader_enforces_hard_limit(isolated_proxy): + proxy, _ = isolated_proxy + response = types.SimpleNamespace( + content=FragmentedContent([b" " * proxy._MAX_TEST_RESPONSE, b" "]) + ) + + with pytest.raises(proxy._ConnectionTestError, match="unexpectedly large"): + await proxy._read_small_json(response) + + +@pytest.mark.asyncio +async def test_proxy_loop_header_is_rejected(isolated_proxy): + proxy, config = isolated_proxy + state = config.as_dict() + config.save( + {"remote_url": "http://manager.local", "timeout": 30, "path_mappings": {}}, + expected_revision=state["revision"], + ) + + async def local_handler(request): + return web.Response(text="local") + + response = await proxy.lm_remote_proxy_middleware( + DummyRequest( + path="/api/lm/health-check", + headers={proxy._PROXY_HOP_HEADER: "1"}, + ), + local_handler, + ) + assert response.status == 508 + + +def test_register_proxy_is_available_while_unconfigured(isolated_proxy): + proxy, config = isolated_proxy + assert not config.is_configured + app = web.Application() + proxy.register_proxy(app) + assert proxy.lm_remote_proxy_middleware in app.middlewares + assert proxy._cleanup_proxy_session in app.on_shutdown + + +def test_register_proxy_keeps_existing_security_middleware_first(isolated_proxy): + proxy, _ = isolated_proxy + + @web.middleware + async def security_guard(request, handler): + return await handler(request) + + app = web.Application(middlewares=[security_guard]) + proxy.register_proxy(app) + + assert list(app.middlewares) == [security_guard, proxy.lm_remote_proxy_middleware] + + +@pytest.mark.asyncio +async def test_chunked_config_body_is_hard_limited(isolated_proxy): + proxy, _ = isolated_proxy + response = await proxy._handle_config( + DummyRequest( + "PUT", + proxy._CONFIG_ROUTE, + raw_body=b" " * (proxy._MAX_CONFIG_BODY + 1), + content_length=None, + ) + ) + + assert response.status == 413 + + +@pytest.mark.asyncio +async def test_remote_client_fetches_all_pages_beyond_server_cap( + isolated_proxy, modules, monkeypatch +): + _, config = isolated_proxy + client = modules.remote_client.RemoteLoraClient() + calls = [] + + async def fake_get_json(path, params=None, *, snapshot=None): + page = int(params["page"]) + calls.append((page, int(params["page_size"]), snapshot.generation)) + start = (page - 1) * 100 + count = 100 if page == 1 else 55 + return { + "items": [ + {"file_name": f"model-{index}"} for index in range(start, start + count) + ], + "total_pages": 2, + } + + monkeypatch.setattr(client, "_get_json", fake_get_json) + items = await client._get_lora_list_cached(snapshot=config.snapshot) + + assert len(items) == 155 + assert items[-1]["file_name"] == "model-154" + assert [call[:2] for call in calls] == [(1, 100), (2, 100)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method_name", "expected_path"), + [ + ("_get_lora_list_cached", "/api/lm/loras/list"), + ("_get_checkpoint_list_cached", "/api/lm/checkpoints/list"), + ], +) +async def test_successful_empty_listing_is_cached( + isolated_proxy, modules, monkeypatch, method_name, expected_path +): + _, config = isolated_proxy + client = modules.remote_client.RemoteLoraClient() + calls = [] + + async def fake_get_all_pages(path, *, snapshot): + calls.append((path, snapshot.generation)) + return [] + + monkeypatch.setattr(client, "_get_all_pages", fake_get_all_pages) + cached_method = getattr(client, method_name) + + assert await cached_method(snapshot=config.snapshot) == [] + assert await cached_method(snapshot=config.snapshot) == [] + assert calls == [(expected_path, config.generation)] + + +@pytest.mark.asyncio +async def test_stale_lookup_never_returns_another_generations_cache( + isolated_proxy, modules, monkeypatch +): + _, config = isolated_proxy + client = modules.remote_client.RemoteLoraClient() + client._lora_cache = [{"file_name": "new-generation"}] + client._lora_cache_generation = config.generation + old_snapshot = modules.config.ConfigSnapshot( + config.generation - 1, + "http://old.local", + 30, + (("/remote", "/old-local"),), + ) + + async def fail_fetch(path, *, snapshot): + raise OSError("old server unavailable") + + monkeypatch.setattr(client, "_get_all_pages", fail_fetch) + + assert await client._get_lora_list_cached(snapshot=old_snapshot) == [] + + +@pytest.mark.asyncio +async def test_lora_info_maps_with_the_fetch_generation( + isolated_proxy, modules, monkeypatch +): + _, config = isolated_proxy + initial = config.as_dict() + config.save( + { + "remote_url": "http://one.local", + "timeout": 30, + "path_mappings": {"/remote": "/local-one"}, + }, + expected_revision=initial["revision"], + ) + client = modules.remote_client.RemoteLoraClient() + + async def fake_list(*, snapshot=None): + latest = config.as_dict() + config.save( + { + "remote_url": "http://two.local", + "timeout": 30, + "path_mappings": {"/remote": "/local-two"}, + }, + expected_revision=latest["revision"], + ) + return [ + { + "file_name": "portrait", + "file_path": "/remote/portrait.safetensors", + "folder": "", + "civitai": {}, + } + ] + + monkeypatch.setattr(client, "_get_lora_list_cached", fake_list) + monkeypatch.setattr(client, "_relative_lora_path", lambda path, folder: path) + + relative, _ = await client.get_lora_info("portrait") + + assert relative == "/local-one/portrait.safetensors" + + +class FakeSession: + def __init__(self, *args, **kwargs): + self.closed = False + + async def close(self): + self.closed = True + + +@pytest.mark.asyncio +async def test_proxy_session_does_not_retain_remote_cookies( + isolated_proxy, monkeypatch +): + proxy, config = isolated_proxy + captured = {} + + def create_session(*args, **kwargs): + captured.update(kwargs) + return FakeSession() + + monkeypatch.setattr(proxy.aiohttp, "ClientSession", create_session) + lease = proxy._proxy_session_lease(config.snapshot) + await lease.__aenter__() + await lease.__aexit__(None, None, None) + + assert isinstance(captured["cookie_jar"], proxy.aiohttp.DummyCookieJar) + await proxy._close_all_proxy_sessions() + + +@pytest.mark.asyncio +async def test_proxy_session_rotation_waits_for_inflight_request( + isolated_proxy, monkeypatch +): + proxy, config = isolated_proxy + created = [] + + def create_session(*args, **kwargs): + session = FakeSession() + created.append(session) + return session + + monkeypatch.setattr(proxy.aiohttp, "ClientSession", create_session) + old_snapshot = config.snapshot + old_lease = proxy._proxy_session_lease(old_snapshot) + old_session = await old_lease.__aenter__() + + state = config.as_dict() + new_snapshot = config.save( + {"remote_url": "http://new.local", "timeout": 30, "path_mappings": {}}, + expected_revision=state["revision"], + ) + await proxy._retire_proxy_sessions(new_snapshot.generation) + assert old_session.closed is False + + new_lease = proxy._proxy_session_lease(new_snapshot) + new_session = await new_lease.__aenter__() + assert new_session is not old_session + await new_lease.__aexit__(None, None, None) + assert new_session.closed is False + + await old_lease.__aexit__(None, None, None) + assert old_session.closed is True + assert new_session.closed is False + + await proxy._close_all_proxy_sessions() + assert new_session.closed is True + + +class FakeWebSocket: + def __init__(self): + self.closed = False + self.close_code = None + + async def close(self, *, code=None, message=None): + self.closed = True + self.close_code = code + + +@pytest.mark.asyncio +async def test_websocket_rotation_closes_only_retired_generation(isolated_proxy): + proxy, config = isolated_proxy + current_generation = config.generation + old_bridge = proxy._ActiveWebSocket( + current_generation - 1, FakeWebSocket(), FakeSession() + ) + current_bridge = proxy._ActiveWebSocket( + current_generation, FakeWebSocket(), FakeSession() + ) + proxy._active_proxy_websockets.update({old_bridge, current_bridge}) + + await proxy._rotate_active_websockets(current_generation) + + assert old_bridge.local_ws.closed is True + assert old_bridge.local_ws.close_code == 1012 + assert old_bridge.session.closed is True + assert current_bridge.local_ws.closed is False + assert proxy._active_proxy_websockets == {current_bridge} + + await proxy._rotate_active_websockets(None) + assert current_bridge.local_ws.closed is True + assert current_bridge.session.closed is True + + +def test_mapped_local_path_is_resolved_against_comfy_lora_roots( + modules, tmp_path, monkeypatch +): + root = tmp_path / "models" / "loras" + fake_folder_paths = types.SimpleNamespace( + get_folder_paths=lambda model_type: [str(root)] + ) + monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths) + relative = modules.remote_client.RemoteLoraClient._relative_lora_path( + str(root / "styles" / "portrait.safetensors"), + "wrong-remote-folder", + ) + assert relative == "styles/portrait.safetensors" diff --git a/web/comfyui/lora_manager_sidebar.css b/web/comfyui/lora_manager_sidebar.css index 33911eb..0d8269c 100644 --- a/web/comfyui/lora_manager_sidebar.css +++ b/web/comfyui/lora_manager_sidebar.css @@ -58,6 +58,17 @@ white-space: nowrap; } +.lmri-header-actions, +.lmri-inline-actions { + display: flex; + align-items: center; + gap: 6px; +} + +.lmri-inline-actions { + margin-top: 10px; +} + .lmri-icon-button, .lmri-name, .lmri-candidate, @@ -347,3 +358,313 @@ grid-template-columns: 1fr; } } + +.lmrc-overlay { + --lmrc-bg: var(--comfy-menu-bg, #18181b); + --lmrc-panel: var(--comfy-input-bg, #242428); + --lmrc-border: var(--border-color, #3a3a40); + --lmrc-text: var(--fg-color, #f4f4f5); + --lmrc-muted: var(--descrip-text, #a1a1aa); + --lmrc-accent: var(--primary-color, #7c3aed); + position: fixed; + z-index: 12000; + inset: 0; + display: grid; + padding: 20px; + overflow-y: auto; + place-items: center; + background: rgb(0 0 0 / 66%); + color: var(--lmrc-text); + font: 13px/1.45 Inter, system-ui, sans-serif; + backdrop-filter: blur(3px); +} + +.lmrc-overlay *, +.lmrc-overlay *::before, +.lmrc-overlay *::after { + box-sizing: border-box; +} + +.lmrc-dialog { + display: flex; + width: min(720px, calc(100vw - 32px)); + max-height: min(820px, calc(100vh - 40px)); + flex-direction: column; + overflow: hidden; + border: 1px solid var(--lmrc-border); + border-radius: 14px; + outline: none; + background: var(--lmrc-bg); + box-shadow: 0 24px 80px rgb(0 0 0 / 55%); +} + +.lmrc-header { + display: flex; + flex: 0 0 auto; + align-items: flex-start; + justify-content: space-between; + gap: 18px; + padding: 18px 20px 15px; + border-bottom: 1px solid var(--lmrc-border); +} + +.lmrc-header h2, +.lmrc-section h3 { + margin: 0; + color: var(--lmrc-text); +} + +.lmrc-header h2 { + font-size: 18px; +} + +.lmrc-header p, +.lmrc-section-heading p { + margin: 4px 0 0; + color: var(--lmrc-muted); +} + +.lmrc-form { + min-height: 0; + flex: 1 1 auto; + padding: 16px 20px 18px; + overflow-y: auto; +} + +.lmrc-section { + margin-bottom: 18px; +} + +.lmrc-section h3 { + font-size: 13px; + font-weight: 650; +} + +.lmrc-connection-grid { + display: grid; + grid-template-columns: minmax(0, 2fr) minmax(130px, 0.7fr); + gap: 12px; + margin-top: 9px; +} + +.lmrc-field { + display: flex; + min-width: 0; + flex-direction: column; + gap: 5px; +} + +.lmrc-label { + color: var(--lmrc-muted); + font-size: 11px; + font-weight: 600; +} + +.lmrc-input { + width: 100%; + min-width: 0; + height: 36px; + padding: 7px 9px; + border: 1px solid var(--lmrc-border); + border-radius: 7px; + outline: none; + background: var(--lmrc-panel); + color: var(--lmrc-text); + font: inherit; +} + +.lmrc-input:focus { + border-color: color-mix(in srgb, var(--lmrc-accent) 78%, white 8%); + box-shadow: 0 0 0 2px color-mix(in srgb, var(--lmrc-accent) 22%, transparent); +} + +.lmrc-input:disabled { + cursor: not-allowed; + opacity: 0.58; +} + +.lmrc-hint { + min-height: 16px; + color: var(--lmrc-muted); + font-size: 10px; + overflow-wrap: anywhere; +} + +.lmrc-section-heading { + display: flex; + align-items: flex-start; + justify-content: space-between; + gap: 14px; +} + +.lmrc-section-heading p { + max-width: 500px; + font-size: 11px; +} + +.lmrc-mapping-rows { + display: flex; + flex-direction: column; + gap: 7px; + margin-top: 10px; +} + +.lmrc-mapping-row { + display: grid; + grid-template-columns: minmax(0, 1fr) 20px minmax(0, 1fr) 32px; + align-items: center; + gap: 7px; +} + +.lmrc-mapping-arrow { + color: var(--lmrc-muted); + text-align: center; +} + +.lmrc-icon-button, +.lmrc-button, +.lmrc-settings-button { + border: 1px solid var(--lmrc-border, var(--border-color, #3a3a40)); + border-radius: 7px; + background: var(--lmrc-panel, var(--comfy-input-bg, #242428)); + color: var(--lmrc-text, var(--fg-color, #f4f4f5)); + font: inherit; + cursor: pointer; +} + +.lmrc-icon-button { + display: grid; + flex: 0 0 32px; + width: 32px; + height: 32px; + padding: 0; + place-items: center; +} + +.lmrc-button, +.lmrc-settings-button { + min-height: 34px; + padding: 7px 11px; +} + +.lmrc-settings-button { + width: 100%; + text-align: left; +} + +.lmrc-icon-button:hover, +.lmrc-button:hover:not(:disabled), +.lmrc-settings-button:hover { + border-color: color-mix(in srgb, var(--lmrc-accent, #7c3aed) 75%, var(--lmrc-border, #3a3a40)); +} + +.lmrc-button:disabled, +.lmrc-icon-button:disabled { + cursor: default; + opacity: 0.5; +} + +.lmrc-primary { + border-color: color-mix(in srgb, var(--lmrc-accent) 72%, white 7%); + background: color-mix(in srgb, var(--lmrc-accent) 54%, var(--lmrc-panel)); +} + +.lmrc-notices { + display: flex; + flex-direction: column; + gap: 6px; +} + +.lmrc-notice, +.lmrc-status { + padding: 8px 10px; + border: 1px solid var(--lmrc-border); + border-radius: 7px; + background: color-mix(in srgb, var(--lmrc-panel) 72%, transparent); + color: var(--lmrc-muted); + font-size: 11px; +} + +.lmrc-notice.warning, +.lmrc-status.error { + border-color: #7f1d1d; + color: #fecaca; +} + +.lmrc-status { + min-height: 35px; + margin-top: 10px; +} + +.lmrc-status-action { + margin-left: 10px; +} + +.lmrc-status:empty { + display: none; +} + +.lmrc-status.success { + border-color: #166534; + color: #bbf7d0; +} + +.lmrc-status.busy { + color: var(--lmrc-text); +} + +.lmrc-footer { + display: flex; + align-items: center; + justify-content: space-between; + gap: 12px; + margin-top: 16px; + padding-top: 14px; + border-top: 1px solid var(--lmrc-border); +} + +.lmrc-footer-group { + display: flex; + gap: 7px; +} + +@media (max-width: 620px) { + .lmrc-overlay { + padding: 8px; + place-items: start center; + } + + .lmrc-dialog { + width: calc(100vw - 16px); + max-height: calc(100vh - 16px); + } + + .lmrc-connection-grid, + .lmrc-mapping-row { + grid-template-columns: 1fr; + } + + .lmrc-mapping-arrow { + transform: rotate(90deg); + } + + .lmrc-mapping-row .lmrc-icon-button { + justify-self: end; + } + + .lmrc-footer { + align-items: stretch; + flex-direction: column; + } + + .lmrc-footer-group { + display: grid; + grid-template-columns: repeat(2, minmax(0, 1fr)); + } + + .lmrc-status-action { + display: block; + width: 100%; + margin: 8px 0 0; + } +} diff --git a/web/comfyui/lora_manager_sidebar.js b/web/comfyui/lora_manager_sidebar.js index 9df0998..6a601d4 100644 --- a/web/comfyui/lora_manager_sidebar.js +++ b/web/comfyui/lora_manager_sidebar.js @@ -11,9 +11,12 @@ import { normalizeLoraIdentifier, normalizeUsageTips, } from "./lora_manager_sidebar_utils.js"; +import { openRemoteConfigDialog } from "./remote_config_dialog.js"; const TAB_ID = "lm-remote-lora-info"; const COMMAND_ID = "LMRemote.OpenLoraInfo"; +const CONFIG_COMMAND_ID = "LMRemote.Configure"; +const CONFIG_SETTING_ID = "LMRemote.Connection.Configure"; const AUTO_OPEN_SETTING = "LMRemote.LoraInfo.AutoOpen"; const STYLE_ID = "lm-remote-lora-info-style"; const NODE_SELECTION_HOOK = Symbol.for("lmRemote.loraInfo.nodeSelectionHook"); @@ -45,6 +48,26 @@ function ensureStyles() { document.head.appendChild(link); } +function openConfiguration() { + ensureStyles(); + openRemoteConfigDialog({ + onSaved: () => { + if (activeName) lookupActiveName(); + }, + }); +} + +function configureSettingControl() { + const button = createElement( + "button", + "lmrc-settings-button", + "Configure LM Remote…" + ); + button.type = "button"; + button.addEventListener("click", openConfiguration); + return button; +} + function selectedNodeLabel() { if (!selectedNode) return ""; return ( @@ -197,7 +220,12 @@ function renderError(content) { const retry = createElement("button", "lmri-button", "Try again"); retry.type = "button"; retry.addEventListener("click", () => lookupActiveName()); - panel.appendChild(retry); + const configure = createElement("button", "lmri-button", "Configure"); + configure.type = "button"; + configure.addEventListener("click", openConfiguration); + const actions = createElement("div", "lmri-inline-actions"); + actions.append(retry, configure); + panel.appendChild(actions); appendSearchActions(panel, activeName); content.appendChild(panel); } @@ -361,7 +389,15 @@ function renderSidebar() { refresh.appendChild(createElement("i", "pi pi-refresh")); refresh.disabled = !activeName || lookupState.status === "loading"; refresh.addEventListener("click", () => lookupActiveName()); - header.append(heading, refresh); + const configure = createElement("button", "lmri-icon-button"); + configure.type = "button"; + configure.title = "Configure LM Remote"; + configure.setAttribute("aria-label", configure.title); + configure.appendChild(createElement("i", "pi pi-cog")); + configure.addEventListener("click", openConfiguration); + const headerActions = createElement("div", "lmri-header-actions"); + headerActions.append(configure, refresh); + header.append(heading, headerActions); sidebarRoot.appendChild(header); if (selectedNames.length > 1) { @@ -424,7 +460,9 @@ async function fallbackListLookup(name, signal) { search: term, fuzzy_search: "true", }); - const response = await api.fetchApi(`/lm/loras/list?${params}`, { signal }); + const response = await api.fetchApi(`/api/lm/loras/list?${params}`, { + signal, + }); if (!response.ok) throw new Error(await responseError(response)); const payload = await response.json(); const result = matchModelItems(name, payload.items); @@ -437,7 +475,7 @@ async function fallbackListLookup(name, signal) { async function resolveManagerCard(name, signal) { const response = await api.fetchApi( - `/lm/loras/resolve?name=${encodeURIComponent(name)}`, + `/api/lm/loras/resolve?name=${encodeURIComponent(name)}`, { signal } ); if (response.status === 404) { @@ -657,6 +695,14 @@ function registerSidebarTab() { app.registerExtension({ name: "LoraManager.RemoteLoraInfoSidebar", settings: [ + { + id: CONFIG_SETTING_ID, + name: "Remote LoRA Manager", + type: configureSettingControl, + defaultValue: "", + category: ["LM Remote", "Connection", "Configure"], + tooltip: "Set the remote URL, timeout, and path mappings.", + }, { id: AUTO_OPEN_SETTING, name: "Open LoRA Info when selecting a LoRA loader", @@ -666,6 +712,12 @@ app.registerExtension({ }, ], commands: [ + { + id: CONFIG_COMMAND_ID, + label: "Configure LM Remote", + icon: "pi pi-cog", + function: openConfiguration, + }, { id: COMMAND_ID, label: "Open LoRA Info", diff --git a/web/comfyui/remote_config_dialog.js b/web/comfyui/remote_config_dialog.js new file mode 100644 index 0000000..c82d65f --- /dev/null +++ b/web/comfyui/remote_config_dialog.js @@ -0,0 +1,461 @@ +import { api } from "../../scripts/api.js"; + +import { + buildConfigDraft, + buildConnectionDraft, + didEffectiveRemoteUrlChange, + isConfigWritable, + mappingsToRows, +} from "./remote_config_utils.js"; + +// Explicit /api paths work with both current and legacy ComfyUI apiURL helpers. +const CONFIG_ENDPOINT = "/api/lm-remote/config"; +const TEST_ENDPOINT = "/api/lm-remote/test-connection"; + +let activeDialog = null; +let pageEffectiveRemoteUrl; + +function element(tag, className, text) { + const value = document.createElement(tag); + if (className) value.className = className; + if (text != null) value.textContent = String(text); + return value; +} + +function labeledInput(labelText, input) { + const label = element("label", "lmrc-field"); + label.append(element("span", "lmrc-label", labelText), input); + return label; +} + +function parseError(response, payload) { + const error = new Error( + payload?.error || `Request failed with HTTP ${response.status}.` + ); + error.field = payload?.field || ""; + error.status = response.status; + error.payload = payload; + return error; +} + +async function requestJson(path, options = {}) { + const response = await api.fetchApi(path, options); + let payload = null; + try { + payload = await response.json(); + } catch { + // The HTTP status below still gives the user a useful failure. + } + if (!response.ok || payload?.success === false) { + throw parseError(response, payload); + } + return payload; +} + +function trapFocus(event, panel) { + if (event.key !== "Tab") return; + const focusable = Array.from( + panel.querySelectorAll( + 'button:not([disabled]), input:not([disabled]), [tabindex]:not([tabindex="-1"])' + ) + ).filter((item) => item.offsetParent !== null); + if (!focusable.length) return; + const first = focusable[0]; + const last = focusable[focusable.length - 1]; + if (event.shiftKey && document.activeElement === first) { + event.preventDefault(); + last.focus(); + } else if (!event.shiftKey && document.activeElement === last) { + event.preventDefault(); + first.focus(); + } +} + +export function openRemoteConfigDialog({ onSaved } = {}) { + if (activeDialog) { + activeDialog.panel.focus(); + return activeDialog; + } + + const opener = document.activeElement; + + const overlay = element("div", "lmrc-overlay"); + overlay.setAttribute("role", "dialog"); + overlay.setAttribute("aria-modal", "true"); + overlay.setAttribute("aria-labelledby", "lmrc-dialog-title"); + const panel = element("section", "lmrc-dialog"); + panel.tabIndex = -1; + overlay.appendChild(panel); + + const header = element("header", "lmrc-header"); + const heading = element("div"); + const title = element("h2", "", "Configure LM Remote"); + title.id = "lmrc-dialog-title"; + heading.append( + title, + element( + "p", + "", + "Connect this ComfyUI instance to LoRA Manager. New requests use saved changes immediately." + ) + ); + const closeButton = element("button", "lmrc-icon-button"); + closeButton.type = "button"; + closeButton.title = "Close"; + closeButton.setAttribute("aria-label", "Close configuration"); + closeButton.appendChild(element("i", "pi pi-times")); + header.append(heading, closeButton); + + const form = element("form", "lmrc-form"); + form.noValidate = true; + const connectionSection = element("section", "lmrc-section"); + connectionSection.appendChild(element("h3", "", "Connection")); + const connectionGrid = element("div", "lmrc-connection-grid"); + + const remoteUrlInput = element("input", "lmrc-input"); + remoteUrlInput.type = "url"; + remoteUrlInput.placeholder = "http://manager.local:8188"; + remoteUrlInput.autocomplete = "url"; + remoteUrlInput.spellcheck = false; + const remoteField = labeledInput("Remote URL", remoteUrlInput); + const remoteHint = element("small", "lmrc-hint"); + remoteField.appendChild(remoteHint); + + const timeoutInput = element("input", "lmrc-input"); + timeoutInput.type = "number"; + timeoutInput.min = "1"; + timeoutInput.max = "300"; + timeoutInput.step = "1"; + const timeoutField = labeledInput("Timeout (seconds)", timeoutInput); + const timeoutHint = element("small", "lmrc-hint"); + timeoutField.appendChild(timeoutHint); + connectionGrid.append(remoteField, timeoutField); + connectionSection.appendChild(connectionGrid); + + const mappingSection = element("section", "lmrc-section"); + const mappingHeading = element("div", "lmrc-section-heading"); + const mappingCopy = element("div"); + mappingCopy.append( + element("h3", "", "Path mappings"), + element( + "p", + "", + "Optional remote-to-local prefixes when both machines mount models at different paths." + ) + ); + const addMappingButton = element("button", "lmrc-button lmrc-secondary", "Add mapping"); + addMappingButton.type = "button"; + mappingHeading.append(mappingCopy, addMappingButton); + const mappingRows = element("div", "lmrc-mapping-rows"); + mappingSection.append(mappingHeading, mappingRows); + + const notices = element("div", "lmrc-notices"); + const status = element("div", "lmrc-status"); + status.setAttribute("role", "status"); + status.setAttribute("aria-live", "polite"); + + const footer = element("footer", "lmrc-footer"); + const leftActions = element("div", "lmrc-footer-group"); + const testButton = element("button", "lmrc-button lmrc-secondary", "Test connection"); + testButton.type = "button"; + const reloadButton = element("button", "lmrc-button lmrc-secondary", "Reload saved"); + reloadButton.type = "button"; + leftActions.append(testButton, reloadButton); + const rightActions = element("div", "lmrc-footer-group"); + const cancelButton = element("button", "lmrc-button lmrc-secondary", "Close"); + cancelButton.type = "button"; + const saveButton = element("button", "lmrc-button lmrc-primary", "Save"); + saveButton.type = "submit"; + rightActions.append(cancelButton, saveButton); + footer.append(leftActions, rightActions); + + form.append(connectionSection, mappingSection, notices, status, footer); + panel.append(header, form); + document.body.appendChild(overlay); + + let loaded = null; + let busy = false; + let formLocked = false; + let dismissLocked = false; + + function updateDisabledState() { + const overrides = loaded?.overrides || {}; + remoteUrlInput.disabled = formLocked || Boolean(overrides.remote_url); + timeoutInput.disabled = formLocked || Boolean(overrides.timeout); + for (const control of mappingRows.querySelectorAll("input, button")) { + control.disabled = formLocked; + } + closeButton.disabled = dismissLocked; + cancelButton.disabled = dismissLocked; + } + + function close() { + if (dismissLocked) return; + if (activeDialog?.overlay !== overlay) return; + overlay.remove(); + activeDialog = null; + if (opener?.isConnected && typeof opener.focus === "function") { + try { + opener.focus({ preventScroll: true }); + } catch { + opener.focus(); + } + } + } + + function setStatus(message = "", kind = "", action = null) { + status.replaceChildren(); + if (message) status.appendChild(element("span", "", message)); + if (action) { + const actionButton = element( + "button", + "lmrc-button lmrc-secondary lmrc-status-action", + action.label + ); + actionButton.type = "button"; + actionButton.addEventListener("click", action.onClick); + status.appendChild(actionButton); + } + status.className = `lmrc-status${kind ? ` ${kind}` : ""}`; + } + + function setBusy( + value, + message = "", + { lockForm = false, lockDismiss = false } = {} + ) { + busy = value; + formLocked = Boolean(value && lockForm); + dismissLocked = Boolean(value && lockDismiss); + testButton.disabled = value || !loaded; + reloadButton.disabled = value; + saveButton.disabled = value || !loaded || !isConfigWritable(loaded); + addMappingButton.disabled = value || !loaded; + form.setAttribute("aria-busy", String(value)); + updateDisabledState(); + if (message) setStatus(message, "busy"); + } + + function addMappingRow(remote = "", local = "") { + const row = element("div", "lmrc-mapping-row"); + const remoteInput = element("input", "lmrc-input"); + remoteInput.type = "text"; + remoteInput.placeholder = "/data/models/loras"; + remoteInput.value = remote; + remoteInput.setAttribute("aria-label", "Remote path prefix"); + const arrow = element("i", "pi pi-arrow-right lmrc-mapping-arrow"); + arrow.setAttribute("aria-hidden", "true"); + const localInput = element("input", "lmrc-input"); + localInput.type = "text"; + localInput.placeholder = "/mnt/nas/models/loras"; + localInput.value = local; + localInput.setAttribute("aria-label", "Local path prefix"); + const remove = element("button", "lmrc-icon-button"); + remove.type = "button"; + remove.title = "Remove path mapping"; + remove.setAttribute("aria-label", remove.title); + remove.appendChild(element("i", "pi pi-trash")); + remove.addEventListener("click", () => row.remove()); + row.append(remoteInput, arrow, localInput, remove); + mappingRows.appendChild(row); + updateDisabledState(); + return row; + } + + function readMappingRows() { + return Array.from(mappingRows.querySelectorAll(".lmrc-mapping-row")).map( + (row) => { + const inputs = row.querySelectorAll("input"); + return { remote: inputs[0]?.value || "", local: inputs[1]?.value || "" }; + } + ); + } + + function renderNotices(payload) { + notices.replaceChildren(); + const source = payload.storage?.source; + if (source === "legacy") { + notices.appendChild( + element( + "div", + "lmrc-notice", + "Loaded the package config. Saving migrates it to ComfyUI user data." + ) + ); + } else if (source === "user") { + notices.appendChild( + element("div", "lmrc-notice", "Stored in ComfyUI user data.") + ); + } else if (source === "explicit") { + notices.appendChild( + element( + "div", + "lmrc-notice", + "Stored in the file selected by LM_REMOTE_CONFIG." + ) + ); + } + if (!isConfigWritable(payload)) { + notices.appendChild( + element( + "div", + "lmrc-notice warning", + "Configuration storage is not writable. Fix its file or directory permissions, then use Reload saved to check again." + ) + ); + } + for (const warning of payload.warnings || []) { + notices.appendChild(element("div", "lmrc-notice warning", warning)); + } + } + + function applyPayload(payload) { + loaded = payload; + const configured = payload.configured || {}; + const effective = payload.effective || configured; + const overrides = payload.overrides || {}; + remoteUrlInput.value = configured.remote_url || ""; + timeoutInput.value = String(configured.timeout ?? 30); + remoteHint.textContent = overrides.remote_url + ? `Managed by ${overrides.remote_url}. Effective: ${effective.remote_url}` + : "The base URL of the standalone LoRA Manager."; + timeoutHint.textContent = overrides.timeout + ? `Managed by ${overrides.timeout}. Effective: ${effective.timeout} seconds.` + : "Used for Manager API and proxy requests."; + mappingRows.replaceChildren(); + const rows = mappingsToRows(configured.path_mappings); + rows.forEach((row) => addMappingRow(row.remote, row.local)); + renderNotices(payload); + setBusy(false); + } + + async function load() { + setBusy(true, "Loading configuration…", { lockForm: true }); + try { + const payload = await requestJson(CONFIG_ENDPOINT); + if (pageEffectiveRemoteUrl === undefined) { + pageEffectiveRemoteUrl = String(payload?.effective?.remote_url || ""); + } + applyPayload(payload); + setStatus(""); + (form.querySelector("input:not([disabled])") || panel).focus(); + } catch (error) { + loaded = null; + setBusy(false); + setStatus(error.message || "Could not load LM Remote configuration.", "error"); + } + } + + function showFormError(error) { + setStatus(error.message, "error"); + const target = + error.field === "remote_url" + ? remoteUrlInput + : error.field === "timeout" + ? timeoutInput + : null; + if (target) { + queueMicrotask(() => { + if (target.isConnected && !target.disabled) target.focus(); + }); + } + } + + addMappingButton.addEventListener("click", () => { + addMappingRow().querySelector("input")?.focus(); + }); + reloadButton.addEventListener("click", () => { + if (!busy) load(); + }); + testButton.addEventListener("click", async () => { + if (busy || !loaded) return; + try { + const draft = buildConnectionDraft({ + remoteUrl: remoteUrlInput.value, + timeout: timeoutInput.value, + effective: loaded.effective, + overrides: loaded.overrides, + }); + setBusy(true, "Testing connection…", { lockForm: true }); + const payload = await requestJson(TEST_ENDPOINT, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(draft), + }); + setStatus(`Connected in ${payload.latency_ms} ms.`, "success"); + } catch (error) { + showFormError(error); + } finally { + setBusy(false); + } + }); + form.addEventListener("submit", async (event) => { + event.preventDefault(); + if (busy || !loaded || !isConfigWritable(loaded)) return; + try { + const config = buildConfigDraft({ + remoteUrl: remoteUrlInput.value, + timeout: timeoutInput.value, + mappingRows: readMappingRows(), + }); + setBusy(true, "Saving configuration…", { + lockForm: true, + lockDismiss: true, + }); + const payload = await requestJson(CONFIG_ENDPOINT, { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ revision: loaded.revision, config }), + }); + applyPayload(payload); + if ( + didEffectiveRemoteUrlChange( + { effective: { remote_url: pageEffectiveRemoteUrl } }, + payload + ) + ) { + setStatus( + "Saved. Reload the ComfyUI page to load Manager assets and reconnect live updates. Save workflow changes first.", + "success", + { + label: "Reload ComfyUI page", + onClick: () => window.location.reload(), + } + ); + } else { + setStatus( + "Saved and applied to new requests. No ComfyUI restart required.", + "success" + ); + } + onSaved?.(payload); + } catch (error) { + if (error.status === 409 && error.payload?.latest) { + setStatus(`${error.message} Use Reload saved to continue.`, "error"); + } else { + showFormError(error); + } + } finally { + setBusy(false); + } + }); + + closeButton.addEventListener("click", close); + cancelButton.addEventListener("click", close); + overlay.addEventListener("mousedown", (event) => { + if (event.target === overlay) close(); + }); + overlay.addEventListener("keydown", (event) => { + if (event.key === "Escape") { + event.preventDefault(); + close(); + } else { + trapFocus(event, panel); + } + }); + + activeDialog = { overlay, panel, close }; + panel.focus(); + load(); + return activeDialog; +} diff --git a/web/comfyui/remote_config_utils.js b/web/comfyui/remote_config_utils.js new file mode 100644 index 0000000..b883700 --- /dev/null +++ b/web/comfyui/remote_config_utils.js @@ -0,0 +1,132 @@ +export class RemoteConfigFormError extends Error { + constructor(field, message) { + super(message); + this.name = "RemoteConfigFormError"; + this.field = field; + } +} + +function normalizeTimeout(value) { + const text = String(value ?? "").trim(); + if (!/^\d+$/.test(text)) { + throw new RemoteConfigFormError("timeout", "Timeout must be a whole number."); + } + const timeout = Number(text); + if (!Number.isSafeInteger(timeout) || timeout < 1 || timeout > 300) { + throw new RemoteConfigFormError( + "timeout", + "Timeout must be between 1 and 300 seconds." + ); + } + return timeout; +} + +function normalizeRemoteUrl(value, allowEmpty) { + const remoteUrl = String(value ?? "").trim(); + if (!remoteUrl) { + if (allowEmpty) return ""; + throw new RemoteConfigFormError( + "remote_url", + "Enter a remote LoRA Manager URL." + ); + } + + let parsed; + try { + parsed = new URL(remoteUrl); + } catch { + throw new RemoteConfigFormError("remote_url", "Enter a valid remote URL."); + } + if (!["http:", "https:"].includes(parsed.protocol)) { + throw new RemoteConfigFormError( + "remote_url", + "Remote URL must use http:// or https://." + ); + } + if (!parsed.hostname) { + throw new RemoteConfigFormError( + "remote_url", + "Remote URL must include a host name." + ); + } + if (parsed.username || parsed.password) { + throw new RemoteConfigFormError( + "remote_url", + "Credentials are not allowed in the remote URL." + ); + } + if (parsed.search || parsed.hash) { + throw new RemoteConfigFormError( + "remote_url", + "Remote URL cannot contain a query or fragment." + ); + } + return remoteUrl.replace(/\/+$/, ""); +} + +function buildMappings(rows) { + const pathMappings = new Map(); + for (const row of rows || []) { + const remote = String(row?.remote ?? "").trim(); + const local = String(row?.local ?? "").trim(); + if (!remote && !local) continue; + if (!remote || !local) { + throw new RemoteConfigFormError( + "path_mappings", + "Each path mapping needs both a remote and local path." + ); + } + const normalizedRemote = remote.replace(/\\/g, "/").replace(/\/+$/, "") || "/"; + if (pathMappings.has(normalizedRemote)) { + throw new RemoteConfigFormError( + "path_mappings", + `Duplicate remote path prefix: ${normalizedRemote}` + ); + } + pathMappings.set(normalizedRemote, local); + } + return Object.fromEntries(pathMappings); +} + +export function mappingsToRows(pathMappings) { + if (!pathMappings || typeof pathMappings !== "object") return []; + return Object.entries(pathMappings).map(([remote, local]) => ({ + remote: String(remote), + local: String(local), + })); +} + +export function buildConfigDraft({ remoteUrl, timeout, mappingRows }) { + return { + remote_url: normalizeRemoteUrl(remoteUrl, true), + timeout: normalizeTimeout(timeout), + path_mappings: buildMappings(mappingRows), + }; +} + +export function buildConnectionDraft({ + remoteUrl, + timeout, + effective, + overrides, +}) { + return { + remote_url: normalizeRemoteUrl( + overrides?.remote_url ? effective?.remote_url : remoteUrl, + false + ), + timeout: normalizeTimeout( + overrides?.timeout ? effective?.timeout : timeout + ), + }; +} + +export function didEffectiveRemoteUrlChange(previousPayload, nextPayload) { + const previous = String(previousPayload?.effective?.remote_url || ""); + const next = String(nextPayload?.effective?.remote_url || ""); + return previous !== next; +} + +export function isConfigWritable(payload) { + return payload?.storage?.writable !== false; +}