diff --git a/api/route_registrars.py b/api/route_registrars.py index b0ff84d..d0a852a 100644 --- a/api/route_registrars.py +++ b/api/route_registrars.py @@ -18,12 +18,16 @@ class RouteSpec: handler: Any -def register_route_family(server, register_route_fn, specs: Iterable[RouteSpec]) -> None: +def register_route_family( + server, register_route_fn, specs: Iterable[RouteSpec] +) -> None: for spec in specs: register_route_fn(server, spec.method, spec.path, spec.handler) -def build_core_route_specs(prefix: str, handlers: dict[str, Any]) -> tuple[RouteSpec, ...]: +def build_core_route_specs( + prefix: str, handlers: dict[str, Any] +) -> tuple[RouteSpec, ...]: return ( RouteSpec("GET", f"{prefix}/admin", handlers["remote_admin_page_handler"]), RouteSpec("GET", f"{prefix}/health", handlers["health_handler"]), diff --git a/services/model_manager.py b/services/model_manager.py index fe7049c..0b7adec 100644 --- a/services/model_manager.py +++ b/services/model_manager.py @@ -8,12 +8,9 @@ import hashlib import json import logging import os -import shutil import tempfile import threading import time -import urllib.request -import uuid from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field from pathlib import Path @@ -21,6 +18,48 @@ from typing import Any, Dict, List, Optional from urllib.parse import urlparse from .job_events import JobEventType, get_job_event_store +from .model_manager_catalog import ( + collect_catalog_entries as _collect_catalog_entries_impl, +) +from .model_manager_catalog import ( + collect_install_entries as _collect_install_entries_impl, +) +from .model_manager_catalog import list_installations as _list_installations_impl +from .model_manager_catalog import load_installations as _load_installations_impl +from .model_manager_catalog import save_installations as _save_installations_impl +from .model_manager_catalog import search_models as _search_models_impl +from .model_manager_tasks import cancel_download_task as _cancel_download_task_impl +from .model_manager_tasks import ( + checkpoint_matches_task as _checkpoint_matches_task_impl, +) +from .model_manager_tasks import checkpoint_path as _checkpoint_path_impl +from .model_manager_tasks import emit as _emit_impl +from .model_manager_tasks import get_download_task as _get_download_task_impl +from .model_manager_tasks import list_download_tasks as _list_download_tasks_impl +from .model_manager_tasks import load_checkpoint as _load_checkpoint_impl +from .model_manager_tasks import load_tasks_from_disk as _load_tasks_from_disk_impl +from .model_manager_tasks import persist_tasks_locked as _persist_tasks_locked_impl +from .model_manager_tasks import progress as _progress_impl +from .model_manager_tasks import ( + recover_incomplete_tasks as _recover_incomplete_tasks_impl, +) +from .model_manager_tasks import save_checkpoint as _save_checkpoint_impl +from .model_manager_tasks import set_resume_status as _set_resume_status_impl +from .model_manager_tasks import tenant_ok as _tenant_ok_impl +from .model_manager_tasks import validators_match as _validators_match_impl +from .model_manager_transfer import assert_budget as _assert_budget_impl +from .model_manager_transfer import create_download_task as _create_download_task_impl +from .model_manager_transfer import download as _download_impl +from .model_manager_transfer import ( + import_downloaded_model as _import_downloaded_model_impl, +) +from .model_manager_transfer import normalize_tenant as _normalize_tenant_impl +from .model_manager_transfer import run_task as _run_task_impl +from .model_manager_transfer import ( + stream_response_to_part as _stream_response_to_part_impl, +) +from .model_manager_transfer import validate_provenance as _validate_provenance_impl +from .model_manager_transfer import validate_url_policy as _validate_url_policy_impl from .safe_io import ( STANDARD_OUTBOUND_POLICY, SSRFError, @@ -395,6 +434,12 @@ class ModelManager: self._futures: Dict[str, Future] = {} self._cancel_events: Dict[str, threading.Event] = {} self._last_tasks_persist_at = 0.0 + self._download_task_cls = DownloadTask + self._download_cancelled_cls = DownloadCancelled + self._ssrf_error_cls = SSRFError + self._default_tenant_id = DEFAULT_TENANT_ID + self._model_type_to_subdir = MODEL_TYPE_TO_SUBDIR + self._threading_event_factory = threading.Event self._executor = ThreadPoolExecutor( max_workers=self.max_workers, thread_name_prefix="openclaw-model-download" ) @@ -418,125 +463,120 @@ class ModelManager: return val return default + @staticmethod + def _error(code: str, detail: str, status: int = 400) -> ModelManagerError: + return ModelManagerError(code, detail, status) + + @staticmethod + def _norm_model_type(value: str) -> str: + return _norm_model_type(value) + + @staticmethod + def _norm_source(value: str) -> str: + return _norm_source(value) + + @staticmethod + def _is_sha256(value: str) -> bool: + return _is_sha256(value) + + @staticmethod + def _sanitize_subdir(value: str) -> str: + return _sanitize_subdir(value) + + @staticmethod + def _sanitize_filename(value: str) -> str: + return _sanitize_filename(value) + + @staticmethod + def _filename_from_url(value: str) -> str: + return _filename_from_url(value) + + @staticmethod + def _file_sha256(path: Path) -> str: + return _file_sha256(path) + + @staticmethod + def _safe_unlink(path: Path) -> None: + _safe_unlink(path) + + @staticmethod + def _parse_content_range(value: str) -> tuple[int, int, int]: + return _parse_content_range(value) + + @staticmethod + def _is_multi_tenant_enabled() -> bool: + return is_multi_tenant_enabled() + + @staticmethod + def _normalize_tenant_id(value: str) -> str: + return normalize_tenant_id(value) + + @staticmethod + def _validate_outbound_download_url( + url: str, + *, + allow_hosts: Optional[set[str]], + allow_any_public_host: bool, + allow_loopback_hosts: Optional[set[str]], + ) -> tuple[str, str, int, list[str]]: + return validate_outbound_url( + url, + allow_hosts=allow_hosts, + allow_any_public_host=allow_any_public_host, + allow_loopback_hosts=allow_loopback_hosts, + policy=STANDARD_OUTBOUND_POLICY, + ) + + @staticmethod + def _build_pinned_download_opener(pinned_ips: list[str]) -> Any: + return _build_pinned_opener(pinned_ips) + + @staticmethod + def _resolve_install_target(root: str, rel_target: str) -> str: + return resolve_under_root(root, rel_target) + def _persist_tasks_locked(self, *, force: bool = False) -> None: - now = time.time() - if not force and (now - self._last_tasks_persist_at) < 0.3: - return - rows = [task.to_dict() for task in self._tasks.values()] - rows.sort(key=lambda row: float(row.get("created_at") or 0.0)) - _atomic_json_write(self.tasks_path, rows) - self._last_tasks_persist_at = now + _persist_tasks_locked_impl( + manager=self, force=force, atomic_json_write=_atomic_json_write + ) def _load_tasks_from_disk(self) -> None: - if not self.tasks_path.exists(): - return - try: - data = json.loads(self.tasks_path.read_text(encoding="utf-8")) - except Exception: - logger.warning( - "F65: failed to parse download task state, ignoring", exc_info=True - ) - return - if not isinstance(data, list): - return - with self._lock: - for item in data: - try: - task = DownloadTask.from_dict(item) - except Exception: - continue - if not task.task_id: - continue - self._tasks[task.task_id] = task - if not task.is_terminal(): - self._cancel_events[task.task_id] = threading.Event() + _load_tasks_from_disk_impl( + manager=self, task_from_dict=DownloadTask.from_dict, logger=logger + ) def _recover_incomplete_tasks(self) -> None: - now = time.time() - with self._lock: - recoverable = sorted( - [t for t in self._tasks.values() if not t.is_terminal()], - key=lambda t: t.created_at, - ) - if not recoverable: - return - for task in recoverable: - task.state = "recovering" - task.updated_at = now - task.error = "restart_recovery_pending" - task.recovery_attempts += 1 - task.resume_status = "restart_recovering" - replayable = recoverable[: self.recovery_replay_limit] - overflow = recoverable[self.recovery_replay_limit :] - for task in overflow: - task.state = "failed" - task.error = "recovery_replay_limit_exceeded" - task.resume_status = "recovery_replay_limit_exceeded" - task.finished_at = now - task.updated_at = now - self._emit(task) - for task in replayable: - event = self._cancel_events.setdefault(task.task_id, threading.Event()) - event.clear() - task.state = "queued" - task.cancel_requested = False - task.error = "" - task.updated_at = now - task.resume_status = "restart_replay_queued" - self._futures[task.task_id] = self._executor.submit( - self._run_task, task.task_id - ) - self._emit(task) - self._persist_tasks_locked(force=True) + _recover_incomplete_tasks_impl(manager=self) @staticmethod def _checkpoint_path(part_path: Path) -> Path: - return Path(f"{part_path}{CHECKPOINT_SUFFIX}") + return _checkpoint_path_impl( + part_path=part_path, checkpoint_suffix=CHECKPOINT_SUFFIX + ) def _load_checkpoint(self, checkpoint_path: Path) -> Dict[str, Any]: - if not checkpoint_path.exists(): - return {} - try: - payload = json.loads(checkpoint_path.read_text(encoding="utf-8")) - except Exception: - return {} - return payload if isinstance(payload, dict) else {} + return _load_checkpoint_impl(checkpoint_path=checkpoint_path) @staticmethod def _checkpoint_matches_task( task: DownloadTask, checkpoint: Dict[str, Any], partial_bytes: int ) -> bool: - if not checkpoint: - return False - if int(checkpoint.get("version") or -1) != CHECKPOINT_VERSION: - return False - if str(checkpoint.get("task_id") or "") != task.task_id: - return False - if str(checkpoint.get("download_url") or "") != task.download_url: - return False - if str(checkpoint.get("expected_sha256") or "") != task.expected_sha256: - return False - if str(checkpoint.get("filename") or "") != task.filename: - return False - if int(checkpoint.get("bytes_downloaded") or -1) != int(partial_bytes): - return False - return True + return _checkpoint_matches_task_impl( + task=task, + checkpoint=checkpoint, + partial_bytes=partial_bytes, + checkpoint_version=CHECKPOINT_VERSION, + ) @staticmethod def _validators_match( checkpoint: Dict[str, Any], response_etag: str, response_last_modified: str ) -> bool: - expected_etag = str(checkpoint.get("etag") or "").strip() - expected_last_modified = str(checkpoint.get("last_modified") or "").strip() - if expected_etag and response_etag and expected_etag != response_etag: - return False - if ( - expected_last_modified - and response_last_modified - and expected_last_modified != response_last_modified - ): - return False - return True + return _validators_match_impl( + checkpoint=checkpoint, + response_etag=response_etag, + response_last_modified=response_last_modified, + ) def _save_checkpoint( self, @@ -548,168 +588,68 @@ class ModelManager: etag: str, last_modified: str, ) -> None: - payload = { - "version": CHECKPOINT_VERSION, - "task_id": task.task_id, - "download_url": task.download_url, - "expected_sha256": task.expected_sha256, - "filename": task.filename, - "bytes_downloaded": max(0, int(bytes_downloaded)), - "total_bytes": max(0, int(total_bytes)), - "etag": str(etag or ""), - "last_modified": str(last_modified or ""), - "updated_at": time.time(), - } - _atomic_json_write(checkpoint_path, payload) - with self._lock: - current = self._tasks.get(task.task_id) - if current is not None: - current.last_checkpoint_at = payload["updated_at"] - self._persist_tasks_locked(force=False) + _save_checkpoint_impl( + manager=self, + checkpoint_path=checkpoint_path, + task=task, + bytes_downloaded=bytes_downloaded, + total_bytes=total_bytes, + etag=etag, + last_modified=last_modified, + checkpoint_version=CHECKPOINT_VERSION, + atomic_json_write=_atomic_json_write, + ) def _set_resume_status(self, task_id: str, status: str) -> None: - with self._lock: - current = self._tasks.get(task_id) - if current is None: - return - current.resume_status = str(status or "not_started")[:120] - current.updated_at = time.time() - self._persist_tasks_locked(force=True) + _set_resume_status_impl(manager=self, task_id=task_id, status=status) def _tenant_ok(self, record_tenant: str, request_tenant: Optional[str]) -> bool: - if not is_multi_tenant_enabled(): - return True - try: - expect = normalize_tenant_id(request_tenant or DEFAULT_TENANT_ID) - except Exception: - expect = DEFAULT_TENANT_ID - try: - got = normalize_tenant_id(record_tenant or DEFAULT_TENANT_ID) - except Exception: - got = DEFAULT_TENANT_ID - return got == expect + return _tenant_ok_impl( + record_tenant=record_tenant, + request_tenant=request_tenant, + default_tenant_id=DEFAULT_TENANT_ID, + is_multi_tenant_enabled=is_multi_tenant_enabled, + normalize_tenant_id=normalize_tenant_id, + ) def _emit(self, task: DownloadTask) -> None: - event_type = { - "queued": JobEventType.QUEUED, - "running": JobEventType.RUNNING, - "completed": JobEventType.COMPLETED, - "failed": JobEventType.FAILED, - "cancelled": JobEventType.CANCELLED, - }.get(task.state) - if event_type is None: - return - get_job_event_store().emit( - event_type=event_type, - prompt_id=f"model_download:{task.task_id}", - trace_id="", - data={ - "channel": "model_download", - "task_id": task.task_id, - "model_id": task.model_id, - "state": task.state, - "progress": task.progress, - "bytes_downloaded": task.bytes_downloaded, - "total_bytes": task.total_bytes, - "error": task.error, - "source": task.source, - "source_label": task.source_label, - "resume_status": task.resume_status, - }, + _emit_impl( + task=task, + event_type_cls=JobEventType, + event_store_getter=get_job_event_store, ) def _load_installations(self) -> List[Dict[str, Any]]: - if not self.installations_path.exists(): - return [] - try: - data = json.loads(self.installations_path.read_text(encoding="utf-8")) - except Exception: - return [] - if not isinstance(data, list): - return [] - return [item for item in data if isinstance(item, dict)] + return _load_installations_impl(installations_path=self.installations_path) def _save_installations(self, rows: List[Dict[str, Any]]) -> None: - _atomic_json_write(self.installations_path, rows) + _save_installations_impl( + installations_path=self.installations_path, + atomic_json_write=_atomic_json_write, + rows=rows, + ) def _collect_install_entries( self, tenant_id: Optional[str] ) -> List[Dict[str, Any]]: - rows = [] - for rec in self._load_installations(): - if not self._tenant_ok( - str(rec.get("tenant_id") or DEFAULT_TENANT_ID), tenant_id - ): - continue - rows.append( - { - "id": str(rec.get("model_id") or rec.get("id") or ""), - "name": str(rec.get("name") or ""), - "model_type": _norm_model_type(str(rec.get("model_type") or "")), - "source": _norm_source(str(rec.get("source") or "managed_install")), - "source_label": str(rec.get("source_label") or "Managed Install"), - "installed": True, - "download_url": str(rec.get("download_url") or ""), - "sha256": str(rec.get("sha256") or "").lower(), - "size_bytes": rec.get("size_bytes"), - "tags": list(rec.get("tags") or []), - "provenance": dict(rec.get("provenance") or {}), - "installation_path": str(rec.get("installation_path") or ""), - "tenant_id": str(rec.get("tenant_id") or DEFAULT_TENANT_ID), - "updated_at": float( - rec.get("installed_at") or rec.get("updated_at") or 0.0 - ), - } - ) - return [row for row in rows if row["id"] and row["name"]] + return _collect_install_entries_impl( + manager=self, + tenant_id=tenant_id, + default_tenant_id=DEFAULT_TENANT_ID, + norm_model_type=_norm_model_type, + norm_source=_norm_source, + ) def _collect_catalog_entries( self, tenant_id: Optional[str] ) -> List[Dict[str, Any]]: - rows = [] - for path in sorted(self.catalog_dir.glob("*.json")): - try: - payload = json.loads(path.read_text(encoding="utf-8")) - except Exception: - continue - if not isinstance(payload, dict): - continue - src = _norm_source(str(payload.get("source") or path.stem)) - src_label = str(payload.get("source_label") or src) - items = payload.get("items") - if not isinstance(items, list): - continue - for item in items: - if not isinstance(item, dict): - continue - tid = str(item.get("tenant_id") or DEFAULT_TENANT_ID) - if not self._tenant_ok(tid, tenant_id): - continue - model_id = str(item.get("id") or item.get("model_id") or "").strip() - name = str(item.get("name") or model_id).strip() - if not model_id or not name: - continue - rows.append( - { - "id": model_id, - "name": name, - "model_type": _norm_model_type( - str(item.get("model_type") or "") - ), - "source": src, - "source_label": src_label, - "installed": False, - "download_url": str(item.get("download_url") or ""), - "sha256": str(item.get("sha256") or "").lower(), - "size_bytes": item.get("size_bytes"), - "tags": list(item.get("tags") or []), - "provenance": dict(item.get("provenance") or {}), - "installation_path": "", - "tenant_id": tid, - "updated_at": float(item.get("updated_at") or 0.0), - } - ) - return rows + return _collect_catalog_entries_impl( + manager=self, + tenant_id=tenant_id, + default_tenant_id=DEFAULT_TENANT_ID, + norm_model_type=_norm_model_type, + norm_source=_norm_source, + ) def search_models( self, @@ -722,112 +662,31 @@ class ModelManager: offset: int = 0, tenant_id: Optional[str] = None, ) -> Dict[str, Any]: - limit = max(1, min(200, int(limit))) - offset = max(0, int(offset)) - q = str(query or "").strip().lower() - src_filter = _norm_source(source) if str(source or "").strip() else "" - type_filter = ( - _norm_model_type(model_type) if str(model_type or "").strip() else "" + return _search_models_impl( + manager=self, + query=query, + source=source, + model_type=model_type, + installed=installed, + limit=limit, + offset=offset, + tenant_id=tenant_id, + norm_source=_norm_source, + norm_model_type=_norm_model_type, + default_tenant_id=DEFAULT_TENANT_ID, ) - rows = self._collect_install_entries(tenant_id) + self._collect_catalog_entries( - tenant_id - ) - out = [] - for row in rows: - if src_filter and row["source"] != src_filter: - continue - if type_filter and row["model_type"] != type_filter: - continue - if installed is not None and bool(row["installed"]) != bool(installed): - continue - if q: - hay = " ".join( - [ - str(row["id"]).lower(), - str(row["name"]).lower(), - " ".join(str(x).lower() for x in (row.get("tags") or [])), - ] - ) - if q not in hay: - continue - out.append(row) - # IMPORTANT: deterministic order is part of the search contract. - out.sort( - key=lambda row: ( - 0 if row["installed"] else 1, - str(row["name"]).lower(), - str(row["id"]).lower(), - str(row["source"]).lower(), - ) - ) - total = len(out) - page = out[offset : offset + limit] - return { - "items": page, - "pagination": {"limit": limit, "offset": offset, "total": total}, - "filters": { - "query": q, - "source": src_filter or None, - "model_type": type_filter or None, - "installed": installed, - }, - } def _validate_url_policy(self, url: str) -> None: - if not str(url or "").strip(): - raise ModelManagerError("invalid_url", "download_url is required") - if not self.allow_hosts and not self.allow_any_public: - raise ModelManagerError( - "download_host_policy_missing", - "set OPENCLAW_MODEL_DOWNLOAD_ALLOW_HOSTS or OPENCLAW_MODEL_DOWNLOAD_ALLOW_ANY_PUBLIC=1", - ) - try: - validate_outbound_url( - str(url).strip(), - allow_hosts=self.allow_hosts or None, - allow_any_public_host=self.allow_any_public, - allow_loopback_hosts=self.allow_loopback_hosts or None, - policy=STANDARD_OUTBOUND_POLICY, - ) - except SSRFError as exc: - raise ModelManagerError("ssrf_blocked", str(exc)) + _validate_url_policy_impl(manager=self, url=url) - @staticmethod - def _validate_provenance(provenance: Dict[str, Any]) -> Dict[str, Any]: - if not isinstance(provenance, dict): - raise ModelManagerError( - "invalid_provenance", "provenance must be an object" - ) - out = { - "publisher": str(provenance.get("publisher") or "").strip(), - "license": str(provenance.get("license") or "").strip(), - "source_url": str(provenance.get("source_url") or "").strip(), - "note": str(provenance.get("note") or "").strip()[:500], - } - if not out["publisher"] or not out["license"] or not out["source_url"]: - raise ModelManagerError( - "invalid_provenance", - "provenance.publisher, provenance.license, provenance.source_url are required", - ) - return out + def _validate_provenance(self, provenance: Dict[str, Any]) -> Dict[str, Any]: + return _validate_provenance_impl(manager=self, provenance=provenance) def _normalize_tenant(self, tenant_id: Optional[str]) -> str: - if not is_multi_tenant_enabled(): - return DEFAULT_TENANT_ID - return normalize_tenant_id(tenant_id or DEFAULT_TENANT_ID) + return _normalize_tenant_impl(manager=self, tenant_id=tenant_id) def _assert_budget(self) -> None: - active = 0 - with self._lock: - for task in self._tasks.values(): - if task.state in {"queued", "running"}: - active += 1 - if active >= self.max_active: - raise ModelManagerError( - "download_queue_full", - f"download queue full (limit={self.max_active})", - 429, - ) + _assert_budget_impl(manager=self) def create_download_task( self, @@ -844,207 +703,28 @@ class ModelManager: filename: Optional[str] = None, tenant_id: Optional[str] = None, ) -> Dict[str, Any]: - self._assert_budget() - model_id = str(model_id or "").strip() - if not model_id: - raise ModelManagerError("validation_error", "model_id is required") - name = str(name or "").strip() - if not name: - raise ModelManagerError("validation_error", "name is required") - digest = str(expected_sha256 or "").strip().lower() - if not _is_sha256(digest): - raise ModelManagerError( - "validation_error", "expected_sha256 must be a 64-char hex string" - ) - self._validate_url_policy(download_url) - provenance = self._validate_provenance(provenance) - mtype = _norm_model_type(model_type) - dest_subdir = _sanitize_subdir( - destination_subdir or MODEL_TYPE_TO_SUBDIR.get(mtype, "misc") - ) - fname = ( - _sanitize_filename(filename) - if filename - else _filename_from_url(download_url) - ) - task = DownloadTask( - task_id=str(uuid.uuid4()), + return _create_download_task_impl( + manager=self, model_id=model_id, name=name, - model_type=mtype, - source=_norm_source(source), - source_label=str(source_label or _norm_source(source))[:80], - download_url=str(download_url).strip(), - destination_subdir=dest_subdir, - filename=fname, - expected_sha256=digest, + model_type=model_type, + source=source, + source_label=source_label, + download_url=download_url, + expected_sha256=expected_sha256, provenance=provenance, - tenant_id=self._normalize_tenant(tenant_id), - resume_status="queued_new", + destination_subdir=destination_subdir, + filename=filename, + tenant_id=tenant_id, ) - with self._lock: - self._tasks[task.task_id] = task - self._cancel_events[task.task_id] = threading.Event() - self._futures[task.task_id] = self._executor.submit( - self._run_task, task.task_id - ) - self._persist_tasks_locked(force=True) - self._emit(task) - return task.to_dict() def _run_task(self, task_id: str) -> None: - with self._lock: - task = self._tasks.get(task_id) - cancel_event = self._cancel_events.get(task_id) - if task is None or cancel_event is None: - return - task.state = "running" - task.started_at = time.time() - task.updated_at = task.started_at - task.resume_status = task.resume_status or "running" - self._emit(task) - self._persist_tasks_locked(force=True) - try: - staged_path, digest = self._download(task, cancel_event) - with self._lock: - current = self._tasks.get(task_id) - if current is None: - return - current.state = "completed" - current.updated_at = time.time() - current.finished_at = current.updated_at - current.progress = 1.0 - current.staged_path = staged_path - current.computed_sha256 = digest - self._emit(current) - self._persist_tasks_locked(force=True) - except DownloadCancelled: - with self._lock: - current = self._tasks.get(task_id) - if current is None: - return - current.state = "cancelled" - current.error = "cancelled" - current.updated_at = time.time() - current.finished_at = current.updated_at - self._emit(current) - self._persist_tasks_locked(force=True) - except Exception as exc: - with self._lock: - current = self._tasks.get(task_id) - if current is None: - return - current.state = "failed" - current.error = str(exc) - current.updated_at = time.time() - current.finished_at = current.updated_at - self._emit(current) - self._persist_tasks_locked(force=True) + _run_task_impl(manager=self, task_id=task_id) def _download( self, task: DownloadTask, cancel_event: threading.Event ) -> tuple[str, str]: - _scheme, _host, _port, pinned_ips = validate_outbound_url( - task.download_url, - allow_hosts=self.allow_hosts or None, - allow_any_public_host=self.allow_any_public, - allow_loopback_hosts=self.allow_loopback_hosts or None, - policy=STANDARD_OUTBOUND_POLICY, - ) - opener = _build_pinned_opener(pinned_ips) - stage_dir = self.staging_dir / task.task_id - stage_dir.mkdir(parents=True, exist_ok=True) - part = stage_dir / f"{task.filename}.part" - checkpoint = self._checkpoint_path(part) - final = stage_dir / task.filename - if final.exists(): - _safe_unlink(final) - - resume_bytes = part.stat().st_size if part.exists() else 0 - checkpoint_data = self._load_checkpoint(checkpoint) - - if resume_bytes > 0 and self._checkpoint_matches_task( - task, checkpoint_data, resume_bytes - ): - digest = hashlib.sha256() - with open(part, "rb") as fh: - while True: - chunk = fh.read(1024 * 1024) - if not chunk: - break - digest.update(chunk) - self._set_resume_status(task.task_id, "resume_attempt") - ( - downloaded, - total, - _etag, - _last_modified, - fallback_reason, - ) = self._stream_response_to_part( - opener=opener, - task=task, - cancel_event=cancel_event, - part=part, - checkpoint=checkpoint, - digest=digest, - resume_from=resume_bytes, - checkpoint_data=checkpoint_data, - ) - if not fallback_reason: - got = digest.hexdigest() - if got != task.expected_sha256: - _safe_unlink(part) - _safe_unlink(checkpoint) - raise ModelManagerError( - "sha256_mismatch", f"expected {task.expected_sha256}, got {got}" - ) - os.replace(part, final) - _safe_unlink(checkpoint) - self._set_resume_status(task.task_id, "resumed_partial") - self._progress(task.task_id, downloaded, total or downloaded) - return str(final), got - self._set_resume_status(task.task_id, fallback_reason) - - elif resume_bytes > 0: - # IMPORTANT: resume only when checkpoint metadata matches this task. - # Blindly appending without metadata validation can corrupt artifacts. - self._set_resume_status(task.task_id, "resume_fallback_checkpoint_mismatch") - - _safe_unlink(part) - _safe_unlink(checkpoint) - digest = hashlib.sha256() - ( - downloaded, - total, - _etag, - _last_modified, - fallback_reason, - ) = self._stream_response_to_part( - opener=opener, - task=task, - cancel_event=cancel_event, - part=part, - checkpoint=checkpoint, - digest=digest, - resume_from=0, - checkpoint_data={}, - ) - if fallback_reason: - raise ModelManagerError("download_resume_failed", fallback_reason) - - got = digest.hexdigest() - if got != task.expected_sha256: - _safe_unlink(part) - _safe_unlink(checkpoint) - raise ModelManagerError( - "sha256_mismatch", f"expected {task.expected_sha256}, got {got}" - ) - os.replace(part, final) - _safe_unlink(checkpoint) - if resume_bytes <= 0: - self._set_resume_status(task.task_id, "started_fresh") - self._progress(task.task_id, downloaded, total or downloaded) - return str(final), got + return _download_impl(manager=self, task=task, cancel_event=cancel_event) def _stream_response_to_part( self, @@ -1058,127 +738,22 @@ class ModelManager: resume_from: int, checkpoint_data: Dict[str, Any], ) -> tuple[int, int, str, str, str]: - req = urllib.request.Request(task.download_url, method="GET") - req.add_header("User-Agent", "ComfyUI-OpenClaw/F65") - if resume_from > 0: - req.add_header("Range", f"bytes={resume_from}-") - - with opener.open(req, timeout=self.timeout_sec) as resp: - code = int(resp.getcode() or 0) - if code in (301, 302, 303, 307, 308): - raise ModelManagerError( - "download_redirect_blocked", - "redirect blocked for managed downloads", - ) - if code >= 400: - raise ModelManagerError("download_http_error", f"HTTP {code}") - - etag = str(resp.headers.get("ETag") or "").strip() - last_modified = str(resp.headers.get("Last-Modified") or "").strip() - content_length = 0 - try: - content_length = max( - 0, - int(str(resp.headers.get("Content-Length") or "0")), - ) - except Exception: - content_length = 0 - - if resume_from > 0: - if code != 206: - return ( - resume_from, - 0, - etag, - last_modified, - "resume_fallback_range_not_supported", - ) - if not self._validators_match(checkpoint_data, etag, last_modified): - return ( - resume_from, - 0, - etag, - last_modified, - "resume_fallback_validator_mismatch", - ) - range_start, _range_end, range_total = _parse_content_range( - str(resp.headers.get("Content-Range") or "") - ) - if range_start != resume_from: - return ( - resume_from, - 0, - etag, - last_modified, - "resume_fallback_content_range_mismatch", - ) - total = ( - range_total if range_total > 0 else (resume_from + content_length) - ) - mode = "ab" - downloaded = resume_from - else: - total = content_length - mode = "wb" - downloaded = 0 - - last_emit = 0.0 - with open(part, mode) as fh: - while True: - if cancel_event.is_set(): - self._save_checkpoint( - checkpoint, - task, - bytes_downloaded=downloaded, - total_bytes=total, - etag=etag, - last_modified=last_modified, - ) - raise DownloadCancelled() - chunk = resp.read(64 * 1024) - if not chunk: - break - fh.write(chunk) - digest.update(chunk) - downloaded += len(chunk) - now = time.time() - if now - last_emit >= 0.35: - self._progress(task.task_id, downloaded, total) - self._save_checkpoint( - checkpoint, - task, - bytes_downloaded=downloaded, - total_bytes=total, - etag=etag, - last_modified=last_modified, - ) - last_emit = now - self._progress(task.task_id, downloaded, total or downloaded) - self._save_checkpoint( - checkpoint, - task, - bytes_downloaded=downloaded, - total_bytes=total, - etag=etag, - last_modified=last_modified, - ) - return downloaded, total, etag, last_modified, "" + return _stream_response_to_part_impl( + manager=self, + opener=opener, + task=task, + cancel_event=cancel_event, + part=part, + checkpoint=checkpoint, + digest=digest, + resume_from=resume_from, + checkpoint_data=checkpoint_data, + ) def _progress(self, task_id: str, downloaded: int, total: int) -> None: - with self._lock: - task = self._tasks.get(task_id) - if task is None: - return - task.bytes_downloaded = max(0, int(downloaded)) - task.total_bytes = max(0, int(total)) - task.progress = ( - min(1.0, (task.bytes_downloaded / task.total_bytes)) - if task.total_bytes - else 0.0 - ) - task.updated_at = time.time() - self._emit(task) - self._persist_tasks_locked(force=False) + _progress_impl( + manager=self, task_id=task_id, downloaded=downloaded, total=total + ) def list_download_tasks( self, @@ -1188,62 +763,27 @@ class ModelManager: limit: int = 100, offset: int = 0, ) -> Dict[str, Any]: - limit = max(1, min(200, int(limit))) - offset = max(0, int(offset)) - state_filter = str(state or "").strip().lower() - with self._lock: - tasks = list(self._tasks.values()) - out = [] - for task in tasks: - if not self._tenant_ok(task.tenant_id, tenant_id): - continue - if state_filter and task.state != state_filter: - continue - out.append(task) - out.sort(key=lambda x: x.created_at, reverse=True) - total = len(out) - page = [item.to_dict() for item in out[offset : offset + limit]] - return { - "tasks": page, - "pagination": {"limit": limit, "offset": offset, "total": total}, - "filters": {"state": state_filter or None}, - } + return _list_download_tasks_impl( + manager=self, + tenant_id=tenant_id, + state=state, + limit=limit, + offset=offset, + ) def get_download_task( self, task_id: str, *, tenant_id: Optional[str] = None ) -> Dict[str, Any]: - with self._lock: - task = self._tasks.get(task_id) - if task is None or not self._tenant_ok(task.tenant_id, tenant_id): - raise ModelManagerError("not_found", "download task not found", 404) - return task.to_dict() + return _get_download_task_impl( + manager=self, task_id=task_id, tenant_id=tenant_id + ) def cancel_download_task( self, task_id: str, *, tenant_id: Optional[str] = None ) -> Dict[str, Any]: - with self._lock: - task = self._tasks.get(task_id) - future = self._futures.get(task_id) - event = self._cancel_events.get(task_id) - if ( - task is None - or event is None - or not self._tenant_ok(task.tenant_id, tenant_id) - ): - raise ModelManagerError("not_found", "download task not found", 404) - if task.is_terminal(): - return task.to_dict() - task.cancel_requested = True - task.updated_at = time.time() - event.set() - if task.state == "queued" and future is not None and future.cancel(): - task.state = "cancelled" - task.error = "cancelled_before_start" - task.finished_at = time.time() - task.updated_at = task.finished_at - self._emit(task) - self._persist_tasks_locked(force=True) - return task.to_dict() + return _cancel_download_task_impl( + manager=self, task_id=task_id, tenant_id=tenant_id + ) def import_downloaded_model( self, @@ -1254,89 +794,14 @@ class ModelManager: filename: Optional[str] = None, tags: Optional[List[str]] = None, ) -> Dict[str, Any]: - with self._lock: - task = self._tasks.get(task_id) - if task is None or not self._tenant_ok(task.tenant_id, tenant_id): - raise ModelManagerError("not_found", "download task not found", 404) - if task.state != "completed": - raise ModelManagerError( - "task_not_ready", "task must be completed before import" - ) - if task.imported: - raise ModelManagerError("already_imported", "task already imported") - staged_path = Path(task.staged_path) - expected = task.expected_sha256 - computed = task.computed_sha256 - if not staged_path.exists(): - raise ModelManagerError("staging_missing", "staged file missing") - # CRITICAL: keep import-time hash verification. Removing this reopens - # tamper window between download completion and activation/import. - actual = _file_sha256(staged_path) - if actual != expected or computed != expected: - raise ModelManagerError( - "sha256_mismatch", f"expected {expected}, got {actual}" - ) - self._validate_provenance(task.provenance) - subdir = _sanitize_subdir(destination_subdir or task.destination_subdir) - fname = _sanitize_filename(filename or task.filename) - rel_target = f"{subdir}/{fname}" - # IMPORTANT: keep root-bounded resolution; plain joins re-enable traversal risks. - abs_target = Path(resolve_under_root(str(self.install_root), rel_target)) - abs_target.parent.mkdir(parents=True, exist_ok=True) - fd, tmp = tempfile.mkstemp( - prefix=f".{abs_target.name}.tmp.", dir=str(abs_target.parent), text=False + return _import_downloaded_model_impl( + manager=self, + task_id=task_id, + tenant_id=tenant_id, + destination_subdir=destination_subdir, + filename=filename, + tags=tags, ) - os.close(fd) - try: - shutil.copy2(staged_path, tmp) - os.replace(tmp, abs_target) - except Exception: - try: - os.remove(tmp) - except OSError: - pass - raise - safe_tags: List[str] = [] - for item in tags or []: - if not isinstance(item, str): - continue - clean = item.strip().lower() - if not clean or clean in safe_tags: - continue - safe_tags.append(clean) - if len(safe_tags) >= 24: - break - rec = { - "id": str(uuid.uuid4()), - "task_id": task.task_id, - "model_id": task.model_id, - "name": task.name, - "model_type": task.model_type, - "source": task.source, - "source_label": task.source_label, - "download_url": task.download_url, - "sha256": expected, - "size_bytes": abs_target.stat().st_size if abs_target.exists() else None, - "provenance": dict(task.provenance), - "installation_path": rel_target.replace("\\", "/"), - "tenant_id": task.tenant_id, - "installed_at": time.time(), - "tags": safe_tags, - } - rows = self._load_installations() - rows.append(rec) - rows.sort(key=lambda x: float(x.get("installed_at") or 0.0), reverse=True) - self._save_installations(rows) - with self._lock: - current = self._tasks.get(task.task_id) - if current is not None: - current.imported = True - current.installation_path = rec["installation_path"] - current.installation_record_id = rec["id"] - current.updated_at = time.time() - self._emit(current) - self._persist_tasks_locked(force=True) - return rec def list_installations( self, @@ -1346,30 +811,15 @@ class ModelManager: limit: int = 100, offset: int = 0, ) -> Dict[str, Any]: - limit = max(1, min(200, int(limit))) - offset = max(0, int(offset)) - type_filter = ( - _norm_model_type(model_type) if str(model_type or "").strip() else "" + return _list_installations_impl( + manager=self, + tenant_id=tenant_id, + model_type=model_type, + limit=limit, + offset=offset, + norm_model_type=_norm_model_type, + default_tenant_id=DEFAULT_TENANT_ID, ) - rows = [] - for rec in self._load_installations(): - if not self._tenant_ok( - str(rec.get("tenant_id") or DEFAULT_TENANT_ID), tenant_id - ): - continue - if ( - type_filter - and _norm_model_type(str(rec.get("model_type") or "")) != type_filter - ): - continue - rows.append(rec) - rows.sort(key=lambda x: float(x.get("installed_at") or 0.0), reverse=True) - total = len(rows) - return { - "installations": rows[offset : offset + limit], - "pagination": {"limit": limit, "offset": offset, "total": total}, - "filters": {"model_type": type_filter or None}, - } model_manager = ModelManager() diff --git a/services/model_manager_catalog.py b/services/model_manager_catalog.py new file mode 100644 index 0000000..595b327 --- /dev/null +++ b/services/model_manager_catalog.py @@ -0,0 +1,227 @@ +""" +Internal catalog/installations helpers for the model manager facade. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + + +def load_installations(*, installations_path: Path) -> List[Dict[str, Any]]: + if not installations_path.exists(): + return [] + try: + data = json.loads(installations_path.read_text(encoding="utf-8")) + except Exception: + return [] + if not isinstance(data, list): + return [] + return [item for item in data if isinstance(item, dict)] + + +def save_installations( + *, + installations_path: Path, + atomic_json_write: Callable[[Path, Any], None], + rows: List[Dict[str, Any]], +) -> None: + atomic_json_write(installations_path, rows) + + +def collect_install_entries( + *, + manager: Any, + tenant_id: Optional[str], + default_tenant_id: str, + norm_model_type: Callable[[str], str], + norm_source: Callable[[str], str], +) -> List[Dict[str, Any]]: + rows = [] + for rec in load_installations(installations_path=manager.installations_path): + if not manager._tenant_ok( + str(rec.get("tenant_id") or default_tenant_id), tenant_id + ): + continue + rows.append( + { + "id": str(rec.get("model_id") or rec.get("id") or ""), + "name": str(rec.get("name") or ""), + "model_type": norm_model_type(str(rec.get("model_type") or "")), + "source": norm_source(str(rec.get("source") or "managed_install")), + "source_label": str(rec.get("source_label") or "Managed Install"), + "installed": True, + "download_url": str(rec.get("download_url") or ""), + "sha256": str(rec.get("sha256") or "").lower(), + "size_bytes": rec.get("size_bytes"), + "tags": list(rec.get("tags") or []), + "provenance": dict(rec.get("provenance") or {}), + "installation_path": str(rec.get("installation_path") or ""), + "tenant_id": str(rec.get("tenant_id") or default_tenant_id), + "updated_at": float( + rec.get("installed_at") or rec.get("updated_at") or 0.0 + ), + } + ) + return [row for row in rows if row["id"] and row["name"]] + + +def collect_catalog_entries( + *, + manager: Any, + tenant_id: Optional[str], + default_tenant_id: str, + norm_model_type: Callable[[str], str], + norm_source: Callable[[str], str], +) -> List[Dict[str, Any]]: + rows = [] + for path in sorted(manager.catalog_dir.glob("*.json")): + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except Exception: + continue + if not isinstance(payload, dict): + continue + src = norm_source(str(payload.get("source") or path.stem)) + src_label = str(payload.get("source_label") or src) + items = payload.get("items") + if not isinstance(items, list): + continue + for item in items: + if not isinstance(item, dict): + continue + tid = str(item.get("tenant_id") or default_tenant_id) + if not manager._tenant_ok(tid, tenant_id): + continue + model_id = str(item.get("id") or item.get("model_id") or "").strip() + name = str(item.get("name") or model_id).strip() + if not model_id or not name: + continue + rows.append( + { + "id": model_id, + "name": name, + "model_type": norm_model_type(str(item.get("model_type") or "")), + "source": src, + "source_label": src_label, + "installed": False, + "download_url": str(item.get("download_url") or ""), + "sha256": str(item.get("sha256") or "").lower(), + "size_bytes": item.get("size_bytes"), + "tags": list(item.get("tags") or []), + "provenance": dict(item.get("provenance") or {}), + "installation_path": "", + "tenant_id": tid, + "updated_at": float(item.get("updated_at") or 0.0), + } + ) + return rows + + +def search_models( + *, + manager: Any, + query: str = "", + source: str = "", + model_type: str = "", + installed: Optional[bool] = None, + limit: int = 50, + offset: int = 0, + tenant_id: Optional[str] = None, + norm_source: Callable[[str], str], + norm_model_type: Callable[[str], str], + default_tenant_id: str, +) -> Dict[str, Any]: + limit = max(1, min(200, int(limit))) + offset = max(0, int(offset)) + q = str(query or "").strip().lower() + src_filter = norm_source(source) if str(source or "").strip() else "" + type_filter = norm_model_type(model_type) if str(model_type or "").strip() else "" + rows = collect_install_entries( + manager=manager, + tenant_id=tenant_id, + default_tenant_id=default_tenant_id, + norm_model_type=norm_model_type, + norm_source=norm_source, + ) + collect_catalog_entries( + manager=manager, + tenant_id=tenant_id, + default_tenant_id=default_tenant_id, + norm_model_type=norm_model_type, + norm_source=norm_source, + ) + out = [] + for row in rows: + if src_filter and row["source"] != src_filter: + continue + if type_filter and row["model_type"] != type_filter: + continue + if installed is not None and bool(row["installed"]) != bool(installed): + continue + if q: + hay = " ".join( + [ + str(row["id"]).lower(), + str(row["name"]).lower(), + " ".join(str(x).lower() for x in (row.get("tags") or [])), + ] + ) + if q not in hay: + continue + out.append(row) + # IMPORTANT: deterministic order is part of the search contract. + out.sort( + key=lambda row: ( + 0 if row["installed"] else 1, + str(row["name"]).lower(), + str(row["id"]).lower(), + str(row["source"]).lower(), + ) + ) + total = len(out) + page = out[offset : offset + limit] + return { + "items": page, + "pagination": {"limit": limit, "offset": offset, "total": total}, + "filters": { + "query": q, + "source": src_filter or None, + "model_type": type_filter or None, + "installed": installed, + }, + } + + +def list_installations( + *, + manager: Any, + tenant_id: Optional[str] = None, + model_type: str = "", + limit: int = 100, + offset: int = 0, + norm_model_type: Callable[[str], str], + default_tenant_id: str, +) -> Dict[str, Any]: + limit = max(1, min(200, int(limit))) + offset = max(0, int(offset)) + type_filter = norm_model_type(model_type) if str(model_type or "").strip() else "" + rows = [] + for rec in load_installations(installations_path=manager.installations_path): + if not manager._tenant_ok( + str(rec.get("tenant_id") or default_tenant_id), tenant_id + ): + continue + if ( + type_filter + and norm_model_type(str(rec.get("model_type") or "")) != type_filter + ): + continue + rows.append(rec) + rows.sort(key=lambda x: float(x.get("installed_at") or 0.0), reverse=True) + total = len(rows) + return { + "installations": rows[offset : offset + limit], + "pagination": {"limit": limit, "offset": offset, "total": total}, + "filters": {"model_type": type_filter or None}, + } diff --git a/services/model_manager_tasks.py b/services/model_manager_tasks.py new file mode 100644 index 0000000..441c990 --- /dev/null +++ b/services/model_manager_tasks.py @@ -0,0 +1,344 @@ +""" +Internal task persistence/recovery helpers for the model manager facade. +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any, Callable, Dict, Optional + + +def persist_tasks_locked( + *, + manager: Any, + force: bool = False, + atomic_json_write: Callable[[Path, Any], None], +) -> None: + now = time.time() + if not force and (now - manager._last_tasks_persist_at) < 0.3: + return + rows = [task.to_dict() for task in manager._tasks.values()] + rows.sort(key=lambda row: float(row.get("created_at") or 0.0)) + atomic_json_write(manager.tasks_path, rows) + manager._last_tasks_persist_at = now + + +def load_tasks_from_disk( + *, + manager: Any, + task_from_dict: Callable[[Dict[str, Any]], Any], + logger: Any, +) -> None: + if not manager.tasks_path.exists(): + return + try: + data = json.loads(manager.tasks_path.read_text(encoding="utf-8")) + except Exception: + logger.warning( + "F65: failed to parse download task state, ignoring", exc_info=True + ) + return + if not isinstance(data, list): + return + with manager._lock: + for item in data: + try: + task = task_from_dict(item) + except Exception: + continue + if not task.task_id: + continue + manager._tasks[task.task_id] = task + if not task.is_terminal(): + manager._cancel_events[task.task_id] = ( + manager._threading_event_factory() + ) + + +def recover_incomplete_tasks(*, manager: Any) -> None: + now = time.time() + with manager._lock: + recoverable = sorted( + [t for t in manager._tasks.values() if not t.is_terminal()], + key=lambda t: t.created_at, + ) + if not recoverable: + return + for task in recoverable: + task.state = "recovering" + task.updated_at = now + task.error = "restart_recovery_pending" + task.recovery_attempts += 1 + task.resume_status = "restart_recovering" + replayable = recoverable[: manager.recovery_replay_limit] + overflow = recoverable[manager.recovery_replay_limit :] + for task in overflow: + task.state = "failed" + task.error = "recovery_replay_limit_exceeded" + task.resume_status = "recovery_replay_limit_exceeded" + task.finished_at = now + task.updated_at = now + manager._emit(task) + for task in replayable: + event = manager._cancel_events.setdefault( + task.task_id, manager._threading_event_factory() + ) + event.clear() + task.state = "queued" + task.cancel_requested = False + task.error = "" + task.updated_at = now + task.resume_status = "restart_replay_queued" + manager._futures[task.task_id] = manager._executor.submit( + manager._run_task, task.task_id + ) + manager._emit(task) + manager._persist_tasks_locked(force=True) + + +def checkpoint_path(*, part_path: Path, checkpoint_suffix: str) -> Path: + return Path(f"{part_path}{checkpoint_suffix}") + + +def load_checkpoint(*, checkpoint_path: Path) -> Dict[str, Any]: + if not checkpoint_path.exists(): + return {} + try: + payload = json.loads(checkpoint_path.read_text(encoding="utf-8")) + except Exception: + return {} + return payload if isinstance(payload, dict) else {} + + +def checkpoint_matches_task( + *, + task: Any, + checkpoint: Dict[str, Any], + partial_bytes: int, + checkpoint_version: int, +) -> bool: + if not checkpoint: + return False + if int(checkpoint.get("version") or -1) != checkpoint_version: + return False + if str(checkpoint.get("task_id") or "") != task.task_id: + return False + if str(checkpoint.get("download_url") or "") != task.download_url: + return False + if str(checkpoint.get("expected_sha256") or "") != task.expected_sha256: + return False + if str(checkpoint.get("filename") or "") != task.filename: + return False + if int(checkpoint.get("bytes_downloaded") or -1) != int(partial_bytes): + return False + return True + + +def validators_match( + *, + checkpoint: Dict[str, Any], + response_etag: str, + response_last_modified: str, +) -> bool: + expected_etag = str(checkpoint.get("etag") or "").strip() + expected_last_modified = str(checkpoint.get("last_modified") or "").strip() + if expected_etag and response_etag and expected_etag != response_etag: + return False + if ( + expected_last_modified + and response_last_modified + and expected_last_modified != response_last_modified + ): + return False + return True + + +def save_checkpoint( + *, + manager: Any, + checkpoint_path: Path, + task: Any, + bytes_downloaded: int, + total_bytes: int, + etag: str, + last_modified: str, + checkpoint_version: int, + atomic_json_write: Callable[[Path, Any], None], +) -> None: + payload = { + "version": checkpoint_version, + "task_id": task.task_id, + "download_url": task.download_url, + "expected_sha256": task.expected_sha256, + "filename": task.filename, + "bytes_downloaded": max(0, int(bytes_downloaded)), + "total_bytes": max(0, int(total_bytes)), + "etag": str(etag or ""), + "last_modified": str(last_modified or ""), + "updated_at": time.time(), + } + atomic_json_write(checkpoint_path, payload) + with manager._lock: + current = manager._tasks.get(task.task_id) + if current is not None: + current.last_checkpoint_at = payload["updated_at"] + manager._persist_tasks_locked(force=False) + + +def set_resume_status(*, manager: Any, task_id: str, status: str) -> None: + with manager._lock: + current = manager._tasks.get(task_id) + if current is None: + return + current.resume_status = str(status or "not_started")[:120] + current.updated_at = time.time() + manager._persist_tasks_locked(force=True) + + +def tenant_ok( + *, + record_tenant: str, + request_tenant: Optional[str], + default_tenant_id: str, + is_multi_tenant_enabled: Callable[[], bool], + normalize_tenant_id: Callable[[str], str], +) -> bool: + if not is_multi_tenant_enabled(): + return True + try: + expect = normalize_tenant_id(request_tenant or default_tenant_id) + except Exception: + expect = default_tenant_id + try: + got = normalize_tenant_id(record_tenant or default_tenant_id) + except Exception: + got = default_tenant_id + return got == expect + + +def emit( + *, + task: Any, + event_type_cls: Any, + event_store_getter: Callable[[], Any], +) -> None: + event_type = { + "queued": event_type_cls.QUEUED, + "running": event_type_cls.RUNNING, + "completed": event_type_cls.COMPLETED, + "failed": event_type_cls.FAILED, + "cancelled": event_type_cls.CANCELLED, + }.get(task.state) + if event_type is None: + return + event_store_getter().emit( + event_type=event_type, + prompt_id=f"model_download:{task.task_id}", + trace_id="", + data={ + "channel": "model_download", + "task_id": task.task_id, + "model_id": task.model_id, + "state": task.state, + "progress": task.progress, + "bytes_downloaded": task.bytes_downloaded, + "total_bytes": task.total_bytes, + "error": task.error, + "source": task.source, + "source_label": task.source_label, + "resume_status": task.resume_status, + }, + ) + + +def progress(*, manager: Any, task_id: str, downloaded: int, total: int) -> None: + with manager._lock: + task = manager._tasks.get(task_id) + if task is None: + return + task.bytes_downloaded = max(0, int(downloaded)) + task.total_bytes = max(0, int(total)) + task.progress = ( + min(1.0, (task.bytes_downloaded / task.total_bytes)) + if task.total_bytes + else 0.0 + ) + task.updated_at = time.time() + manager._emit(task) + manager._persist_tasks_locked(force=False) + + +def list_download_tasks( + *, + manager: Any, + tenant_id: Optional[str] = None, + state: str = "", + limit: int = 100, + offset: int = 0, +) -> Dict[str, Any]: + limit = max(1, min(200, int(limit))) + offset = max(0, int(offset)) + state_filter = str(state or "").strip().lower() + with manager._lock: + tasks = list(manager._tasks.values()) + out = [] + for task in tasks: + if not manager._tenant_ok(task.tenant_id, tenant_id): + continue + if state_filter and task.state != state_filter: + continue + out.append(task) + out.sort(key=lambda x: x.created_at, reverse=True) + total = len(out) + page = [item.to_dict() for item in out[offset : offset + limit]] + return { + "tasks": page, + "pagination": {"limit": limit, "offset": offset, "total": total}, + "filters": {"state": state_filter or None}, + } + + +def get_download_task( + *, + manager: Any, + task_id: str, + tenant_id: Optional[str] = None, +) -> Dict[str, Any]: + with manager._lock: + task = manager._tasks.get(task_id) + if task is None or not manager._tenant_ok(task.tenant_id, tenant_id): + raise manager._error("not_found", "download task not found", 404) + return task.to_dict() + + +def cancel_download_task( + *, + manager: Any, + task_id: str, + tenant_id: Optional[str] = None, +) -> Dict[str, Any]: + with manager._lock: + task = manager._tasks.get(task_id) + future = manager._futures.get(task_id) + event = manager._cancel_events.get(task_id) + if ( + task is None + or event is None + or not manager._tenant_ok(task.tenant_id, tenant_id) + ): + raise manager._error("not_found", "download task not found", 404) + if task.is_terminal(): + return task.to_dict() + task.cancel_requested = True + task.updated_at = time.time() + event.set() + if task.state == "queued" and future is not None and future.cancel(): + task.state = "cancelled" + task.error = "cancelled_before_start" + task.finished_at = time.time() + task.updated_at = task.finished_at + manager._emit(task) + manager._persist_tasks_locked(force=True) + return task.to_dict() diff --git a/services/model_manager_transfer.py b/services/model_manager_transfer.py new file mode 100644 index 0000000..cf01eea --- /dev/null +++ b/services/model_manager_transfer.py @@ -0,0 +1,487 @@ +""" +Internal download/import lifecycle helpers for the model manager facade. +""" + +from __future__ import annotations + +import hashlib +import os +import shutil +import tempfile +import time +import urllib.request +import uuid +from pathlib import Path +from typing import Any, Dict, List, Optional + + +def validate_url_policy(*, manager: Any, url: str) -> None: + if not str(url or "").strip(): + raise manager._error("invalid_url", "download_url is required") + if not manager.allow_hosts and not manager.allow_any_public: + raise manager._error( + "download_host_policy_missing", + "set OPENCLAW_MODEL_DOWNLOAD_ALLOW_HOSTS or OPENCLAW_MODEL_DOWNLOAD_ALLOW_ANY_PUBLIC=1", + ) + try: + manager._validate_outbound_download_url( + str(url).strip(), + allow_hosts=manager.allow_hosts or None, + allow_any_public_host=manager.allow_any_public, + allow_loopback_hosts=manager.allow_loopback_hosts or None, + ) + except manager._ssrf_error_cls as exc: + raise manager._error("ssrf_blocked", str(exc)) + + +def validate_provenance(*, manager: Any, provenance: Dict[str, Any]) -> Dict[str, Any]: + if not isinstance(provenance, dict): + raise manager._error("invalid_provenance", "provenance must be an object") + out = { + "publisher": str(provenance.get("publisher") or "").strip(), + "license": str(provenance.get("license") or "").strip(), + "source_url": str(provenance.get("source_url") or "").strip(), + "note": str(provenance.get("note") or "").strip()[:500], + } + if not out["publisher"] or not out["license"] or not out["source_url"]: + raise manager._error( + "invalid_provenance", + "provenance.publisher, provenance.license, provenance.source_url are required", + ) + return out + + +def normalize_tenant(*, manager: Any, tenant_id: Optional[str]) -> str: + if not manager._is_multi_tenant_enabled(): + return manager._default_tenant_id + return manager._normalize_tenant_id(tenant_id or manager._default_tenant_id) + + +def assert_budget(*, manager: Any) -> None: + active = 0 + with manager._lock: + for task in manager._tasks.values(): + if task.state in {"queued", "running"}: + active += 1 + if active >= manager.max_active: + raise manager._error( + "download_queue_full", + f"download queue full (limit={manager.max_active})", + 429, + ) + + +def create_download_task( + *, + manager: Any, + model_id: str, + name: str, + model_type: str, + source: str, + source_label: str, + download_url: str, + expected_sha256: str, + provenance: Dict[str, Any], + destination_subdir: Optional[str] = None, + filename: Optional[str] = None, + tenant_id: Optional[str] = None, +) -> Dict[str, Any]: + manager._assert_budget() + model_id = str(model_id or "").strip() + if not model_id: + raise manager._error("validation_error", "model_id is required") + name = str(name or "").strip() + if not name: + raise manager._error("validation_error", "name is required") + digest = str(expected_sha256 or "").strip().lower() + if not manager._is_sha256(digest): + raise manager._error( + "validation_error", "expected_sha256 must be a 64-char hex string" + ) + manager._validate_url_policy(download_url) + provenance = manager._validate_provenance(provenance) + mtype = manager._norm_model_type(model_type) + dest_subdir = manager._sanitize_subdir( + destination_subdir or manager._model_type_to_subdir.get(mtype, "misc") + ) + fname = ( + manager._sanitize_filename(filename) + if filename + else manager._filename_from_url(download_url) + ) + task = manager._download_task_cls( + task_id=str(uuid.uuid4()), + model_id=model_id, + name=name, + model_type=mtype, + source=manager._norm_source(source), + source_label=str(source_label or manager._norm_source(source))[:80], + download_url=str(download_url).strip(), + destination_subdir=dest_subdir, + filename=fname, + expected_sha256=digest, + provenance=provenance, + tenant_id=manager._normalize_tenant(tenant_id), + resume_status="queued_new", + ) + with manager._lock: + manager._tasks[task.task_id] = task + manager._cancel_events[task.task_id] = manager._threading_event_factory() + manager._futures[task.task_id] = manager._executor.submit( + manager._run_task, task.task_id + ) + manager._persist_tasks_locked(force=True) + manager._emit(task) + return task.to_dict() + + +def run_task(*, manager: Any, task_id: str) -> None: + with manager._lock: + task = manager._tasks.get(task_id) + cancel_event = manager._cancel_events.get(task_id) + if task is None or cancel_event is None: + return + task.state = "running" + task.started_at = time.time() + task.updated_at = task.started_at + task.resume_status = task.resume_status or "running" + manager._emit(task) + manager._persist_tasks_locked(force=True) + try: + staged_path, digest = manager._download(task, cancel_event) + with manager._lock: + current = manager._tasks.get(task_id) + if current is None: + return + current.state = "completed" + current.updated_at = time.time() + current.finished_at = current.updated_at + current.progress = 1.0 + current.staged_path = staged_path + current.computed_sha256 = digest + manager._emit(current) + manager._persist_tasks_locked(force=True) + except manager._download_cancelled_cls: + with manager._lock: + current = manager._tasks.get(task_id) + if current is None: + return + current.state = "cancelled" + current.error = "cancelled" + current.updated_at = time.time() + current.finished_at = current.updated_at + manager._emit(current) + manager._persist_tasks_locked(force=True) + except Exception as exc: + with manager._lock: + current = manager._tasks.get(task_id) + if current is None: + return + current.state = "failed" + current.error = str(exc) + current.updated_at = time.time() + current.finished_at = current.updated_at + manager._emit(current) + manager._persist_tasks_locked(force=True) + + +def download(*, manager: Any, task: Any, cancel_event: Any) -> tuple[str, str]: + _scheme, _host, _port, pinned_ips = manager._validate_outbound_download_url( + task.download_url, + allow_hosts=manager.allow_hosts or None, + allow_any_public_host=manager.allow_any_public, + allow_loopback_hosts=manager.allow_loopback_hosts or None, + ) + opener = manager._build_pinned_download_opener(pinned_ips) + stage_dir = manager.staging_dir / task.task_id + stage_dir.mkdir(parents=True, exist_ok=True) + part = stage_dir / f"{task.filename}.part" + checkpoint = manager._checkpoint_path(part) + final = stage_dir / task.filename + if final.exists(): + manager._safe_unlink(final) + + resume_bytes = part.stat().st_size if part.exists() else 0 + checkpoint_data = manager._load_checkpoint(checkpoint) + + if resume_bytes > 0 and manager._checkpoint_matches_task( + task, checkpoint_data, resume_bytes + ): + digest = hashlib.sha256() + with open(part, "rb") as fh: + while True: + chunk = fh.read(1024 * 1024) + if not chunk: + break + digest.update(chunk) + manager._set_resume_status(task.task_id, "resume_attempt") + downloaded, total, _etag, _last_modified, fallback_reason = ( + manager._stream_response_to_part( + opener=opener, + task=task, + cancel_event=cancel_event, + part=part, + checkpoint=checkpoint, + digest=digest, + resume_from=resume_bytes, + checkpoint_data=checkpoint_data, + ) + ) + if not fallback_reason: + got = digest.hexdigest() + if got != task.expected_sha256: + manager._safe_unlink(part) + manager._safe_unlink(checkpoint) + raise manager._error( + "sha256_mismatch", f"expected {task.expected_sha256}, got {got}" + ) + os.replace(part, final) + manager._safe_unlink(checkpoint) + manager._set_resume_status(task.task_id, "resumed_partial") + manager._progress(task.task_id, downloaded, total or downloaded) + return str(final), got + manager._set_resume_status(task.task_id, fallback_reason) + elif resume_bytes > 0: + # IMPORTANT: resume only when checkpoint metadata matches this task. + # Blindly appending without metadata validation can corrupt artifacts. + manager._set_resume_status(task.task_id, "resume_fallback_checkpoint_mismatch") + + manager._safe_unlink(part) + manager._safe_unlink(checkpoint) + digest = hashlib.sha256() + downloaded, total, _etag, _last_modified, fallback_reason = ( + manager._stream_response_to_part( + opener=opener, + task=task, + cancel_event=cancel_event, + part=part, + checkpoint=checkpoint, + digest=digest, + resume_from=0, + checkpoint_data={}, + ) + ) + if fallback_reason: + raise manager._error("download_resume_failed", fallback_reason) + + got = digest.hexdigest() + if got != task.expected_sha256: + manager._safe_unlink(part) + manager._safe_unlink(checkpoint) + raise manager._error( + "sha256_mismatch", f"expected {task.expected_sha256}, got {got}" + ) + os.replace(part, final) + manager._safe_unlink(checkpoint) + if resume_bytes <= 0: + manager._set_resume_status(task.task_id, "started_fresh") + manager._progress(task.task_id, downloaded, total or downloaded) + return str(final), got + + +def stream_response_to_part( + *, + manager: Any, + opener: Any, + task: Any, + cancel_event: Any, + part: Path, + checkpoint: Path, + digest: Any, + resume_from: int, + checkpoint_data: Dict[str, Any], +) -> tuple[int, int, str, str, str]: + req = urllib.request.Request(task.download_url, method="GET") + req.add_header("User-Agent", "ComfyUI-OpenClaw/F65") + if resume_from > 0: + req.add_header("Range", f"bytes={resume_from}-") + + with opener.open(req, timeout=manager.timeout_sec) as resp: + code = int(resp.getcode() or 0) + if code in (301, 302, 303, 307, 308): + raise manager._error( + "download_redirect_blocked", + "redirect blocked for managed downloads", + ) + if code >= 400: + raise manager._error("download_http_error", f"HTTP {code}") + + etag = str(resp.headers.get("ETag") or "").strip() + last_modified = str(resp.headers.get("Last-Modified") or "").strip() + content_length = 0 + try: + content_length = max(0, int(str(resp.headers.get("Content-Length") or "0"))) + except Exception: + content_length = 0 + + if resume_from > 0: + if code != 206: + return ( + resume_from, + 0, + etag, + last_modified, + "resume_fallback_range_not_supported", + ) + if not manager._validators_match(checkpoint_data, etag, last_modified): + return ( + resume_from, + 0, + etag, + last_modified, + "resume_fallback_validator_mismatch", + ) + range_start, _range_end, range_total = manager._parse_content_range( + str(resp.headers.get("Content-Range") or "") + ) + if range_start != resume_from: + return ( + resume_from, + 0, + etag, + last_modified, + "resume_fallback_content_range_mismatch", + ) + total = range_total if range_total > 0 else (resume_from + content_length) + mode = "ab" + downloaded = resume_from + else: + total = content_length + mode = "wb" + downloaded = 0 + + last_emit = 0.0 + with open(part, mode) as fh: + while True: + if cancel_event.is_set(): + manager._save_checkpoint( + checkpoint, + task, + bytes_downloaded=downloaded, + total_bytes=total, + etag=etag, + last_modified=last_modified, + ) + raise manager._download_cancelled_cls() + chunk = resp.read(64 * 1024) + if not chunk: + break + fh.write(chunk) + digest.update(chunk) + downloaded += len(chunk) + now = time.time() + if now - last_emit >= 0.35: + manager._progress(task.task_id, downloaded, total) + manager._save_checkpoint( + checkpoint, + task, + bytes_downloaded=downloaded, + total_bytes=total, + etag=etag, + last_modified=last_modified, + ) + last_emit = now + manager._progress(task.task_id, downloaded, total or downloaded) + manager._save_checkpoint( + checkpoint, + task, + bytes_downloaded=downloaded, + total_bytes=total, + etag=etag, + last_modified=last_modified, + ) + return downloaded, total, etag, last_modified, "" + + +def import_downloaded_model( + *, + manager: Any, + task_id: str, + tenant_id: Optional[str] = None, + destination_subdir: Optional[str] = None, + filename: Optional[str] = None, + tags: Optional[List[str]] = None, +) -> Dict[str, Any]: + with manager._lock: + task = manager._tasks.get(task_id) + if task is None or not manager._tenant_ok(task.tenant_id, tenant_id): + raise manager._error("not_found", "download task not found", 404) + if task.state != "completed": + raise manager._error( + "task_not_ready", "task must be completed before import" + ) + if task.imported: + raise manager._error("already_imported", "task already imported") + staged_path = Path(task.staged_path) + expected = task.expected_sha256 + computed = task.computed_sha256 + if not staged_path.exists(): + raise manager._error("staging_missing", "staged file missing") + # CRITICAL: keep import-time hash verification. Removing this reopens + # tamper window between download completion and activation/import. + actual = manager._file_sha256(staged_path) + if actual != expected or computed != expected: + raise manager._error("sha256_mismatch", f"expected {expected}, got {actual}") + manager._validate_provenance(task.provenance) + subdir = manager._sanitize_subdir(destination_subdir or task.destination_subdir) + fname = manager._sanitize_filename(filename or task.filename) + rel_target = f"{subdir}/{fname}" + # IMPORTANT: keep root-bounded resolution; plain joins re-enable traversal risks. + abs_target = Path( + manager._resolve_install_target(str(manager.install_root), rel_target) + ) + abs_target.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp( + prefix=f".{abs_target.name}.tmp.", dir=str(abs_target.parent), text=False + ) + os.close(fd) + try: + shutil.copy2(staged_path, tmp) + os.replace(tmp, abs_target) + except Exception: + try: + os.remove(tmp) + except OSError: + pass + raise + safe_tags: List[str] = [] + for item in tags or []: + if not isinstance(item, str): + continue + clean = item.strip().lower() + if not clean or clean in safe_tags: + continue + safe_tags.append(clean) + if len(safe_tags) >= 24: + break + rec = { + "id": str(uuid.uuid4()), + "task_id": task.task_id, + "model_id": task.model_id, + "name": task.name, + "model_type": task.model_type, + "source": task.source, + "source_label": task.source_label, + "download_url": task.download_url, + "sha256": expected, + "size_bytes": abs_target.stat().st_size if abs_target.exists() else None, + "provenance": dict(task.provenance), + "installation_path": rel_target.replace("\\", "/"), + "tenant_id": task.tenant_id, + "installed_at": time.time(), + "tags": safe_tags, + } + rows = manager._load_installations() + rows.append(rec) + rows.sort(key=lambda x: float(x.get("installed_at") or 0.0), reverse=True) + manager._save_installations(rows) + with manager._lock: + current = manager._tasks.get(task.task_id) + if current is not None: + current.imported = True + current.installation_path = rec["installation_path"] + current.installation_record_id = rec["id"] + current.updated_at = time.time() + manager._emit(current) + manager._persist_tasks_locked(force=True) + return rec diff --git a/tests/e2e/specs/notifications.spec.js b/tests/e2e/specs/notifications.spec.js index 5cf25e7..ec27bb5 100644 --- a/tests/e2e/specs/notifications.spec.js +++ b/tests/e2e/specs/notifications.spec.js @@ -14,6 +14,18 @@ function isPath(pathname, suffix) { test.describe("Notification Center", () => { test("persists model-manager failures across reload until dismissed", async ({ page }) => { + await page.addInitScript(() => { + try { + if (!window.name.includes("__openclaw_notifications_storage_reset__")) { + window.localStorage.clear(); + window.sessionStorage.clear(); + window.name = `${window.name}__openclaw_notifications_storage_reset__`; + } + } catch { + // ignore storage reset failures in restrictive browser contexts + } + }); + await mockComfyUiCore(page); await page.route("**/models/search**", async (route) => { @@ -63,23 +75,81 @@ test.describe("Notification Center", () => { }); }); + const okJson = JSON.stringify({ ok: true, entries: [], config: {}, stats: {} }); + await page.route("**/events/stream**", async (route) => { + await route.fulfill({ + status: 200, + contentType: "text/event-stream", + body: "", + }); + }); + await page.route("**/logs/tail**", async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: okJson, + }); + }); + await page.route("**/config**", async (route) => { + const url = new URL(route.request().url()); + if (isPath(url.pathname, "/config")) { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: okJson, + }); + return; + } + await route.fallback(); + }); + await page.route("**/system_stats**", async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: okJson, + }); + }); + await page.route("**/system_info**", async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: okJson, + }); + }); + await page.route("**/version**", async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: okJson, + }); + }); + await page.goto("test-harness.html"); await waitForOpenClawReady(page); await clickTab(page, "Model Manager"); + await page.locator("#mm-refresh-btn").click(); const toggle = page.locator("#openclaw-notification-toggle"); - await expect(toggle.locator(".openclaw-notification-badge")).toHaveText("1"); - - await toggle.click(); - await expect(page.locator("#openclaw-notification-panel")).toContainText("search: search_failed"); - await expect(page.locator("#openclaw-notification-panel")).toContainText("Open Model Manager"); + await toggle.dispatchEvent("click"); + const targetNotification = page + .locator("#openclaw-notification-panel .openclaw-notification-item") + .filter({ hasText: "search: search_failed" }) + .first(); + await expect(targetNotification).toContainText("search: search_failed", { timeout: 15000 }); + await expect(targetNotification).toContainText("Open Model Manager"); await page.reload(); await waitForOpenClawReady(page); - await page.locator("#openclaw-notification-toggle").click(); - await expect(page.locator("#openclaw-notification-panel")).toContainText("search: search_failed"); + await page.locator("#openclaw-notification-toggle").dispatchEvent("click"); + const reloadedNotification = page + .locator("#openclaw-notification-panel .openclaw-notification-item") + .filter({ hasText: "search: search_failed" }) + .first(); + await expect(reloadedNotification).toContainText("search: search_failed", { timeout: 15000 }); - await page.getByRole("button", { name: "Dismiss" }).first().click(); + await reloadedNotification + .getByRole("button", { name: "Dismiss notification: search: search_failed" }) + .click(); await expect(page.locator("#openclaw-notification-panel")).not.toContainText("search: search_failed"); }); }); diff --git a/web/openclaw_ui.js b/web/openclaw_ui.js index e1db958..a6bc7de 100644 --- a/web/openclaw_ui.js +++ b/web/openclaw_ui.js @@ -314,8 +314,9 @@ export class OpenClawUI { } list.innerHTML = activeEntries.map((entry) => { + const escapedMessage = String(entry.message || "").replace(/"/g, """); const actionHtml = entry.action?.type && entry.action?.payload - ? `` + ? `` : ""; const countHtml = entry.count > 1 ? `x${entry.count}` @@ -336,8 +337,8 @@ export class OpenClawUI {
${actionHtml} - - + +