diff --git a/core/routes/metadata.py b/core/routes/metadata.py index a1d237f..16dd1fe 100644 --- a/core/routes/metadata.py +++ b/core/routes/metadata.py @@ -1,25 +1,15 @@ """Metadata route registration.""" +from ..services.metadata_service import MetadataService from ..services.scanner_service import ScannerService from .context import RouteContext def register_metadata_routes(context: RouteContext): - asyncio = context.get('asyncio') - audit_metadata_sizes = context.get('audit_metadata_sizes') - build_missing_local_metadata = context.get('build_missing_local_metadata') - cancel_progress_response = context.get('cancel_progress_response') - get_metadata_build_capabilities = context.get('get_metadata_build_capabilities') - get_progress_response = context.get('get_progress_response') json_api_endpoint = context.get('json_api_endpoint') - normalize_metadata_build_mode = context.get('normalize_metadata_build_mode') routes = context.get('routes') - run_in_background_thread = context.get('run_in_background_thread') - self = context.get('self') - to_bool = context.get('to_bool') - to_int = context.get('to_int') - web = context.get('web') scanner_service = ScannerService(context) + metadata_service = MetadataService(context) @routes.get("/model_resolver/models") @json_api_endpoint("get_models") @@ -31,165 +21,28 @@ def register_metadata_routes(context: RouteContext): @json_api_endpoint("metadata-size-audit") async def metadata_size_audit_route(request): """Check local model .metadata.json sidecars for stale size values.""" - payload = {} - try: - if request.can_read_body: - payload = await request.json() - except Exception: - payload = {} - if not isinstance(payload, dict): - payload = {} - - force_rescan = to_bool(payload.get("force_rescan"), True) - result = await asyncio.to_thread( - audit_metadata_sizes, - force_rescan=force_rescan, - worker_count=payload.get("worker_count"), - batch_size=payload.get("batch_size"), - ) - return web.json_response(result) + return await metadata_service.metadata_size_audit(request) @routes.get("/model_resolver/metadata-build/capabilities") @json_api_endpoint("metadata-build-capabilities") async def metadata_build_capabilities_route(request): """Return local CPU and worker limits for the metadata builder.""" - return web.json_response(get_metadata_build_capabilities()) + return await metadata_service.metadata_build_capabilities(request) @routes.post("/model_resolver/metadata-build/start") @json_api_endpoint("metadata-build-start") async def metadata_build_start_route(request): """Create missing local metadata sidecars and fill missing SHA256 values.""" - import uuid - - payload = {} - try: - if request.can_read_body: - payload = await request.json() - except Exception: - payload = {} - if not isinstance(payload, dict): - payload = {} - - force_rescan = to_bool(payload.get("force_rescan"), True) - worker_count = to_int(payload.get("worker_count"), 0) - metadata_mode = normalize_metadata_build_mode( - payload.get("metadata_mode") - ) - self.metadata_builder_progress.cleanup() - progress_id = f"metadata_build_{uuid.uuid4().hex}" - self.metadata_builder_progress.update( - progress_id, - status="queued", - stage="queued", - message="Preparing local metadata build...", - percent=0, - current=0, - total=0, - requested_worker_count=worker_count, - metadata_mode=metadata_mode, - ) - - def update_metadata_build_progress(progress_payload): - self._update_metadata_build_progress( - progress_id, progress_payload - ) - - def is_metadata_build_cancelled(): - return self.metadata_builder_progress.is_cancelled(progress_id) - - def run_metadata_build_task(): - result = build_missing_local_metadata( - force_rescan=force_rescan, - worker_count=worker_count, - metadata_mode=metadata_mode, - progress_callback=update_metadata_build_progress, - is_cancelled=is_metadata_build_cancelled, - ) - return result - - def on_success(result): - self.metadata_builder_progress.update( - progress_id, - status="done", - stage="done", - message="Local metadata build completed.", - percent=100, - active_models=[], - active_worker_count=0, - current_model="", - current_path="", - bytes_read=0, - total_bytes=0, - result=result, - **result, - ) - - def on_cancel(result=None): - self.metadata_builder_progress.update( - progress_id, - status="cancelled", - stage="cancelled", - message="Metadata build cancelled", - percent=100, - active_models=[], - active_worker_count=0, - current_model="", - current_path="", - bytes_read=0, - total_bytes=0, - result=result or {}, - **(result or {}), - ) - - def on_error(exc): - self.metadata_builder_progress.update( - progress_id, - status="error", - stage="error", - message=str(exc) or "Metadata build failed", - percent=100, - active_models=[], - active_worker_count=0, - current_model="", - current_path="", - bytes_read=0, - total_bytes=0, - error=str(exc) or "Metadata build failed", - ) - - run_in_background_thread( - self.metadata_builder_progress, - progress_id, - run_metadata_build_task, - on_success, - on_cancel, - on_error, - error_log_msg="Metadata build failed", - ) - return web.json_response( - { - "success": True, - "progress_id": progress_id, - "metadata_mode": metadata_mode, - } - ) + return await metadata_service.metadata_build_start(request) @routes.get("/model_resolver/metadata-build/progress/{progress_id}") @json_api_endpoint("metadata-build-progress") async def metadata_build_progress_route(request): """Return progress for local metadata sidecar creation.""" - return get_progress_response( - self.metadata_builder_progress, - request, - not_found_status=404, - ) + return await metadata_service.metadata_build_progress(request) @routes.post("/model_resolver/metadata-build/cancel/{progress_id}") @json_api_endpoint("metadata-build-cancel") async def metadata_build_cancel_route(request): """Cancel local metadata sidecar creation.""" - return cancel_progress_response( - self.metadata_builder_progress, - request, - cancel_message="Stopping metadata build...", - ) + return await metadata_service.metadata_build_cancel(request) diff --git a/core/services/metadata_service.py b/core/services/metadata_service.py new file mode 100644 index 0000000..2e431b6 --- /dev/null +++ b/core/services/metadata_service.py @@ -0,0 +1,192 @@ +"""Metadata audit and build orchestration used by HTTP route adapters.""" + +import uuid + +from ..routes.context import RouteContext + + +class MetadataService: + """Coordinate local metadata audit and background build operations.""" + + def __init__(self, context: RouteContext): + self.asyncio = context.get("asyncio") + self.audit_metadata_sizes_fn = context.get("audit_metadata_sizes") + self.build_missing_local_metadata_fn = context.get( + "build_missing_local_metadata" + ) + self.cancel_progress_response_fn = context.get( + "cancel_progress_response" + ) + self.get_metadata_build_capabilities_fn = context.get( + "get_metadata_build_capabilities" + ) + self.get_progress_response_fn = context.get("get_progress_response") + self.normalize_metadata_build_mode = context.get( + "normalize_metadata_build_mode" + ) + self.run_in_background_thread = context.get("run_in_background_thread") + self.extension = context.get("self") + self.to_bool = context.get("to_bool") + self.to_int = context.get("to_int") + self.web = context.get("web") + + @property + def metadata_builder_progress(self): + return self.extension.metadata_builder_progress + + async def metadata_size_audit(self, request): + """Check local metadata sidecars for stale file-size values.""" + payload = {} + try: + if request.can_read_body: + payload = await request.json() + except Exception: + payload = {} + if not isinstance(payload, dict): + payload = {} + + force_rescan = self.to_bool(payload.get("force_rescan"), True) + result = await self.asyncio.to_thread( + self.audit_metadata_sizes_fn, + force_rescan=force_rescan, + worker_count=payload.get("worker_count"), + batch_size=payload.get("batch_size"), + ) + return self.web.json_response(result) + + async def metadata_build_capabilities(self, request): + """Return local CPU and worker limits for the metadata builder.""" + return self.web.json_response( + self.get_metadata_build_capabilities_fn() + ) + + async def metadata_build_start(self, request): + """Start building missing local metadata sidecars in the background.""" + payload = {} + try: + if request.can_read_body: + payload = await request.json() + except Exception: + payload = {} + if not isinstance(payload, dict): + payload = {} + + force_rescan = self.to_bool(payload.get("force_rescan"), True) + worker_count = self.to_int(payload.get("worker_count"), 0) + metadata_mode = self.normalize_metadata_build_mode( + payload.get("metadata_mode") + ) + self.metadata_builder_progress.cleanup() + progress_id = f"metadata_build_{uuid.uuid4().hex}" + self.metadata_builder_progress.update( + progress_id, + status="queued", + stage="queued", + message="Preparing local metadata build...", + percent=0, + current=0, + total=0, + requested_worker_count=worker_count, + metadata_mode=metadata_mode, + ) + + def update_metadata_build_progress(progress_payload): + self.extension._update_metadata_build_progress( + progress_id, + progress_payload, + ) + + def is_metadata_build_cancelled(): + return self.metadata_builder_progress.is_cancelled(progress_id) + + def run_metadata_build_task(): + return self.build_missing_local_metadata_fn( + force_rescan=force_rescan, + worker_count=worker_count, + metadata_mode=metadata_mode, + progress_callback=update_metadata_build_progress, + is_cancelled=is_metadata_build_cancelled, + ) + + def on_success(result): + self.metadata_builder_progress.update( + progress_id, + status="done", + stage="done", + message="Local metadata build completed.", + percent=100, + active_models=[], + active_worker_count=0, + current_model="", + current_path="", + bytes_read=0, + total_bytes=0, + result=result, + **result, + ) + + def on_cancel(result=None): + self.metadata_builder_progress.update( + progress_id, + status="cancelled", + stage="cancelled", + message="Metadata build cancelled", + percent=100, + active_models=[], + active_worker_count=0, + current_model="", + current_path="", + bytes_read=0, + total_bytes=0, + result=result or {}, + **(result or {}), + ) + + def on_error(exc): + self.metadata_builder_progress.update( + progress_id, + status="error", + stage="error", + message=str(exc) or "Metadata build failed", + percent=100, + active_models=[], + active_worker_count=0, + current_model="", + current_path="", + bytes_read=0, + total_bytes=0, + error=str(exc) or "Metadata build failed", + ) + + self.run_in_background_thread( + self.metadata_builder_progress, + progress_id, + run_metadata_build_task, + on_success, + on_cancel, + on_error, + error_log_msg="Metadata build failed", + ) + return self.web.json_response( + { + "success": True, + "progress_id": progress_id, + "metadata_mode": metadata_mode, + } + ) + + async def metadata_build_progress(self, request): + """Return progress for local metadata sidecar creation.""" + return self.get_progress_response_fn( + self.metadata_builder_progress, + request, + not_found_status=404, + ) + + async def metadata_build_cancel(self, request): + """Cancel local metadata sidecar creation.""" + return self.cancel_progress_response_fn( + self.metadata_builder_progress, + request, + cancel_message="Stopping metadata build...", + ) diff --git a/tests/test_metadata_routes.py b/tests/test_metadata_routes.py new file mode 100644 index 0000000..3e5a511 --- /dev/null +++ b/tests/test_metadata_routes.py @@ -0,0 +1,273 @@ +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from aiohttp import web + +from core.routes.context import RouteContext +from core.routes.metadata import register_metadata_routes +from core.type_utils import to_bool, to_int + + +class _Routes: + def __init__(self): + self.handlers = {} + + def get(self, path): + def decorator(handler): + self.handlers[("GET", path)] = handler + return handler + + return decorator + + def post(self, path): + def decorator(handler): + self.handlers[("POST", path)] = handler + return handler + + return decorator + + +def _json_api_endpoint(_name, **_kwargs): + return lambda handler: handler + + +def _build_routes(): + routes = _Routes() + metadata_progress = MagicMock() + metadata_progress.is_cancelled.return_value = False + extension = SimpleNamespace( + metadata_builder_progress=metadata_progress, + _update_metadata_build_progress=MagicMock(), + ) + values = { + "asyncio": asyncio, + "audit_metadata_sizes": MagicMock( + return_value={"stale": 1, "checked": 3} + ), + "build_missing_local_metadata": MagicMock( + return_value={"created_metadata": 1} + ), + "cancel_progress_response": MagicMock( + return_value={"success": True, "cancelled": True} + ), + "get_metadata_build_capabilities": MagicMock( + return_value={"max_workers": 4} + ), + "get_model_files": MagicMock(return_value=[]), + "get_progress_response": MagicMock( + return_value={"status": "running"} + ), + "invalidate_local_hash_match_cache": MagicMock(), + "infer_download_path_templates": MagicMock(), + "json_api_endpoint": _json_api_endpoint, + "normalize_metadata_build_mode": MagicMock( + return_value="import_existing" + ), + "routes": routes, + "run_in_background_thread": MagicMock(), + "self": extension, + "to_bool": to_bool, + "to_int": to_int, + "web": web, + } + register_metadata_routes(RouteContext(values)) + return routes.handlers, values + + +def _request(payload=None, *, can_read_body=True, progress_id="build-1"): + return SimpleNamespace( + can_read_body=can_read_body, + json=AsyncMock(return_value=payload), + match_info={"progress_id": progress_id}, + query={}, + ) + + +@pytest.mark.asyncio +async def test_metadata_size_audit_forwards_payload_options(): + handlers, values = _build_routes() + handler = handlers[("POST", "/model_resolver/metadata-size-audit")] + + response = await handler( + _request( + { + "force_rescan": False, + "worker_count": 2, + "batch_size": 25, + } + ) + ) + + assert json.loads(response.text) == {"stale": 1, "checked": 3} + values["audit_metadata_sizes"].assert_called_once_with( + force_rescan=False, + worker_count=2, + batch_size=25, + ) + + +@pytest.mark.asyncio +async def test_metadata_size_audit_uses_defaults_for_unreadable_body(): + handlers, values = _build_routes() + handler = handlers[("POST", "/model_resolver/metadata-size-audit")] + request = _request(can_read_body=False) + + await handler(request) + + request.json.assert_not_awaited() + values["audit_metadata_sizes"].assert_called_once_with( + force_rescan=True, + worker_count=None, + batch_size=None, + ) + + +@pytest.mark.asyncio +async def test_metadata_routes_handle_json_parse_failures(): + handlers, values = _build_routes() + + audit_request = _request() + audit_request.json = AsyncMock(side_effect=ValueError("invalid json")) + await handlers[("POST", "/model_resolver/metadata-size-audit")]( + audit_request + ) + values["audit_metadata_sizes"].assert_called_once_with( + force_rescan=True, + worker_count=None, + batch_size=None, + ) + + empty_body_request = _request(can_read_body=False) + await handlers[("POST", "/model_resolver/metadata-build/start")]( + empty_body_request + ) + values["normalize_metadata_build_mode"].assert_called_once_with(None) + + build_request = _request() + build_request.json = AsyncMock(side_effect=ValueError("invalid json")) + await handlers[("POST", "/model_resolver/metadata-build/start")]( + build_request + ) + assert values["normalize_metadata_build_mode"].call_count == 2 + + +@pytest.mark.asyncio +async def test_metadata_build_capabilities_progress_and_cancel_routes(): + handlers, values = _build_routes() + + capabilities = await handlers[ + ("GET", "/model_resolver/metadata-build/capabilities") + ](_request()) + assert json.loads(capabilities.text) == {"max_workers": 4} + + progress_request = _request(progress_id="progress-1") + progress = await handlers[ + ("GET", "/model_resolver/metadata-build/progress/{progress_id}") + ](progress_request) + assert progress == {"status": "running"} + values["get_progress_response"].assert_called_once_with( + values["self"].metadata_builder_progress, + progress_request, + not_found_status=404, + ) + + cancel_request = _request(progress_id="cancel-1") + cancel = await handlers[ + ("POST", "/model_resolver/metadata-build/cancel/{progress_id}") + ](cancel_request) + assert cancel == {"success": True, "cancelled": True} + values["cancel_progress_response"].assert_called_once_with( + values["self"].metadata_builder_progress, + cancel_request, + cancel_message="Stopping metadata build...", + ) + + +@pytest.mark.asyncio +async def test_metadata_build_start_registers_task_and_updates_callbacks(): + handlers, values = _build_routes() + captured = {} + + def capture_background_task( + progress, + progress_id, + task, + on_success, + on_cancel, + on_error, + **kwargs, + ): + captured.update( + progress=progress, + progress_id=progress_id, + task=task, + on_success=on_success, + on_cancel=on_cancel, + on_error=on_error, + kwargs=kwargs, + ) + + values["run_in_background_thread"].side_effect = capture_background_task + handler = handlers[("POST", "/model_resolver/metadata-build/start")] + + response = await handler( + _request( + { + "force_rescan": False, + "worker_count": "3", + "metadata_mode": "import_existing", + } + ) + ) + + body = json.loads(response.text) + progress_id = body["progress_id"] + assert body == { + "success": True, + "progress_id": progress_id, + "metadata_mode": "import_existing", + } + assert progress_id.startswith("metadata_build_") + values["self"].metadata_builder_progress.cleanup.assert_called_once_with() + assert captured["progress_id"] == progress_id + assert captured["kwargs"] == {"error_log_msg": "Metadata build failed"} + + task_result = captured["task"]() + assert task_result == {"created_metadata": 1} + build_call = values["build_missing_local_metadata"].call_args.kwargs + assert build_call["force_rescan"] is False + assert build_call["worker_count"] == 3 + assert build_call["metadata_mode"] == "import_existing" + assert captured["progress"] is values["self"].metadata_builder_progress + + progress_callback = build_call["progress_callback"] + progress_callback({"percent": 45}) + values["self"]._update_metadata_build_progress.assert_called_once_with( + progress_id, + {"percent": 45}, + ) + assert build_call["is_cancelled"]() is False + + captured["on_success"]({"created_metadata": 2}) + captured["on_cancel"]({"cancelled": True}) + captured["on_error"](ValueError()) + updates = values["self"].metadata_builder_progress.update.call_args_list + statuses = [call.kwargs["status"] for call in updates] + assert statuses == ["queued", "done", "cancelled", "error"] + assert updates[-1].kwargs["message"] == "Metadata build failed" + + +@pytest.mark.asyncio +async def test_metadata_build_start_normalizes_invalid_payload(): + handlers, values = _build_routes() + handler = handlers[("POST", "/model_resolver/metadata-build/start")] + + response = await handler(_request(["not", "a", "mapping"])) + + assert json.loads(response.text)["success"] is True + values["normalize_metadata_build_mode"].assert_called_once_with(None) + build_call = values["build_missing_local_metadata"].call_args + assert build_call is None