diff --git a/configure_model_root.py b/configure_model_root.py new file mode 100644 index 0000000..31c646a --- /dev/null +++ b/configure_model_root.py @@ -0,0 +1,48 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Configure BackEnd model-root state before a newly installed Comfy can start.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_EXTENSION_ROOT = Path(__file__).resolve().parent +if str(_EXTENSION_ROOT) not in sys.path: + sys.path.insert(0, str(_EXTENSION_ROOT)) + +from substitute_backend.features.environment_management.infrastructure.model_root_store import ( # noqa: E402 + ModelRootStore, +) + + +def main(argv: list[str] | None = None) -> int: + """Persist a model-root selection for an offline Comfy installation.""" + + parser = argparse.ArgumentParser() + parser.add_argument("--comfy-root", required=True, type=Path) + selection = parser.add_mutually_exclusive_group(required=True) + selection.add_argument("--default", action="store_true") + selection.add_argument("--path", type=Path) + arguments = parser.parse_args(argv) + store = ModelRootStore(arguments.comfy_root) + store.save(None if arguments.default else arguments.path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/package-lock.json b/package-lock.json index 98e9933..3f6bedf 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "substitute-backend", - "version": "1.6.2", + "version": "1.7.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "substitute-backend", - "version": "1.6.2", + "version": "1.7.0", "license": "AGPL-3.0-or-later", "devDependencies": { "@semantic-release/changelog": "^6.0.3", diff --git a/package.json b/package.json index 9587af9..598c973 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "substitute-backend", "private": true, - "version": "1.6.2", + "version": "1.7.0", "license": "AGPL-3.0-or-later", "type": "module", "scripts": { diff --git a/prestartup_script.py b/prestartup_script.py new file mode 100644 index 0000000..d955197 --- /dev/null +++ b/prestartup_script.py @@ -0,0 +1,45 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Run BackEnd model-root configuration in ComfyUI's prestartup phase.""" + +from __future__ import annotations + +import importlib +import importlib.util +from pathlib import Path +from types import ModuleType + + +def _load_prestartup_support(extension_root: Path) -> ModuleType: + """Load the isolated support module without importing the node package.""" + + module_path = extension_root / "substitute_backend_prestartup.py" + spec = importlib.util.spec_from_file_location( + "substitute_backend_prestartup_support", module_path + ) + if spec is None or spec.loader is None: + raise RuntimeError("Substitute BackEnd prestartup support could not be loaded.") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_EXTENSION_ROOT = Path(__file__).resolve().parent +_SUPPORT = _load_prestartup_support(_EXTENSION_ROOT) +_SUPPORT.apply_model_root( + _EXTENSION_ROOT.parents[1], + importlib.import_module("folder_paths"), +) diff --git a/pyproject.toml b/pyproject.toml index d6bfb1a..03afb83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "substitute-backend" description = "Backend liaison services for SugarSubstitute and ComfyUI." -version = "1.6.2" +version = "1.7.0" license = "AGPL-3.0-or-later" license-files = ["LICENSE"] requires-python = ">=3.12" @@ -31,6 +31,9 @@ dev = [ Repository = "https://github.com/Artificial-Sweetener/Substitute-BackEnd" "Bug Tracker" = "https://github.com/Artificial-Sweetener/Substitute-BackEnd/issues" +[tool.setuptools] +py-modules = ["substitute_backend_prestartup"] + [tool.comfy] PublisherId = "artificialsweetener" DisplayName = "Substitute BackEnd" diff --git a/substitute_backend/__init__.py b/substitute_backend/__init__.py index 653c944..ff1a5c8 100644 --- a/substitute_backend/__init__.py +++ b/substitute_backend/__init__.py @@ -17,7 +17,7 @@ from __future__ import annotations -__version__ = "1.6.2" +__version__ = "1.7.0" EXTENSION_NAME = "Substitute BackEnd" API_VERSION = 1 diff --git a/substitute_backend/features/environment_management/api/routes.py b/substitute_backend/features/environment_management/api/routes.py index 8f3e936..8a4533a 100644 --- a/substitute_backend/features/environment_management/api/routes.py +++ b/substitute_backend/features/environment_management/api/routes.py @@ -27,6 +27,7 @@ from substitute_backend.api.errors import BackendHttpError, json_error from substitute_backend.features.environment_management.application.services import ( EnvironmentManagementServices, ) +from substitute_backend.features.environment_management.domain.model_root import ModelRootMode RouteHandler = Callable[[web.Request], Awaitable[web.StreamResponse]] @@ -37,6 +38,8 @@ class EnvironmentRouteHandlers: capabilities: RouteHandler status: RouteHandler + get_model_root: RouteHandler + update_model_root: RouteHandler list_packages: RouteHandler list_components: RouteHandler plan_operation: RouteHandler @@ -85,6 +88,44 @@ def build_environment_route_handlers( ) ) + async def get_model_root(request: web.Request) -> web.Response: + """Return persisted and active Comfy model-root state.""" + + _ = request + try: + return web.json_response(services.model_root.get_status().to_payload()) + except (OSError, ValueError) as exc: + logger.warning( + "model-root status unavailable", + extra={"operation": "model-root-status", "error": repr(exc)}, + ) + return json_error( + BackendHttpError( + message="Model-root configuration could not be read.", + status=500, + code="model-root-status-unavailable", + ) + ) + + async def update_model_root(request: web.Request) -> web.Response: + """Persist the model root that ComfyUI will use after restart.""" + + try: + body = await _json_object_body(request) + mode, raw_path = _parse_model_root_update(body) + status = services.model_root.configure(mode, raw_path) + return web.json_response(status.to_payload()) + except BackendHttpError as exc: + return json_error(exc) + except (OSError, ValueError) as exc: + return json_error( + BackendHttpError( + message=str(exc), + status=400, + code="invalid-model-root", + ) + ) + async def restart(request: web.Request) -> web.Response: """Queue a Comfy restart operation.""" @@ -302,6 +343,8 @@ def build_environment_route_handlers( return EnvironmentRouteHandlers( capabilities=capabilities, status=status, + get_model_root=get_model_root, + update_model_root=update_model_root, list_packages=list_packages, list_components=list_components, plan_operation=plan_operation, @@ -330,6 +373,36 @@ async def _json_object_body(request: web.Request) -> dict[str, object]: return body +def _parse_model_root_update( + body: dict[str, object], +) -> tuple[ModelRootMode, str | None]: + """Parse one model-root update without leaking untyped request data.""" + + raw_mode = body.get("mode") + if not isinstance(raw_mode, str): + raise BackendHttpError( + message="Model-root mode is required.", + status=400, + code="invalid-model-root", + ) + try: + mode = ModelRootMode(raw_mode) + except ValueError as exc: + raise BackendHttpError( + message="Model-root mode must be 'default' or 'custom'.", + status=400, + code="invalid-model-root", + ) from exc + raw_path = body.get("path") + if raw_path is not None and not isinstance(raw_path, str): + raise BackendHttpError( + message="Model-root path must be a string.", + status=400, + code="invalid-model-root", + ) + return mode, raw_path + + def _required_int(data: dict[str, object], key: str) -> int: """Read one required integer request field.""" diff --git a/substitute_backend/features/environment_management/application/environment_service.py b/substitute_backend/features/environment_management/application/environment_service.py index 982e44d..0e57a43 100644 --- a/substitute_backend/features/environment_management/application/environment_service.py +++ b/substitute_backend/features/environment_management/application/environment_service.py @@ -50,6 +50,7 @@ class EnvironmentService: EnvironmentFeature.PACKAGE_INVENTORY, EnvironmentFeature.COMPONENT_INVENTORY, EnvironmentFeature.OPERATION_PLANNING, + EnvironmentFeature.MODEL_ROOT_MANAGEMENT, ] if restart_support.supported: supported_features.append(EnvironmentFeature.RESTART) @@ -59,6 +60,7 @@ class EnvironmentService: restart_supported=restart_support.supported, package_mutation_supported=False, operation_planning_supported=True, + model_root_management_supported=True, restart_unavailable_reason=restart_support.unavailable_reason, ) diff --git a/substitute_backend/features/environment_management/application/model_root_service.py b/substitute_backend/features/environment_management/application/model_root_service.py new file mode 100644 index 0000000..63c8758 --- /dev/null +++ b/substitute_backend/features/environment_management/application/model_root_service.py @@ -0,0 +1,74 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Coordinate persisted and active Comfy model-root state.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from substitute_backend.features.environment_management.domain.model_root import ( + ModelRootMode, + ModelRootStatus, +) +from substitute_backend.features.environment_management.infrastructure.model_root_runtime import ( + ModelRootRuntime, +) +from substitute_backend.features.environment_management.infrastructure.model_root_store import ( + ModelRootStore, +) + +_MODEL_ROOT_SCHEMA_VERSION = 1 + + +@dataclass(frozen=True) +class ModelRootService: + """Own model-root queries and configuration changes for one Comfy host.""" + + comfy_root: Path + store: ModelRootStore + runtime: ModelRootRuntime + + def get_status(self) -> ModelRootStatus: + """Return the persisted selection alongside the running process state.""" + + default_root = (self.comfy_root / "models").resolve() + configured_root = self.store.load() + active_root = self.runtime.active_model_root() + desired_root = configured_root or default_root + return ModelRootStatus( + schema_version=_MODEL_ROOT_SCHEMA_VERSION, + default_model_root=default_root, + configured_model_root=configured_root, + active_model_root=active_root, + restart_required=desired_root != active_root, + ) + + def configure(self, mode: ModelRootMode, path: str | None) -> ModelRootStatus: + """Persist an explicit default or custom selection for the next launch.""" + + if mode is ModelRootMode.DEFAULT: + if path is not None: + raise ValueError("Default model-root mode does not accept a path.") + self.store.save(None) + else: + if path is None or not path.strip(): + raise ValueError("Custom model-root mode requires a path.") + self.store.save(ModelRootStore.resolve_custom_root(path)) + return self.get_status() + + +__all__ = ["ModelRootService"] diff --git a/substitute_backend/features/environment_management/application/services.py b/substitute_backend/features/environment_management/application/services.py index 28ba3a0..b0e12e1 100644 --- a/substitute_backend/features/environment_management/application/services.py +++ b/substitute_backend/features/environment_management/application/services.py @@ -32,6 +32,9 @@ from substitute_backend.features.environment_management.application.inventory_se from substitute_backend.features.environment_management.application.job_service import ( JobService, ) +from substitute_backend.features.environment_management.application.model_root_service import ( + ModelRootService, +) from substitute_backend.features.environment_management.application.restart_service import ( RestartService, ) @@ -47,3 +50,4 @@ class EnvironmentManagementServices: maintenance_plan: MaintenancePlanService operation_planning: OperationPlanningService restart: RestartService + model_root: ModelRootService diff --git a/substitute_backend/features/environment_management/domain/capabilities.py b/substitute_backend/features/environment_management/domain/capabilities.py index 7ba7edc..ac308ba 100644 --- a/substitute_backend/features/environment_management/domain/capabilities.py +++ b/substitute_backend/features/environment_management/domain/capabilities.py @@ -31,6 +31,7 @@ class EnvironmentFeature(StrEnum): RESTART = "restart" OPERATION_PLANNING = "operation-planning" PACKAGE_MUTATION = "package-mutation" + MODEL_ROOT_MANAGEMENT = "model-root-management" @dataclass(frozen=True) @@ -42,6 +43,7 @@ class EnvironmentCapabilities: restart_supported: bool package_mutation_supported: bool operation_planning_supported: bool + model_root_management_supported: bool restart_unavailable_reason: str | None = None def to_payload(self) -> JsonObject: @@ -53,6 +55,7 @@ class EnvironmentCapabilities: "restartSupported": self.restart_supported, "packageMutationSupported": self.package_mutation_supported, "operationPlanningSupported": self.operation_planning_supported, + "modelRootManagementSupported": self.model_root_management_supported, } if self.restart_unavailable_reason is not None: payload["restartUnavailableReason"] = self.restart_unavailable_reason diff --git a/substitute_backend/features/environment_management/domain/model_root.py b/substitute_backend/features/environment_management/domain/model_root.py new file mode 100644 index 0000000..22c0f23 --- /dev/null +++ b/substitute_backend/features/environment_management/domain/model_root.py @@ -0,0 +1,65 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Model-root contracts owned by the Comfy environment.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path + +from substitute_backend.api.serialization import JsonObject + + +class ModelRootMode(StrEnum): + """Identify whether Comfy uses its default or a configured model root.""" + + DEFAULT = "default" + CUSTOM = "custom" + + +@dataclass(frozen=True) +class ModelRootStatus: + """Describe persisted and active model-root state for one Comfy host.""" + + schema_version: int + default_model_root: Path + configured_model_root: Path | None + active_model_root: Path + restart_required: bool + + @property + def uses_default(self) -> bool: + """Return whether the persisted configuration selects Comfy's default.""" + + return self.configured_model_root is None + + def to_payload(self) -> JsonObject: + """Return a stable API payload for external environment clients.""" + + return { + "schemaVersion": self.schema_version, + "defaultModelRoot": str(self.default_model_root), + "configuredModelRoot": ( + str(self.configured_model_root) if self.configured_model_root is not None else None + ), + "activeModelRoot": str(self.active_model_root), + "usesDefault": self.uses_default, + "restartRequired": self.restart_required, + } + + +__all__ = ["ModelRootMode", "ModelRootStatus"] diff --git a/substitute_backend/features/environment_management/infrastructure/model_root_runtime.py b/substitute_backend/features/environment_management/infrastructure/model_root_runtime.py new file mode 100644 index 0000000..bc02012 --- /dev/null +++ b/substitute_backend/features/environment_management/infrastructure/model_root_runtime.py @@ -0,0 +1,48 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Read active model-root state from ComfyUI at the host boundary.""" + +from __future__ import annotations + +import importlib +from pathlib import Path +from typing import Protocol, cast + + +class ModelRootRuntime(Protocol): + """Expose the model root active in the current Comfy process.""" + + def active_model_root(self) -> Path: + """Return the model root currently used by ComfyUI.""" + + +class _FolderPathsModule(Protocol): + """Describe the ComfyUI state required by the runtime adapter.""" + + models_dir: str + + +class ComfyModelRootRuntime: + """Read the active root from ComfyUI's initialized folder registry.""" + + def active_model_root(self) -> Path: + """Return ComfyUI's normalized active model root.""" + + folder_paths = cast("_FolderPathsModule", importlib.import_module("folder_paths")) + return Path(folder_paths.models_dir).resolve() + + +__all__ = ["ComfyModelRootRuntime", "ModelRootRuntime"] diff --git a/substitute_backend/features/environment_management/infrastructure/model_root_store.py b/substitute_backend/features/environment_management/infrastructure/model_root_store.py new file mode 100644 index 0000000..f38fc4a --- /dev/null +++ b/substitute_backend/features/environment_management/infrastructure/model_root_store.py @@ -0,0 +1,110 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Persist the Comfy installation's authoritative model-root selection.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from uuid import uuid4 + +_CONFIG_DIRECTORY = ".substitute" +_CONFIG_FILE = "model_root.json" +_SCHEMA_VERSION = 1 + + +class ModelRootStore: + """Read and atomically replace model-root configuration for one Comfy root.""" + + def __init__(self, comfy_root: Path) -> None: + """Initialize persistence beneath the supplied Comfy installation.""" + + self._comfy_root = comfy_root.resolve() + self._config_path = self._comfy_root / _CONFIG_DIRECTORY / _CONFIG_FILE + + @property + def config_path(self) -> Path: + """Return the canonical persisted configuration path.""" + + return self._config_path + + def load(self) -> Path | None: + """Return the configured custom root, or ``None`` for Comfy's default.""" + + if not self._config_path.exists(): + return None + payload = json.loads(self._config_path.read_text(encoding="utf-8")) + if payload.get("schemaVersion") != _SCHEMA_VERSION: + raise ValueError("Unsupported model-root configuration schema.") + value = payload.get("modelRoot") + if not isinstance(value, str) or not value.strip(): + raise ValueError("Model-root configuration does not contain a path.") + return self.resolve_custom_root(value) + + def save(self, model_root: Path | None) -> Path | None: + """Persist a custom root, or remove configuration for Comfy's default.""" + + if model_root is None: + self._remove_config() + return None + resolved_root = self.resolve_custom_root(str(model_root)) + resolved_root.mkdir(parents=True, exist_ok=True) + if not resolved_root.is_dir(): + raise ValueError("Model root must be a directory.") + self._config_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = self._config_path.with_name(f".{self._config_path.name}.{uuid4().hex}.tmp") + try: + temporary_path.write_text( + json.dumps( + { + "schemaVersion": _SCHEMA_VERSION, + "modelRoot": str(resolved_root), + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + os.replace(temporary_path, self._config_path) + finally: + if temporary_path.exists(): + temporary_path.unlink() + return resolved_root + + @staticmethod + def resolve_custom_root(value: str) -> Path: + """Validate and normalize one host-side custom model root.""" + + expanded = Path(os.path.expandvars(value)).expanduser() + if not expanded.is_absolute(): + raise ValueError("Model root must be an absolute path.") + resolved = expanded.resolve() + if resolved.exists() and not resolved.is_dir(): + raise ValueError("Model root must be a directory.") + return resolved + + def _remove_config(self) -> None: + """Remove the custom selection while preserving unrelated host state.""" + + if self._config_path.exists(): + self._config_path.unlink() + directory = self._config_path.parent + if directory.exists() and not any(directory.iterdir()): + directory.rmdir() + + +__all__ = ["ModelRootStore"] diff --git a/substitute_backend/host/extension.py b/substitute_backend/host/extension.py index 0339a87..a4de828 100644 --- a/substitute_backend/host/extension.py +++ b/substitute_backend/host/extension.py @@ -62,6 +62,9 @@ from substitute_backend.features.environment_management.application.inventory_se from substitute_backend.features.environment_management.application.job_service import ( JobService, ) +from substitute_backend.features.environment_management.application.model_root_service import ( + ModelRootService, +) from substitute_backend.features.environment_management.application.restart_service import ( RestartService, ) @@ -75,6 +78,12 @@ from substitute_backend.features.environment_management.infrastructure import ( PipInspector, ) from substitute_backend.features.environment_management.infrastructure.job_store import JobStore +from substitute_backend.features.environment_management.infrastructure.model_root_runtime import ( + ComfyModelRootRuntime, +) +from substitute_backend.features.environment_management.infrastructure.model_root_store import ( + ModelRootStore, +) from substitute_backend.features.environment_management.infrastructure.python_environment import ( PythonEnvironmentInspector, ) @@ -367,6 +376,11 @@ def build_environment_management_services(extension_root: Path) -> EnvironmentMa coordinator=restart_coordinator, logger=get_logger("environment.restart"), ), + model_root=ModelRootService( + comfy_root=comfy_root, + store=ModelRootStore(comfy_root), + runtime=ComfyModelRootRuntime(), + ), ) diff --git a/substitute_backend/host/routes.py b/substitute_backend/host/routes.py index ee03f0e..9c654f1 100644 --- a/substitute_backend/host/routes.py +++ b/substitute_backend/host/routes.py @@ -91,6 +91,9 @@ class RouteRegistrar(Protocol): def post(self, path: str) -> Callable[[_RouteHandler], _RouteHandler]: """Return a decorator registering a POST handler.""" + def put(self, path: str) -> Callable[[_RouteHandler], _RouteHandler]: + """Return a decorator registering a PUT handler.""" + def delete(self, path: str) -> Callable[[_RouteHandler], _RouteHandler]: """Return a decorator registering a DELETE handler.""" @@ -239,6 +242,8 @@ def register_routes( routes.post("/substitute/v1/cube-library/sync-and-check")(cube_library_handlers.sync_and_check) routes.get("/substitute/v1/environment/capabilities")(environment_handlers.capabilities) routes.get("/substitute/v1/environment/status")(environment_handlers.status) + routes.get("/substitute/v1/environment/model-root")(environment_handlers.get_model_root) + routes.put("/substitute/v1/environment/model-root")(environment_handlers.update_model_root) routes.get("/substitute/v1/environment/packages")(environment_handlers.list_packages) routes.get("/substitute/v1/environment/components")(environment_handlers.list_components) routes.post("/substitute/v1/environment/operations/plan")(environment_handlers.plan_operation) diff --git a/substitute_backend_prestartup.py b/substitute_backend_prestartup.py new file mode 100644 index 0000000..d9eac50 --- /dev/null +++ b/substitute_backend_prestartup.py @@ -0,0 +1,165 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Apply BackEnd-owned model configuration before normal custom-node imports.""" + +from __future__ import annotations + +import json +import logging +import os +import shutil +from pathlib import Path +from typing import Any +from uuid import uuid4 + +_CONFIG_RELATIVE_PATH = Path(".substitute") / "model_root.json" +_LEGACY_CONFIG_RELATIVE_PATH = Path(".substitute") / "managed_model_root.json" +_LEGACY_HOOK_DIRECTORY = "SubstituteManagedModelRoot" +_SCHEMA_VERSION = 1 +_LOGGER = logging.getLogger("substitute_backend.prestartup") + + +def apply_model_root(comfy_root: Path, folder_paths: Any) -> Path: + """Apply the persisted root and return the model root active after startup.""" + + _migrate_legacy_configuration(comfy_root) + old_root = Path(str(folder_paths.models_dir)).resolve() + configured_root = _read_configured_root(comfy_root) + new_root = configured_root or (comfy_root / "models").resolve() + new_root.mkdir(parents=True, exist_ok=True) + if new_root != old_root: + _redirect_registered_paths(folder_paths, old_root, new_root) + folder_paths.models_dir = str(new_root) + _clear_caches(folder_paths) + _LOGGER.info("Substitute BackEnd configured ComfyUI model root: %s", new_root) + _remove_legacy_hook(comfy_root) + return new_root + + +def _migrate_legacy_configuration(comfy_root: Path) -> None: + """Adopt desktop-owned persisted state before the normal startup read.""" + + canonical_path = comfy_root / _CONFIG_RELATIVE_PATH + legacy_path = comfy_root / _LEGACY_CONFIG_RELATIVE_PATH + if not legacy_path.exists(): + return + if canonical_path.exists(): + legacy_path.unlink() + _LOGGER.info("Removed superseded legacy Substitute model-root configuration.") + return + payload = json.loads(legacy_path.read_text(encoding="utf-8")) + value = payload.get("model_root") + if not isinstance(value, str) or not value.strip(): + raise ValueError("Legacy Substitute model-root configuration is invalid.") + model_root = Path(value).expanduser() + if not model_root.is_absolute(): + raise ValueError("Legacy Substitute model root must be absolute.") + _write_configuration(canonical_path, model_root.resolve()) + legacy_path.unlink() + _LOGGER.info("Migrated legacy Substitute model-root configuration.") + + +def _write_configuration(config_path: Path, model_root: Path) -> None: + """Atomically write the lightweight prestartup configuration shape.""" + + config_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = config_path.with_name(f".{config_path.name}.{uuid4().hex}.tmp") + try: + temporary_path.write_text( + json.dumps( + {"schemaVersion": _SCHEMA_VERSION, "modelRoot": str(model_root)}, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + os.replace(temporary_path, config_path) + finally: + if temporary_path.exists(): + temporary_path.unlink() + + +def _remove_legacy_hook(comfy_root: Path) -> None: + """Remove the superseded desktop-generated custom node after migration.""" + + hook_directory = comfy_root / "custom_nodes" / _LEGACY_HOOK_DIRECTORY + if not hook_directory.exists(): + return + try: + shutil.rmtree(hook_directory) + except OSError as exc: + _LOGGER.warning( + "Could not remove legacy Substitute model-root hook: %s", + exc, + ) + return + _LOGGER.info("Removed legacy Substitute model-root startup hook.") + + +def _read_configured_root(comfy_root: Path) -> Path | None: + """Read the lightweight configuration without importing BackEnd services.""" + + config_path = comfy_root / _CONFIG_RELATIVE_PATH + if not config_path.exists(): + return None + payload = json.loads(config_path.read_text(encoding="utf-8")) + if payload.get("schemaVersion") != _SCHEMA_VERSION: + raise ValueError("Unsupported Substitute BackEnd model-root schema.") + value = payload.get("modelRoot") + if not isinstance(value, str) or not value.strip(): + raise ValueError("Substitute BackEnd model-root configuration is invalid.") + configured_root = Path(value).expanduser() + if not configured_root.is_absolute(): + raise ValueError("Substitute BackEnd model root must be absolute.") + return configured_root.resolve() + + +def _redirect_registered_paths( + folder_paths: Any, + old_root: Path, + new_root: Path, +) -> None: + """Move default-root registry entries while preserving external additions.""" + + registry = getattr(folder_paths, "folder_names_and_paths", {}) + for folder_name, values in list(registry.items()): + paths, extensions = values + rewritten: list[str] = [] + changed = False + for raw_path in paths: + path = Path(str(raw_path)).resolve() + try: + relative = path.relative_to(old_root) + except ValueError: + rewritten.append(str(raw_path)) + continue + rewritten.append(str(new_root / relative)) + changed = True + if changed: + registry[folder_name] = (rewritten, extensions) + + +def _clear_caches(folder_paths: Any) -> None: + """Clear supported Comfy filename caches after registry redirection.""" + + for attribute in ("filename_list_cache", "cache_helper"): + cache = getattr(folder_paths, attribute, None) + clear = getattr(cache, "clear", None) + if callable(clear): + clear() + + +__all__ = ["apply_model_root"] diff --git a/tests/test_capabilities_and_routes.py b/tests/test_capabilities_and_routes.py index f3a81aa..58a181a 100644 --- a/tests/test_capabilities_and_routes.py +++ b/tests/test_capabilities_and_routes.py @@ -68,6 +68,11 @@ class FakeRoutes: return self._record("POST", path) + def put(self, path: str) -> Callable[[_RouteHandler], _RouteHandler]: + """Record a PUT route registration.""" + + return self._record("PUT", path) + def delete(self, path: str) -> Callable[[_RouteHandler], _RouteHandler]: """Record a DELETE route registration.""" @@ -152,6 +157,8 @@ def test_register_routes_uses_expected_surface(tmp_path: Path) -> None: ("POST", "/substitute/v1/cube-library/sync-and-check"), ("GET", "/substitute/v1/environment/capabilities"), ("GET", "/substitute/v1/environment/status"), + ("GET", "/substitute/v1/environment/model-root"), + ("PUT", "/substitute/v1/environment/model-root"), ("GET", "/substitute/v1/environment/packages"), ("GET", "/substitute/v1/environment/components"), ("POST", "/substitute/v1/environment/operations/plan"), diff --git a/tests/test_download_telemetry.py b/tests/test_download_telemetry.py index 21afa91..3f1106c 100644 --- a/tests/test_download_telemetry.py +++ b/tests/test_download_telemetry.py @@ -637,6 +637,11 @@ class _FakeRoutes: return self._record("POST", path) + def put(self, path: str) -> Any: + """Record a PUT handler.""" + + return self._record("PUT", path) + def delete(self, path: str) -> Any: """Record a DELETE handler.""" diff --git a/tests/test_environment_management.py b/tests/test_environment_management.py index d52ead3..52cf3b1 100644 --- a/tests/test_environment_management.py +++ b/tests/test_environment_management.py @@ -46,6 +46,9 @@ from substitute_backend.features.environment_management.application.inventory_se from substitute_backend.features.environment_management.application.job_service import ( JobService, ) +from substitute_backend.features.environment_management.application.model_root_service import ( + ModelRootService, +) from substitute_backend.features.environment_management.application.restart_service import ( RestartService, ) @@ -79,6 +82,12 @@ from substitute_backend.features.environment_management.infrastructure import ( from substitute_backend.features.environment_management.infrastructure.job_store import ( JobStore, ) +from substitute_backend.features.environment_management.infrastructure.model_root_runtime import ( + ModelRootRuntime, +) +from substitute_backend.features.environment_management.infrastructure.model_root_store import ( + ModelRootStore, +) from substitute_backend.features.environment_management.infrastructure.python_environment import ( PythonEnvironmentInspector, ) @@ -129,6 +138,20 @@ class UnsupportedRestartCoordinator(RestartCoordinator): ) +class StaticModelRootRuntime(ModelRootRuntime): + """Return a fixed active model root for service tests.""" + + def __init__(self, active_root: Path) -> None: + """Store the active root.""" + + self._active_root = active_root + + def active_model_root(self) -> Path: + """Return the configured active root.""" + + return self._active_root.resolve() + + class StaticPipInspector(PipInspector): """Return fixed pip packages for inventory tests.""" @@ -1007,6 +1030,47 @@ def _maintenance_plan_service( ) +def test_model_root_routes_report_and_persist_host_state(tmp_path: Path) -> None: + """Environment routes expose a typed model-root mutation contract.""" + + services = _environment_services(tmp_path, StaticPipInspector()) + handlers = build_environment_route_handlers( + services, + get_logger("tests.environment.routes"), + ) + custom_root = tmp_path / "shared-models" + + async def run_routes() -> None: + """Exercise model-root GET and update handlers.""" + + initial = await handlers.get_model_root(cast("web.Request", object())) + initial_payload = _response_payload(initial) + assert initial_payload["usesDefault"] is True + assert initial_payload["restartRequired"] is False + + updated = await handlers.update_model_root( + cast( + "web.Request", + FakeJsonRequest({"mode": "custom", "path": str(custom_root)}), + ) + ) + updated_payload = _response_payload(updated) + assert updated_payload["configuredModelRoot"] == str(custom_root.resolve()) + assert updated_payload["restartRequired"] is True + + invalid = await handlers.update_model_root( + cast( + "web.Request", + FakeJsonRequest({"mode": "custom", "path": "relative/models"}), + ) + ) + assert invalid.status == 400 + error = cast("JsonObject", _response_payload(invalid)["error"]) + assert error["code"] == "invalid-model-root" + + asyncio.run(run_routes()) + + def _expected_triton_requirement() -> str: """Return the platform-specific Triton requirement expected by the service.""" @@ -1048,6 +1112,11 @@ def _environment_services( coordinator=SupportedRestartCoordinator(), logger=get_logger("tests.environment.restart"), ), + model_root=ModelRootService( + comfy_root=tmp_path, + store=ModelRootStore(tmp_path), + runtime=StaticModelRootRuntime(tmp_path / "models"), + ), ) diff --git a/tests/test_model_root_management.py b/tests/test_model_root_management.py new file mode 100644 index 0000000..41203e1 --- /dev/null +++ b/tests/test_model_root_management.py @@ -0,0 +1,227 @@ +# Substitute BackEnd - backend liaison services for SugarSubstitute and ComfyUI +# Copyright (C) 2026 Artificial Sweetener and contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +"""Tests for authoritative Comfy model-root management.""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import ModuleType +from typing import Any, cast + +import pytest + +from substitute_backend.features.environment_management.application.model_root_service import ( + ModelRootService, +) +from substitute_backend.features.environment_management.domain.model_root import ModelRootMode +from substitute_backend.features.environment_management.infrastructure.model_root_runtime import ( + ModelRootRuntime, +) +from substitute_backend.features.environment_management.infrastructure.model_root_store import ( + ModelRootStore, +) +from substitute_backend_prestartup import apply_model_root + + +class MutableModelRootRuntime(ModelRootRuntime): + """Expose a mutable active root for restart-transition tests.""" + + def __init__(self, active_root: Path) -> None: + """Initialize the active root.""" + + self.active_root = active_root + + def active_model_root(self) -> Path: + """Return the current active root.""" + + return self.active_root.resolve() + + +class ClearTracker: + """Record Comfy cache invalidation calls.""" + + def __init__(self) -> None: + """Initialize an uncleared tracker.""" + + self.cleared = False + + def clear(self) -> None: + """Record cache invalidation.""" + + self.cleared = True + + +def test_store_defaults_to_comfy_models_without_configuration(tmp_path: Path) -> None: + """An unconfigured host uses its own models directory.""" + + runtime = MutableModelRootRuntime(tmp_path / "models") + service = ModelRootService(tmp_path, ModelRootStore(tmp_path), runtime) + + status = service.get_status() + + assert status.uses_default is True + assert status.default_model_root == (tmp_path / "models").resolve() + assert status.active_model_root == (tmp_path / "models").resolve() + assert status.restart_required is False + + +def test_custom_root_is_persisted_atomically_and_requires_restart(tmp_path: Path) -> None: + """A custom selection is durable while active state remains process-owned.""" + + runtime = MutableModelRootRuntime(tmp_path / "models") + store = ModelRootStore(tmp_path) + service = ModelRootService(tmp_path, store, runtime) + custom_root = tmp_path / "shared models" + + status = service.configure(ModelRootMode.CUSTOM, str(custom_root)) + + assert status.configured_model_root == custom_root.resolve() + assert status.active_model_root == (tmp_path / "models").resolve() + assert status.restart_required is True + assert store.load() == custom_root.resolve() + assert not tuple(store.config_path.parent.glob("*.tmp")) + + runtime.active_root = custom_root + assert service.get_status().restart_required is False + + +def test_default_mode_removes_only_owned_configuration(tmp_path: Path) -> None: + """Resetting to default preserves unrelated Substitute host state.""" + + store = ModelRootStore(tmp_path) + store.save(tmp_path / "shared") + unrelated = store.config_path.parent / "other.json" + unrelated.write_text("{}", encoding="utf-8") + service = ModelRootService( + tmp_path, + store, + MutableModelRootRuntime(tmp_path / "shared"), + ) + + status = service.configure(ModelRootMode.DEFAULT, None) + + assert status.uses_default is True + assert status.restart_required is True + assert not store.config_path.exists() + assert unrelated.exists() + + +@pytest.mark.parametrize("value", ["models", "", "."]) +def test_custom_root_rejects_relative_paths(tmp_path: Path, value: str) -> None: + """Host configuration never interprets relative paths ambiguously.""" + + service = ModelRootService( + tmp_path, + ModelRootStore(tmp_path), + MutableModelRootRuntime(tmp_path / "models"), + ) + + with pytest.raises(ValueError, match=r"absolute|requires"): + service.configure(ModelRootMode.CUSTOM, value) + + +def test_prestartup_redirects_default_categories_and_preserves_external_paths( + tmp_path: Path, +) -> None: + """Prestartup moves Comfy defaults without consuming extra model paths.""" + + comfy_root = tmp_path / "ComfyUI" + old_root = comfy_root / "models" + new_root = tmp_path / "shared-models" + external_root = tmp_path / "external-loras" + config_path = comfy_root / ".substitute" / "model_root.json" + config_path.parent.mkdir(parents=True) + config_path.write_text( + json.dumps({"schemaVersion": 1, "modelRoot": str(new_root)}), + encoding="utf-8", + ) + folder_paths = cast(Any, ModuleType("folder_paths")) + folder_paths.models_dir = str(old_root) + folder_paths.folder_names_and_paths = { + "checkpoints": ([str(old_root / "checkpoints")], {".safetensors"}), + "loras": ( + [str(old_root / "loras"), str(external_root)], + {".safetensors"}, + ), + } + filename_cache = ClearTracker() + helper_cache = ClearTracker() + folder_paths.filename_list_cache = filename_cache + folder_paths.cache_helper = helper_cache + + active_root = apply_model_root(comfy_root, folder_paths) + + assert active_root == new_root.resolve() + assert Path(folder_paths.models_dir) == new_root.resolve() + assert folder_paths.folder_names_and_paths["checkpoints"][0] == [str(new_root / "checkpoints")] + assert folder_paths.folder_names_and_paths["loras"][0] == [ + str(new_root / "loras"), + str(external_root), + ] + assert filename_cache.cleared is True + assert helper_cache.cleared is True + + +def test_prestartup_uses_default_without_configuration(tmp_path: Path) -> None: + """The BackEnd hook does not require SugarSubstitute launch state.""" + + comfy_root = tmp_path / "ComfyUI" + folder_paths = cast(Any, ModuleType("folder_paths")) + folder_paths.models_dir = str(comfy_root / "models") + folder_paths.folder_names_and_paths = {} + + assert apply_model_root(comfy_root, folder_paths) == (comfy_root / "models").resolve() + + +def test_prestartup_migrates_legacy_desktop_state_and_hook(tmp_path: Path) -> None: + """BackEnd adopts the previous owner before removing its generated hook.""" + + comfy_root = tmp_path / "ComfyUI" + legacy_config = comfy_root / ".substitute" / "managed_model_root.json" + legacy_hook = comfy_root / "custom_nodes" / "SubstituteManagedModelRoot" + custom_root = tmp_path / "legacy-models" + legacy_config.parent.mkdir(parents=True) + legacy_config.write_text( + json.dumps({"schema_version": 1, "model_root": str(custom_root)}), + encoding="utf-8", + ) + legacy_hook.mkdir(parents=True) + (legacy_hook / "prestartup_script.py").write_text("", encoding="utf-8") + folder_paths = cast(Any, ModuleType("folder_paths")) + folder_paths.models_dir = str(comfy_root / "models") + folder_paths.folder_names_and_paths = {} + + assert apply_model_root(comfy_root, folder_paths) == custom_root.resolve() + assert not legacy_config.exists() + assert not legacy_hook.exists() + assert ModelRootStore(comfy_root).load() == custom_root.resolve() + + +def test_prestartup_keeps_canonical_selection_over_legacy_state(tmp_path: Path) -> None: + """Canonical BackEnd state should win and remove a leftover desktop file.""" + + comfy_root = tmp_path / "ComfyUI" + canonical_root = tmp_path / "canonical-models" + ModelRootStore(comfy_root).save(canonical_root) + legacy_config = comfy_root / ".substitute" / "managed_model_root.json" + legacy_config.write_text("not valid anymore", encoding="utf-8") + folder_paths = cast(Any, ModuleType("folder_paths")) + folder_paths.models_dir = str(comfy_root / "models") + folder_paths.folder_names_and_paths = {} + + assert apply_model_root(comfy_root, folder_paths) == canonical_root.resolve() + assert not legacy_config.exists()