import contextlib import logging import os import urllib.parse import uuid from typing import Optional from aiohttp import web from pydantic import ValidationError import folder_paths from ... import user_manager from .. import manager, scanner from . import schemas_in, schemas_out ROUTES = web.RouteTableDef() USER_MANAGER: Optional[user_manager.UserManager] = None LOGGER = logging.getLogger(__name__) # UUID regex (canonical hyphenated form, case-insensitive) UUID_RE = r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}" @ROUTES.head("/api/assets/hash/{hash}") async def head_asset_by_hash(request: web.Request) -> web.Response: hash_str = request.match_info.get("hash", "").strip().lower() if not hash_str or ":" not in hash_str: return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:'") algo, digest = hash_str.split(":", 1) if algo != "blake3" or not digest or any(c for c in digest if c not in "0123456789abcdef"): return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:'") exists = await manager.asset_exists(asset_hash=hash_str) return web.Response(status=200 if exists else 404) @ROUTES.get("/api/assets") async def list_assets(request: web.Request) -> web.Response: qp = request.rel_url.query query_dict = {} if "include_tags" in qp: query_dict["include_tags"] = qp.getall("include_tags") if "exclude_tags" in qp: query_dict["exclude_tags"] = qp.getall("exclude_tags") for k in ("name_contains", "metadata_filter", "limit", "offset", "sort", "order"): v = qp.get(k) if v is not None: query_dict[k] = v try: q = schemas_in.ListAssetsQuery.model_validate(query_dict) except ValidationError as ve: return _validation_error_response("INVALID_QUERY", ve) payload = await manager.list_assets( include_tags=q.include_tags, exclude_tags=q.exclude_tags, name_contains=q.name_contains, metadata_filter=q.metadata_filter, limit=q.limit, offset=q.offset, sort=q.sort, order=q.order, owner_id=USER_MANAGER.get_request_user_id(request), ) return web.json_response(payload.model_dump(mode="json")) @ROUTES.get(f"/api/assets/{{id:{UUID_RE}}}/content") async def download_asset_content(request: web.Request) -> web.Response: disposition = request.query.get("disposition", "attachment").lower().strip() if disposition not in {"inline", "attachment"}: disposition = "attachment" try: abs_path, content_type, filename = await manager.resolve_asset_content_for_download( asset_info_id=str(uuid.UUID(request.match_info["id"])), owner_id=USER_MANAGER.get_request_user_id(request), ) except ValueError as ve: return _error_response(404, "ASSET_NOT_FOUND", str(ve)) except NotImplementedError as nie: return _error_response(501, "BACKEND_UNSUPPORTED", str(nie)) except FileNotFoundError: return _error_response(404, "FILE_NOT_FOUND", "Underlying file not found on disk.") quoted = (filename or "").replace("\r", "").replace("\n", "").replace('"', "'") cd = f'{disposition}; filename="{quoted}"; filename*=UTF-8\'\'{urllib.parse.quote(filename)}' resp = web.FileResponse(abs_path) resp.content_type = content_type resp.headers["Content-Disposition"] = cd return resp @ROUTES.post("/api/assets/from-hash") async def create_asset_from_hash(request: web.Request) -> web.Response: try: payload = await request.json() body = schemas_in.CreateFromHashBody.model_validate(payload) except ValidationError as ve: return _validation_error_response("INVALID_BODY", ve) except Exception: return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.") result = await manager.create_asset_from_hash( hash_str=body.hash, name=body.name, tags=body.tags, user_metadata=body.user_metadata, owner_id=USER_MANAGER.get_request_user_id(request), ) if result is None: return _error_response(404, "ASSET_NOT_FOUND", f"Asset content {body.hash} does not exist") return web.json_response(result.model_dump(mode="json"), status=201) @ROUTES.post("/api/assets") async def upload_asset(request: web.Request) -> web.Response: """Multipart/form-data endpoint for Asset uploads.""" if not (request.content_type or "").lower().startswith("multipart/"): return _error_response(415, "UNSUPPORTED_MEDIA_TYPE", "Use multipart/form-data for uploads.") reader = await request.multipart() file_present = False file_client_name: Optional[str] = None tags_raw: list[str] = [] provided_name: Optional[str] = None user_metadata_raw: Optional[str] = None provided_hash: Optional[str] = None provided_hash_exists: Optional[bool] = None file_written = 0 tmp_path: Optional[str] = None while True: field = await reader.next() if field is None: break fname = getattr(field, "name", "") or "" if fname == "hash": try: s = ((await field.text()) or "").strip().lower() except Exception: return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:'") if s: if ":" not in s: return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:'") algo, digest = s.split(":", 1) if algo != "blake3" or not digest or any(c for c in digest if c not in "0123456789abcdef"): return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:'") provided_hash = f"{algo}:{digest}" try: provided_hash_exists = await manager.asset_exists(asset_hash=provided_hash) except Exception: provided_hash_exists = None # do not fail the whole request here elif fname == "file": file_present = True file_client_name = (field.filename or "").strip() if provided_hash and provided_hash_exists is True: # If client supplied a hash that we know exists, drain but do not write to disk try: while True: chunk = await field.read_chunk(8 * 1024 * 1024) if not chunk: break file_written += len(chunk) except Exception: return _error_response(500, "UPLOAD_IO_ERROR", "Failed to receive uploaded file.") continue # Do not create temp file; we will create AssetInfo from the existing content # Otherwise, store to temp for hashing/ingest uploads_root = os.path.join(folder_paths.get_temp_directory(), "uploads") unique_dir = os.path.join(uploads_root, uuid.uuid4().hex) os.makedirs(unique_dir, exist_ok=True) tmp_path = os.path.join(unique_dir, ".upload.part") try: with open(tmp_path, "wb") as f: while True: chunk = await field.read_chunk(8 * 1024 * 1024) if not chunk: break f.write(chunk) file_written += len(chunk) except Exception: try: if os.path.exists(tmp_path or ""): os.remove(tmp_path) finally: return _error_response(500, "UPLOAD_IO_ERROR", "Failed to receive and store uploaded file.") elif fname == "tags": tags_raw.append((await field.text()) or "") elif fname == "name": provided_name = (await field.text()) or None elif fname == "user_metadata": user_metadata_raw = (await field.text()) or None # If client did not send file, and we are not doing a from-hash fast path -> error if not file_present and not (provided_hash and provided_hash_exists): return _error_response(400, "MISSING_FILE", "Form must include a 'file' part or a known 'hash'.") if file_present and file_written == 0 and not (provided_hash and provided_hash_exists): # Empty upload is only acceptable if we are fast-pathing from existing hash try: if tmp_path and os.path.exists(tmp_path): os.remove(tmp_path) finally: return _error_response(400, "EMPTY_UPLOAD", "Uploaded file is empty.") try: spec = schemas_in.UploadAssetSpec.model_validate({ "tags": tags_raw, "name": provided_name, "user_metadata": user_metadata_raw, "hash": provided_hash, }) except ValidationError as ve: try: if tmp_path and os.path.exists(tmp_path): os.remove(tmp_path) finally: return _validation_error_response("INVALID_BODY", ve) # Validate models category against configured folders (consistent with previous behavior) if spec.tags and spec.tags[0] == "models": if len(spec.tags) < 2 or spec.tags[1] not in folder_paths.folder_names_and_paths: if tmp_path and os.path.exists(tmp_path): os.remove(tmp_path) return _error_response( 400, "INVALID_BODY", f"unknown models category '{spec.tags[1] if len(spec.tags) >= 2 else ''}'" ) owner_id = USER_MANAGER.get_request_user_id(request) # Fast path: if a valid provided hash exists, create AssetInfo without writing anything if spec.hash and provided_hash_exists is True: try: result = await manager.create_asset_from_hash( hash_str=spec.hash, name=spec.name or (spec.hash.split(":", 1)[1]), tags=spec.tags, user_metadata=spec.user_metadata or {}, owner_id=owner_id, ) except Exception: LOGGER.exception("create_asset_from_hash failed for hash=%s, owner_id=%s", spec.hash, owner_id) return _error_response(500, "INTERNAL", "Unexpected server error.") if result is None: return _error_response(404, "ASSET_NOT_FOUND", f"Asset content {spec.hash} does not exist") # Drain temp if we accidentally saved (e.g., hash field came after file) if tmp_path and os.path.exists(tmp_path): with contextlib.suppress(Exception): os.remove(tmp_path) status = 200 if (not result.created_new) else 201 return web.json_response(result.model_dump(mode="json"), status=status) # Otherwise, we must have a temp file path to ingest if not tmp_path or not os.path.exists(tmp_path): # The only case we reach here without a temp file is: client sent a hash that does not exist and no file return _error_response(404, "ASSET_NOT_FOUND", "Provided hash not found and no file uploaded.") try: created = await manager.upload_asset_from_temp_path( spec, temp_path=tmp_path, client_filename=file_client_name, owner_id=owner_id, expected_asset_hash=spec.hash, ) status = 201 if created.created_new else 200 return web.json_response(created.model_dump(mode="json"), status=status) except ValueError as e: if tmp_path and os.path.exists(tmp_path): os.remove(tmp_path) msg = str(e) if "HASH_MISMATCH" in msg or msg.strip().upper() == "HASH_MISMATCH": return _error_response( 400, "HASH_MISMATCH", "Uploaded file hash does not match provided hash.", ) return _error_response(400, "BAD_REQUEST", "Invalid inputs.") except Exception: if tmp_path and os.path.exists(tmp_path): os.remove(tmp_path) LOGGER.exception("upload_asset_from_temp_path failed for tmp_path=%s, owner_id=%s", tmp_path, owner_id) return _error_response(500, "INTERNAL", "Unexpected server error.") @ROUTES.get(f"/api/assets/{{id:{UUID_RE}}}") async def get_asset(request: web.Request) -> web.Response: asset_info_id = str(uuid.UUID(request.match_info["id"])) try: result = await manager.get_asset( asset_info_id=asset_info_id, owner_id=USER_MANAGER.get_request_user_id(request), ) except ValueError as ve: return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) except Exception: LOGGER.exception( "get_asset failed for asset_info_id=%s, owner_id=%s", asset_info_id, USER_MANAGER.get_request_user_id(request), ) return _error_response(500, "INTERNAL", "Unexpected server error.") return web.json_response(result.model_dump(mode="json"), status=200) @ROUTES.put(f"/api/assets/{{id:{UUID_RE}}}") async def update_asset(request: web.Request) -> web.Response: asset_info_id = str(uuid.UUID(request.match_info["id"])) try: body = schemas_in.UpdateAssetBody.model_validate(await request.json()) except ValidationError as ve: return _validation_error_response("INVALID_BODY", ve) except Exception: return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.") try: result = await manager.update_asset( asset_info_id=asset_info_id, name=body.name, tags=body.tags, user_metadata=body.user_metadata, owner_id=USER_MANAGER.get_request_user_id(request), ) except (ValueError, PermissionError) as ve: return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) except Exception: LOGGER.exception( "update_asset failed for asset_info_id=%s, owner_id=%s", asset_info_id, USER_MANAGER.get_request_user_id(request), ) return _error_response(500, "INTERNAL", "Unexpected server error.") return web.json_response(result.model_dump(mode="json"), status=200) @ROUTES.put(f"/api/assets/{{id:{UUID_RE}}}/preview") async def set_asset_preview(request: web.Request) -> web.Response: asset_info_id = str(uuid.UUID(request.match_info["id"])) try: body = schemas_in.SetPreviewBody.model_validate(await request.json()) except ValidationError as ve: return _validation_error_response("INVALID_BODY", ve) except Exception: return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.") try: result = await manager.set_asset_preview( asset_info_id=asset_info_id, preview_asset_id=body.preview_id, owner_id=USER_MANAGER.get_request_user_id(request), ) except (PermissionError, ValueError) as ve: return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) except Exception: LOGGER.exception( "set_asset_preview failed for asset_info_id=%s, owner_id=%s", asset_info_id, USER_MANAGER.get_request_user_id(request), ) return _error_response(500, "INTERNAL", "Unexpected server error.") return web.json_response(result.model_dump(mode="json"), status=200) @ROUTES.delete(f"/api/assets/{{id:{UUID_RE}}}") async def delete_asset(request: web.Request) -> web.Response: asset_info_id = str(uuid.UUID(request.match_info["id"])) delete_content = request.query.get("delete_content") delete_content = True if delete_content is None else delete_content.lower() not in {"0", "false", "no"} try: deleted = await manager.delete_asset_reference( asset_info_id=asset_info_id, owner_id=USER_MANAGER.get_request_user_id(request), delete_content_if_orphan=delete_content, ) except Exception: LOGGER.exception( "delete_asset_reference failed for asset_info_id=%s, owner_id=%s", asset_info_id, USER_MANAGER.get_request_user_id(request), ) return _error_response(500, "INTERNAL", "Unexpected server error.") if not deleted: return _error_response(404, "ASSET_NOT_FOUND", f"AssetInfo {asset_info_id} not found.") return web.Response(status=204) @ROUTES.get("/api/tags") async def get_tags(request: web.Request) -> web.Response: query_map = dict(request.rel_url.query) try: query = schemas_in.TagsListQuery.model_validate(query_map) except ValidationError as ve: return web.json_response( {"error": {"code": "INVALID_QUERY", "message": "Invalid query parameters", "details": ve.errors()}}, status=400, ) result = await manager.list_tags( prefix=query.prefix, limit=query.limit, offset=query.offset, order=query.order, include_zero=query.include_zero, owner_id=USER_MANAGER.get_request_user_id(request), ) return web.json_response(result.model_dump(mode="json")) @ROUTES.post(f"/api/assets/{{id:{UUID_RE}}}/tags") async def add_asset_tags(request: web.Request) -> web.Response: asset_info_id = str(uuid.UUID(request.match_info["id"])) try: payload = await request.json() data = schemas_in.TagsAdd.model_validate(payload) except ValidationError as ve: return _error_response(400, "INVALID_BODY", "Invalid JSON body for tags add.", {"errors": ve.errors()}) except Exception: return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.") try: result = await manager.add_tags_to_asset( asset_info_id=asset_info_id, tags=data.tags, origin="manual", owner_id=USER_MANAGER.get_request_user_id(request), ) except (ValueError, PermissionError) as ve: return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) except Exception: LOGGER.exception( "add_tags_to_asset failed for asset_info_id=%s, owner_id=%s", asset_info_id, USER_MANAGER.get_request_user_id(request), ) return _error_response(500, "INTERNAL", "Unexpected server error.") return web.json_response(result.model_dump(mode="json"), status=200) @ROUTES.delete(f"/api/assets/{{id:{UUID_RE}}}/tags") async def delete_asset_tags(request: web.Request) -> web.Response: asset_info_id = str(uuid.UUID(request.match_info["id"])) try: payload = await request.json() data = schemas_in.TagsRemove.model_validate(payload) except ValidationError as ve: return _error_response(400, "INVALID_BODY", "Invalid JSON body for tags remove.", {"errors": ve.errors()}) except Exception: return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.") try: result = await manager.remove_tags_from_asset( asset_info_id=asset_info_id, tags=data.tags, owner_id=USER_MANAGER.get_request_user_id(request), ) except ValueError as ve: return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) except Exception: LOGGER.exception( "remove_tags_from_asset failed for asset_info_id=%s, owner_id=%s", asset_info_id, USER_MANAGER.get_request_user_id(request), ) return _error_response(500, "INTERNAL", "Unexpected server error.") return web.json_response(result.model_dump(mode="json"), status=200) @ROUTES.post("/api/assets/scan/seed") async def seed_assets(request: web.Request) -> web.Response: try: payload = await request.json() except Exception: payload = {} try: body = schemas_in.ScheduleAssetScanBody.model_validate(payload) except ValidationError as ve: return _validation_error_response("INVALID_BODY", ve) try: await scanner.sync_seed_assets(body.roots) except Exception: LOGGER.exception("sync_seed_assets failed for roots=%s", body.roots) return _error_response(500, "INTERNAL", "Unexpected server error.") return web.json_response({"synced": True, "roots": body.roots}, status=200) @ROUTES.post("/api/assets/scan/schedule") async def schedule_asset_scan(request: web.Request) -> web.Response: try: payload = await request.json() except Exception: payload = {} try: body = schemas_in.ScheduleAssetScanBody.model_validate(payload) except ValidationError as ve: return _validation_error_response("INVALID_BODY", ve) states = await scanner.schedule_scans(body.roots) return web.json_response(states.model_dump(mode="json"), status=202) @ROUTES.get("/api/assets/scan") async def get_asset_scan_status(request: web.Request) -> web.Response: root = request.query.get("root", "").strip().lower() states = scanner.current_statuses() if root in {"models", "input", "output"}: states = [s for s in states.scans if s.root == root] # type: ignore states = schemas_out.AssetScanStatusResponse(scans=states) return web.json_response(states.model_dump(mode="json"), status=200) def register_assets_system(app: web.Application, user_manager_instance: user_manager.UserManager) -> None: global USER_MANAGER USER_MANAGER = user_manager_instance app.add_routes(ROUTES) def _error_response(status: int, code: str, message: str, details: Optional[dict] = None) -> web.Response: return web.json_response({"error": {"code": code, "message": message, "details": details or {}}}, status=status) def _validation_error_response(code: str, ve: ValidationError) -> web.Response: return _error_response(400, code, "Validation failed.", {"errors": ve.json()})