From 2320b0b9ee59b9a59062cbabd9c97b8cdd42c4e4 Mon Sep 17 00:00:00 2001 From: DrHepa Date: Mon, 21 Sep 2026 19:50:54 +0200 Subject: [PATCH 1/2] feat(workflows): support scene model artifacts --- README.md | 12 + api/README.md | 1 + api/routers/generation.py | 107 +++++++-- api/runner.py | 50 ++++- api/schemas/generation.py | 13 +- api/services/artifact_input.py | 29 +++ api/services/extension_process.py | 55 ++++- api/services/generator_registry.py | 53 ++++- api/services/generators/base.py | 21 +- api/services/scene_input.py | 138 ++++++++++++ api/tests/test_extension_process.py | 28 +++ api/tests/test_generation_router.py | 3 + api/tests/test_generator_registry.py | 47 ++++ api/tests/test_runner.py | 25 +++ api/tests/test_scene_generation.py | 115 ++++++++++ api/tests/test_scene_input.py | 67 ++++++ .../main/artifact-registry-service.test.ts | 12 + electron/main/artifact-registry-service.ts | 3 +- .../main/extension-install-utils.test.mjs | 44 ++++ electron/main/extension-install-utils.ts | 34 ++- electron/main/ipc-handlers.ts | 15 +- src/areas/workflows/WorkflowsPage.tsx | 12 +- src/areas/workflows/mockExtensions.ts | 6 +- src/areas/workflows/nodes/ExtensionNode.tsx | 2 + src/areas/workflows/nodes/LoadSceneNode.tsx | 134 ++++++++++++ src/areas/workflows/preflight.test.mjs | 35 +++ src/areas/workflows/preflight.ts | 26 ++- src/areas/workflows/workflowRunStore.ts | 64 ++++-- src/areas/workflows/workflowSceneRun.test.mjs | 60 +++++ .../workflows/workflowSceneSource.test.mjs | 29 +++ src/areas/workflows/workflowSceneSource.ts | 205 ++++++++++++++++++ src/shared/stores/workflowsStore.ts | 2 +- src/shared/types/artifacts.ts | 14 ++ src/shared/types/electron.d.ts | 6 +- 34 files changed, 1390 insertions(+), 77 deletions(-) create mode 100644 api/services/artifact_input.py create mode 100644 api/services/scene_input.py create mode 100644 api/tests/test_scene_generation.py create mode 100644 api/tests/test_scene_input.py create mode 100644 src/areas/workflows/nodes/LoadSceneNode.tsx create mode 100644 src/areas/workflows/workflowSceneRun.test.mjs create mode 100644 src/areas/workflows/workflowSceneSource.test.mjs create mode 100644 src/areas/workflows/workflowSceneSource.ts diff --git a/README.md b/README.md index b162cf23..95a0b6ff 100644 --- a/README.md +++ b/README.md @@ -153,6 +153,18 @@ original behavior. ## Workflows Start with a basic workflow first. For example, on the "Workflows" tab, try: Image -> Generate Mesh -> Add to Scene. Make sure there is a connection between each of the steps. Go to the "Generate" tab, make sure the workflow is selected, then click on "Generate 3D Model". Click on "Settings/Logs/Errors" to see any issues. +Model extensions may also declare `scene` as a node input or output. A scene is +a workspace directory containing `scene-manifest.json` with schema +`modly.scene-manifest.v1`; it is not an arbitrary JSON file. Use the **Load +Scene** workflow node to select and validate an existing scene directory. +Scene-capable generators implement `generate_artifact(input_kind, +artifact_path, ...)`; legacy image generators and `POST /generate/from-image` +remain unchanged. The generic `POST /generate/from-artifact` boundary currently +accepts only `scene`, leaving future artifact kinds to separate reviewed changes. +For this first contract, `scene` is model-only and must be declared as the single +`input` value (not inside `inputs`); process and mixed-input scene nodes are rejected. +Model nodes may still accept multiple images and produce a scene. + ## Modly CLI diff --git a/api/README.md b/api/README.md index ee45bc3f..cdcd4d6b 100644 --- a/api/README.md +++ b/api/README.md @@ -29,6 +29,7 @@ uvicorn main:app --host 127.0.0.1 --port 8765 --reload | GET | `/model/status` | Model download / load status | | GET | `/model/download` | SSE stream of download progress | | POST | `/generate/from-image` | Start image-to-3D job | +| POST | `/generate/from-artifact` | Start a typed-artifact model job (`scene` only) | | GET | `/generate/status/{job_id}` | Poll job status | ## Model diff --git a/api/routers/generation.py b/api/routers/generation.py index 8481deb4..7c355014 100644 --- a/api/routers/generation.py +++ b/api/routers/generation.py @@ -4,7 +4,8 @@ import time import traceback import uuid -from typing import Dict +from pathlib import Path +from typing import Dict, Optional, Union from fastapi import APIRouter, File, Form, UploadFile, HTTPException, BackgroundTasks from services.generators.base import smooth_progress, GenerationCancelled @@ -14,7 +15,8 @@ # binding captured at import would keep writing output to the old directory. import services.generator_registry as registry from services.generator_registry import generator_registry -from schemas.generation import JobStatus +from schemas.generation import GenerateFromArtifactRequest, JobStatus +from services.artifact_input import TypedArtifactInput, validate_artifact_input router = APIRouter(tags=["generation"]) @@ -105,6 +107,7 @@ async def generate_from_image( # Verify the requested model exists in the registry try: generator_registry.get_generator(model_id) + output_kind = generator_registry.get_manifest(model_id).get("output", "mesh") except ValueError as e: raise HTTPException(400, str(e)) @@ -131,11 +134,49 @@ async def generate_from_image( _jobs[job_id] = job _cancel_events[job_id] = threading.Event() - background_tasks.add_task(_run_generation, job_id, image_bytes, full_params, collection) + background_tasks.add_task( + _run_generation, job_id, image_bytes, full_params, collection, output_kind, model_id + ) return {"job_id": job_id} +_RESERVED_ARTIFACT_PARAMS = { + "artifact_path", "input_kind", "input_path", "scene_path", "scene_manifest_path", +} + + +@router.post("/from-artifact") +async def generate_from_artifact( + request: GenerateFromArtifactRequest, + background_tasks: BackgroundTasks, +): + """Queue a validated typed artifact without serializing it as image bytes.""" + try: + manifest = generator_registry.get_manifest(request.model_id) + except (KeyError, ValueError) as exc: + raise HTTPException(400, str(exc)) from exc + declared = manifest.get("inputs") or [manifest.get("input", "image")] + if request.input_kind not in declared: + raise HTTPException(400, f"Model {request.model_id} does not accept {request.input_kind} input") + try: + artifact = validate_artifact_input(registry.WORKSPACE_DIR, request.input_kind, request.input_path) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + + params = {k: v for k, v in request.params.items() if k not in _RESERVED_ARTIFACT_PARAMS} + params["scene_manifest_path"] = str(artifact.path) + collection = sanitize_collection(request.collection) + job_id = str(uuid.uuid4()) + _purge_old_jobs() + _jobs[job_id] = JobStatus(job_id=job_id, status="pending", progress=0) + _cancel_events[job_id] = threading.Event() + background_tasks.add_task( + _run_generation, job_id, artifact, params, collection, manifest.get("output", "mesh"), request.model_id + ) + return {"job_id": job_id} + + @router.get("/status/{job_id}") async def job_status(job_id: str): @@ -170,7 +211,14 @@ async def cancel_job(job_id: str): return {"cancelled": True} -async def _run_generation(job_id: str, image_bytes: bytes, params: dict, collection: str = "Default") -> None: +async def _run_generation( + job_id: str, + model_input: Union[bytes, TypedArtifactInput], + params: dict, + collection: str = "Default", + output_kind: str = "mesh", + model_id: Optional[str] = None, +) -> None: job = _jobs[job_id] job.status = "running" @@ -189,8 +237,14 @@ def progress_cb(pct: int, step: str = "") -> None: # Check if the model needs to be loaded BEFORE calling get_active(), # because get_active() loads the model in a blocking manner. # active_status() is an instantaneous operation (simple dict lookup). - if not generator_registry.active_status()["loaded"]: - active = generator_registry.active_status() + get_generator = (lambda: generator_registry.get_ready_generator(model_id)) \ + if model_id is not None else generator_registry.get_active + status_reader = getattr(generator_registry, "model_status", None) + status = (status_reader(model_id) if model_id is not None and status_reader + else generator_registry.active_status() if model_id is None + else {"name": model_id, "downloaded": True, "loaded": False}) + if not status["loaded"]: + active = status model_name = active['name'] init_label = f"Downloading {model_name}…" if not active['downloaded'] else f"Loading {model_name}…" progress_cb(0, init_label) @@ -202,11 +256,11 @@ def progress_cb(pct: int, step: str = "") -> None: ) load_thread.start() try: - gen = await loop.run_in_executor(None, generator_registry.get_active) + gen = await loop.run_in_executor(None, get_generator) finally: stop_load_evt.set() else: - gen = await loop.run_in_executor(None, generator_registry.get_active) + gen = await loop.run_in_executor(None, get_generator) if job_id in _cancelled: return @@ -217,18 +271,39 @@ def progress_cb(pct: int, step: str = "") -> None: gen.outputs_dir = coll_dir cancel_event = _cancel_events.get(job_id) - import inspect - supports_cancel = "cancel_event" in inspect.signature(gen.generate).parameters - output_path = await loop.run_in_executor( - None, - lambda: gen.generate(image_bytes, params, progress_cb, cancel_event) - if supports_cancel - else gen.generate(image_bytes, params, progress_cb), - ) + if isinstance(model_input, TypedArtifactInput): + # Revalidate just before crossing the inference boundary. The + # subprocess runner repeats this check inside the worker. + from services.artifact_input import revalidate_artifact_input + model_input = revalidate_artifact_input(registry.WORKSPACE_DIR, model_input) + import inspect + supports_cancel = "cancel_event" in inspect.signature(gen.generate_artifact).parameters + output_path = await loop.run_in_executor( + None, + lambda: gen.generate_artifact(model_input.kind, model_input.path, params, progress_cb, cancel_event) + if supports_cancel else gen.generate_artifact(model_input.kind, model_input.path, params, progress_cb), + ) + else: + import inspect + supports_cancel = "cancel_event" in inspect.signature(gen.generate).parameters + output_path = await loop.run_in_executor( + None, + lambda: gen.generate(model_input, params, progress_cb, cancel_event) + if supports_cancel else gen.generate(model_input, params, progress_cb), + ) if job_id in _cancelled: return + output_path = Path(output_path).resolve(strict=True) + if output_kind == "scene": + from services.scene_input import validate_scene_input + try: + output_relative = output_path.relative_to(registry.WORKSPACE_DIR.resolve()) + output_path = validate_scene_input(registry.WORKSPACE_DIR, output_relative.as_posix()) + except (OSError, ValueError) as exc: + raise ValueError("Generated scene output is not a valid workspace scene") from exc + job.status = "done" job.progress = 100 _completed_at[job_id] = time.monotonic() diff --git a/api/runner.py b/api/runner.py index 0dd21392..0dd17ddb 100644 --- a/api/runner.py +++ b/api/runner.py @@ -151,6 +151,32 @@ def _apply_manifest_metadata(gen, manifest: dict, node: dict) -> None: gen._params_schema = node.get("params_schema") or manifest.get("params_schema", []) +def decode_model_input(msg: dict): + """Decode legacy image bytes or revalidate a typed artifact in the worker.""" + if "input" not in msg: + return base64.b64decode(msg["image_b64"]) + value = msg["input"] + if not isinstance(value, dict) or set(value) != {"kind", "path"}: + raise ValueError("Typed artifact input must contain exactly kind and path") + from services.artifact_input import TypedArtifactInput, revalidate_artifact_input + typed = TypedArtifactInput(kind=value.get("kind"), path=Path(value.get("path", ""))) + return revalidate_artifact_input(WORKSPACE_DIR, typed) + + +def validate_requested_model(msg: dict, manifest: dict, node: dict) -> None: + """Reject requests routed to a worker for a different manifest node.""" + requested = msg.get("model_id") + if requested is None: # Backward compatibility with already-running legacy hosts. + return + expected = manifest["id"] + if node.get("id"): + expected = f"{expected}/{node['id']}" + if requested != expected: + raise ValueError( + f"Generation request model '{requested}' does not match worker model '{expected}'" + ) + + # ------------------------------------------------------------------ # # Main loop # ------------------------------------------------------------------ # @@ -201,10 +227,17 @@ def main() -> None: # ---- generate -------------------------------------------- elif action == "generate": + validate_requested_model(msg, manifest, node) cancel_evt = threading.Event() _cancel[rid] = cancel_evt - image_bytes = base64.b64decode(msg["image_b64"]) + model_input = decode_model_input(msg) params = msg.get("params", {}) + if hasattr(model_input, "kind"): + if not isinstance(params, dict): + raise ValueError("Model params must be an object") + reserved = {"artifact_path", "input_kind", "input_path", "scene_path", "scene_manifest_path"} + params = {key: value for key, value in params.items() if key not in reserved} + params["scene_manifest_path"] = str(model_input.path) if msg.get("outputs_dir"): gen.outputs_dir = Path(msg["outputs_dir"]) gen.outputs_dir.mkdir(parents=True, exist_ok=True) @@ -217,7 +250,20 @@ def progress_cb(pct: int, step: str = "") -> None: send({"type": "log", "level": "warning", "message": ("Model was not loaded (earlier setup failure?); " "reloaded before generating.")}) - output_path = gen.generate(image_bytes, params, progress_cb, cancel_evt) + if hasattr(model_input, "kind"): + output_path = gen.generate_artifact( + model_input.kind, model_input.path, params, progress_cb, cancel_evt + ) + else: + output_path = gen.generate(model_input, params, progress_cb, cancel_evt) + if node.get("output") == "scene": + from services.scene_input import validate_scene_input + resolved_output = Path(output_path).resolve(strict=True) + try: + relative_output = resolved_output.relative_to(WORKSPACE_DIR.resolve(strict=True)) + except (OSError, ValueError) as exc: + raise ValueError("Generated scene output is outside the workspace") from exc + output_path = validate_scene_input(WORKSPACE_DIR, relative_output.as_posix()) send({"type": "done", "id": rid, "output_path": str(output_path)}) except Exception as exc: # Detect GenerationCancelled by name to avoid import issues diff --git a/api/schemas/generation.py b/api/schemas/generation.py index 7ed6ca62..04c18c85 100644 --- a/api/schemas/generation.py +++ b/api/schemas/generation.py @@ -1,5 +1,5 @@ -from typing import Literal, Optional -from pydantic import BaseModel +from typing import Any, Literal, Optional +from pydantic import BaseModel, Field class JobStatus(BaseModel): @@ -9,3 +9,12 @@ class JobStatus(BaseModel): step: Optional[str] = None # Human-readable current step output_url: Optional[str] = None error: Optional[str] = None + + +class GenerateFromArtifactRequest(BaseModel): + """Generic typed-artifact request. Only scene is public in this release.""" + input_kind: Literal["scene"] + input_path: str + model_id: str + collection: str = "Workflows" + params: dict[str, Any] = Field(default_factory=dict) diff --git a/api/services/artifact_input.py b/api/services/artifact_input.py new file mode 100644 index 00000000..165e29ec --- /dev/null +++ b/api/services/artifact_input.py @@ -0,0 +1,29 @@ +"""Typed model artifact inputs shared by the API, worker bridge, and runner.""" +from dataclasses import dataclass +from pathlib import Path + +from services.scene_input import validate_scene_input + +SUPPORTED_ARTIFACT_INPUTS = frozenset({"scene"}) + + +@dataclass(frozen=True) +class TypedArtifactInput: + kind: str + path: Path + + +def validate_artifact_input(workspace: Path, kind: str, input_path: str) -> TypedArtifactInput: + if kind not in SUPPORTED_ARTIFACT_INPUTS: + raise ValueError(f"Unsupported artifact input kind: {kind}") + if kind == "scene": + return TypedArtifactInput(kind="scene", path=validate_scene_input(workspace, input_path)) + raise ValueError(f"Unsupported artifact input kind: {kind}") + + +def revalidate_artifact_input(workspace: Path, value: TypedArtifactInput) -> TypedArtifactInput: + try: + relative = value.path.resolve(strict=True).relative_to(workspace.resolve(strict=True)) + except (OSError, ValueError) as exc: + raise ValueError("Artifact input is outside the workspace") from exc + return validate_artifact_input(workspace, value.kind, relative.as_posix()) diff --git a/api/services/extension_process.py b/api/services/extension_process.py index 67565d36..3a39bfc5 100644 --- a/api/services/extension_process.py +++ b/api/services/extension_process.py @@ -291,16 +291,18 @@ def generate( progress_cb: Optional[Callable[[int, str], None]] = None, cancel_event: Optional[threading.Event] = None, ) -> Path: - from services.generators.base import GenerationCancelled + return self._generate_request( + {"image_b64": base64.b64encode(image_bytes).decode()}, + params, progress_cb, cancel_event, + ) - req_id = str(uuid.uuid4()) - self._send({ - "action": "generate", - "id": req_id, - "image_b64": base64.b64encode(image_bytes).decode(), - "params": params, - "outputs_dir": str(self.outputs_dir) if self.outputs_dir else None, - }) + def _receive_generation( + self, + req_id: str, + progress_cb: Optional[Callable[[int, str], None]], + cancel_event: Optional[threading.Event], + ) -> Path: + from services.generators.base import GenerationCancelled # Grace period after sending a cooperative cancel before hard-killing # the subprocess. Long enough to let generators that check cancel_event @@ -377,6 +379,41 @@ def generate( elif t == "log": print(f"[{self.MODEL_ID}] {msg.get('message', '')}", file=sys.stderr) + def generate_artifact( + self, + input_kind: str, + artifact_path: Path, + params: dict, + progress_cb: Optional[Callable[[int, str], None]] = None, + cancel_event: Optional[threading.Event] = None, + ) -> Path: + """Send a typed artifact envelope to the isolated runner.""" + from services.artifact_input import TypedArtifactInput, revalidate_artifact_input + from services.generator_registry import WORKSPACE_DIR + + validated = revalidate_artifact_input( + WORKSPACE_DIR, TypedArtifactInput(kind=input_kind, path=artifact_path) + ) + return self._generate_request( + {"input": {"kind": validated.kind, "path": str(validated.path)}}, + params, progress_cb, cancel_event, + ) + + def _generate_request( + self, + input_payload: dict, + params: dict, + progress_cb: Optional[Callable[[int, str], None]], + cancel_event: Optional[threading.Event], + ) -> Path: + req_id = str(uuid.uuid4()) + self._send({ + "action": "generate", "id": req_id, "model_id": self.MODEL_ID, + **input_payload, "params": params, + "outputs_dir": str(self.outputs_dir) if self.outputs_dir else None, + }) + return self._receive_generation(req_id, progress_cb, cancel_event) + def params_schema(self) -> list: return self._params_schema diff --git a/api/services/generator_registry.py b/api/services/generator_registry.py index 348a42cb..de1750ca 100644 --- a/api/services/generator_registry.py +++ b/api/services/generator_registry.py @@ -452,6 +452,25 @@ def _discover_extensions( node for node in raw_nodes if isinstance(node, dict) and node.get("id") ] + allowed_io = {"image", "text", "mesh", "audio", "scene"} + for node in nodes: + declared_inputs = node.get("inputs") or [node.get("input", "image")] + if (not isinstance(declared_inputs, list) + or any(value not in allowed_io for value in declared_inputs)): + raise ValueError( + f'model node "{node.get("id", "unknown")}" has an unsupported input type' + ) + if node.get("output", "mesh") not in allowed_io: + raise ValueError( + f'model node "{node.get("id", "unknown")}" has an unsupported output type' + ) + if "scene" in declared_inputs and ( + "inputs" in node or node.get("input", "image") != "scene" + ): + raise ValueError( + f'model node "{node.get("id", "unknown")}" must declare scene ' + 'as its single input field' + ) # Markers left while setup or runtime registration is unfinished: # the folder is not ready to be loaded. The readable manifest lets @@ -540,6 +559,7 @@ def _discover_extensions( "hf_include_prefixes": node.get("hf_include_prefixes", []), "params_schema": node.get("params_schema", manifest.get("params_schema", [])), "input": node.get("input", "image"), + "inputs": node.get("inputs"), "output": node.get("output", "mesh"), } if model_sources is not None: @@ -612,6 +632,9 @@ def initialize( ) # Subprocess mode: wrap in ExtensionProcess gen = ExtensionProcess(ext_dir, manifest) + # Pin the subprocess envelope to the exact registry key; + # multi-node workers must never fall back to an extension ID. + gen.MODEL_ID = model_id gen.model_dir = MODELS_DIR / model_id gen.outputs_dir = WORKSPACE_DIR else: @@ -704,10 +727,13 @@ def _assert_not_quarantined(model_id: str) -> None: def get_active(self) -> BaseGenerator: """Returns the active generator. Downloads and loads if necessary.""" - self._assert_not_quarantined(self._active_id) - gen = self._generators[self._active_id] - downloaded = self._is_downloaded(self._active_id, gen) - if "model_sources" in self._manifests[self._active_id] and not downloaded: + return self.get_ready_generator(self._active_id) + + def get_ready_generator(self, model_id: str) -> BaseGenerator: + """Load and return exactly ``model_id`` without consulting active state.""" + gen = self.get_generator(model_id) + downloaded = self._is_downloaded(model_id, gen) + if "model_sources" in self._manifests[model_id] and not downloaded: raise RuntimeError( "Model sources are incomplete. Download this node's weights " "from the Modly Models page before generation." @@ -724,6 +750,16 @@ def get_active(self) -> BaseGenerator: gen.load() return gen + def model_status(self, model_id: str) -> dict: + gen = self.get_generator(model_id) + manifest = self._manifests[model_id] + return { + "id": model_id, + "name": manifest.get("name", gen.DISPLAY_NAME), + "downloaded": self._is_downloaded(model_id, gen), + "loaded": gen.is_loaded(), + } + def get_generator(self, model_id: str) -> BaseGenerator: self._assert_not_quarantined(model_id) if model_id not in self._generators: @@ -765,14 +801,7 @@ def _is_downloaded(self, model_id: str, gen: BaseGenerator) -> bool: return gen.is_downloaded() def active_status(self) -> dict: - gen = self._generators[self._active_id] - manifest = self._manifests[self._active_id] - return { - "id": self._active_id, - "name": manifest.get("name", gen.DISPLAY_NAME), - "downloaded": self._is_downloaded(self._active_id, gen), - "loaded": gen.is_loaded(), - } + return self.model_status(self._active_id) def all_status(self) -> list: result = [] diff --git a/api/services/generators/base.py b/api/services/generators/base.py index fd62ceef..0b5169a2 100644 --- a/api/services/generators/base.py +++ b/api/services/generators/base.py @@ -143,7 +143,6 @@ def is_loaded(self) -> bool: # Inference # ------------------------------------------------------------------ # - @abstractmethod def generate( self, image_bytes: bytes, @@ -157,7 +156,25 @@ def generate( progress_cb(percent: int, step_label: str) cancel_event: set this to interrupt generation between steps. """ - ... + raise NotImplementedError( + f"{type(self).__name__} does not implement legacy image generation" + ) + + def generate_artifact( + self, + input_kind: str, + artifact_path: Path, + params: dict, + progress_cb: Optional[Callable[[int, str], None]] = None, + cancel_event: Optional[threading.Event] = None, + ) -> Path: + """Generate from a validated typed artifact. + + New extensions should override this method. The default delegates to + ``generate`` with the canonical path so scene-capable extensions built + against the pre-release contract remain compatible. + """ + return self.generate(artifact_path, params, progress_cb, cancel_event) # type: ignore[arg-type] def _check_cancelled(self, cancel_event: Optional[threading.Event]) -> None: """Raises GenerationCancelled if cancel_event is set.""" diff --git a/api/services/scene_input.py b/api/services/scene_input.py new file mode 100644 index 00000000..c5395050 --- /dev/null +++ b/api/services/scene_input.py @@ -0,0 +1,138 @@ +"""Validation shared by the API and isolated model runner for scene inputs.""" +import json +import math +import re +import stat +from pathlib import Path, PurePosixPath, PureWindowsPath + +SCHEMA = "modly.scene-manifest.v1" +MANIFEST = "scene-manifest.json" +MAX_MANIFEST_BYTES = 1024 * 1024 +MAX_REFERENCES = 4096 +MAX_REFERENCED_FILE_BYTES = 16 * 1024**3 +MAX_REFERENCED_TOTAL_BYTES = 64 * 1024**3 + + +def _relative(value: str, *, allow_dot: bool = False) -> Path: + if not isinstance(value, str) or not value or value != value.strip() or "\x00" in value: + raise ValueError("Scene path must be a nonempty workspace-relative path") + value = value.replace("\\", "/") + if value == "." and allow_dot: + return Path(".") + if (PurePosixPath(value).is_absolute() or PureWindowsPath(value).is_absolute() + or re.match(r"^[A-Za-z][A-Za-z0-9+.-]*:", value) + or re.search(r"%(?:25|2e|2f|5c|00)", value, re.I) + or re.search(r"%(?![0-9a-f]{2})", value, re.I) + or any(part in ("", ".", "..") for part in value.split("/"))): + raise ValueError("Scene path must be a safe workspace-relative path") + return Path(*value.split("/")) + + +def _inside(path: Path, root: Path) -> Path: + try: + relative = path.relative_to(root) + except ValueError as exc: + raise ValueError("Scene path escapes its allowed root") from exc + current = root + for part in relative.parts: + current = current / part + try: + info = current.lstat() + except OSError as exc: + raise ValueError("Scene referenced path is missing or unreadable") from exc + is_reparse = bool(getattr(info, "st_file_attributes", 0) & getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0)) + if current.is_symlink() or is_reparse: + raise ValueError("Scene referenced path must not use symlinks or reparse points") + try: + resolved = path.resolve(strict=True) + except OSError as exc: + raise ValueError("Scene referenced path is missing or unreadable") from exc + if not resolved.is_relative_to(root): + raise ValueError("Scene path escapes its allowed root") + return resolved + + +def validate_scene_input(workspace: Path, scene_path: str) -> Path: + """Return a canonical manifest Path, rejecting traversal and symlink escapes. + + V1 sceneRoot is workspace-relative, except '.' denotes the manifest directory. + An asset's path and preview references are sceneRoot-relative; workspacePath + is always workspace-relative. Existence never determines which base applies. + """ + root = workspace.resolve(strict=True) + relative = _relative(scene_path) + requested = root / relative + if requested.name != MANIFEST: + if requested.suffix.lower() == ".json": + raise ValueError(f"Scene input accepts a directory or {MANIFEST}") + requested = requested / MANIFEST + try: + manifest_path = _inside(requested, root) + if not manifest_path.is_file(): + raise ValueError("Scene manifest is not a file") + if manifest_path.stat().st_size > MAX_MANIFEST_BYTES: + raise ValueError("Scene manifest exceeds the 1 MiB limit") + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise ValueError("Scene manifest is missing or invalid JSON") from exc + if not isinstance(manifest, dict) or manifest.get("schema") != SCHEMA: + raise ValueError(f"Scene manifest schema must be {SCHEMA}") + scene_root = _relative(manifest.get("sceneRoot"), allow_dot=True) + # v1 sceneRoot is workspace-relative; only '.' means the manifest's directory. + # Never select a base according to which unrelated path happens to exist. + scene_root_candidate = manifest_path.parent if scene_root == Path(".") else root / scene_root + root_path = _inside(scene_root_candidate, root) + if not root_path.is_dir(): + raise ValueError("Scene root is not a directory") + if not isinstance(manifest.get("assets"), list): + raise ValueError("Scene manifest assets must be an array") + if len(manifest["assets"]) > MAX_REFERENCES: + raise ValueError("Scene manifest contains too many asset references") + total_bytes = 0 + for asset in manifest["assets"]: + # Assets may be opaque extension metadata. Validate only references the + # host understands; never silently resolve a supplied path outside workspace. + if isinstance(asset, dict): + for field, base in (("workspacePath", root), ("path", root_path)): + if field not in asset: + continue + target = _inside(base / _relative(asset[field]), root) + if not target.is_file(): + raise ValueError(f"Scene asset {field} is not a file") + size = target.stat().st_size + if size > MAX_REFERENCED_FILE_BYTES: + raise ValueError("Scene referenced file exceeds the size limit") + total_bytes += size + if total_bytes > MAX_REFERENCED_TOTAL_BYTES: + raise ValueError("Scene referenced files exceed the total size limit") + preview = manifest.get("preview", {}) + if not isinstance(preview, dict): + raise ValueError("Scene preview must be an object") + for name in ("image", "video"): + if name in preview: + relative_preview = _relative(preview[name]) + target = _inside(root_path / relative_preview, root) + if not target.is_file(): + raise ValueError(f"Scene preview {name} is not a file") + view = manifest.get("initialView") + if view is not None: + if not isinstance(view, dict): + raise ValueError("Scene initialView must be an object") + for field in ("position", "target", "up"): + triple = view.get(field) + if triple is None and field == "up": + continue + if (not isinstance(triple, list) or len(triple) != 3 + or any(not isinstance(n, (int, float)) or isinstance(n, bool) + or not math.isfinite(n) for n in triple)): + raise ValueError(f"Scene initialView {field} must be a finite numeric triple") + if view["position"] == view["target"] or view.get("up") == [0, 0, 0]: + raise ValueError("Scene initialView has degenerate camera vectors") + return manifest_path + + +def revalidate_scene_manifest(workspace: Path, manifest_path: Path) -> Path: + """Recheck the file immediately before model invocation in the worker.""" + root = workspace.resolve(strict=True) + candidate = _inside(manifest_path, root) + return validate_scene_input(root, candidate.relative_to(root).as_posix()) diff --git a/api/tests/test_extension_process.py b/api/tests/test_extension_process.py index 348e293f..7db13e4a 100644 --- a/api/tests/test_extension_process.py +++ b/api/tests/test_extension_process.py @@ -2,6 +2,9 @@ import platform import queue import unittest +import json +import tempfile +from unittest.mock import patch from pathlib import Path from services.extension_process import ExtensionProcess, _venv_python @@ -12,6 +15,31 @@ def _make_proc() -> ExtensionProcess: class ExtensionProcessTests(unittest.TestCase): + def test_generation_envelope_pins_worker_model_id(self) -> None: + proc = _make_proc() + sent = [] + proc._send = sent.append + proc._receive_generation = lambda *args: Path("result.glb") + proc._generate_request({"image_b64": ""}, {}, None, None) + self.assertEqual(sent[0]["model_id"], "demo") + + def test_generate_artifact_sends_typed_scene_without_image_bytes(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workspace = Path(tmp) / "workspace" + scene = workspace / "Workflows" / "room" + scene.mkdir(parents=True) + manifest = scene / "scene-manifest.json" + manifest.write_text(json.dumps({ + "schema": "modly.scene-manifest.v1", "sceneRoot": ".", "assets": [], + })) + proc = _make_proc() + calls = [] + proc._generate_request = lambda payload, params, progress, cancel: calls.append((payload, params)) or manifest + with patch("services.generator_registry.WORKSPACE_DIR", workspace): + result = proc.generate_artifact("scene", manifest, {"quality": "high"}) + self.assertEqual(result, manifest) + self.assertEqual(calls, [({"input": {"kind": "scene", "path": str(manifest.resolve())}}, {"quality": "high"})]) + def test_read_loop_writes_sentinel_to_own_queue_only(self) -> None: proc = _make_proc() diff --git a/api/tests/test_generation_router.py b/api/tests/test_generation_router.py index 20fdda94..5e03c241 100644 --- a/api/tests/test_generation_router.py +++ b/api/tests/test_generation_router.py @@ -49,6 +49,9 @@ def get_active(self) -> _FakeGenerator: def get_generator(self, model_id: str) -> _FakeGenerator: return self._gen + def get_manifest(self, model_id: str) -> dict: + return {"output": "mesh"} + def switch_model(self, model_id: str) -> None: pass diff --git a/api/tests/test_generator_registry.py b/api/tests/test_generator_registry.py index ff9d090c..a4814b9e 100644 --- a/api/tests/test_generator_registry.py +++ b/api/tests/test_generator_registry.py @@ -155,6 +155,53 @@ def test_legacy_generator_supports_eager_and_lazy_sibling_imports(self) -> None: self.registry.reload() self.assertNotIn(str(extension.resolve()), sys.path) + def test_scene_io_is_registered_but_capture_and_video_are_rejected(self) -> None: + for extension_id, input_kind in (("scene-io", "scene"), ("capture-io", "capture"), ("video-io", "video")): + extension = self._make_extension(extension_id) + manifest = { + "id": extension_id, "name": extension_id, "type": "model", + "generator_class": "TestGenerator", + "nodes": [{"id": "generate", "input": input_kind, "output": "scene"}], + } + (extension / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") + (extension / "generator.py").write_text( + "from services.generators.base import BaseGenerator\n" + "class TestGenerator(BaseGenerator):\n" + " def load(self): self._model = object()\n" + " def generate(self, value, params, progress_cb=None, cancel_event=None): return self.outputs_dir\n", + encoding="utf-8", + ) + + self.registry.initialize() + self.assertEqual(self.registry.get_manifest("scene-io/generate")["input"], "scene") + self.assertIn("capture-io/generate", self.registry.load_errors()) + self.assertIn("video-io/generate", self.registry.load_errors()) + + def test_scene_input_rejects_multi_input_shapes_but_image_multi_can_output_scene(self) -> None: + cases = { + "scene-mixed": {"input": "scene", "inputs": ["scene", "text"], "output": "mesh"}, + "scene-array": {"input": "scene", "inputs": ["scene"], "output": "mesh"}, + "images-scene": {"input": "image", "inputs": ["image", "image"], "output": "scene"}, + } + for extension_id, node in cases.items(): + extension = self._make_extension(extension_id) + (extension / "manifest.json").write_text(json.dumps({ + "id": extension_id, "name": extension_id, "type": "model", + "generator_class": "TestGenerator", + "nodes": [{"id": "generate", **node}], + }), encoding="utf-8") + (extension / "generator.py").write_text( + "from services.generators.base import BaseGenerator\n" + "class TestGenerator(BaseGenerator):\n" + " def load(self): self._model = object()\n" + " def generate(self, value, params, progress_cb=None, cancel_event=None): return self.outputs_dir\n", + encoding="utf-8", + ) + self.registry.initialize() + self.assertIn("scene-mixed/generate", self.registry.load_errors()) + self.assertIn("scene-array/generate", self.registry.load_errors()) + self.assertIn("images-scene/generate", self.registry._generators) + def test_declared_sources_block_generation_even_when_generator_overrides_readiness(self) -> None: extension = self._make_extension("multi-source") manifest = { diff --git a/api/tests/test_runner.py b/api/tests/test_runner.py index 8fce3d31..a666f87a 100644 --- a/api/tests/test_runner.py +++ b/api/tests/test_runner.py @@ -5,6 +5,7 @@ import json import tempfile import importlib +from unittest.mock import patch from contextlib import redirect_stdout from pathlib import Path @@ -20,6 +21,30 @@ class RunnerTests(unittest.TestCase): + def test_decode_typed_scene_revalidates_worker_workspace_and_keeps_legacy_image(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + workspace = Path(tmp) / "workspace" + scene = workspace / "Workflows" / "room" + scene.mkdir(parents=True) + manifest = scene / "scene-manifest.json" + manifest.write_text(json.dumps({ + "schema": "modly.scene-manifest.v1", "sceneRoot": ".", "assets": [], + })) + with patch.object(runner, "WORKSPACE_DIR", workspace): + typed = runner.decode_model_input({"input": {"kind": "scene", "path": str(manifest)}}) + self.assertEqual(typed.kind, "scene") + self.assertEqual(typed.path, manifest.resolve()) + self.assertEqual(runner.decode_model_input({"image_b64": "aW1hZ2U="}), b"image") + with self.assertRaises(ValueError): + runner.decode_model_input({"input": {"kind": "video", "path": str(manifest)}}) + + def test_runner_model_envelope_rejects_cross_node_dispatch(self) -> None: + manifest = {"id": "pixal3d"} + node = {"id": "worldsculpt"} + runner.validate_requested_model({"model_id": "pixal3d/worldsculpt"}, manifest, node) + with self.assertRaisesRegex(ValueError, "does not match"): + runner.validate_requested_model({"model_id": "pixal3d/generate"}, manifest, node) + def test_select_node_uses_model_dir_override(self) -> None: manifest = { "nodes": [ diff --git a/api/tests/test_scene_generation.py b/api/tests/test_scene_generation.py new file mode 100644 index 00000000..a804171a --- /dev/null +++ b/api/tests/test_scene_generation.py @@ -0,0 +1,115 @@ +import asyncio +import json +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from fastapi import BackgroundTasks, HTTPException +from pydantic import ValidationError + +import routers.generation as generation +import services.generator_registry as registry +from schemas.generation import GenerateFromArtifactRequest + + +class _Registry: + def __init__(self): + self.switched = False + def get_generator(self, model_id): return object() + def get_manifest(self, model_id): return {"input": "scene"} + def switch_model(self, model_id): self.switched = True + + +class SceneGenerationTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.workspace = Path(self.tmp.name) / "workspace" + self.scene = self.workspace / "Workflows" / "room" + self.scene.mkdir(parents=True) + self.manifest = self.scene / "scene-manifest.json" + self.manifest.write_text(json.dumps({"schema": "modly.scene-manifest.v1", "sceneRoot": ".", "assets": []})) + self.registry = _Registry() + self.patches = [patch.object(generation, "generator_registry", self.registry), patch.object(registry, "WORKSPACE_DIR", self.workspace)] + for item in self.patches: item.start() + + def tearDown(self): + for item in reversed(self.patches): item.stop() + generation._jobs.clear(); generation._cancel_events.clear(); generation._cancelled.clear(); generation._completed_at.clear() + self.tmp.cleanup() + + def test_generic_route_queues_typed_scene_and_strips_reserved_params(self): + tasks = BackgroundTasks() + result = asyncio.run(generation.generate_from_artifact(GenerateFromArtifactRequest( + input_kind="scene", input_path="Workflows/room", model_id="demo/scene", + params={"artifact_path": "/etc/passwd", "input_kind": "image", "quality": "high"}, + ), tasks)) + queued = tasks.tasks[0] + self.assertEqual(queued.args[1].kind, "scene") + self.assertEqual(queued.args[1].path, self.manifest.resolve()) + self.assertEqual(queued.args[2]["scene_manifest_path"], str(self.manifest.resolve())) + self.assertNotIn("artifact_path", queued.args[2]) + self.assertNotIn("input_kind", queued.args[2]) + self.assertEqual(queued.args[5], "demo/scene") + self.assertEqual(result["job_id"], queued.args[0]) + + def test_generic_route_rejects_unsupported_kind_and_model_mismatch(self): + for kind in ("video", "capture", "image"): + with self.subTest(kind=kind), self.assertRaises(ValidationError): + GenerateFromArtifactRequest( + input_kind=kind, input_path="Workflows/room", model_id="demo/scene") + self.registry.get_manifest = lambda _model_id: {"input": "image"} + with self.assertRaises(HTTPException) as caught: + asyncio.run(generation.generate_from_artifact(GenerateFromArtifactRequest( + input_kind="scene", input_path="Workflows/room", model_id="demo/image"), BackgroundTasks())) + self.assertEqual(caught.exception.status_code, 400) + + def test_rejects_traversal_before_switch_or_queue(self): + with self.assertRaises(HTTPException): + asyncio.run(generation.generate_from_artifact(GenerateFromArtifactRequest( + input_kind="scene", input_path="../outside", model_id="demo/scene"), BackgroundTasks())) + self.assertFalse(self.registry.switched) + + def test_queued_scene_job_is_pinned_to_requested_model(self): + calls = [] + + class Generator: + outputs_dir = None + def is_loaded(self): return True + def generate_artifact(self, kind, path, params, progress_cb, cancel_event=None): + calls.append(("model-a", kind, path)) + output = Path(self.outputs_dir) / "result.glb" + output.write_bytes(b"glb") + return output + + generator = Generator() + registry_stub = type("Registry", (), { + "get_ready_generator": lambda self, model_id: generator if model_id == "demo/a" else (_ for _ in ()).throw(ValueError(f"Unknown model ID: {model_id}")), + "get_active": lambda self: (_ for _ in ()).throw(AssertionError("mutable active model must not be used")), + })() + job_id = "pinned-scene" + generation._jobs[job_id] = generation.JobStatus(job_id=job_id, status="pending", progress=0) + generation._cancel_events[job_id] = __import__("threading").Event() + with patch.object(generation, "generator_registry", registry_stub): + asyncio.run(generation._run_generation( + job_id, generation.TypedArtifactInput("scene", self.manifest.resolve()), {}, + "Workflows", "mesh", "demo/a", + )) + self.assertEqual(calls[0][0], "model-a") + self.assertEqual(generation._jobs[job_id].status, "done") + + def test_missing_pinned_model_fails_actionably(self): + registry_stub = type("Registry", (), { + "get_ready_generator": lambda self, model_id: (_ for _ in ()).throw(ValueError(f"Unknown model ID: {model_id}")), + "get_active": lambda self: (_ for _ in ()).throw(AssertionError("must not use active model")), + })() + job_id = "missing-scene" + generation._jobs[job_id] = generation.JobStatus(job_id=job_id, status="pending", progress=0) + generation._cancel_events[job_id] = __import__("threading").Event() + with patch.object(generation, "generator_registry", registry_stub): + asyncio.run(generation._run_generation( + job_id, generation.TypedArtifactInput("scene", self.manifest.resolve()), {}, + "Workflows", "mesh", "demo/missing", + )) + self.assertEqual(generation._jobs[job_id].status, "error") + self.assertIn("Unknown model ID: demo/missing", generation._jobs[job_id].error) diff --git a/api/tests/test_scene_input.py b/api/tests/test_scene_input.py new file mode 100644 index 00000000..43bd9946 --- /dev/null +++ b/api/tests/test_scene_input.py @@ -0,0 +1,67 @@ +import json +import tempfile +import unittest +from pathlib import Path + +from services.scene_input import validate_scene_input, revalidate_scene_manifest + + +class SceneInputTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.workspace = Path(self.tmp.name) / "workspace" + self.scene = self.workspace / "Workflows" / "room" + self.scene.mkdir(parents=True) + (self.scene / "model.glb").write_bytes(b"mesh") + self.manifest = self.scene / "scene-manifest.json" + self.manifest.write_text(json.dumps({ + "schema": "modly.scene-manifest.v1", + "sceneRoot": ".", + "assets": [{"path": "model.glb"}], + })) + + def tearDown(self): + self.tmp.cleanup() + + def test_directory_and_manifest_are_canonical_paths(self): + expected = self.manifest.resolve() + self.assertEqual(validate_scene_input(self.workspace, "Workflows/room"), expected) + self.assertEqual(validate_scene_input(self.workspace, "Workflows/room/scene-manifest.json"), expected) + self.assertEqual(revalidate_scene_manifest(self.workspace, expected), expected) + + def test_rejects_traversal_absolute_encoded_and_non_manifest_json(self): + for path in ("../outside", "/etc/passwd", "C:/outside", "Workflows/room/../room", "Workflows/room/other.json", "Workflows/%2e%2e"): + with self.subTest(path=path), self.assertRaises(ValueError): + validate_scene_input(self.workspace, path) + + def test_rejects_symlinks_missing_assets_and_oversized_manifest(self): + outside = Path(self.tmp.name) / "outside" + outside.mkdir() + (outside / "scene-manifest.json").write_text(self.manifest.read_text()) + (self.workspace / "Workflows" / "link").symlink_to(outside, target_is_directory=True) + with self.assertRaises(ValueError): + validate_scene_input(self.workspace, "Workflows/link") + + data = json.loads(self.manifest.read_text()) + data["assets"] = [{"path": "missing.glb"}] + self.manifest.write_text(json.dumps(data)) + with self.assertRaises(ValueError): + validate_scene_input(self.workspace, "Workflows/room") + + self.manifest.write_text(" " * (1024 * 1024 + 1)) + with self.assertRaises(ValueError): + validate_scene_input(self.workspace, "Workflows/room") + + def test_rejects_malformed_manifest_and_asset_paths(self): + original = json.loads(self.manifest.read_text()) + for patch in ( + {"schema": "wrong"}, + {"assets": "bad"}, + {"assets": [{"path": "../escape.glb"}]}, + {"preview": {"image": None}}, + {"initialView": {"position": [0, 0, 0], "target": [0, 0, 0]}}, + ): + with self.subTest(patch=patch): + self.manifest.write_text(json.dumps({**original, **patch})) + with self.assertRaises(ValueError): + validate_scene_input(self.workspace, "Workflows/room") diff --git a/electron/main/artifact-registry-service.test.ts b/electron/main/artifact-registry-service.test.ts index 882120d4..276e90a1 100644 --- a/electron/main/artifact-registry-service.test.ts +++ b/electron/main/artifact-registry-service.test.ts @@ -64,6 +64,18 @@ test('lists Workflows and Exports assets while skipping hidden, cache, and inter assert.equal(result.success && result.entries.find((entry) => entry.workspacePath.endsWith('exported.ply'))?.openable, false) })) +test('registers a generated scene directory through its canonical manifest artifact', () => withWorkspace(async (workspaceDir) => { + await mkdir(path.join(workspaceDir, 'Workflows/world'), { recursive: true }) + await writeFile(path.join(workspaceDir, 'Workflows/world/scene-manifest.json'), JSON.stringify({ + schema: 'modly.scene-manifest.v1', sceneRoot: '.', assets: [], + })) + const result = await listWorkspaceAssetLibrary({ workspaceDir }) + assert.equal(result.success, true) + const scene = result.success && result.entries.find((entry) => entry.workspacePath === 'Workflows/world/scene-manifest.json') + assert.equal(scene && scene.capability, 'scene-manifest') + assert.equal(scene && scene.state, 'ready') +})) + test('reads and opens only safe GLB/GLTF workspace assets', () => withWorkspace(async (workspaceDir) => { await mkdir(path.join(workspaceDir, 'Workflows/checkpoints'), { recursive: true }) await mkdir(path.join(workspaceDir, 'Exports'), { recursive: true }) diff --git a/electron/main/artifact-registry-service.ts b/electron/main/artifact-registry-service.ts index 3c26de73..2acb657f 100644 --- a/electron/main/artifact-registry-service.ts +++ b/electron/main/artifact-registry-service.ts @@ -124,7 +124,7 @@ export function classifyAssetLibraryCandidate(candidate: AssetLibraryClassificat if (candidate.workspacePath.endsWith('.world.json')) { return { capability: 'generated-world', state: 'ready', previewKind: 'text', openable: false, nonOpenableReason: 'Generated worlds are list-only in this release.' } } - if (candidate.workspacePath.endsWith('.scene.json')) { + if (candidate.workspacePath.endsWith('.scene.json') || candidate.workspacePath.endsWith('/scene-manifest.json')) { return { capability: 'scene-manifest', state: 'ready', previewKind: 'text', openable: false, nonOpenableReason: 'Scene manifests are list-only in this release.' } } if (INTRINSIC_MOTION_EXTENSIONS.has(extension)) { @@ -153,6 +153,7 @@ function objectField(value: unknown): Record | undefined { function manifestCapabilityFor(workspacePath: string): 'generated-world' | 'scene-manifest' | undefined { if (workspacePath.endsWith('.world.json')) return 'generated-world' if (workspacePath.endsWith('.scene.json')) return 'scene-manifest' + if (workspacePath.endsWith('/scene-manifest.json')) return 'scene-manifest' return undefined } diff --git a/electron/main/extension-install-utils.test.mjs b/electron/main/extension-install-utils.test.mjs index 84139f9a..4dcf42da 100644 --- a/electron/main/extension-install-utils.test.mjs +++ b/electron/main/extension-install-utils.test.mjs @@ -85,6 +85,50 @@ test('validateInstallManifest accepts multi-source nodes and preserves legacy sh }, { hasEntryFile: () => false, hasGeneratorFile: () => true }, 'repository')) }) +test('validateInstallManifest accepts scene IO and rejects undeclared future artifact kinds', () => { + const mod = loadModule() + const files = { hasEntryFile: () => false, hasGeneratorFile: () => true } + assert.doesNotThrow(() => mod.validateInstallManifest({ + id: 'scene-model', generator_class: 'Generator', + nodes: [{ id: 'normalize', input: 'scene', output: 'scene' }], + }, files, 'repository')) + for (const input of ['capture', 'video']) { + assert.throws(() => mod.validateInstallManifest({ + id: 'future-model', generator_class: 'Generator', + nodes: [{ id: 'future', input, output: 'scene' }], + }, files, 'repository'), /supported artifact type/) + } +}) + +test('scene is model-only, single-input, while image-multi to scene stays valid', () => { + const mod = loadModule() + const modelFiles = { hasEntryFile: () => false, hasGeneratorFile: () => true } + const processFiles = { hasEntryFile: () => true, hasGeneratorFile: () => false } + for (const node of [ + { id: 'mixed', input: 'scene', inputs: ['scene', 'text'], output: 'mesh' }, + { id: 'duplicate', input: 'scene', inputs: ['scene', 'scene'], output: 'mesh' }, + { id: 'hidden', input: 'image', inputs: ['scene'], output: 'mesh' }, + ]) { + assert.throws(() => mod.validateInstallManifest({ id: 'bad', generator_class: 'Generator', nodes: [node] }, modelFiles, 'repository'), /scene.*single|single.*scene/i) + } + for (const node of [ + { id: 'input', input: 'scene', output: 'mesh' }, + { id: 'output', input: 'image', output: 'scene' }, + ]) { + assert.throws(() => mod.validateInstallManifest({ id: 'proc', type: 'process', entry: 'processor.js', nodes: [node] }, processFiles, 'repository'), /scene.*model|model.*scene/i) + } + assert.doesNotThrow(() => mod.validateInstallManifest({ + id: 'images-to-scene', generator_class: 'Generator', + nodes: [{ id: 'prepare', input: 'image', inputs: ['image', 'image'], output: 'scene' }], + }, modelFiles, 'repository')) + for (const output of ['scene', 'mesh']) { + assert.doesNotThrow(() => mod.validateInstallManifest({ + id: `scene-to-${output}`, generator_class: 'Generator', + nodes: [{ id: 'generate', input: 'scene', output }], + }, modelFiles, 'repository')) + } +}) + test('validateInstallManifest rejects malformed or process model_sources', () => { const mod = loadModule() const source = { diff --git a/electron/main/extension-install-utils.ts b/electron/main/extension-install-utils.ts index 05b965b0..41e9798c 100644 --- a/electron/main/extension-install-utils.ts +++ b/electron/main/extension-install-utils.ts @@ -10,7 +10,13 @@ export interface InstallManifest { entry?: string generator_class?: string model_sources?: unknown - nodes?: Array<{ id?: string; model_sources?: unknown } & ModelSourceNode> + nodes?: Array<{ + id?: string + input?: unknown + inputs?: unknown + output?: unknown + model_sources?: unknown + } & ModelSourceNode> } export interface ValidatedInstallManifest { @@ -27,6 +33,21 @@ export interface ExtensionReloadPayload { errors: Record } +export function assertSupportedSceneNodeShape( + kind: 'model' | 'process', + node: { id?: string; input?: unknown; inputs?: unknown; output?: unknown }, + declaredInputs: unknown[], + output: unknown, +): void { + const usesSceneInput = declaredInputs.includes('scene') + if (kind === 'process' && (usesSceneInput || output === 'scene')) { + throw new Error('manifest.json: scene input and output are supported only for model nodes') + } + if (kind === 'model' && usesSceneInput && (node.inputs !== undefined || node.input !== 'scene')) { + throw new Error(`manifest.json: ${node.id ?? 'node'} must declare scene as its single input field`) + } +} + export type IncompleteInstallRecoveryAction = | 'none' | 'remove-incomplete' @@ -45,11 +66,22 @@ export function validateInstallManifest( const isProcess = manifest.type === 'process' const entryFile = manifest.entry ?? 'processor.js' const nodes = Array.isArray(manifest.nodes) ? manifest.nodes.filter((node) => node?.id) : [] + const allowedIo = new Set(['image', 'text', 'mesh', 'audio', 'scene']) if (manifest.model_sources !== undefined) { throw new Error('manifest.json: model_sources must be declared on a model node') } for (const node of Array.isArray(manifest.nodes) ? manifest.nodes : []) { + const declaredInputs = node.inputs === undefined ? [node.input ?? 'image'] : node.inputs + if (!Array.isArray(declaredInputs) || declaredInputs.length === 0 + || declaredInputs.some((value) => typeof value !== 'string' || !allowedIo.has(value))) { + throw new Error(`manifest.json: ${node.id ?? 'node'}.input must use a supported artifact type`) + } + const output = node.output ?? 'mesh' + if (typeof output !== 'string' || !allowedIo.has(output)) { + throw new Error(`manifest.json: ${node.id ?? 'node'}.output must use a supported artifact type`) + } + assertSupportedSceneNodeShape(isProcess ? 'process' : 'model', node, declaredInputs, output) if (node.model_sources === undefined) continue if (isProcess) { throw new Error('manifest.json: model_sources is supported only for model nodes') diff --git a/electron/main/ipc-handlers.ts b/electron/main/ipc-handlers.ts index 4fa24788..9e5131de 100644 --- a/electron/main/ipc-handlers.ts +++ b/electron/main/ipc-handlers.ts @@ -49,6 +49,7 @@ import { validateExtensionReloadPayload, validateExistingExtensionReplacement, validateInstallManifest, + assertSupportedSceneNodeShape, } from './extension-install-utils' import { beginExtensionRegistrationTransaction, @@ -844,10 +845,10 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe nodes?: { id: string name?: string - input?: 'mesh' | 'image' | 'text' | 'audio' - inputs?: ('mesh' | 'image' | 'text' | 'audio')[] + input?: 'mesh' | 'image' | 'text' | 'audio' | 'scene' + inputs?: ('mesh' | 'image' | 'text' | 'audio' | 'scene')[] input_labels?: string[] - output?: 'mesh' | 'image' | 'text' | 'audio' + output?: 'mesh' | 'image' | 'text' | 'audio' | 'scene' params_schema?: unknown[] param_defaults?: Record hf_repo?: string @@ -873,7 +874,15 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe if (parsed.model_sources !== undefined) { throw new Error('manifest.json: model_sources must be declared on a model node') } + const allowedIo = new Set(['image', 'text', 'mesh', 'audio', 'scene']) const nodes = (parsed.nodes ?? []).map(n => { + const declaredInputs = n.inputs ?? [n.input ?? 'image'] + for (const input of declaredInputs) { + if (!allowedIo.has(input)) throw new Error(`manifest.json: unsupported node input type "${input}"`) + } + const output = n.output ?? 'mesh' + if (!allowedIo.has(output)) throw new Error(`manifest.json: unsupported node output type "${output}"`) + assertSupportedSceneNodeShape(parsed.type === 'process' ? 'process' : 'model', n, declaredInputs, output) if (parsed.type === 'process' && n.model_sources !== undefined) { throw new Error('manifest.json: model_sources is supported only for model nodes') } diff --git a/src/areas/workflows/WorkflowsPage.tsx b/src/areas/workflows/WorkflowsPage.tsx index d15c06ff..32452774 100644 --- a/src/areas/workflows/WorkflowsPage.tsx +++ b/src/areas/workflows/WorkflowsPage.tsx @@ -27,6 +27,7 @@ import ImageNode from './nodes/ImageNode' import TextNode from './nodes/TextNode' import AddToSceneNode from './nodes/AddToSceneNode' import Load3DMeshNode from './nodes/Load3DMeshNode' +import LoadSceneNode from './nodes/LoadSceneNode' import PreviewImageNode from './nodes/PreviewImageNode' import ImagePreviewNode from './nodes/ImagePreviewNode' import WaitNode from './nodes/WaitNode' @@ -38,7 +39,7 @@ import WorkflowEdge from './nodes/WorkflowEdge' const DRAG_KEY = 'modly/extension-id' const DRAG_NODE_KEY = 'modly/node-type' -const NODE_TYPES = { extensionNode: ExtensionNode, imageNode: ImageNode, textNode: TextNode, outputNode: AddToSceneNode, meshNode: Load3DMeshNode, previewNode: PreviewImageNode, imagePreviewNode: ImagePreviewNode, waitNode: WaitNode, whileNode: WhileNode, forEachNode: ForEachNode } +const NODE_TYPES = { extensionNode: ExtensionNode, imageNode: ImageNode, textNode: TextNode, outputNode: AddToSceneNode, meshNode: Load3DMeshNode, sceneNode: LoadSceneNode, previewNode: PreviewImageNode, imagePreviewNode: ImagePreviewNode, waitNode: WaitNode, whileNode: WhileNode, forEachNode: ForEachNode } // Loop-container node types: resizable frames whose children form a loop body. // (For Each iterators are plain source nodes, not containers.) @@ -62,14 +63,15 @@ function findWhileContainerAt(nodes: Node[], pos: { x: number; y: number }): Nod // ─── IO badge ───────────────────────────────────────────────────────────────── -const IO_STYLES: Record<'image' | 'text' | 'mesh' | 'audio', string> = { +const IO_STYLES: Record<'image' | 'text' | 'mesh' | 'audio' | 'scene', string> = { audio: 'bg-emerald-500/15 text-emerald-400 border-emerald-500/25', image: 'bg-sky-500/15 text-sky-400 border-sky-500/25', mesh: 'bg-violet-500/15 text-violet-400 border-violet-500/25', text: 'bg-amber-500/15 text-amber-400 border-amber-500/25', + scene: 'bg-emerald-500/15 text-emerald-400 border-emerald-500/25', } -function IoBadge({ type }: { type: 'image' | 'text' | 'mesh' | 'audio' }) { +function IoBadge({ type }: { type: 'image' | 'text' | 'mesh' | 'audio' | 'scene' }) { return ( {type} @@ -99,6 +101,7 @@ const PANEL_BUILTIN_NODES = [ { type: 'imageNode', label: 'Image', color: '#38bdf8', icon: <> }, { type: 'textNode', label: 'Text', color: '#fbbf24', icon: <> }, { type: 'meshNode', label: 'Load 3D Mesh', color: '#a78bfa', icon: <> }, + { type: 'sceneNode', label: 'Load Scene', color: '#34d399', icon: <> }, { type: 'outputNode', label: 'Add to Scene', color: '#a78bfa', icon: <> }, { type: 'previewNode', label: 'Preview Views', color: '#38bdf8', icon: <> }, { type: 'imagePreviewNode', label: 'Preview Image', color: '#38bdf8', icon: <> }, @@ -346,6 +349,7 @@ const BUILTIN_NODES = [ { type: 'imageNode', label: 'Image', color: '#38bdf8', description: 'Image input' }, { type: 'textNode', label: 'Text', color: '#fbbf24', description: 'Text input' }, { type: 'meshNode', label: 'Load 3D Mesh', color: '#a78bfa', description: 'Load a 3D mesh file or use current model' }, + { type: 'sceneNode', label: 'Load Scene', color: '#34d399', description: 'Load and validate a workspace scene directory' }, { type: 'outputNode', label: 'Add to Scene', color: '#a78bfa', description: 'Output node — adds the mesh to the 3D scene' }, { type: 'previewNode', label: 'Preview Views', color: '#38bdf8', description: 'Displays multi-view image outputs in a 2×3 grid' }, { type: 'imagePreviewNode', label: 'Preview Image', color: '#38bdf8', description: 'Displays a single image output in the workflow' }, @@ -703,6 +707,7 @@ function getNodeOutputType(node: Node | undefined, allExts: WorkflowExtension[]) if (!node) return undefined if (node.type === 'imageNode') return 'image' if (node.type === 'meshNode') return 'mesh' + if (node.type === 'sceneNode') return 'scene' if (node.type === 'textNode') return 'text' if (node.type === 'imagePreviewNode') return 'image' return allExts.find((e) => e.id === (node.data as WFNodeData)?.extensionId)?.output @@ -1374,6 +1379,7 @@ const MINI_NODE_TINTS: Record = { imageNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, textNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, meshNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, + sceneNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, extensionNode: { fill: 'rgba(167,139,250,0.24)', stroke: '#a78bfa' }, outputNode: { fill: 'rgba(56,189,248,0.22)', stroke: '#38bdf8' }, previewNode: { fill: 'rgba(56,189,248,0.22)', stroke: '#38bdf8' }, diff --git a/src/areas/workflows/mockExtensions.ts b/src/areas/workflows/mockExtensions.ts index 2bbc8cc6..be727fea 100644 --- a/src/areas/workflows/mockExtensions.ts +++ b/src/areas/workflows/mockExtensions.ts @@ -10,10 +10,10 @@ export interface WorkflowExtension { nodeId: string // "node_id" name: string description: string - input: 'image' | 'text' | 'mesh' | 'audio' - inputs?: ('image' | 'text' | 'mesh' | 'audio')[] // multi-input; overrides input when set + input: 'image' | 'text' | 'mesh' | 'audio' | 'scene' + inputs?: ('image' | 'text' | 'mesh' | 'audio' | 'scene')[] // multi-input; overrides input when set inputLabels?: string[] // display labels per input slot - output: 'image' | 'text' | 'mesh' | 'audio' + output: 'image' | 'text' | 'mesh' | 'audio' | 'scene' params: ParamSchema[] builtin: boolean type: 'model' | 'process' diff --git a/src/areas/workflows/nodes/ExtensionNode.tsx b/src/areas/workflows/nodes/ExtensionNode.tsx index abe2f9f6..4be11a14 100644 --- a/src/areas/workflows/nodes/ExtensionNode.tsx +++ b/src/areas/workflows/nodes/ExtensionNode.tsx @@ -16,6 +16,7 @@ const HANDLE_COLOR: Record = { image: '#38bdf8', mesh: '#a78bfa', text: '#fbbf24', + scene: '#34d399', } const TAG_CLS: Record = { @@ -23,6 +24,7 @@ const TAG_CLS: Record = { image: 'border-sky-500/30 bg-sky-500/10 text-sky-400', mesh: 'border-violet-500/30 bg-violet-500/10 text-violet-400', text: 'border-amber-500/30 bg-amber-500/10 text-amber-400', + scene: 'border-emerald-500/30 bg-emerald-500/10 text-emerald-400', } // ─── Param control ──────────────────────────────────────────────────────────── diff --git a/src/areas/workflows/nodes/LoadSceneNode.tsx b/src/areas/workflows/nodes/LoadSceneNode.tsx new file mode 100644 index 00000000..2ac7dded --- /dev/null +++ b/src/areas/workflows/nodes/LoadSceneNode.tsx @@ -0,0 +1,134 @@ +import { useCallback, useLayoutEffect, useRef, useState } from 'react' +import { Handle, Position, useReactFlow } from '@xyflow/react' +import type { WFNodeData } from '@shared/types/electron.d' + +import BaseNode from './BaseNode' +import { resolveSceneSourceManifest } from '../workflowSceneSource' + +const OUTPUT_COLOR = '#34d399' + +async function validateAndPersistScenePath(args: { + id: string + data: WFNodeData + nextPath: string + updateNodeData: ReturnType['updateNodeData'] +}): Promise { + const settings = await window.electron.settings.get() + const resolution = await resolveSceneSourceManifest({ + scenePath: args.nextPath, + workspaceDir: settings.workspaceDir, + readFileBase64: window.electron.fs.readFileBase64, + }) + + if (!resolution.ok) { + args.updateNodeData(args.id, { + params: { + ...args.data.params, + path: args.nextPath, + manifestPath: undefined, + sceneRoot: undefined, + error: resolution.error, + }, + }) + return + } + + args.updateNodeData(args.id, { + params: { + ...args.data.params, + path: resolution.inputWorkspacePath, + manifestPath: resolution.manifestWorkspacePath, + sceneRoot: resolution.sceneRoot, + sourceKind: resolution.sourceKind, + error: undefined, + }, + }) +} + +export default function LoadSceneNode({ id, data, selected }: { id: string; data: WFNodeData; selected?: boolean }) { + const { updateNodeData } = useReactFlow() + const ioRowRef = useRef(null) + const [handleTop, setHandleTop] = useState('50%') + + useLayoutEffect(() => { + if (ioRowRef.current) { + const center = ioRowRef.current.offsetTop + ioRowRef.current.offsetHeight / 2 + setHandleTop(`${center}px`) + } + }, []) + + const scenePath = typeof data.params.path === 'string' ? data.params.path : '' + const manifestPath = typeof data.params.manifestPath === 'string' ? data.params.manifestPath : undefined + const sceneRoot = typeof data.params.sceneRoot === 'string' ? data.params.sceneRoot : undefined + const error = typeof data.params.error === 'string' ? data.params.error : undefined + + const browseDirectory = useCallback(async () => { + const path = await window.electron.fs.selectDirectory() + if (!path) return + await validateAndPersistScenePath({ id, data, nextPath: path, updateNodeData }) + }, [id, data, updateNodeData]) + + const validatePath = useCallback(async () => { + if (!scenePath.trim()) return + await validateAndPersistScenePath({ id, data, nextPath: scenePath, updateNodeData }) + }, [id, data, scenePath, updateNodeData]) + + return ( + + + + + + + } + subheader={ +
+ scene +
+ } + handles={ + + } + > +
+ updateNodeData(id, { params: { ...data.params, path: event.target.value } })} + className="nodrag w-full rounded-lg border border-zinc-700 bg-zinc-800 px-2.5 py-2 text-[10px] text-zinc-200 placeholder-zinc-600 focus:outline-none focus:border-emerald-500/40" + /> +
+ + +
+ {manifestPath ? ( +
+
Manifest: {manifestPath}
+ {sceneRoot &&
sceneRoot: {sceneRoot}
} +
+ ) : ( +
+ Loads an existing workspace scene manifest for downstream scene nodes. +
+ )} + {error &&
{error}
} +
+ + ) +} diff --git a/src/areas/workflows/preflight.test.mjs b/src/areas/workflows/preflight.test.mjs index 75b09d71..3b4df8aa 100644 --- a/src/areas/workflows/preflight.test.mjs +++ b/src/areas/workflows/preflight.test.mjs @@ -131,3 +131,38 @@ test('multi-input extension requires every declared input type', () => { assert.ok(!issues.some((i) => i.key === 'proc:missing:image')) assert.ok(issues.some((i) => i.key === 'proc:missing:text')) }) + +test('scene input requires a validated Load Scene source and rejects image wiring', () => { + const { validateWorkflowPreflight } = loadModule() + const model = { id: 'model', type: 'extensionNode', position: { x: 0, y: 0 }, data: { extensionId: 'pack/process-node' } } + const scene = { id: 'scene', type: 'sceneNode', position: { x: 0, y: 0 }, data: { params: { manifestPath: 'Workflows/room/scene-manifest.json' } } } + for (const output of ['scene', 'mesh']) { + const extension = ext({ input: 'scene', output, type: 'model' }) + assert.deepEqual(validateWorkflowPreflight(wf([scene, model], [{ id: 'scene-edge', source: 'scene', target: 'model' }]), [extension]), []) + } + + const extension = ext({ input: 'scene', output: 'scene', type: 'model' }) + const issues = validateWorkflowPreflight(wf([imageNode(), model], [{ id: 'image-edge', source: 'img', target: 'model' }]), [extension]) + assert.ok(issues.some((issue) => issue.key === 'model:missing:scene')) + assert.ok(issues.some((issue) => issue.key === 'model:type:image-edge')) +}) + +test('Load Scene must be validated before a workflow can run', () => { + const { validateWorkflowPreflight } = loadModule() + const scene = { id: 'scene', type: 'sceneNode', position: { x: 0, y: 0 }, data: { params: { path: 'Workflows/room' } } } + const issues = validateWorkflowPreflight(wf([scene], []), []) + assert.equal(issues[0].key, 'scene:scene-invalid') +}) + +test('renderer fails closed for unsupported process and mixed scene node shapes', () => { + const { validateWorkflowPreflight } = loadModule() + const scene = { id: 'scene', type: 'sceneNode', position: { x: 0, y: 0 }, data: { params: { manifestPath: 'Workflows/room/scene-manifest.json' } } } + const target = { id: 'target', type: 'extensionNode', position: { x: 0, y: 0 }, data: { extensionId: 'pack/process-node' } } + for (const extension of [ + ext({ input: 'scene', output: 'mesh', type: 'process' }), + ext({ input: 'scene', inputs: ['scene', 'text'], output: 'mesh', type: 'model' }), + ]) { + const issues = validateWorkflowPreflight(wf([scene, target], [{ id: 'e', source: 'scene', target: 'target' }]), [extension]) + assert.ok(issues.some((issue) => issue.key === 'target:unsupported-scene-shape')) + } +}) diff --git a/src/areas/workflows/preflight.ts b/src/areas/workflows/preflight.ts index 3985855c..df7d35ca 100644 --- a/src/areas/workflows/preflight.ts +++ b/src/areas/workflows/preflight.ts @@ -2,7 +2,7 @@ import type { Workflow, WFNode } from '@shared/types/electron.d' import { getWorkflowExtension, type WorkflowExtension } from './mockExtensions' import { isPassthrough, isBranchConsumer, resolveDataSource, nearestUpstreamWaits } from './nodeBehaviors' -type DataType = 'image' | 'text' | 'mesh' | 'audio' +type DataType = 'image' | 'text' | 'mesh' | 'audio' | 'scene' export interface WorkflowPreflightIssue { key: string @@ -14,6 +14,7 @@ function nodeLabel(node: WFNode, allExtensions: WorkflowExtension[]): string { if (node.type === 'imageNode') return 'Image' if (node.type === 'textNode') return 'Text' if (node.type === 'meshNode') return 'Load 3D Mesh' + if (node.type === 'sceneNode') return 'Load Scene' if (node.type === 'outputNode') return 'Add to Scene' if (node.type === 'previewNode') return 'Preview Views' if (node.type === 'imagePreviewNode') return 'Preview Image' @@ -28,6 +29,7 @@ function nodeLabel(node: WFNode, allExtensions: WorkflowExtension[]): string { } function formatType(type: DataType): string { + if (type === 'scene') return 'scene' if (type === 'mesh') return 'mesh' if (type === 'image') return 'image' if (type === 'audio') return 'audio' @@ -44,6 +46,7 @@ function getNodeOutputType(node: WFNode, allExtensions: WorkflowExtension[]): Da if (node.type === 'imageNode') return 'image' if (node.type === 'textNode') return 'text' if (node.type === 'meshNode' || node.type === 'outputNode') return 'mesh' + if (node.type === 'sceneNode') return 'scene' if (node.type === 'previewNode') return 'image' if (node.type === 'imagePreviewNode') return 'image' if (node.type === 'forEachNode') { @@ -96,6 +99,13 @@ export function validateWorkflowPreflight( }) } + if (node.type === 'sceneNode' && !((node.data.params?.manifestPath as string | undefined)?.trim())) { + pushIssue(issues, { + key: `${node.id}:scene-invalid`, nodeId: node.id, + message: 'Load Scene needs a validated scene directory.', + }) + } + // A node fed by two different Wait branches can't be scheduled into a single // branch — it would run before either branch produces its mesh. if ( @@ -121,6 +131,20 @@ export function validateWorkflowPreflight( continue } + const usesSceneInput = ext.input === 'scene' || ext.inputs?.includes('scene') === true + const unsupportedSceneShape = + (ext.type === 'process' && (usesSceneInput || ext.output === 'scene')) + || (ext.type === 'model' && usesSceneInput + && (ext.inputs !== undefined || ext.input !== 'scene')) + if (unsupportedSceneShape) { + pushIssue(issues, { + key: `${node.id}:unsupported-scene-shape`, + nodeId: node.id, + message: `${ext.name} uses an unsupported scene input or output declaration.`, + }) + continue + } + const incomingEdges = workflow.edges.filter((edge) => edge.target === node.id) const requiredTypes = [...new Set((ext.inputs ?? [ext.input]) as DataType[])] diff --git a/src/areas/workflows/workflowRunStore.ts b/src/areas/workflows/workflowRunStore.ts index 15cd9ace..82166635 100644 --- a/src/areas/workflows/workflowRunStore.ts +++ b/src/areas/workflows/workflowRunStore.ts @@ -306,6 +306,14 @@ async function executeExtensionNode( selectedImagePath, selectedImageData } = ctx const ext = getWorkflowExtension(node.data.extensionId ?? '', allExtensions) + if (ext) { + const usesSceneInput = ext.input === 'scene' || ext.inputs?.includes('scene') === true + if ((ext.type === 'process' && (usesSceneInput || ext.output === 'scene')) + || (ext.type === 'model' && usesSceneInput + && (ext.inputs !== undefined || ext.input !== 'scene'))) { + throw new Error(`${ext.name} uses an unsupported scene input or output declaration`) + } + } // Freshest params at the moment the node starts (so loop iterations / Retry pick // up edits made while paused, not the values captured at run start). const liveParams = _liveParams.current.get(node.id) ?? node.data.params ?? {} @@ -318,6 +326,7 @@ async function executeExtensionNode( let nodeInputPath: string | undefined let nodeInputText: string | undefined let nodeInputMeshPath: string | undefined + let nodeInputScenePath: string | undefined // Per-slot texts for multi-text-input nodes (e.g. positive/negative prompts). // Indexed by target handle: input-0 → texts[0], input-1 → texts[1]. const nodeInputTexts: (string | undefined)[] = [] @@ -354,21 +363,24 @@ async function executeExtensionNode( const src = resolveSource(edge.source) if (src?.filePath !== undefined) nodeInputPath = src.filePath if (src?.text !== undefined && src.text.trim().length > 0) nodeInputText = src.text + if (src?.outputType === 'scene') nodeInputScenePath = src.filePath } } const isModelNode = ext?.type === 'model' if (isModelNode) { + const isSceneInput = ext?.inputs ? ext.inputs.includes('scene') : ext?.input === 'scene' const isTextInput = ext?.inputs ? ext.inputs.every((i) => i === 'text') : ext?.input === 'text' - const activeImagePath = isTextInput ? undefined : (nodeInputPath ?? selectedImagePath) - if (!isTextInput && !selectedImageData && (!activeImagePath || activeImagePath.trim().length === 0)) { + if (isSceneInput && !nodeInputScenePath) throw new Error(`${ext?.name ?? 'Model'} needs an incoming scene connection`) + const activeImagePath = (isTextInput || isSceneInput) ? undefined : (nodeInputPath ?? selectedImagePath) + if (!isTextInput && !isSceneInput && !selectedImageData && (!activeImagePath || activeImagePath.trim().length === 0)) { throw new Error('No input image selected for model node') } let blob: Blob let fname: string - if (isTextInput || (selectedImageData && nodeInputPath === undefined)) { + if (isTextInput || isSceneInput || (selectedImageData && nodeInputPath === undefined)) { const base64 = selectedImageData && nodeInputPath === undefined ? selectedImageData : 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==' // 1x1 transparent PNG @@ -401,21 +413,30 @@ async function executeExtensionNode( ) const effectiveParams = { ...schemaDefaults, ...liveParams } - const fd = new FormData() - fd.append('image', blob, fname) - fd.append('model_id', node.data.extensionId ?? '') - fd.append('collection', 'Workflows') - fd.append('remesh', 'none') - fd.append('enable_texture', 'false') - fd.append('texture_resolution', '1024') - fd.append('params', JSON.stringify({ ...effectiveParams, ...extraParams })) - setRunState((s) => ({ ...s, blockProgress: 5, blockStep: 'Submitting to model…' })) - - const { data } = await client.post<{ job_id: string }>( - '/generate/from-image', fd, - { headers: { 'Content-Type': 'multipart/form-data' } }, - ) + let submission: { data: { job_id: string } } + if (isSceneInput) { + const normalized = nodeInputScenePath!.replace(/\\/g, '/') + const inputPath = normalized.startsWith(`${workspaceDir}/`) + ? normalized.slice(workspaceDir.length + 1) + : normalized.replace(/^\/workspace\//, '') + submission = await client.post('/generate/from-artifact', { + input_kind: 'scene', input_path: inputPath, + model_id: node.data.extensionId ?? '', collection: 'Workflows', + params: { ...effectiveParams, ...extraParams }, + }) + } else { + const fd = new FormData() + fd.append('image', blob, fname) + fd.append('model_id', node.data.extensionId ?? '') + fd.append('collection', 'Workflows') + fd.append('remesh', 'none') + fd.append('enable_texture', 'false') + fd.append('texture_resolution', '1024') + fd.append('params', JSON.stringify({ ...effectiveParams, ...extraParams })) + submission = await client.post('/generate/from-image', fd, { headers: { 'Content-Type': 'multipart/form-data' } }) + } + const { data } = submission _activeJobId.current = data.job_id while (true) { @@ -594,7 +615,7 @@ export const useWorkflowRunStore = create((set, get) => { if (!outputUrl) { for (const [, o] of ctx.nodeOutputs) { if (o.filePath) { - if (o.outputType === 'audio') { + if (o.outputType === 'audio' || o.outputType === 'scene') { outputPath = o.filePath continue } @@ -802,6 +823,13 @@ export const useWorkflowRunStore = create((set, get) => { if (fp) nodeOutputs.set(node.id, { filePath: fp, outputType: 'mesh' }) } } + if (node.type === 'sceneNode') { + const manifestPath = node.data.params?.manifestPath as string | undefined + if (manifestPath) nodeOutputs.set(node.id, { + filePath: `${workspaceDir}/${manifestPath.replace(/^\/+/, '')}`, + outputType: 'scene', + }) + } } const ctx: RunContext = { diff --git a/src/areas/workflows/workflowSceneRun.test.mjs b/src/areas/workflows/workflowSceneRun.test.mjs new file mode 100644 index 00000000..a2c2b2d0 --- /dev/null +++ b/src/areas/workflows/workflowSceneRun.test.mjs @@ -0,0 +1,60 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { build } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +const dir = mkdtempSync(join(tmpdir(), 'modly-scene-run-')) +const stub = (name, source) => { const path = join(dir, name); writeFileSync(path, source); return path } +const appStoreStub = stub('app.ts', ` +export const appState: any = { apiUrl: 'http://modly.test', currentJob: null, + setCurrentJob(value: any) { this.currentJob = value }, + updateCurrentJob(value: any) { this.currentJob = { ...(this.currentJob ?? {}), ...value } } } +export const useAppStore: any = (selector: any) => selector(appState) +useAppStore.getState = () => appState +`) +const axiosStub = stub('axios.ts', `const axios: any = { create: () => (globalThis as any).__sceneClient }; export default axios; export type AxiosInstance = any`) +const extStub = stub('ext.ts', `export const getWorkflowExtension = (id: string, all: any[]) => all.find((value) => value.id === id); export type WorkflowExtension = any`) +const notifyStub = stub('notify.ts', `export const showCompletionNotification = async () => {}`) +const aliases = new Map([ + ['axios', axiosStub], ['@shared/stores/appStore', appStoreStub], + ['./mockExtensions', extStub], ['@shared/utils/notification', notifyStub], +]) +const outfile = join(dir, 'store.cjs') +writeFileSync(outfile, (await build({ + entryPoints: [resolve('src/areas/workflows/workflowRunStore.ts')], bundle: true, + platform: 'node', format: 'cjs', write: false, + plugins: [{ name: 'aliases', setup(build) { build.onResolve({ filter: /.*/ }, (args) => aliases.has(args.path) ? { path: aliases.get(args.path) } : null) } }], +})).outputFiles[0].text) +const { useWorkflowRunStore } = createRequire(import.meta.url)(outfile) + +test('scene model uses typed artifact route and preserves scene output', async () => { + const posts = [] + globalThis.window = { electron: { + settings: { get: async () => ({ workspaceDir: '/workspace' }) }, + fs: { deleteDirectory: async () => ({ success: true }), listFiles: async () => [], readFileBase64: async () => { throw new Error('scene must not be read as image bytes') } }, + } } + globalThis.__sceneClient = { + post: async (url, body) => { posts.push({ url, body }); return { data: { job_id: 'scene-job' } } }, + get: async () => ({ data: { status: 'done', progress: 100, output_url: '/workspace/Workflows/result/scene-manifest.json' } }), + } + const workflow = { + id: 'wf', name: 'Scene', description: '', createdAt: '', updatedAt: '', + nodes: [ + { id: 'source', type: 'sceneNode', position: { x: 0, y: 0 }, data: { enabled: true, params: { manifestPath: 'Workflows/input/scene-manifest.json' } } }, + { id: 'model', type: 'extensionNode', position: { x: 1, y: 0 }, data: { enabled: true, extensionId: 'pixal/world', params: {} } }, + ], + edges: [{ id: 'e', source: 'source', target: 'model' }], + } + const extension = { id: 'pixal/world', name: 'World', type: 'model', input: 'scene', output: 'scene', params: [] } + await useWorkflowRunStore.getState().run(workflow, [extension]) + assert.equal(posts[0].url, '/generate/from-artifact') + assert.deepEqual(posts[0].body, { + input_kind: 'scene', input_path: 'Workflows/input/scene-manifest.json', + model_id: 'pixal/world', collection: 'Workflows', params: {}, + }) + assert.equal(useWorkflowRunStore.getState().runState.outputPath, '/workspace/Workflows/result/scene-manifest.json') + assert.equal(useWorkflowRunStore.getState().runState.outputUrl, undefined) +}) diff --git a/src/areas/workflows/workflowSceneSource.test.mjs b/src/areas/workflows/workflowSceneSource.test.mjs new file mode 100644 index 00000000..ad2a0280 --- /dev/null +++ b/src/areas/workflows/workflowSceneSource.test.mjs @@ -0,0 +1,29 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +const outfile = join(mkdtempSync(join(tmpdir(), 'modly-scene-source-')), 'scene.cjs') +writeFileSync(outfile, buildSync({ entryPoints: [resolve('src/areas/workflows/workflowSceneSource.ts')], bundle: true, platform: 'node', format: 'cjs', write: false }).outputFiles[0].text) +const { resolveSceneSourceManifest } = createRequire(import.meta.url)(outfile) +const encoded = Buffer.from(JSON.stringify({ schema: 'modly.scene-manifest.v1', sceneRoot: '.', assets: [] })).toString('base64') + +test('Load Scene resolves directory and manifest without image bytes', async () => { + for (const scenePath of ['Workflows/room', 'Workflows/room/scene-manifest.json']) { + const result = await resolveSceneSourceManifest({ scenePath, workspaceDir: '/workspace', readFileBase64: async () => encoded }) + assert.equal(result.ok, true) + assert.equal(result.manifestWorkspacePath, 'Workflows/room/scene-manifest.json') + } +}) + +test('Load Scene refuses unsafe paths before reading', async () => { + let reads = 0 + for (const scenePath of ['../outside', '/etc/passwd', 'C:/outside', 'Workflows/%2e%2e', 'Workflows/room/']) { + const result = await resolveSceneSourceManifest({ scenePath, workspaceDir: '/workspace', readFileBase64: async () => { reads++; return encoded } }) + assert.equal(result.ok, false, scenePath) + } + assert.equal(reads, 0) +}) diff --git a/src/areas/workflows/workflowSceneSource.ts b/src/areas/workflows/workflowSceneSource.ts new file mode 100644 index 00000000..5873b692 --- /dev/null +++ b/src/areas/workflows/workflowSceneSource.ts @@ -0,0 +1,205 @@ +import type { SceneArtifactManifestInitialView, SceneArtifactManifestPreview, SceneArtifactManifestV1 } from '../../shared/types/artifacts' + +export const SCENE_MANIFEST_FILE_NAME = 'scene-manifest.json' + +export type SceneSourceKind = 'manifest' | 'directory' + +export type ResolveSceneSourceSuccess = { + ok: true + sourceKind: SceneSourceKind + inputWorkspacePath: string + manifestWorkspacePath: string + manifestAbsolutePath: string + sceneRoot: string + manifest: SceneArtifactManifestV1 +} + +export type ResolveSceneSourceFailure = { + ok: false + error: string +} + +export type ResolveSceneSourceResult = ResolveSceneSourceSuccess | ResolveSceneSourceFailure + +type ResolveSceneSourceArgs = { + scenePath: string + workspaceDir: string + readFileBase64: (filePath: string) => Promise +} + +function isAbsolutePath(value: string): boolean { + return value.startsWith('/') || /^[A-Za-z]:\//.test(value) +} + +function trimTrailingSlashes(value: string): string { + return value.replace(/\/+$/, '') +} + +function isSafeRelativePath(value: unknown, allowDot = false): value is string { + if (typeof value !== 'string' || !value || value !== value.trim() || value.includes('\u0000')) return false + const normalized = value.replace(/\\/g, '/') + if (allowDot && normalized === '.') return true + if (isAbsolutePath(normalized) || /^[A-Za-z][A-Za-z0-9+.-]*:/.test(normalized) + || /%(?:25|2e|2f|5c|00)/i.test(normalized) || /%(?![0-9a-f]{2})/i.test(normalized)) return false + return normalized.split('/').every((segment) => segment.length > 0 && segment !== '.' && segment !== '..') +} + +function normalizeWorkspaceRelativePath(value: string | undefined, workspaceDir: string): string | undefined { + const normalizedValue = value?.replace(/\\/g, '/') + if (!normalizedValue) return undefined + + const normalizedWorkspace = trimTrailingSlashes(workspaceDir.replace(/\\/g, '/')) + let relativePath: string | undefined + + if (normalizedValue.startsWith('/workspace/')) { + relativePath = normalizedValue.slice('/workspace/'.length) + } else if (normalizedValue === normalizedWorkspace) { + return undefined + } else if (normalizedValue.startsWith(`${normalizedWorkspace}/`)) { + relativePath = normalizedValue.slice(normalizedWorkspace.length + 1) + } else if (!isAbsolutePath(normalizedValue)) { + relativePath = normalizedValue + } + + if (!relativePath) return undefined + return isSafeRelativePath(relativePath) ? relativePath : undefined +} + +function resolveSceneSourceKind(inputWorkspacePath: string): SceneSourceKind | undefined { + if (inputWorkspacePath === SCENE_MANIFEST_FILE_NAME || inputWorkspacePath.endsWith(`/${SCENE_MANIFEST_FILE_NAME}`)) { + return 'manifest' + } + if (inputWorkspacePath.toLowerCase().endsWith('.json')) return undefined + return 'directory' +} + +function decodeBase64Utf8(base64: string): string { + const bytes = Uint8Array.from(atob(base64), (char) => char.charCodeAt(0)) + return new TextDecoder().decode(bytes) +} + +function isPlainObject(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +function isSafePreviewReference(value: unknown): value is string { + return isSafeRelativePath(value) +} + +function isSceneManifestPreview(value: unknown): value is SceneArtifactManifestPreview { + return isPlainObject(value) + && (!('image' in value) || isSafePreviewReference(value.image)) + && (!('video' in value) || isSafePreviewReference(value.video)) +} + +function isFiniteTriple(value: unknown): value is [number, number, number] { + return Array.isArray(value) && value.length === 3 + && value.every((component) => typeof component === 'number' && Number.isFinite(component)) +} + +function isSceneManifestInitialView(value: unknown): value is SceneArtifactManifestInitialView { + if (!isPlainObject(value)) return false + const { position, target, up } = value + if (!isFiniteTriple(position) || !isFiniteTriple(target)) return false + if (position.every((component, index) => component === target[index])) return false + return up === undefined || (isFiniteTriple(up) && up.some((component) => component !== 0)) +} + +function normalizeSceneRoot(sceneRoot: unknown): string | undefined { + return isSafeRelativePath(sceneRoot, true) ? sceneRoot.replace(/\\/g, '/') : undefined +} + +function validateSceneManifest(manifest: unknown): { ok: true; manifest: SceneArtifactManifestV1; sceneRoot: string } | { ok: false; error: string } { + if (!isPlainObject(manifest)) { + return { ok: false, error: 'Scene manifest must be a JSON object.' } + } + if (manifest.schema !== 'modly.scene-manifest.v1') { + return { ok: false, error: 'Scene manifest schema must be modly.scene-manifest.v1.' } + } + + const rawSceneRoot = manifest.sceneRoot + const sceneRoot = normalizeSceneRoot(rawSceneRoot) + if (typeof rawSceneRoot !== 'string' || !sceneRoot) { + return { ok: false, error: 'Scene manifest sceneRoot must be a safe relative path.' } + } + if (!Array.isArray(manifest.assets)) { + return { ok: false, error: 'Scene manifest assets must be an array.' } + } + if (manifest.assets.some((asset) => isPlainObject(asset) + && (('workspacePath' in asset && !isSafeRelativePath(asset.workspacePath)) + || ('path' in asset && !isSafeRelativePath(asset.path))))) { + return { ok: false, error: 'Scene manifest asset paths must be safe relative file references.' } + } + + const { preview, initialView, ...metadata } = manifest + if (preview !== undefined && !isPlainObject(preview)) { + return { ok: false, error: 'Scene manifest preview must be a JSON object.' } + } + if (preview !== undefined && !isSceneManifestPreview(preview)) { + return { ok: false, error: 'Scene manifest preview image/video must be safe relative file references.' } + } + if (initialView !== undefined && !isPlainObject(initialView)) { + return { ok: false, error: 'Scene manifest initialView must be a JSON object.' } + } + if (initialView !== undefined && !isSceneManifestInitialView(initialView)) { + return { ok: false, error: 'Scene manifest initialView requires distinct finite numeric position/target triples and optional non-zero finite up.' } + } + + return { + ok: true, + sceneRoot, + manifest: { + ...metadata, + schema: 'modly.scene-manifest.v1', + sceneRoot: rawSceneRoot, + assets: manifest.assets, + ...(preview !== undefined ? { preview } : {}), + ...(initialView !== undefined ? { initialView } : {}), + }, + } +} + +export async function resolveSceneSourceManifest(args: ResolveSceneSourceArgs): Promise { + const inputWorkspacePath = normalizeWorkspaceRelativePath(args.scenePath, args.workspaceDir) + if (!inputWorkspacePath) { + return { ok: false, error: 'Load Scene requires a safe workspace-relative scene path.' } + } + + const sourceKind = resolveSceneSourceKind(inputWorkspacePath) + if (!sourceKind) { + return { ok: false, error: `Load Scene accepts ${SCENE_MANIFEST_FILE_NAME} or a scene directory.` } + } + + const manifestWorkspacePath = sourceKind === 'manifest' + ? inputWorkspacePath + : `${inputWorkspacePath}/${SCENE_MANIFEST_FILE_NAME}` + const normalizedWorkspace = trimTrailingSlashes(args.workspaceDir.replace(/\\/g, '/')) + const manifestAbsolutePath = `${normalizedWorkspace}/${manifestWorkspacePath}` + + let manifestRaw: string + try { + manifestRaw = decodeBase64Utf8(await args.readFileBase64(manifestAbsolutePath)) + } catch (error) { + return { ok: false, error: `Unable to read scene manifest: ${String(error)}` } + } + + let parsed: unknown + try { + parsed = JSON.parse(manifestRaw) + } catch (error) { + return { ok: false, error: `Scene manifest is not valid JSON: ${String(error)}` } + } + + const validation = validateSceneManifest(parsed) + if (!validation.ok) return validation + + return { + ok: true, + sourceKind, + inputWorkspacePath, + manifestWorkspacePath, + manifestAbsolutePath, + sceneRoot: validation.sceneRoot, + manifest: validation.manifest, + } +} diff --git a/src/shared/stores/workflowsStore.ts b/src/shared/stores/workflowsStore.ts index b8753375..f6fb829c 100644 --- a/src/shared/stores/workflowsStore.ts +++ b/src/shared/stores/workflowsStore.ts @@ -98,7 +98,7 @@ interface LegacyWorkflow { // Source-only nodes have no target handle; sink-only nodes have no source handle. // An edge into/out of the wrong side can't resolve a handle and makes React Flow // warn ("Couldn't create edge for target handle id: null") on every render. -export const NODE_TYPES_WITHOUT_TARGET = new Set(['imageNode', 'textNode', 'meshNode', 'inputNode', 'forEachNode']) +export const NODE_TYPES_WITHOUT_TARGET = new Set(['imageNode', 'textNode', 'meshNode', 'sceneNode', 'inputNode', 'forEachNode']) export const NODE_TYPES_WITHOUT_SOURCE = new Set(['outputNode', 'previewNode']) function sanitizeEdges(nodes: WFNode[], edges: WFEdge[]): WFEdge[] { diff --git a/src/shared/types/artifacts.ts b/src/shared/types/artifacts.ts index a8dfc6eb..57c6b8d4 100644 --- a/src/shared/types/artifacts.ts +++ b/src/shared/types/artifacts.ts @@ -4,3 +4,17 @@ export interface ArtifactProvenance { source?: string [key: string]: unknown } + +export interface SceneArtifactManifestPreview { image?: string; video?: string } +export interface SceneArtifactManifestInitialView { + position: [number, number, number] + target: [number, number, number] + up?: [number, number, number] +} +export interface SceneArtifactManifestV1 { + schema: 'modly.scene-manifest.v1' + sceneRoot: string + assets: unknown[] + preview?: SceneArtifactManifestPreview + initialView?: SceneArtifactManifestInitialView +} diff --git a/src/shared/types/electron.d.ts b/src/shared/types/electron.d.ts index ffe434ed..c6cca213 100644 --- a/src/shared/types/electron.d.ts +++ b/src/shared/types/electron.d.ts @@ -14,10 +14,10 @@ import type { export interface ExtensionNode { id: string name: string - input: 'image' | 'text' | 'mesh' | 'audio' - inputs?: ('image' | 'text' | 'mesh' | 'audio')[] // multi-input nodes; overrides input when set + input: 'image' | 'text' | 'mesh' | 'audio' | 'scene' + inputs?: ('image' | 'text' | 'mesh' | 'audio' | 'scene')[] // multi-input nodes; overrides input when set inputLabels?: string[] // display labels per input slot (e.g. positive/negative) - output: 'image' | 'text' | 'mesh' | 'audio' + output: 'image' | 'text' | 'mesh' | 'audio' | 'scene' paramsSchema: ParamSchema[] paramDefaults?: Record hfRepo?: string From 10cdc15ce1a6c717f1e6e99d5427c44a7c30dcd6 Mon Sep 17 00:00:00 2001 From: DrHepa Date: Thu, 1 Oct 2026 22:52:43 +0200 Subject: [PATCH 2/2] fix(workflows): address scene artifact review --- api/routers/generation.py | 104 ++++++-- api/routers/model.py | 8 +- api/routers/workflow_runs.py | 38 +-- api/services/generator_registry.py | 148 ++++++------ api/services/generators/base.py | 10 +- api/tests/test_base_generator.py | 8 +- api/tests/test_generator_registry.py | 37 ++- api/tests/test_model_router.py | 45 ++++ api/tests/test_scene_generation.py | 224 +++++++++++++++++- api/tests/test_scene_input.py | 8 +- api/tests/test_workflow_runs_lifecycle.py | 107 +++++++++ api/tests/test_workflow_runs_router.py | 5 +- electron/main/bounded-file-reader.test.mjs | 48 ++++ electron/main/bounded-file-reader.ts | 40 ++++ .../main/extension-install-utils.test.mjs | 6 +- electron/main/extension-install-utils.ts | 11 +- electron/main/ipc-handlers.ts | 22 +- src/areas/workflows/WorkflowsPage.tsx | 8 +- src/areas/workflows/nodes/ExtensionNode.tsx | 4 +- src/areas/workflows/nodes/LoadSceneNode.tsx | 67 +++--- src/areas/workflows/workflowRunStore.ts | 6 +- .../workflows/workflowSceneSource.test.mjs | 43 +++- src/areas/workflows/workflowSceneSource.ts | 41 ++++ 23 files changed, 817 insertions(+), 221 deletions(-) create mode 100644 electron/main/bounded-file-reader.test.mjs create mode 100644 electron/main/bounded-file-reader.ts diff --git a/api/routers/generation.py b/api/routers/generation.py index 7c355014..02879218 100644 --- a/api/routers/generation.py +++ b/api/routers/generation.py @@ -4,6 +4,7 @@ import time import traceback import uuid +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Dict, Optional, Union from fastapi import APIRouter, File, Form, UploadFile, HTTPException, BackgroundTasks @@ -28,6 +29,14 @@ _cancelled: set = set() _cancel_events: Dict[str, threading.Event] = {} _completed_at: Dict[str, float] = {} +_job_generators: Dict[str, object] = {} +# A pinned generation owns the complete switch/load/generate lifecycle. Keeping +# that lifecycle on one dedicated worker provides process-wide serialization +# without parking default-executor workers on a blocking lock acquisition. +_pinned_generation_executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="modly-pinned-generation", +) _JOB_TTL = 1800 # purge terminal jobs after 30 minutes @@ -39,6 +48,7 @@ def _purge_old_jobs() -> None: _jobs.pop(jid, None) _cancelled.discard(jid) _cancel_events.pop(jid, None) + _job_generators.pop(jid, None) _completed_at.pop(jid, None) @@ -111,8 +121,6 @@ async def generate_from_image( except ValueError as e: raise HTTPException(400, str(e)) - generator_registry.switch_model(model_id) - # Parse model-specific params from JSON and merge with common fields try: model_params = json.loads(params) @@ -197,11 +205,10 @@ async def cancel_job(job_id: str): if job.status in ("pending", "running"): job.status = "cancelled" _completed_at[job_id] = time.monotonic() - # Kill the active generator subprocess immediately so inference stops now. - # _run_generation will catch the resulting exception, see job_id in _cancelled, - # and return cleanly without setting an error status. + # Kill only the subprocess bound to this job. A queued cancellation must not + # terminate whichever earlier job currently owns the active generator. try: - gen = generator_registry._generators.get(generator_registry._active_id) + gen = _job_generators.get(job_id) if gen is not None and hasattr(gen, "_proc") and gen._proc and gen._proc.poll() is None: gen._proc.kill() gen._loaded = False @@ -219,6 +226,50 @@ async def _run_generation( output_kind: str = "mesh", model_id: Optional[str] = None, ) -> None: + # Pinned jobs share one model lifecycle. Switching is deliberately deferred + # until this job runs on the dedicated worker: request-time switches can + # otherwise unload a running job or leave A loading beside B. Crucially, + # queued jobs are executor work items rather than default-executor threads + # blocked on a lock, so cancelling a waiter cannot orphan queue ownership or + # starve the worker that performs generation. + loop = asyncio.get_running_loop() + executor = _pinned_generation_executor if model_id is not None else None + future = loop.run_in_executor( + executor, + _run_generation_impl, + job_id, + model_input, + params, + collection, + output_kind, + model_id, + ) + try: + await future + except asyncio.CancelledError: + # asyncio cancellation attempts to cancel a queued concurrent future. + # If it has already begun, the event lets the generator stop safely. + _cancelled.add(job_id) + cancel_event = _cancel_events.get(job_id) + if cancel_event is not None: + cancel_event.set() + job = _jobs.get(job_id) + if job is not None and job.status in ("pending", "running"): + job.status = "cancelled" + _completed_at[job_id] = time.monotonic() + raise + + +def _run_generation_impl( + job_id: str, + model_input: Union[bytes, TypedArtifactInput], + params: dict, + collection: str, + output_kind: str, + model_id: Optional[str], +) -> None: + if job_id in _cancelled: + return job = _jobs[job_id] job.status = "running" @@ -232,17 +283,15 @@ def progress_cb(pct: int, step: str = "") -> None: job.step = step try: - loop = asyncio.get_running_loop() - - # Check if the model needs to be loaded BEFORE calling get_active(), - # because get_active() loads the model in a blocking manner. - # active_status() is an instantaneous operation (simple dict lookup). - get_generator = (lambda: generator_registry.get_ready_generator(model_id)) \ + # Check if the model needs to be loaded BEFORE calling the generator + # getter, because that call can load the model in a blocking manner. + get_generator = (lambda: generator_registry.activate_ready_generator(model_id)) \ if model_id is not None else generator_registry.get_active - status_reader = getattr(generator_registry, "model_status", None) - status = (status_reader(model_id) if model_id is not None and status_reader - else generator_registry.active_status() if model_id is None - else {"name": model_id, "downloaded": True, "loaded": False}) + if model_id is not None: + _job_generators[job_id] = generator_registry.get_generator(model_id) + status_reader = (lambda: generator_registry.model_status(model_id)) \ + if model_id is not None else generator_registry.active_status + status = status_reader() if not status["loaded"]: active = status model_name = active['name'] @@ -256,12 +305,13 @@ def progress_cb(pct: int, step: str = "") -> None: ) load_thread.start() try: - gen = await loop.run_in_executor(None, get_generator) + gen = get_generator() finally: stop_load_evt.set() else: - gen = await loop.run_in_executor(None, get_generator) + gen = get_generator() + _job_generators[job_id] = gen if job_id in _cancelled: return @@ -278,18 +328,18 @@ def progress_cb(pct: int, step: str = "") -> None: model_input = revalidate_artifact_input(registry.WORKSPACE_DIR, model_input) import inspect supports_cancel = "cancel_event" in inspect.signature(gen.generate_artifact).parameters - output_path = await loop.run_in_executor( - None, - lambda: gen.generate_artifact(model_input.kind, model_input.path, params, progress_cb, cancel_event) - if supports_cancel else gen.generate_artifact(model_input.kind, model_input.path, params, progress_cb), + output_path = ( + gen.generate_artifact(model_input.kind, model_input.path, params, progress_cb, cancel_event) + if supports_cancel + else gen.generate_artifact(model_input.kind, model_input.path, params, progress_cb) ) else: import inspect supports_cancel = "cancel_event" in inspect.signature(gen.generate).parameters - output_path = await loop.run_in_executor( - None, - lambda: gen.generate(model_input, params, progress_cb, cancel_event) - if supports_cancel else gen.generate(model_input, params, progress_cb), + output_path = ( + gen.generate(model_input, params, progress_cb, cancel_event) + if supports_cancel + else gen.generate(model_input, params, progress_cb) ) if job_id in _cancelled: @@ -328,3 +378,5 @@ def progress_cb(pct: int, step: str = "") -> None: job.status = "error" job.error = tb.strip() _completed_at[job_id] = time.monotonic() + finally: + _job_generators.pop(job_id, None) diff --git a/api/routers/model.py b/api/routers/model.py index 0b40d155..7d34f6cd 100644 --- a/api/routers/model.py +++ b/api/routers/model.py @@ -59,20 +59,20 @@ def _check_download_control(control: dict[str, threading.Event]) -> None: @router.get("/status") async def model_status(): """Status of the active model.""" - return generator_registry.active_status() + return await asyncio.to_thread(generator_registry.active_status) @router.get("/all") async def all_models_status(): """Status of all known models (downloaded, loaded, required VRAM).""" - return generator_registry.all_status() + return await asyncio.to_thread(generator_registry.all_status) @router.get("/params") async def model_params(model_id: Optional[str] = None): """Parameter schema of the active model (or a specified model).""" try: - return generator_registry.params_schema(model_id) + return await asyncio.to_thread(generator_registry.params_schema, model_id) except KeyError: raise HTTPException(404, f"Unknown model ID: {model_id}") @@ -81,7 +81,7 @@ async def model_params(model_id: Optional[str] = None): async def switch_model(model_id: str): """Switch the active model.""" try: - generator_registry.switch_model(model_id) + await asyncio.to_thread(generator_registry.switch_model, model_id) return {"active": model_id} except ValueError as e: raise HTTPException(400, str(e)) diff --git a/api/routers/workflow_runs.py b/api/routers/workflow_runs.py index 9b78a205..ba654a53 100644 --- a/api/routers/workflow_runs.py +++ b/api/routers/workflow_runs.py @@ -1,6 +1,5 @@ import json import threading -import time import uuid from typing import Optional from fastapi import APIRouter, BackgroundTasks, File, Form, HTTPException, UploadFile @@ -9,11 +8,10 @@ from routers.generation import ( VALID_REMESH_MODES, _cancel_events, - _cancelled, - _completed_at, _jobs, _purge_old_jobs, _run_generation, + cancel_job, sanitize_collection, ) from schemas.generation import JobStatus @@ -59,10 +57,7 @@ async def create_run_from_image( **model_params, } - # Same constraint /generate/from-image enforces on this field, checked before touching - # the registry below for the same reason that endpoint checks it first: switch_model() - # unloads whatever generator is currently active, and a request rejected for a bad - # remesh value should not pay for -- or force a reload after -- evicting it. + # Keep the same request validation as /generate/from-image before filing a job. if full_params["remesh"] not in VALID_REMESH_MODES: raise HTTPException(400, "remesh must be 'quad', 'triangle', or 'none'") @@ -70,11 +65,10 @@ async def create_run_from_image( try: generator_registry.get_generator(model_id) + output_kind = generator_registry.get_manifest(model_id).get("output", "mesh") except ValueError as e: raise HTTPException(400, str(e)) - generator_registry.switch_model(model_id) - job_id = str(uuid.uuid4()) image_bytes = await image.read() @@ -83,7 +77,9 @@ async def create_run_from_image( _jobs[job_id] = JobStatus(job_id=job_id, status="pending", progress=0) _cancel_events[job_id] = threading.Event() - background_tasks.add_task(_run_generation, job_id, image_bytes, full_params, collection) + background_tasks.add_task( + _run_generation, job_id, image_bytes, full_params, collection, output_kind, model_id + ) return {"run_id": job_id, "status": "pending"} @@ -111,24 +107,4 @@ async def get_run(run_id: str): @router.post("/{run_id}/cancel") async def cancel_run(run_id: str): - job = _jobs.get(run_id) - if not job: - raise HTTPException(404, f"Run {run_id} not found") - - _cancelled.add(run_id) - if run_id in _cancel_events: - _cancel_events[run_id].set() - if job.status in ("pending", "running"): - job.status = "cancelled" - _completed_at[run_id] = time.monotonic() - - try: - gen = generator_registry._generators.get(generator_registry._active_id) - if gen is not None and hasattr(gen, "_proc") and gen._proc and gen._proc.poll() is None: - gen._proc.kill() - gen._loaded = False - gen._proc = None - except Exception: - pass - - return {"cancelled": True} + return await cancel_job(run_id) diff --git a/api/services/generator_registry.py b/api/services/generator_registry.py index de1750ca..40c08db7 100644 --- a/api/services/generator_registry.py +++ b/api/services/generator_registry.py @@ -452,19 +452,13 @@ def _discover_extensions( node for node in raw_nodes if isinstance(node, dict) and node.get("id") ] - allowed_io = {"image", "text", "mesh", "audio", "scene"} for node in nodes: - declared_inputs = node.get("inputs") or [node.get("input", "image")] - if (not isinstance(declared_inputs, list) - or any(value not in allowed_io for value in declared_inputs)): - raise ValueError( - f'model node "{node.get("id", "unknown")}" has an unsupported input type' - ) - if node.get("output", "mesh") not in allowed_io: - raise ValueError( - f'model node "{node.get("id", "unknown")}" has an unsupported output type' - ) - if "scene" in declared_inputs and ( + declared_inputs = node.get("inputs") + uses_scene_input = ( + node.get("input", "image") == "scene" + or (isinstance(declared_inputs, list) and "scene" in declared_inputs) + ) + if uses_scene_input and ( "inputs" in node or node.get("input", "image") != "scene" ): raise ValueError( @@ -606,6 +600,7 @@ def __init__(self) -> None: self._generators: Dict[str, BaseGenerator] = {} self._manifests: Dict[str, dict] = {} self._errors: Dict[str, str] = {} + self._lifecycle_lock = threading.RLock() self._legacy_imports = _LegacyImportManager() self._active_id: str = os.environ.get("SELECTED_MODEL_ID", "sf3d") @@ -727,38 +722,47 @@ def _assert_not_quarantined(model_id: str) -> None: def get_active(self) -> BaseGenerator: """Returns the active generator. Downloads and loads if necessary.""" - return self.get_ready_generator(self._active_id) + with self._lifecycle_lock: + return self.get_ready_generator(self._active_id) def get_ready_generator(self, model_id: str) -> BaseGenerator: """Load and return exactly ``model_id`` without consulting active state.""" - gen = self.get_generator(model_id) - downloaded = self._is_downloaded(model_id, gen) - if "model_sources" in self._manifests[model_id] and not downloaded: - raise RuntimeError( - "Model sources are incomplete. Download this node's weights " - "from the Modly Models page before generation." - ) - if not gen.is_loaded(): - if not downloaded: - if isinstance(gen, ExtensionProcess): - # Let the subprocess handle its own download logic during - # load() — some extensions (e.g. mv-adapter) need custom - # multi-repo downloads that the standard HF endpoint can't do. - pass - else: - gen._auto_download() - gen.load() - return gen + with self._lifecycle_lock: + gen = self.get_generator(model_id) + downloaded = self._is_downloaded(model_id, gen) + if "model_sources" in self._manifests[model_id] and not downloaded: + raise RuntimeError( + "Model sources are incomplete. Download this node's weights " + "from the Modly Models page before generation." + ) + if not gen.is_loaded(): + if not downloaded: + if isinstance(gen, ExtensionProcess): + # Let the subprocess handle its own download logic during + # load() — some extensions (e.g. mv-adapter) need custom + # multi-repo downloads that the standard HF endpoint can't do. + pass + else: + gen._auto_download() + gen.load() + return gen + + def activate_ready_generator(self, model_id: str) -> BaseGenerator: + """Atomically make ``model_id`` active and return it ready for inference.""" + with self._lifecycle_lock: + self.switch_model(model_id) + return self.get_ready_generator(model_id) def model_status(self, model_id: str) -> dict: - gen = self.get_generator(model_id) - manifest = self._manifests[model_id] - return { - "id": model_id, - "name": manifest.get("name", gen.DISPLAY_NAME), - "downloaded": self._is_downloaded(model_id, gen), - "loaded": gen.is_loaded(), - } + with self._lifecycle_lock: + gen = self.get_generator(model_id) + manifest = self._manifests[model_id] + return { + "id": model_id, + "name": manifest.get("name", gen.DISPLAY_NAME), + "downloaded": self._is_downloaded(model_id, gen), + "loaded": gen.is_loaded(), + } def get_generator(self, model_id: str) -> BaseGenerator: self._assert_not_quarantined(model_id) @@ -777,16 +781,17 @@ def get_manifest(self, model_id: str) -> dict: def switch_model(self, model_id: str) -> None: """Switches the active model. Unloads the previous one if different.""" - self._assert_not_quarantined(model_id) - if model_id not in self._generators: - raise ValueError( - f"Unknown model ID: '{model_id}'. " - f"Available: {list(self._generators.keys())}" - ) - if model_id != self._active_id: - if self._active_id in self._generators: - self._generators[self._active_id].unload() - self._active_id = model_id + with self._lifecycle_lock: + self._assert_not_quarantined(model_id) + if model_id not in self._generators: + raise ValueError( + f"Unknown model ID: '{model_id}'. " + f"Available: {list(self._generators.keys())}" + ) + if model_id != self._active_id: + if self._active_id in self._generators: + self._generators[self._active_id].unload() + self._active_id = model_id # ------------------------------------------------------------------ # # Status @@ -801,31 +806,34 @@ def _is_downloaded(self, model_id: str, gen: BaseGenerator) -> bool: return gen.is_downloaded() def active_status(self) -> dict: - return self.model_status(self._active_id) + with self._lifecycle_lock: + return self.model_status(self._active_id) def all_status(self) -> list: - result = [] - for model_id, gen in self._generators.items(): - manifest = self._manifests[model_id] - result.append({ - "id": model_id, - "name": manifest.get("name", gen.DISPLAY_NAME), - "description": manifest.get("description", ""), - "version": manifest.get("version", ""), - "vram_gb": manifest.get("vram_gb", gen.VRAM_GB), - "hf_repo": manifest.get("hf_repo", ""), - "tags": manifest.get("tags", []), - "downloaded": self._is_downloaded(model_id, gen), - "loaded": gen.is_loaded(), - "active": model_id == self._active_id, - }) - return result + with self._lifecycle_lock: + result = [] + for model_id, gen in self._generators.items(): + manifest = self._manifests[model_id] + result.append({ + "id": model_id, + "name": manifest.get("name", gen.DISPLAY_NAME), + "description": manifest.get("description", ""), + "version": manifest.get("version", ""), + "vram_gb": manifest.get("vram_gb", gen.VRAM_GB), + "hf_repo": manifest.get("hf_repo", ""), + "tags": manifest.get("tags", []), + "downloaded": self._is_downloaded(model_id, gen), + "loaded": gen.is_loaded(), + "active": model_id == self._active_id, + }) + return result def params_schema(self, model_id: Optional[str] = None) -> list: - target_id = model_id or self._active_id - if target_id not in self._generators: - raise KeyError(target_id) - return self._generators[target_id].params_schema() + with self._lifecycle_lock: + target_id = model_id or self._active_id + if target_id not in self._generators: + raise KeyError(target_id) + return self._generators[target_id].params_schema() # ------------------------------------------------------------------ # # Paths update & shutdown diff --git a/api/services/generators/base.py b/api/services/generators/base.py index 0b5169a2..adeb9802 100644 --- a/api/services/generators/base.py +++ b/api/services/generators/base.py @@ -170,11 +170,13 @@ def generate_artifact( ) -> Path: """Generate from a validated typed artifact. - New extensions should override this method. The default delegates to - ``generate`` with the canonical path so scene-capable extensions built - against the pre-release contract remain compatible. + Typed-artifact extensions must implement this explicitly. Falling back + to ``generate`` would pass a filesystem ``Path`` to the legacy + image-bytes ABI and fail far from the actual contract violation. """ - return self.generate(artifact_path, params, progress_cb, cancel_event) # type: ignore[arg-type] + raise NotImplementedError( + f"{type(self).__name__} does not implement {input_kind} artifact generation" + ) def _check_cancelled(self, cancel_event: Optional[threading.Event]) -> None: """Raises GenerationCancelled if cancel_event is set.""" diff --git a/api/tests/test_base_generator.py b/api/tests/test_base_generator.py index 580b3d25..9adfe00c 100644 --- a/api/tests/test_base_generator.py +++ b/api/tests/test_base_generator.py @@ -43,6 +43,12 @@ def test_download_progress_print_is_ascii_safe(self) -> None: self.assertTrue(first_line.endswith("...")) self.assertTrue(all(ord(ch) < 128 for ch in printed)) + def test_typed_artifact_generation_requires_an_explicit_override(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + gen = _DownloadGen(Path(tmp) / "model", Path(tmp) / "out") + with self.assertRaisesRegex(NotImplementedError, "scene artifact generation"): + gen.generate_artifact("scene", Path(tmp) / "scene-manifest.json", {}) + if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/api/tests/test_generator_registry.py b/api/tests/test_generator_registry.py index a4814b9e..76691e1b 100644 --- a/api/tests/test_generator_registry.py +++ b/api/tests/test_generator_registry.py @@ -155,7 +155,7 @@ def test_legacy_generator_supports_eager_and_lazy_sibling_imports(self) -> None: self.registry.reload() self.assertNotIn(str(extension.resolve()), sys.path) - def test_scene_io_is_registered_but_capture_and_video_are_rejected(self) -> None: + def test_scene_and_existing_custom_io_types_are_registered(self) -> None: for extension_id, input_kind in (("scene-io", "scene"), ("capture-io", "capture"), ("video-io", "video")): extension = self._make_extension(extension_id) manifest = { @@ -174,8 +174,8 @@ def test_scene_io_is_registered_but_capture_and_video_are_rejected(self) -> None self.registry.initialize() self.assertEqual(self.registry.get_manifest("scene-io/generate")["input"], "scene") - self.assertIn("capture-io/generate", self.registry.load_errors()) - self.assertIn("video-io/generate", self.registry.load_errors()) + self.assertEqual(self.registry.get_manifest("capture-io/generate")["input"], "capture") + self.assertEqual(self.registry.get_manifest("video-io/generate")["input"], "video") def test_scene_input_rejects_multi_input_shapes_but_image_multi_can_output_scene(self) -> None: cases = { @@ -258,6 +258,37 @@ def test_declared_sources_block_generation_even_when_generator_overrides_readine with self.assertRaisesRegex(RuntimeError, "Model sources are incomplete"): self.registry.get_active() + def test_activate_ready_generator_switches_before_loading_exact_model(self) -> None: + class Generator: + DISPLAY_NAME = "test" + def __init__(self): + self.loaded = False + self.unloads = 0 + def is_downloaded(self): return True + def is_loaded(self): return self.loaded + def load(self): self.loaded = True + def unload(self): + self.loaded = False + self.unloads += 1 + + first = Generator() + second = Generator() + first.loaded = True + self.registry._generators = {"demo/a": first, "demo/b": second} + self.registry._manifests = { + "demo/a": {"name": "A"}, + "demo/b": {"name": "B"}, + } + self.registry._active_id = "demo/a" + + selected = self.registry.activate_ready_generator("demo/b") + + self.assertIs(selected, second) + self.assertEqual(self.registry._active_id, "demo/b") + self.assertFalse(first.loaded) + self.assertEqual(first.unloads, 1) + self.assertTrue(second.loaded) + def test_reload_preserves_legacy_path_owned_by_the_host(self) -> None: extension = self._make_extension("host-owned-path") self._write_manifest(extension, extension_id="host-owned-path") diff --git a/api/tests/test_model_router.py b/api/tests/test_model_router.py index 3abaef0b..6196f3cb 100644 --- a/api/tests/test_model_router.py +++ b/api/tests/test_model_router.py @@ -2,6 +2,7 @@ import json import sys import tempfile +import threading import types import unittest from pathlib import Path @@ -197,5 +198,49 @@ def test_composite_model_unload_route_uses_path_converter(self) -> None: self.assertEqual(model_router.Request.__module__, "urllib.request") +class ModelLifecycleEndpointTests(unittest.TestCase): + def test_status_and_switch_do_not_block_the_event_loop_while_lifecycle_is_contended(self) -> None: + class ContendedRegistry: + def __init__(self): + self.lock = threading.Lock() + def active_status(self): + with self.lock: + return {"id": "demo/a", "loaded": True} + def switch_model(self, model_id): + with self.lock: + return None + + registry = ContendedRegistry() + previous = model_router.generator_registry + model_router.generator_registry = registry + + async def assert_responsive(call): + registry.lock.acquire() + release = threading.Timer(0.25, registry.lock.release) + release.start() + try: + start = asyncio.get_running_loop().time() + task = asyncio.create_task(call()) + await asyncio.sleep(0.02) + elapsed = asyncio.get_running_loop().time() - start + self.assertLess(elapsed, 0.15) + return await task + finally: + release.join() + + async def run(): + status = await assert_responsive(model_router.model_status) + switched = await assert_responsive(lambda: model_router.switch_model("demo/b")) + return status, switched + + try: + status, switched = asyncio.run(run()) + finally: + model_router.generator_registry = previous + + self.assertEqual(status["id"], "demo/a") + self.assertEqual(switched, {"active": "demo/b"}) + + if __name__ == "__main__": unittest.main() diff --git a/api/tests/test_scene_generation.py b/api/tests/test_scene_generation.py index a804171a..14288e6e 100644 --- a/api/tests/test_scene_generation.py +++ b/api/tests/test_scene_generation.py @@ -1,7 +1,9 @@ import asyncio import json import tempfile +import threading import unittest +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from unittest.mock import patch @@ -35,7 +37,7 @@ def setUp(self): def tearDown(self): for item in reversed(self.patches): item.stop() - generation._jobs.clear(); generation._cancel_events.clear(); generation._cancelled.clear(); generation._completed_at.clear() + generation._jobs.clear(); generation._cancel_events.clear(); generation._cancelled.clear(); generation._completed_at.clear(); generation._job_generators.clear() self.tmp.cleanup() def test_generic_route_queues_typed_scene_and_strips_reserved_params(self): @@ -52,6 +54,7 @@ def test_generic_route_queues_typed_scene_and_strips_reserved_params(self): self.assertNotIn("input_kind", queued.args[2]) self.assertEqual(queued.args[5], "demo/scene") self.assertEqual(result["job_id"], queued.args[0]) + self.assertFalse(self.registry.switched) def test_generic_route_rejects_unsupported_kind_and_model_mismatch(self): for kind in ("video", "capture", "image"): @@ -84,7 +87,9 @@ def generate_artifact(self, kind, path, params, progress_cb, cancel_event=None): generator = Generator() registry_stub = type("Registry", (), { - "get_ready_generator": lambda self, model_id: generator if model_id == "demo/a" else (_ for _ in ()).throw(ValueError(f"Unknown model ID: {model_id}")), + "model_status": lambda self, model_id: {"name": model_id, "downloaded": True, "loaded": True}, + "get_generator": lambda self, model_id: generator if model_id == "demo/a" else (_ for _ in ()).throw(ValueError(f"Unknown model ID: {model_id}")), + "activate_ready_generator": lambda self, model_id: generator if model_id == "demo/a" else (_ for _ in ()).throw(ValueError(f"Unknown model ID: {model_id}")), "get_active": lambda self: (_ for _ in ()).throw(AssertionError("mutable active model must not be used")), })() job_id = "pinned-scene" @@ -100,7 +105,9 @@ def generate_artifact(self, kind, path, params, progress_cb, cancel_event=None): def test_missing_pinned_model_fails_actionably(self): registry_stub = type("Registry", (), { - "get_ready_generator": lambda self, model_id: (_ for _ in ()).throw(ValueError(f"Unknown model ID: {model_id}")), + "model_status": lambda self, model_id: {"name": model_id, "downloaded": True, "loaded": False}, + "get_generator": lambda self, model_id: object(), + "activate_ready_generator": lambda self, model_id: (_ for _ in ()).throw(ValueError(f"Unknown model ID: {model_id}")), "get_active": lambda self: (_ for _ in ()).throw(AssertionError("must not use active model")), })() job_id = "missing-scene" @@ -113,3 +120,214 @@ def test_missing_pinned_model_fails_actionably(self): )) self.assertEqual(generation._jobs[job_id].status, "error") self.assertIn("Unknown model ID: demo/missing", generation._jobs[job_id].error) + + def test_interleaved_pinned_jobs_serialize_model_lifecycle_and_cancel_exact_job(self): + class Proc: + def __init__(self, owner): + self.owner = owner + self.alive = True + def poll(self): + return None if self.alive else 0 + def kill(self): + self.alive = False + self.owner.killed = True + + class Generator: + def __init__(self, model_id): + self.model_id = model_id + self.outputs_dir = None + self.loaded = False + self.started = threading.Event() + self.killed = False + self._loaded = True + self._proc = Proc(self) + def is_loaded(self): return self.loaded + def is_downloaded(self): return True + def load(self): self.loaded = True + def unload(self): self.loaded = False + def generate_artifact(self, kind, path, params, progress_cb, cancel_event=None): + self.started.set() + if self.model_id == "demo/a": + if cancel_event is None or not cancel_event.wait(2): + raise AssertionError("first job was not cancelled") + raise generation.GenerationCancelled() + output = Path(self.outputs_dir) / "result.glb" + output.write_bytes(b"glb") + return output + + class Registry: + def __init__(self): + self.generators = {model_id: Generator(model_id) for model_id in ("demo/a", "demo/b")} + self.active_id = "demo/a" + def get_generator(self, model_id): return self.generators[model_id] + def model_status(self, model_id): + gen = self.generators[model_id] + return {"name": model_id, "downloaded": True, "loaded": gen.loaded} + def switch_model(self, model_id): + if model_id != self.active_id: + self.generators[self.active_id].unload() + self.active_id = model_id + def activate_ready_generator(self, model_id): + self.switch_model(model_id) + gen = self.generators[model_id] + gen.load() + if sum(candidate.loaded for candidate in self.generators.values()) != 1: + raise AssertionError("more than one generator is resident") + return gen + def get_active(self): raise AssertionError("pinned jobs must not consult mutable active state") + + registry_stub = Registry() + # Reproduce the request interleaving: A switches first, then B switches + # before either background task has started. + registry_stub.switch_model("demo/a") + registry_stub.switch_model("demo/b") + for job_id in ("job-a", "job-b"): + generation._jobs[job_id] = generation.JobStatus(job_id=job_id, status="pending", progress=0) + generation._cancel_events[job_id] = threading.Event() + + async def run_both(): + first = asyncio.create_task(generation._run_generation( + "job-a", generation.TypedArtifactInput("scene", self.manifest.resolve()), {}, + "Workflows", "mesh", "demo/a", + )) + second = asyncio.create_task(generation._run_generation( + "job-b", generation.TypedArtifactInput("scene", self.manifest.resolve()), {}, + "Workflows", "mesh", "demo/b", + )) + started = await asyncio.to_thread(registry_stub.generators["demo/a"].started.wait, 1) + self.assertTrue(started) + await asyncio.sleep(0.05) + self.assertFalse(registry_stub.generators["demo/b"].started.is_set()) + await generation.cancel_job("job-a") + await asyncio.gather(first, second) + + with patch.object(generation, "generator_registry", registry_stub): + asyncio.run(run_both()) + + self.assertTrue(registry_stub.generators["demo/a"].killed) + self.assertFalse(registry_stub.generators["demo/b"].killed) + self.assertEqual(generation._jobs["job-a"].status, "cancelled") + self.assertEqual(generation._jobs["job-b"].status, "done") + self.assertEqual(registry_stub.active_id, "demo/b") + self.assertFalse(registry_stub.generators["demo/a"].loaded) + self.assertTrue(registry_stub.generators["demo/b"].loaded) + + def test_cancelling_queued_generation_does_not_wedge_later_jobs(self): + release_first = threading.Event() + first_started = threading.Event() + + class Generator: + outputs_dir = None + + def __init__(self): + self.calls = 0 + + def generate(self, image_bytes, params, progress_cb, cancel_event=None): + self.calls += 1 + call_number = self.calls + if call_number == 1: + first_started.set() + if not release_first.wait(2): + raise AssertionError("first generation was not released") + output = Path(self.outputs_dir) / f"result-{call_number}.glb" + output.write_bytes(b"glb") + return output + + generator = Generator() + registry_stub = type("Registry", (), { + "get_generator": lambda self, model_id: generator, + "model_status": lambda self, model_id: { + "name": model_id, "downloaded": True, "loaded": True, + }, + "activate_ready_generator": lambda self, model_id: generator, + "get_active": lambda self: (_ for _ in ()).throw( + AssertionError("pinned jobs must not use mutable active state") + ), + })() + for job_id in ("blocking", "cancelled-waiter", "after-cancel"): + generation._jobs[job_id] = generation.JobStatus( + job_id=job_id, status="pending", progress=0, + ) + generation._cancel_events[job_id] = threading.Event() + + async def run_scenario(): + first = asyncio.create_task(generation._run_generation( + "blocking", b"image", {}, "Workflows", "mesh", "demo/a", + )) + started = await asyncio.to_thread(first_started.wait, 1) + self.assertTrue(started) + + queued = asyncio.create_task(generation._run_generation( + "cancelled-waiter", b"image", {}, "Workflows", "mesh", "demo/b", + )) + await asyncio.sleep(0) + queued.cancel() + with self.assertRaises(asyncio.CancelledError): + await queued + + release_first.set() + await asyncio.wait_for(first, 1) + await asyncio.wait_for(generation._run_generation( + "after-cancel", b"image", {}, "Workflows", "mesh", "demo/c", + ), 1) + + with patch.object(generation, "generator_registry", registry_stub): + asyncio.run(run_scenario()) + + self.assertEqual(generation._jobs["cancelled-waiter"].status, "cancelled") + self.assertEqual(generation._jobs["after-cancel"].status, "done") + self.assertEqual(generator.calls, 2) + + def test_pinned_jobs_do_not_starve_small_default_executor(self): + class Generator: + outputs_dir = None + + def __init__(self): + self.calls = 0 + + def generate(self, image_bytes, params, progress_cb, cancel_event=None): + self.calls += 1 + output = Path(self.outputs_dir) / f"result-{self.calls}.glb" + output.write_bytes(b"glb") + return output + + generator = Generator() + registry_stub = type("Registry", (), { + "get_generator": lambda self, model_id: generator, + "model_status": lambda self, model_id: { + "name": model_id, "downloaded": True, "loaded": True, + }, + "activate_ready_generator": lambda self, model_id: generator, + "get_active": lambda self: (_ for _ in ()).throw( + AssertionError("pinned jobs must not use mutable active state") + ), + })() + job_ids = ("small-pool-a", "small-pool-b", "small-pool-c") + for job_id in job_ids: + generation._jobs[job_id] = generation.JobStatus( + job_id=job_id, status="pending", progress=0, + ) + generation._cancel_events[job_id] = threading.Event() + + default_executor = ThreadPoolExecutor(max_workers=2) + + async def run_scenario(): + asyncio.get_running_loop().set_default_executor(default_executor) + await asyncio.wait_for(asyncio.gather(*( + generation._run_generation( + job_id, b"image", {}, "Workflows", "mesh", f"demo/{job_id}", + ) + for job_id in job_ids + )), 2) + + try: + with patch.object(generation, "generator_registry", registry_stub): + asyncio.run(run_scenario()) + finally: + default_executor.shutdown(wait=True) + + self.assertEqual(generator.calls, 3) + self.assertEqual( + [generation._jobs[job_id].status for job_id in job_ids], + ["done", "done", "done"], + ) diff --git a/api/tests/test_scene_input.py b/api/tests/test_scene_input.py index 43bd9946..7d570d10 100644 --- a/api/tests/test_scene_input.py +++ b/api/tests/test_scene_input.py @@ -34,14 +34,18 @@ def test_rejects_traversal_absolute_encoded_and_non_manifest_json(self): with self.subTest(path=path), self.assertRaises(ValueError): validate_scene_input(self.workspace, path) - def test_rejects_symlinks_missing_assets_and_oversized_manifest(self): + def test_rejects_symlink_escape_when_supported(self): outside = Path(self.tmp.name) / "outside" outside.mkdir() (outside / "scene-manifest.json").write_text(self.manifest.read_text()) - (self.workspace / "Workflows" / "link").symlink_to(outside, target_is_directory=True) + try: + (self.workspace / "Workflows" / "link").symlink_to(outside, target_is_directory=True) + except OSError as exc: + self.skipTest(f"directory symlinks are unavailable on this host: {exc}") with self.assertRaises(ValueError): validate_scene_input(self.workspace, "Workflows/link") + def test_rejects_missing_assets_and_oversized_manifest(self): data = json.loads(self.manifest.read_text()) data["assets"] = [{"path": "missing.glb"}] self.manifest.write_text(json.dumps(data)) diff --git a/api/tests/test_workflow_runs_lifecycle.py b/api/tests/test_workflow_runs_lifecycle.py index f8fe7371..7360cc03 100644 --- a/api/tests/test_workflow_runs_lifecycle.py +++ b/api/tests/test_workflow_runs_lifecycle.py @@ -1,12 +1,15 @@ import asyncio +import tempfile import threading import time import unittest +from pathlib import Path from fastapi import BackgroundTasks import routers.generation as generation import routers.workflow_runs as workflow_runs +import services.generator_registry as registry_module from schemas.generation import JobStatus @@ -30,6 +33,9 @@ class _FakeRegistry: def get_generator(self, model_id: str) -> object: return object() + def get_manifest(self, model_id: str) -> dict: + return {"output": "mesh"} + def switch_model(self, model_id: str) -> None: pass @@ -40,6 +46,7 @@ def _clear_job_stores() -> None: generation._cancel_events, generation._cancelled, generation._completed_at, + generation._job_generators, ): store.clear() @@ -91,6 +98,106 @@ def test_cancel_run_records_completion_so_it_can_be_purged(self) -> None: # Without a _completed_at stamp the purge sweep can never evict a cancelled run. self.assertIn(run_id, generation._completed_at) + def test_workflow_and_generate_routes_share_exact_model_queue_and_cancellation(self) -> None: + class Proc: + def __init__(self, owner): + self.owner = owner + self.alive = True + def poll(self): return None if self.alive else 0 + def kill(self): + self.alive = False + self.owner.killed = True + + class Generator: + def __init__(self, model_id): + self.model_id = model_id + self.outputs_dir = None + self.loaded = False + self.started = threading.Event() + self.killed = False + self._loaded = True + self._proc = Proc(self) + def is_loaded(self): return self.loaded + def load(self): self.loaded = True + def unload(self): self.loaded = False + def generate(self, image_bytes, params, progress_cb, cancel_event=None): + self.started.set() + if self.model_id == "demo/workflow": + if cancel_event is None or not cancel_event.wait(2): + raise AssertionError("workflow job was not cancelled") + raise generation.GenerationCancelled() + output = Path(self.outputs_dir) / "result.glb" + output.write_bytes(b"glb") + return output + + class Registry: + def __init__(self): + ids = ("demo/workflow", "demo/generate") + self.generators = {model_id: Generator(model_id) for model_id in ids} + self.active_id = "demo/generate" + def get_generator(self, model_id): return self.generators[model_id] + def get_manifest(self, model_id): return {"output": "mesh"} + def model_status(self, model_id): + gen = self.generators[model_id] + return {"name": model_id, "downloaded": True, "loaded": gen.loaded} + def activate_ready_generator(self, model_id): + if model_id != self.active_id: + self.generators[self.active_id].unload() + self.active_id = model_id + gen = self.generators[model_id] + gen.load() + if sum(candidate.loaded for candidate in self.generators.values()) != 1: + raise AssertionError("more than one generator is resident") + return gen + def get_active(self): raise AssertionError("queued routes must use their pinned model id") + + fake_registry = Registry() + previous_generation_registry = generation.generator_registry + previous_workflow_registry = workflow_runs.generator_registry + previous_workspace = registry_module.WORKSPACE_DIR + with tempfile.TemporaryDirectory() as tmp: + registry_module.WORKSPACE_DIR = Path(tmp) + generation.generator_registry = fake_registry + workflow_runs.generator_registry = fake_registry + + async def run_cross_route_jobs(): + workflow_background = BackgroundTasks() + workflow_response = await workflow_runs.create_run_from_image( + workflow_background, _FakeUpload(), "demo/workflow", "Workflows", "{}", + ) + generate_background = BackgroundTasks() + generate_response = await generation.generate_from_image( + generate_background, _FakeUpload(), "demo/generate", "Workflows", + "quad", False, 1024, "{}", + ) + self.assertEqual(workflow_background.tasks[0].args[5], "demo/workflow") + self.assertEqual(generate_background.tasks[0].args[5], "demo/generate") + + first = asyncio.create_task(workflow_background.tasks[0]()) + second = asyncio.create_task(generate_background.tasks[0]()) + started = await asyncio.to_thread( + fake_registry.generators["demo/workflow"].started.wait, 1, + ) + self.assertTrue(started) + await asyncio.sleep(0.05) + self.assertFalse(fake_registry.generators["demo/generate"].started.is_set()) + await workflow_runs.cancel_run(workflow_response["run_id"]) + await asyncio.gather(first, second) + return workflow_response["run_id"], generate_response["job_id"] + + try: + workflow_job_id, generate_job_id = asyncio.run(run_cross_route_jobs()) + finally: + generation.generator_registry = previous_generation_registry + workflow_runs.generator_registry = previous_workflow_registry + registry_module.WORKSPACE_DIR = previous_workspace + + self.assertTrue(fake_registry.generators["demo/workflow"].killed) + self.assertFalse(fake_registry.generators["demo/generate"].killed) + self.assertEqual(generation._jobs[workflow_job_id].status, "cancelled") + self.assertEqual(generation._jobs[generate_job_id].status, "done") + self.assertEqual(fake_registry.active_id, "demo/generate") + if __name__ == "__main__": unittest.main() diff --git a/api/tests/test_workflow_runs_router.py b/api/tests/test_workflow_runs_router.py index f045bef8..7537a484 100644 --- a/api/tests/test_workflow_runs_router.py +++ b/api/tests/test_workflow_runs_router.py @@ -24,6 +24,9 @@ class _FakeRegistry: def get_generator(self, model_id: str) -> object: return object() + def get_manifest(self, model_id: str) -> dict: + return {"output": "mesh"} + def switch_model(self, model_id: str) -> None: pass @@ -79,7 +82,7 @@ def _collection_forwarded(self, collection: str) -> str: params="{}", ) ) - # add_task(_run_generation, job_id, image_bytes, full_params, collection) + # add_task(_run_generation, ..., collection, output_kind, model_id) return background.tasks[0].args[3] def test_collection_is_forwarded_to_the_run(self) -> None: diff --git a/electron/main/bounded-file-reader.test.mjs b/electron/main/bounded-file-reader.test.mjs new file mode 100644 index 00000000..03071180 --- /dev/null +++ b/electron/main/bounded-file-reader.test.mjs @@ -0,0 +1,48 @@ +import assert from 'node:assert/strict' +import { readFileSync } from 'node:fs' +import { resolve } from 'node:path' +import test from 'node:test' + +const { + MAX_SCENE_MANIFEST_BYTES, + readLocalFileBase64, +} = await import(new URL('./bounded-file-reader.ts', import.meta.url).href) + +test('the authoritative readFileBase64 IPC delegates to the bounded reader', () => { + const handlers = readFileSync(resolve('electron/main/ipc-handlers.ts'), 'utf8') + assert.match( + handlers, + /ipcMain\.handle\('fs:readFileBase64',[\s\S]*?readLocalFileBase64\(filePath\)/, + ) +}) + +test('oversized scene manifests are rejected before their bytes are read', async () => { + let reads = 0 + await assert.rejects( + readLocalFileBase64('/workspace/room/scene-manifest.json', { + statFile: async () => ({ size: MAX_SCENE_MANIFEST_BYTES + 1, isFile: () => true }), + readBytes: async () => { reads++; return new Uint8Array() }, + }), + /exceeds the 1 MiB limit/, + ) + assert.equal(reads, 0) +}) + +test('scene manifests at the bound are read and encoded', async () => { + const bytes = new TextEncoder().encode('{"schema":"modly.scene-manifest.v1"}') + const encoded = await readLocalFileBase64('/workspace/room/scene-manifest.json', { + statFile: async () => ({ size: MAX_SCENE_MANIFEST_BYTES, isFile: () => true }), + readBytes: async () => bytes, + }) + assert.equal(encoded, Buffer.from(bytes).toString('base64')) +}) + +test('a manifest that grows after stat is rejected before base64 transfer', async () => { + await assert.rejects( + readLocalFileBase64('/workspace/room/scene-manifest.json', { + statFile: async () => ({ size: 1, isFile: () => true }), + readBytes: async () => new Uint8Array(MAX_SCENE_MANIFEST_BYTES + 1), + }), + /exceeds the 1 MiB limit/, + ) +}) diff --git a/electron/main/bounded-file-reader.ts b/electron/main/bounded-file-reader.ts new file mode 100644 index 00000000..f7b18ba4 --- /dev/null +++ b/electron/main/bounded-file-reader.ts @@ -0,0 +1,40 @@ +import { basename } from 'node:path' +import { readFile, stat } from 'node:fs/promises' + +export const SCENE_MANIFEST_FILE_NAME = 'scene-manifest.json' +export const MAX_SCENE_MANIFEST_BYTES = 1024 * 1024 + +type FileReadDependencies = { + statFile: (filePath: string) => Promise<{ size: number; isFile: () => boolean }> + readBytes: (filePath: string) => Promise +} + +const DEFAULT_DEPENDENCIES: FileReadDependencies = { + statFile: stat, + readBytes: readFile, +} + +export async function readLocalFileBase64( + filePath: string, + dependencies: FileReadDependencies = DEFAULT_DEPENDENCIES, +): Promise { + if (typeof filePath !== 'string' || filePath.trim().length === 0) { + throw new Error('fs:readFileBase64 requires a non-empty file path') + } + + const isSceneManifest = basename(filePath) === SCENE_MANIFEST_FILE_NAME + if (isSceneManifest) { + const fileInfo = await dependencies.statFile(filePath) + if (!fileInfo.isFile()) throw new Error('Scene manifest is not a file') + if (fileInfo.size > MAX_SCENE_MANIFEST_BYTES) { + throw new Error('Scene manifest exceeds the 1 MiB limit') + } + } + + const bytes = await dependencies.readBytes(filePath) + // Recheck the transferred bytes in case the file grew between stat and read. + if (isSceneManifest && bytes.byteLength > MAX_SCENE_MANIFEST_BYTES) { + throw new Error('Scene manifest exceeds the 1 MiB limit') + } + return Buffer.from(bytes).toString('base64') +} diff --git a/electron/main/extension-install-utils.test.mjs b/electron/main/extension-install-utils.test.mjs index 4dcf42da..3c5ac9bb 100644 --- a/electron/main/extension-install-utils.test.mjs +++ b/electron/main/extension-install-utils.test.mjs @@ -85,7 +85,7 @@ test('validateInstallManifest accepts multi-source nodes and preserves legacy sh }, { hasEntryFile: () => false, hasGeneratorFile: () => true }, 'repository')) }) -test('validateInstallManifest accepts scene IO and rejects undeclared future artifact kinds', () => { +test('validateInstallManifest accepts scene IO without rejecting third-party artifact kinds', () => { const mod = loadModule() const files = { hasEntryFile: () => false, hasGeneratorFile: () => true } assert.doesNotThrow(() => mod.validateInstallManifest({ @@ -93,10 +93,10 @@ test('validateInstallManifest accepts scene IO and rejects undeclared future art nodes: [{ id: 'normalize', input: 'scene', output: 'scene' }], }, files, 'repository')) for (const input of ['capture', 'video']) { - assert.throws(() => mod.validateInstallManifest({ + assert.doesNotThrow(() => mod.validateInstallManifest({ id: 'future-model', generator_class: 'Generator', nodes: [{ id: 'future', input, output: 'scene' }], - }, files, 'repository'), /supported artifact type/) + }, files, 'repository')) } }) diff --git a/electron/main/extension-install-utils.ts b/electron/main/extension-install-utils.ts index 41e9798c..c50e85e0 100644 --- a/electron/main/extension-install-utils.ts +++ b/electron/main/extension-install-utils.ts @@ -66,21 +66,12 @@ export function validateInstallManifest( const isProcess = manifest.type === 'process' const entryFile = manifest.entry ?? 'processor.js' const nodes = Array.isArray(manifest.nodes) ? manifest.nodes.filter((node) => node?.id) : [] - const allowedIo = new Set(['image', 'text', 'mesh', 'audio', 'scene']) - if (manifest.model_sources !== undefined) { throw new Error('manifest.json: model_sources must be declared on a model node') } for (const node of Array.isArray(manifest.nodes) ? manifest.nodes : []) { - const declaredInputs = node.inputs === undefined ? [node.input ?? 'image'] : node.inputs - if (!Array.isArray(declaredInputs) || declaredInputs.length === 0 - || declaredInputs.some((value) => typeof value !== 'string' || !allowedIo.has(value))) { - throw new Error(`manifest.json: ${node.id ?? 'node'}.input must use a supported artifact type`) - } + const declaredInputs = Array.isArray(node.inputs) ? node.inputs : [node.input ?? 'image'] const output = node.output ?? 'mesh' - if (typeof output !== 'string' || !allowedIo.has(output)) { - throw new Error(`manifest.json: ${node.id ?? 'node'}.output must use a supported artifact type`) - } assertSupportedSceneNodeShape(isProcess ? 'process' : 'model', node, declaredInputs, output) if (node.model_sources === undefined) continue if (isProcess) { diff --git a/electron/main/ipc-handlers.ts b/electron/main/ipc-handlers.ts index 9e5131de..09f1b649 100644 --- a/electron/main/ipc-handlers.ts +++ b/electron/main/ipc-handlers.ts @@ -68,6 +68,7 @@ import { } from './extension-install-recovery' import { registerWorkspaceAssetLibraryIpcHandlers } from './artifact-registry-service' import { updatesSupported } from './updater' +import { readLocalFileBase64 } from './bounded-file-reader' type WindowGetter = () => BrowserWindow | null const pExecFile = promisify(execFile) @@ -379,13 +380,7 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe }) // Read local file → base64 (bypasses file:// restrictions in the renderer) - ipcMain.handle('fs:readFileBase64', async (_, filePath: string) => { - if (typeof filePath !== 'string' || filePath.trim().length === 0) { - throw new Error('fs:readFileBase64 requires a non-empty file path') - } - const buffer = await readFile(filePath) - return buffer.toString('base64') - }) + ipcMain.handle('fs:readFileBase64', (_, filePath: string) => readLocalFileBase64(filePath)) ipcMain.handle('fs:readScreenshotDataUrl', async (_, filename: string) => { const filePath = app.isPackaged @@ -845,10 +840,10 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe nodes?: { id: string name?: string - input?: 'mesh' | 'image' | 'text' | 'audio' | 'scene' - inputs?: ('mesh' | 'image' | 'text' | 'audio' | 'scene')[] + input?: string + inputs?: string[] input_labels?: string[] - output?: 'mesh' | 'image' | 'text' | 'audio' | 'scene' + output?: string params_schema?: unknown[] param_defaults?: Record hf_repo?: string @@ -874,14 +869,9 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe if (parsed.model_sources !== undefined) { throw new Error('manifest.json: model_sources must be declared on a model node') } - const allowedIo = new Set(['image', 'text', 'mesh', 'audio', 'scene']) const nodes = (parsed.nodes ?? []).map(n => { - const declaredInputs = n.inputs ?? [n.input ?? 'image'] - for (const input of declaredInputs) { - if (!allowedIo.has(input)) throw new Error(`manifest.json: unsupported node input type "${input}"`) - } + const declaredInputs = Array.isArray(n.inputs) ? n.inputs : [n.input ?? 'image'] const output = n.output ?? 'mesh' - if (!allowedIo.has(output)) throw new Error(`manifest.json: unsupported node output type "${output}"`) assertSupportedSceneNodeShape(parsed.type === 'process' ? 'process' : 'model', n, declaredInputs, output) if (parsed.type === 'process' && n.model_sources !== undefined) { throw new Error('manifest.json: model_sources is supported only for model nodes') diff --git a/src/areas/workflows/WorkflowsPage.tsx b/src/areas/workflows/WorkflowsPage.tsx index 32452774..8fc4cea6 100644 --- a/src/areas/workflows/WorkflowsPage.tsx +++ b/src/areas/workflows/WorkflowsPage.tsx @@ -68,7 +68,7 @@ const IO_STYLES: Record<'image' | 'text' | 'mesh' | 'audio' | 'scene', string> = image: 'bg-sky-500/15 text-sky-400 border-sky-500/25', mesh: 'bg-violet-500/15 text-violet-400 border-violet-500/25', text: 'bg-amber-500/15 text-amber-400 border-amber-500/25', - scene: 'bg-emerald-500/15 text-emerald-400 border-emerald-500/25', + scene: 'bg-pink-500/15 text-pink-400 border-pink-500/25', } function IoBadge({ type }: { type: 'image' | 'text' | 'mesh' | 'audio' | 'scene' }) { @@ -101,7 +101,7 @@ const PANEL_BUILTIN_NODES = [ { type: 'imageNode', label: 'Image', color: '#38bdf8', icon: <> }, { type: 'textNode', label: 'Text', color: '#fbbf24', icon: <> }, { type: 'meshNode', label: 'Load 3D Mesh', color: '#a78bfa', icon: <> }, - { type: 'sceneNode', label: 'Load Scene', color: '#34d399', icon: <> }, + { type: 'sceneNode', label: 'Load Scene', color: '#f472b6', icon: <> }, { type: 'outputNode', label: 'Add to Scene', color: '#a78bfa', icon: <> }, { type: 'previewNode', label: 'Preview Views', color: '#38bdf8', icon: <> }, { type: 'imagePreviewNode', label: 'Preview Image', color: '#38bdf8', icon: <> }, @@ -349,7 +349,7 @@ const BUILTIN_NODES = [ { type: 'imageNode', label: 'Image', color: '#38bdf8', description: 'Image input' }, { type: 'textNode', label: 'Text', color: '#fbbf24', description: 'Text input' }, { type: 'meshNode', label: 'Load 3D Mesh', color: '#a78bfa', description: 'Load a 3D mesh file or use current model' }, - { type: 'sceneNode', label: 'Load Scene', color: '#34d399', description: 'Load and validate a workspace scene directory' }, + { type: 'sceneNode', label: 'Load Scene', color: '#f472b6', description: 'Load and validate a workspace scene directory' }, { type: 'outputNode', label: 'Add to Scene', color: '#a78bfa', description: 'Output node — adds the mesh to the 3D scene' }, { type: 'previewNode', label: 'Preview Views', color: '#38bdf8', description: 'Displays multi-view image outputs in a 2×3 grid' }, { type: 'imagePreviewNode', label: 'Preview Image', color: '#38bdf8', description: 'Displays a single image output in the workflow' }, @@ -1379,7 +1379,7 @@ const MINI_NODE_TINTS: Record = { imageNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, textNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, meshNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, - sceneNode: { fill: 'rgba(52,211,153,0.22)', stroke: '#34d399' }, + sceneNode: { fill: 'rgba(244,114,182,0.22)', stroke: '#f472b6' }, extensionNode: { fill: 'rgba(167,139,250,0.24)', stroke: '#a78bfa' }, outputNode: { fill: 'rgba(56,189,248,0.22)', stroke: '#38bdf8' }, previewNode: { fill: 'rgba(56,189,248,0.22)', stroke: '#38bdf8' }, diff --git a/src/areas/workflows/nodes/ExtensionNode.tsx b/src/areas/workflows/nodes/ExtensionNode.tsx index 4be11a14..e5dc3143 100644 --- a/src/areas/workflows/nodes/ExtensionNode.tsx +++ b/src/areas/workflows/nodes/ExtensionNode.tsx @@ -16,7 +16,7 @@ const HANDLE_COLOR: Record = { image: '#38bdf8', mesh: '#a78bfa', text: '#fbbf24', - scene: '#34d399', + scene: '#f472b6', } const TAG_CLS: Record = { @@ -24,7 +24,7 @@ const TAG_CLS: Record = { image: 'border-sky-500/30 bg-sky-500/10 text-sky-400', mesh: 'border-violet-500/30 bg-violet-500/10 text-violet-400', text: 'border-amber-500/30 bg-amber-500/10 text-amber-400', - scene: 'border-emerald-500/30 bg-emerald-500/10 text-emerald-400', + scene: 'border-pink-500/30 bg-pink-500/10 text-pink-400', } // ─── Param control ──────────────────────────────────────────────────────────── diff --git a/src/areas/workflows/nodes/LoadSceneNode.tsx b/src/areas/workflows/nodes/LoadSceneNode.tsx index 2ac7dded..40f3af1c 100644 --- a/src/areas/workflows/nodes/LoadSceneNode.tsx +++ b/src/areas/workflows/nodes/LoadSceneNode.tsx @@ -1,18 +1,26 @@ import { useCallback, useLayoutEffect, useRef, useState } from 'react' import { Handle, Position, useReactFlow } from '@xyflow/react' -import type { WFNodeData } from '@shared/types/electron.d' +import type { ReactFlowInstance } from '@xyflow/react' +import type { WFNode, WFNodeData } from '@shared/types/electron.d' import BaseNode from './BaseNode' -import { resolveSceneSourceManifest } from '../workflowSceneSource' +import { + applySceneValidationResult, + invalidateValidatedScenePath, + resolveSceneSourceManifest, +} from '../workflowSceneSource' -const OUTPUT_COLOR = '#34d399' +const OUTPUT_COLOR = '#f472b6' async function validateAndPersistScenePath(args: { id: string - data: WFNodeData nextPath: string - updateNodeData: ReturnType['updateNodeData'] + updateNodeData: ReactFlowInstance['updateNodeData'] }): Promise { + args.updateNodeData(args.id, (node) => ({ + params: invalidateValidatedScenePath(node.data.params, args.nextPath), + })) + const settings = await window.electron.settings.get() const resolution = await resolveSceneSourceManifest({ scenePath: args.nextPath, @@ -20,33 +28,14 @@ async function validateAndPersistScenePath(args: { readFileBase64: window.electron.fs.readFileBase64, }) - if (!resolution.ok) { - args.updateNodeData(args.id, { - params: { - ...args.data.params, - path: args.nextPath, - manifestPath: undefined, - sceneRoot: undefined, - error: resolution.error, - }, - }) - return - } - - args.updateNodeData(args.id, { - params: { - ...args.data.params, - path: resolution.inputWorkspacePath, - manifestPath: resolution.manifestWorkspacePath, - sceneRoot: resolution.sceneRoot, - sourceKind: resolution.sourceKind, - error: undefined, - }, + args.updateNodeData(args.id, (node) => { + const params = applySceneValidationResult(node.data.params, args.nextPath, resolution) + return params ? { params } : {} }) } export default function LoadSceneNode({ id, data, selected }: { id: string; data: WFNodeData; selected?: boolean }) { - const { updateNodeData } = useReactFlow() + const { updateNodeData } = useReactFlow() const ioRowRef = useRef(null) const [handleTop, setHandleTop] = useState('50%') @@ -65,13 +54,13 @@ export default function LoadSceneNode({ id, data, selected }: { id: string; data const browseDirectory = useCallback(async () => { const path = await window.electron.fs.selectDirectory() if (!path) return - await validateAndPersistScenePath({ id, data, nextPath: path, updateNodeData }) - }, [id, data, updateNodeData]) + await validateAndPersistScenePath({ id, nextPath: path, updateNodeData }) + }, [id, updateNodeData]) const validatePath = useCallback(async () => { if (!scenePath.trim()) return - await validateAndPersistScenePath({ id, data, nextPath: scenePath, updateNodeData }) - }, [id, data, scenePath, updateNodeData]) + await validateAndPersistScenePath({ id, nextPath: scenePath, updateNodeData }) + }, [id, scenePath, updateNodeData]) return ( - scene + scene } handles={ @@ -106,20 +95,20 @@ export default function LoadSceneNode({ id, data, selected }: { id: string; data type="text" value={scenePath} placeholder="Scenes/castle or Scenes/castle/scene-manifest.json" - onChange={(event) => updateNodeData(id, { params: { ...data.params, path: event.target.value } })} - className="nodrag w-full rounded-lg border border-zinc-700 bg-zinc-800 px-2.5 py-2 text-[10px] text-zinc-200 placeholder-zinc-600 focus:outline-none focus:border-emerald-500/40" + onChange={(event) => updateNodeData(id, { params: invalidateValidatedScenePath(data.params, event.target.value) })} + className="nodrag w-full rounded-lg border border-zinc-700 bg-zinc-800 px-2.5 py-2 text-[10px] text-zinc-200 placeholder-zinc-600 focus:outline-none focus:border-pink-500/40" />
- -
{manifestPath ? ( -
-
Manifest: {manifestPath}
+
+
Manifest: {manifestPath}
{sceneRoot &&
sceneRoot: {sceneRoot}
}
) : ( diff --git a/src/areas/workflows/workflowRunStore.ts b/src/areas/workflows/workflowRunStore.ts index 82166635..ba6faf70 100644 --- a/src/areas/workflows/workflowRunStore.ts +++ b/src/areas/workflows/workflowRunStore.ts @@ -7,6 +7,7 @@ import type { WorkflowExtension } from './mockExtensions' import type { Workflow, WFNode, WFEdge } from '@shared/types/electron.d' import { isBranchStarter, isSceneOutput, resolveDataSource, reachesSceneOutput, nearestUpstreamWaits } from './nodeBehaviors' import { assignSlotFilePaths } from './slotInputs' +import type { SlotInputType } from './slotInputs' // ─── Types ──────────────────────────────────────────────────────────────────── @@ -337,7 +338,10 @@ async function executeExtensionNode( const incomingEdges = workflow.edges.filter((e) => e.target === node.id) if (ext?.inputs && ext.inputs.length > 1) { - const inputTypes = ext.inputs + // Scene is intentionally a single-input-only model contract. The guard + // above rejects it before this multi-slot path, and filtering it here also + // narrows the remaining declarations to assignSlotFilePaths' exact ABI. + const inputTypes: SlotInputType[] = ext.inputs.filter((input) => input !== 'scene') // Resolved by target handle first, then typed by that slot's declared input -- // not by the arrival order of `incomingEdges`, which does not match slot order. const inputPaths = new Array(inputTypes.length).fill(undefined) diff --git a/src/areas/workflows/workflowSceneSource.test.mjs b/src/areas/workflows/workflowSceneSource.test.mjs index ad2a0280..778e13d3 100644 --- a/src/areas/workflows/workflowSceneSource.test.mjs +++ b/src/areas/workflows/workflowSceneSource.test.mjs @@ -8,7 +8,11 @@ import { join, resolve } from 'node:path' const outfile = join(mkdtempSync(join(tmpdir(), 'modly-scene-source-')), 'scene.cjs') writeFileSync(outfile, buildSync({ entryPoints: [resolve('src/areas/workflows/workflowSceneSource.ts')], bundle: true, platform: 'node', format: 'cjs', write: false }).outputFiles[0].text) -const { resolveSceneSourceManifest } = createRequire(import.meta.url)(outfile) +const { + applySceneValidationResult, + invalidateValidatedScenePath, + resolveSceneSourceManifest, +} = createRequire(import.meta.url)(outfile) const encoded = Buffer.from(JSON.stringify({ schema: 'modly.scene-manifest.v1', sceneRoot: '.', assets: [] })).toString('base64') test('Load Scene resolves directory and manifest without image bytes', async () => { @@ -27,3 +31,40 @@ test('Load Scene refuses unsafe paths before reading', async () => { } assert.equal(reads, 0) }) + +test('editing a validated Load Scene path clears every derived scene reference', () => { + const params = invalidateValidatedScenePath({ + path: 'Workflows/old', + manifestPath: 'Workflows/old/scene-manifest.json', + sceneRoot: '.', + sourceKind: 'directory', + }, 'Workflows/new') + + assert.equal(params.path, 'Workflows/new') + assert.equal(params.manifestPath, undefined) + assert.equal(params.sceneRoot, undefined) + assert.equal(params.sourceKind, undefined) +}) + +test('an async validation result cannot overwrite a subsequently edited path', async () => { + const oldResolution = await resolveSceneSourceManifest({ + scenePath: 'Workflows/old', + workspaceDir: '/workspace', + readFileBase64: async () => encoded, + }) + assert.equal(oldResolution.ok, true) + + const currentParams = invalidateValidatedScenePath({ + path: 'Workflows/old', + manifestPath: 'Workflows/old/scene-manifest.json', + sceneRoot: '.', + sourceKind: 'directory', + }, 'Workflows/new') + + assert.equal( + applySceneValidationResult(currentParams, 'Workflows/old', oldResolution), + undefined, + ) + assert.equal(currentParams.path, 'Workflows/new') + assert.equal(currentParams.manifestPath, undefined) +}) diff --git a/src/areas/workflows/workflowSceneSource.ts b/src/areas/workflows/workflowSceneSource.ts index 5873b692..efe2b3ed 100644 --- a/src/areas/workflows/workflowSceneSource.ts +++ b/src/areas/workflows/workflowSceneSource.ts @@ -21,6 +21,47 @@ export type ResolveSceneSourceFailure = { export type ResolveSceneSourceResult = ResolveSceneSourceSuccess | ResolveSceneSourceFailure +export function applySceneValidationResult( + currentParams: Record, + validatedPath: string, + resolution: ResolveSceneSourceResult, +): Record | undefined { + if (currentParams.path !== validatedPath) return undefined + + if (!resolution.ok) { + return { + ...currentParams, + manifestPath: undefined, + sceneRoot: undefined, + sourceKind: undefined, + error: resolution.error, + } + } + + return { + ...currentParams, + path: resolution.inputWorkspacePath, + manifestPath: resolution.manifestWorkspacePath, + sceneRoot: resolution.sceneRoot, + sourceKind: resolution.sourceKind, + error: undefined, + } +} + +export function invalidateValidatedScenePath( + params: Record, + nextPath: string, +): Record { + return { + ...params, + path: nextPath, + manifestPath: undefined, + sceneRoot: undefined, + sourceKind: undefined, + error: undefined, + } +} + type ResolveSceneSourceArgs = { scenePath: string workspaceDir: string