Add explicit asset upload subfolder handling

Amp-Thread-ID: https://ampcode.com/threads/T-019ecf39-2e6f-747d-ae80-addba6b8e4f5
Co-authored-by: Amp <amp@ampcode.com>

app/assets/api/routes.py

@@ -409,6 +409,7 @@ async def upload_asset(request: web.Request) -> web.Response:
                 "hash": parsed.provided_hash,
                 "mime_type": parsed.provided_mime_type,
                 "preview_id": parsed.provided_preview_id,
+                "subfolder": parsed.provided_subfolder,
             }
         )
     except ValidationError as ve:
@@ -454,6 +455,7 @@ async def upload_asset(request: web.Request) -> web.Response:
                 expected_hash=spec.hash,
                 mime_type=spec.mime_type,
                 preview_id=spec.preview_id,
+                subfolder=spec.subfolder,
             )
     except AssetValidationError as e:
         delete_temp_file_if_exists(parsed.tmp_path)

app/assets/api/schemas_in.py

@@ -47,6 +47,7 @@ class ParsedUpload:
     provided_hash_exists: bool | None
     provided_mime_type: str | None = None
     provided_preview_id: str | None = None
+    provided_subfolder: str | None = None
 
 
 class ListAssetsQuery(BaseModel):
@@ -239,8 +240,9 @@ class TagsRemove(TagsAdd):
 class UploadAssetSpec(BaseModel):
     """Upload Asset operation.
 
-    - tags: optional list; if provided, first is root ('models'|'input'|'output');
+    - tags: labels plus one destination role ('models'|'input'|'output') for new bytes;
-            if root == 'models', second must be a valid category
+            if role == 'models', exactly one model_type:<folder_name> tag is required
+    - subfolder: optional destination subfolder for new bytes
     - name: display name
     - user_metadata: arbitrary JSON object (optional)
     - hash: optional canonical 'blake3:<hex>' for validation / fast-path
@@ -258,6 +260,7 @@ class UploadAssetSpec(BaseModel):
     hash: str | None = Field(default=None)
     mime_type: str | None = Field(default=None)
     preview_id: str | None = Field(default=None)  # references an asset_reference id
+    subfolder: str | None = Field(default=None, max_length=1024)
 
     @field_validator("hash", mode="before")
     @classmethod
@@ -315,6 +318,14 @@ class UploadAssetSpec(BaseModel):
                 norm.append(tnorm)
         return norm
 
+    @field_validator("subfolder", mode="before")
+    @classmethod
+    def _parse_subfolder(cls, v):
+        if v is None:
+            return None
+        s = str(v).strip()
+        return s or None
+
     @field_validator("user_metadata", mode="before")
     @classmethod
     def _parse_metadata_json(cls, v):

app/assets/api/upload.py

@@ -54,6 +54,7 @@ async def parse_multipart_upload(
     provided_hash_exists: bool | None = None
     provided_mime_type: str | None = None
     provided_preview_id: str | None = None
+    provided_subfolder: str | None = None
 
     file_written = 0
     tmp_path: str | None = None
@@ -140,6 +141,8 @@ async def parse_multipart_upload(
             provided_mime_type = ((await field.text()) or "").strip() or None
         elif fname == "preview_id":
             provided_preview_id = ((await field.text()) or "").strip() or None
+        elif fname == "subfolder":
+            provided_subfolder = ((await field.text()) or "").strip() or None
 
     if not file_present and not (provided_hash and provided_hash_exists):
         raise UploadError(
@@ -166,6 +169,7 @@ async def parse_multipart_upload(
         provided_hash_exists=provided_hash_exists,
         provided_mime_type=provided_mime_type,
         provided_preview_id=provided_preview_id,
+        provided_subfolder=provided_subfolder,
     )
 
 

app/assets/services/ingest.py

@@ -463,6 +463,7 @@ def upload_from_temp_path(
     expected_hash: str | None = None,
     mime_type: str | None = None,
     preview_id: str | None = None,
+    subfolder: str | None = None,
 ) -> UploadResult:
     try:
         digest, _ = hashing.compute_blake3_hash(temp_path)
@@ -507,7 +508,7 @@ def upload_from_temp_path(
 
     if not tags:
         raise ValueError("tags are required for new asset uploads")
-    base_dir, subdirs = resolve_destination_from_tags(tags)
+    base_dir, subdirs = resolve_destination_from_tags(tags, subfolder=subfolder)
     dest_dir = os.path.join(base_dir, *subdirs) if subdirs else base_dir
     os.makedirs(dest_dir, exist_ok=True)
 

app/assets/services/path_utils.py

@@ -25,11 +25,27 @@ def get_comfy_models_folders() -> list[tuple[str, list[str]]]:
     return targets
 
 
-def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]:
+def _validate_subfolder(subfolder: str | None) -> list[str]:
+    if not subfolder:
+        return []
+
+    parts = Path(subfolder).parts
+    invalid = {"", ".", ".."}
+    if Path(subfolder).is_absolute() or any(part in invalid for part in parts):
+        raise ValueError("invalid subfolder path")
+    if any("/" in part or "\\" in part for part in parts):
+        raise ValueError("invalid subfolder path")
+    return list(parts)
+
+
+def resolve_destination_from_tags(
+    tags: list[str], subfolder: str | None = None
+) -> tuple[str, list[str]]:
     """Validates and maps upload routing tags -> (base_dir, subdirs_for_fs).
 
     The request tags are only used to choose the write destination. Extra tags
     remain labels; they do not become path components or trusted classification.
+    Explicit subfolder is the only request field that can add path components.
     """
     destination_roles = [t for t in tags if t in {"input", "models", "output"}]
     if len(destination_roles) != 1:
@@ -56,7 +72,7 @@ def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]:
     else:
         base_dir = os.path.abspath(folder_paths.get_output_directory())
 
-    return base_dir, []
+    return base_dir, _validate_subfolder(subfolder)
 
 
 def validate_path_within_base(candidate: str, base: str) -> None:

server.py

@@ -440,7 +440,10 @@ class PromptServer():
                 if args.enable_assets:
                     try:
                         tag = image_upload_type if image_upload_type in ("input", "output") else "input"
-                        result = register_file_in_place(abs_path=filepath, name=filename, tags=[tag])
+                        tags = [tag]
+                        if subfolder in {"3d", "pasted", "painter", "threed", "webcam"}:
+                            tags.append(subfolder)
+                        result = register_file_in_place(abs_path=filepath, name=filename, tags=tags)
                         resp["asset"] = {
                             "id": result.ref.id,
                             "name": result.ref.name,

tests-unit/assets_test/services/test_path_utils.py

@@ -171,6 +171,19 @@ class TestGetAssetCategoryAndRelativePath:
 
 
 class TestResolveDestinationFromTags:
+    def test_explicit_subfolder_is_path_component(self, fake_dirs):
+        base_dir, subdirs = resolve_destination_from_tags(
+            ["input", "unit-tests", "foo"], subfolder="foo/bar"
+        )
+
+        assert base_dir == os.path.abspath(fake_dirs["input"])
+        assert subdirs == ["foo", "bar"]
+
+    @pytest.mark.parametrize("subfolder", ["../escape", "foo/../bar", "/abs", "foo\\bar"])
+    def test_explicit_subfolder_rejects_unsafe_paths(self, fake_dirs, subfolder: str):
+        with pytest.raises(ValueError, match="invalid subfolder"):
+            resolve_destination_from_tags(["input", "unit-tests"], subfolder=subfolder)
+
     def test_model_upload_rejects_non_writable_registered_folders(self):
         with tempfile.TemporaryDirectory() as root:
             root_path = Path(root)

tests-unit/assets_test/test_uploads.py

@@ -366,6 +366,39 @@ def test_upload_extra_tags_are_labels_not_path_components(http: requests.Session
     assert "model_type:checkpoints" in body["tags"]
 
 
+def test_upload_subfolder_is_explicit_path_component(
+    http: requests.Session, api_base: str, comfy_tmp_base_dir: Path
+):
+    files = {"file": ("nested.bin", b"nested" * 64, "application/octet-stream")}
+    form = {
+        "tags": json.dumps(["input", "unit-tests", "foo"]),
+        "subfolder": "foo/bar",
+        "name": "nested.bin",
+    }
+    r = http.post(api_base + "/api/assets", data=form, files=files, timeout=120)
+    body = r.json()
+
+    assert r.status_code == 201, body
+    stored_name = get_asset_filename(body["asset_hash"], ".bin")
+    assert (comfy_tmp_base_dir / "input" / "foo" / "bar" / stored_name).exists()
+    assert body["file_path"] == f"input/foo/bar/{stored_name}"
+    assert "foo" in body["tags"]
+
+
+def test_upload_rejects_unsafe_subfolder(http: requests.Session, api_base: str):
+    files = {"file": ("escape.bin", b"escape" * 64, "application/octet-stream")}
+    form = {
+        "tags": json.dumps(["input", "unit-tests"]),
+        "subfolder": "../escape",
+        "name": "escape.bin",
+    }
+    r = http.post(api_base + "/api/assets", data=form, files=files, timeout=120)
+    body = r.json()
+
+    assert r.status_code == 400, body
+    assert body["error"]["code"] == "INVALID_BODY"
+
+
 def test_multipart_upload_accepts_system_looking_extra_labels(
     http: requests.Session, api_base: str
 ):