From e5ddbcfb8aa499f1d2a2750925c34d8fd8b73dc4 Mon Sep 17 00:00:00 2001 From: Berry Wahlberg <40695099+BerryUIKI@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:50:57 +0800 Subject: [PATCH] feat(backend): implement M1-M2 core runtime, isolated supervisors, and SQLite persistence --- backend/app/__init__.py | 2 +- backend/app/core/cache.py | 122 +++- backend/app/core/dag.py | 11 + backend/app/main.py | 729 ++++++++++++++++++++-- backend/app/runners/api_runner.py | 52 +- backend/app/runners/comfy_runner.py | 80 ++- backend/app/runners/creative_runner.py | 315 ++++++++++ backend/app/runners/macro_compiler.py | 218 ++++++- backend/app/runners/webui_runner.py | 174 ++++++ backend/app/runtime/credentials.py | 184 ++++++ backend/app/runtime/engine_manager.py | 168 +++++ backend/app/runtime/hardware.py | 143 +++++ backend/app/runtime/installer.py | 358 +++++++++++ backend/app/runtime/supervisor.py | 94 ++- backend/app/runtime/webui_supervisor.py | 241 +++++++ backend/app/schemas/cloud.py | 40 ++ backend/app/schemas/creative.py | 59 ++ backend/app/schemas/engine.py | 126 ++++ backend/app/schemas/events.py | 16 +- backend/app/schemas/hardware.py | 33 + backend/app/schemas/model.py | 59 ++ backend/app/schemas/project.py | 38 ++ backend/app/schemas/task.py | 47 ++ backend/app/storage/__init__.py | 1 + backend/app/storage/asset_store.py | 172 +++++ backend/app/storage/db.py | 125 ++++ backend/app/storage/model_store.py | 279 +++++++++ backend/app/storage/project_store.py | 112 ++++ backend/requirements.txt | 2 + backend/tests/__init__.py | 2 +- backend/tests/test_health.py | 2 +- backend/tests/test_m1_core.py | 271 ++++++++ backend/tests/test_m2_engines.py | 246 ++++++++ backend/tests/test_m3_creative.py | 206 ++++++ backend/tests/test_m4_cloud.py | 174 ++++++ backend/tests/test_m5_release.py | 206 ++++++ backend/tests/test_m6_launcher_manager.py | 204 ++++++ 37 files changed, 5177 insertions(+), 134 deletions(-) create mode 100644 backend/app/runners/creative_runner.py create mode 100644 backend/app/runners/webui_runner.py create mode 100644 backend/app/runtime/credentials.py create mode 100644 backend/app/runtime/engine_manager.py create mode 100644 backend/app/runtime/hardware.py create mode 100644 backend/app/runtime/installer.py create mode 100644 backend/app/runtime/webui_supervisor.py create mode 100644 backend/app/schemas/cloud.py create mode 100644 backend/app/schemas/creative.py create mode 100644 backend/app/schemas/engine.py create mode 100644 backend/app/schemas/hardware.py create mode 100644 backend/app/schemas/model.py create mode 100644 backend/app/schemas/project.py create mode 100644 backend/app/schemas/task.py create mode 100644 backend/app/storage/__init__.py create mode 100644 backend/app/storage/asset_store.py create mode 100644 backend/app/storage/db.py create mode 100644 backend/app/storage/model_store.py create mode 100644 backend/app/storage/project_store.py create mode 100644 backend/tests/test_m1_core.py create mode 100644 backend/tests/test_m2_engines.py create mode 100644 backend/tests/test_m3_creative.py create mode 100644 backend/tests/test_m4_cloud.py create mode 100644 backend/tests/test_m5_release.py create mode 100644 backend/tests/test_m6_launcher_manager.py diff --git a/backend/app/__init__.py b/backend/app/__init__.py index fae8917..ec2fbae 100644 --- a/backend/app/__init__.py +++ b/backend/app/__init__.py @@ -1 +1 @@ -"""AI-Workflow backend application package.""" +"""Berry AI Studio backend application package.""" diff --git a/backend/app/core/cache.py b/backend/app/core/cache.py index 82f376e..afdf626 100644 --- a/backend/app/core/cache.py +++ b/backend/app/core/cache.py @@ -1,23 +1,67 @@ -"""Deterministic hashing and output result caching engine.""" +"""Deterministic hashing and output result caching engine with SQLite persistence.""" import hashlib import json -from typing import Any, Dict, Optional +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional, Tuple + +from app.storage.db import DatabaseManager, db_manager + + +def compute_content_hash(value: Any) -> str: + """Compute a deterministic SHA-256 hash of output content.""" + if value is None: + return hashlib.sha256(b"null").hexdigest() + if isinstance(value, (int, float, bool)): + return hashlib.sha256(str(value).encode("utf-8")).hexdigest() + if isinstance(value, str): + return hashlib.sha256(value.encode("utf-8")).hexdigest() + if isinstance(value, (dict, list)): + serialized = json.dumps(value, sort_keys=True, default=str) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() + return hashlib.sha256(str(value).encode("utf-8")).hexdigest() + + +def compute_semantic_node_hash( + node_type: str, + params: Dict[str, Any], + input_bindings: List[Tuple[str, str, str]], # (target_handle, upstream_content_hash, source_handle) +) -> str: + """ + Compute a deterministic port-aware SHA-256 hash representing a node's exact semantic state. + NodeHash = SHA256(NodeType + SerializedCanonicalParams + OrderedPortBindings) + + OrderedPortBindings strictly associates input ports to upstream outputs: + target_handle:source_handle:content_hash + """ + hasher = hashlib.sha256() + hasher.update(node_type.encode("utf-8")) + + # Canonicalize params: exclude transient secrets from hash + clean_params = {k: v for k, v in params.items() if not k.lower().endswith("key")} + params_str = json.dumps(clean_params, sort_keys=True, default=str) + hasher.update(params_str.encode("utf-8")) + + # Deterministic binding sort by target_handle + sorted_bindings = sorted(input_bindings, key=lambda b: (b[0], b[2])) + for target_handle, content_hash, source_handle in sorted_bindings: + binding_repr = f"|in:{target_handle}->out:{source_handle}#{content_hash}" + hasher.update(binding_repr.encode("utf-8")) + + return hasher.hexdigest() def compute_node_hash(node_type: str, params: Dict[str, Any], parent_hashes: list[str]) -> str: """ - Compute a deterministic SHA-256 hash representing a node's exact state. - NodeHash = SHA256(NodeType + SerializedParams + SortedParentHashes) + Legacy helper: Compute SHA-256 representing a node's state with parent hashes. + Kept for backwards compatibility with baseline tests. """ hasher = hashlib.sha256() hasher.update(node_type.encode("utf-8")) - # Serialize parameters with sorted keys for deterministic encoding params_str = json.dumps(params, sort_keys=True, default=str) hasher.update(params_str.encode("utf-8")) - # Concatenate upstream parent hashes in deterministic sorted order for parent_hash in sorted(parent_hashes): hasher.update(parent_hash.encode("utf-8")) @@ -25,29 +69,81 @@ def compute_node_hash(node_type: str, params: Dict[str, Any], parent_hashes: lis class CacheStore: - """Thread-safe in-memory cache store for node execution outputs.""" + """Thread-safe cache store with in-memory cache and SQLite persistence.""" - def __init__(self) -> None: + def __init__(self, manager: Optional[DatabaseManager] = None) -> None: self._store: Dict[str, Dict[str, Any]] = {} + self.manager = manager or db_manager def get(self, node_hash: str) -> Optional[Dict[str, Any]]: - """Retrieve cached output data by node hash.""" + """Retrieve cached output data by node hash synchronously from memory.""" return self._store.get(node_hash) def set(self, node_hash: str, output: Dict[str, Any]) -> None: - """Store output data for a node hash.""" + """Store output data for a node hash in memory.""" self._store[node_hash] = output def has(self, node_hash: str) -> bool: - """Check if output for node hash is cached.""" + """Check if output for node hash is cached in memory.""" return node_hash in self._store + async def get_async(self, node_hash: str) -> Optional[Dict[str, Any]]: + """Retrieve cached output, checking memory first then SQLite.""" + if node_hash in self._store: + return self._store[node_hash] + + try: + conn = await self.manager.get_connection() + async with conn.execute( + "SELECT output_json FROM cache_entries WHERE node_hash = ?", + (node_hash,), + ) as cursor: + row = await cursor.fetchone() + if row: + output = json.loads(row["output_json"]) + self._store[node_hash] = output + return output + except Exception: + pass + return None + + async def set_async(self, node_hash: str, output: Dict[str, Any]) -> None: + """Store output data in both memory and SQLite.""" + self._store[node_hash] = output + try: + conn = await self.manager.get_connection() + now = datetime.now(timezone.utc).isoformat() + await conn.execute( + """ + INSERT OR REPLACE INTO cache_entries (node_hash, output_json, asset_ids_json, created_at) + VALUES (?, ?, ?, ?) + """, + (node_hash, json.dumps(output), json.dumps([]), now), + ) + await conn.commit() + except Exception: + pass + + async def has_async(self, node_hash: str) -> bool: + """Check if output exists in memory or SQLite.""" + return (await self.get_async(node_hash)) is not None + def clear(self) -> None: - """Purge all cached results.""" + """Purge in-memory cached results.""" + self._store.clear() + + async def clear_all_async(self) -> None: + """Purge both in-memory and persistent SQLite cache.""" self._store.clear() + try: + conn = await self.manager.get_connection() + await conn.execute("DELETE FROM cache_entries") + await conn.commit() + except Exception: + pass def size(self) -> int: - """Return number of cached node states.""" + """Return number of in-memory cached node states.""" return len(self._store) diff --git a/backend/app/core/dag.py b/backend/app/core/dag.py index d6af2c1..869885d 100644 --- a/backend/app/core/dag.py +++ b/backend/app/core/dag.py @@ -77,6 +77,17 @@ def _get_ancestors(self, node_id: str) -> Set[str]: queue.extend(self.parents[curr]) return ancestors + def get_descendants(self, node_id: str) -> Set[str]: + """Recursively gather all downstream descendant nodes.""" + descendants: Set[str] = set() + queue = deque(self.adjacency[node_id]) + while queue: + curr = queue.popleft() + if curr not in descendants: + descendants.add(curr) + queue.extend(self.adjacency[curr]) + return descendants + def get_parent_ids(self, node_id: str) -> List[str]: """Return direct upstream parent node IDs.""" return self.parents.get(node_id, []) diff --git a/backend/app/main.py b/backend/app/main.py index 42658d7..459ffab 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -1,36 +1,86 @@ -"""FastAPI application entrypoint for AI-Workflow.""" +"""FastAPI application entrypoint for Berry AI Studio.""" +import asyncio import json +import logging import time -from typing import Any, Dict, List +import uuid +import os +from pathlib import Path +from typing import Any, Dict, List, Optional -from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect +from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, UploadFile, File from fastapi.middleware.cors import CORSMiddleware - -from app.core.cache import cache_store, compute_node_hash +from fastapi.responses import FileResponse, HTMLResponse +from pydantic import BaseModel + +from app.core.cache import ( + cache_store, + compute_content_hash, + compute_node_hash, + compute_semantic_node_hash, +) from app.core.dag import CyclicDependencyError, DAGResolver from app.nodes.registry import registry from app.runners.api_runner import NODE_RUNNERS, run_input_text_node from app.runners.comfy_runner import comfy_client, run_comfy_txt2img_node +from app.runners.creative_runner import creative_runner from app.runtime.supervisor import supervisor - -# Register ComfyUI node runners -NODE_RUNNERS["image.comfy.txt2img"] = run_comfy_txt2img_node +from app.runtime.webui_supervisor import webui_supervisor +from app.runtime.engine_manager import engine_manager +from app.runtime.hardware import check_hardware_readiness +from app.runtime.installer import installer +from app.runtime.credentials import credentials_manager +from app.storage.model_store import model_store +from app.schemas.creative import CreativeActionRequest, CreativeActionResult +from app.schemas.cloud import ( + CloudProviderId, + CloudProviderInfo, + SetCredentialRequest, + TestKeyRequest, + TestKeyResult, +) +from app.schemas.engine import ( + EngineConnection, + EngineConnectRequest, + EngineInstallManifest, + EngineOwnership, + EngineType, + EngineUpdateManifest, + LauncherConfig, + ManagerStatusResponse, + ShutdownRequest, + ShutdownResponse, +) +from app.schemas.hardware import HardwareReadiness +from app.schemas.model import ModelRecord, ModelRoot, ModelRootCreate from app.schemas.events import ( GraphFinishedEvent, GraphStartedEvent, + NodeErrorEvent, NodeOutputEvent, NodeStatusEvent, + RunCancelledEvent, ) from app.schemas.node import NodeDefinition +from app.schemas.project import Project, ProjectCreate, ProjectUpdate, AssetRecord +from app.schemas.task import WorkflowRunRequest from app.schemas.workflow import ExecutionPlan, PlannedNodeStep, WorkflowGraph +from app.storage.asset_store import asset_store +from app.storage.db import db_manager +from app.storage.project_store import project_store + +# Register ComfyUI node runners +NODE_RUNNERS["image.comfy.txt2img"] = run_comfy_txt2img_node # Import builtin nodes to trigger auto-registration import app.nodes.builtin # noqa: F401 +logger = logging.getLogger(__name__) + app = FastAPI( - title="AI-Workflow Engine API", - description="Lightweight, API-first execution engine for multimodal workflows", + title="Berry AI Studio Engine API", + description="Lightweight, API-first creative engine with persistent projects and deterministic caching", version="0.1.0", ) @@ -43,6 +93,9 @@ allow_headers=["*"], ) +# Active run cancellation tracker: run_id -> asyncio.Event +active_cancellations: Dict[str, asyncio.Event] = {} + @app.get("/health") async def health_check() -> dict[str, str]: @@ -54,7 +107,7 @@ async def health_check() -> dict[str, str]: async def system_info() -> dict[str, object]: installed = supervisor.is_installed() return { - "name": "AI-Workflow", + "name": "Berry AI Studio", "version": "0.1.0", "runners": { "api": {"status": "ready", "type": "cloud"}, @@ -64,6 +117,115 @@ async def system_info() -> dict[str, object]: } +app_start_time = time.time() +launcher_config = LauncherConfig() + + +# --------------------------------------------------------------------------- +# Launcher & Environment Manager Endpoints (L01-L12) +# --------------------------------------------------------------------------- + +@app.get("/api/v1/manager/status", response_model=ManagerStatusResponse) +async def manager_status() -> ManagerStatusResponse: + """Return unified environment and process status for Berry, engines, and models.""" + frontend_dir = Path(__file__).resolve().parent.parent.parent / "frontend" / "dist" + frontend_packaged = (frontend_dir / "index.html").is_file() + + managed_comfy = engine_manager.get_engine("managed_comfyui") + managed_web = engine_manager.get_engine("managed_webui") + external_list = [e for e in engine_manager.list_engines() if e.ownership == EngineOwnership.EXTERNAL] + + configured_clouds = sum(1 for p in credentials_manager.list_providers() if p.is_configured) + models_count = len(model_store.list_models()) + + return ManagerStatusResponse( + app_name="Berry AI Studio", + version="0.1.0", + pid=os.getpid(), + uptime_seconds=round(time.time() - app_start_time, 2), + port=launcher_config.port, + frontend_packaged=frontend_packaged, + managed_comfyui=managed_comfy, + managed_webui=managed_web, + external_engines=external_list, + cloud_providers_configured=configured_clouds, + models_indexed=models_count, + active_tasks=len(active_cancellations), + launcher_config=launcher_config, + ) + + +@app.get("/api/v1/manager/config", response_model=LauncherConfig) +async def get_launcher_config() -> LauncherConfig: + """Get current launcher configuration.""" + return launcher_config + + +@app.post("/api/v1/manager/config", response_model=LauncherConfig) +async def set_launcher_config(config: LauncherConfig) -> LauncherConfig: + """Update launcher configuration.""" + global launcher_config + launcher_config = config + return launcher_config + + +@app.post("/api/v1/manager/shutdown", response_model=ShutdownResponse) +async def manager_shutdown(request: ShutdownRequest) -> ShutdownResponse: + """ + Explicit controlled shutdown of Berry AI Studio (L07). + Guards active generation tasks and predictably terminates managed engines if configured. + Never terminates external user engines. + """ + active_count = len(active_cancellations) + if active_count > 0 and not request.force: + raise HTTPException( + status_code=409, + detail=f"Cannot shutdown: {active_count} active generation task(s) running. Provide force=true to abort tasks.", + ) + + # Abort active tasks if forced + if active_count > 0: + for run_id, cancel_evt in list(active_cancellations.items()): + cancel_evt.set() + + # Determine whether to stop managed engines + stop_managed = request.stop_managed_engines + if stop_managed is None: + stop_managed = launcher_config.stop_managed_engines_on_exit + + stopped_engines = [] + if stop_managed: + if supervisor.is_running(): + supervisor.stop() + stopped_engines.append("managed_comfyui") + if webui_supervisor.is_running(): + webui_supervisor.stop() + stopped_engines.append("managed_webui") + + # Schedule self-termination + def _delayed_exit(): + logger.info("Berry AI Studio server exiting upon manager request.") + os._exit(0) + + try: + loop = asyncio.get_running_loop() + loop.call_later(0.5, _delayed_exit) + except Exception: + pass + + return ShutdownResponse( + status="shutting_down", + message="Berry AI Studio shutdown initiated.", + active_tasks_cancelled=active_count, + managed_engines_stopped=stopped_engines, + ) + + + +# --------------------------------------------------------------------------- +# Engine & Supervisor Endpoints +# --------------------------------------------------------------------------- + @app.get("/api/v1/comfy/status") async def comfy_status() -> dict[str, Any]: """Check connectivity and hardware stats of the local ComfyUI instance.""" @@ -94,12 +256,344 @@ async def runtime_stop() -> dict[str, Any]: return supervisor.stop() +@app.get("/api/v1/runtime/webui/status") +async def runtime_webui_status() -> dict[str, Any]: + """Check the status of the sandboxed WebUI runtime and supervisor.""" + return webui_supervisor.get_status() + + +@app.post("/api/v1/runtime/webui/start") +async def runtime_webui_start() -> dict[str, Any]: + """Launch the isolated WebUI subprocess via supervisor.""" + return webui_supervisor.start() + + +@app.post("/api/v1/runtime/webui/stop") +async def runtime_webui_stop() -> dict[str, Any]: + """Stop the isolated WebUI subprocess via supervisor.""" + return webui_supervisor.stop() + + +@app.get("/api/v1/runtime/{engine_type}/manifest", response_model=EngineInstallManifest) +async def get_engine_manifest(engine_type: EngineType) -> EngineInstallManifest: + """Retrieve the installation manifest for an isolated engine.""" + return installer.read_manifest(engine_type) + + +@app.post("/api/v1/runtime/{engine_type}/install", response_model=EngineInstallManifest) +async def trigger_engine_install(engine_type: EngineType) -> EngineInstallManifest: + """Start or check isolated installation of an engine without host pollution.""" + # Spawn in background task to avoid blocking HTTP call + asyncio.create_task(installer.install_engine(engine_type)) + return installer.read_manifest(engine_type) + + +@app.get("/api/v1/runtime/{engine_type}/update/manifest", response_model=EngineUpdateManifest) +async def get_engine_update_manifest(engine_type: EngineType) -> EngineUpdateManifest: + """Retrieve the update manifest and rollback state for a managed engine.""" + return installer.read_update_manifest(engine_type) + + +@app.post("/api/v1/runtime/{engine_type}/update", response_model=EngineUpdateManifest) +async def trigger_engine_update(engine_type: EngineType) -> EngineUpdateManifest: + """ + Safely update a managed engine with active job checking and rollback protection (L08, L12). + """ + has_active = len(active_cancellations) > 0 + return await installer.update_engine(engine_type, has_active_tasks_fn=lambda: has_active) + + +@app.get("/api/v1/updates/check") +async def check_all_updates() -> dict[str, Any]: + """Check update availability for Berry AI Studio and managed engines separately (L08).""" + comfy_manifest = installer.read_update_manifest(EngineType.COMFYUI) + webui_manifest = installer.read_update_manifest(EngineType.WEBUI) + + return { + "app": { + "name": "Berry AI Studio", + "current_version": "0.1.0", + "latest_version": "0.1.0", + "update_available": False, + "release_notes_url": "https://github.com/BerryUIKI/AI-Studio/releases", + }, + "engines": { + "comfyui": { + "installed": supervisor.is_installed(), + "update_manifest": comfy_manifest, + }, + "webui": { + "installed": webui_supervisor.is_installed(), + "update_manifest": webui_manifest, + }, + }, + } + + +@app.post("/api/v1/updates/app") +async def trigger_app_update() -> dict[str, Any]: + """ + Check and report or trigger Berry application updates (L08). + Decoupled from inference engine updates. + """ + repo_root = Path(__file__).resolve().parent.parent.parent + is_git_repo = (repo_root / ".git").is_dir() + + if is_git_repo: + try: + proc = await asyncio.create_subprocess_exec( + "git", "pull", "--ff-only", + cwd=str(repo_root), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate() + if proc.returncode == 0: + return { + "mode": "git", + "status": "updated", + "message": "Application code updated successfully. Please restart Berry AI Studio.", + "details": stdout.decode().strip(), + } + else: + return { + "mode": "git", + "status": "failed", + "message": f"Git update failed: {stderr.decode().strip()}", + } + except Exception as e: + return {"mode": "git", "status": "error", "message": str(e)} + else: + return { + "mode": "packaged", + "status": "guidance", + "message": "In packaged Windows release mode, download the latest package to update.", + "download_url": "https://github.com/BerryUIKI/AI-Studio/releases", + "notes": "User projects, credentials, and models in %LOCALAPPDATA% are strictly preserved during updates.", + } + + + + +# --------------------------------------------------------------------------- +# Hardware & Engine Management Endpoints (M2) +# --------------------------------------------------------------------------- + +@app.get("/api/v1/hardware/readiness", response_model=HardwareReadiness) +async def get_hardware_readiness() -> HardwareReadiness: + """Diagnose GPU VRAM, drivers, and filesystem storage readiness for local inference.""" + return await check_hardware_readiness() + + +@app.get("/api/v1/engines", response_model=List[EngineConnection]) +async def list_engines() -> List[EngineConnection]: + """List all available managed and external engine connections.""" + return engine_manager.list_engines() + + +@app.post("/api/v1/engines/connect", response_model=EngineConnection) +async def connect_external_engine(request: EngineConnectRequest) -> EngineConnection: + """Connect to a running user engine without process ownership.""" + return await engine_manager.connect_external_engine( + engine_type=request.engine_type, + endpoint_url=request.endpoint_url, + name=request.name, + ) + + +@app.post("/api/v1/engines/{engine_id}/test", response_model=EngineConnection) +async def test_engine(engine_id: str) -> EngineConnection: + """Test connectivity and readiness of a registered engine.""" + conn = await engine_manager.test_connection_by_id(engine_id) + if not conn: + raise HTTPException(status_code=404, detail=f"Engine '{engine_id}' not found") + return conn + + +# --------------------------------------------------------------------------- +# Model Catalog & Roots Endpoints (M2) +# --------------------------------------------------------------------------- + +@app.get("/api/v1/models", response_model=List[ModelRecord]) +async def list_models() -> List[ModelRecord]: + """List all models currently discovered across registered roots.""" + return await model_store.scan_all_roots_async() + + +@app.post("/api/v1/models/scan", response_model=List[ModelRecord]) +async def trigger_model_scan() -> List[ModelRecord]: + """Force re-scan of all model root directories.""" + return await model_store.scan_all_roots_async() + + +@app.get("/api/v1/models/roots", response_model=List[ModelRoot]) +async def list_model_roots() -> List[ModelRoot]: + """List all configured model root directories.""" + return model_store.list_roots() + + +class AddModelRootRequest(BaseModel): + path: str + label: str + engine_type: Optional[str] = None + + +@app.post("/api/v1/models/roots", response_model=ModelRoot) +async def add_model_root(req: AddModelRootRequest) -> ModelRoot: + """Register a new user-specified directory for model discovery.""" + root_id = f"root_{abs(hash(req.path)) % 10000}" + return model_store.add_root(root_id, req.path, req.label, req.engine_type) + + +@app.post("/api/v1/models/rescan", response_model=List[ModelRecord]) +async def trigger_model_rescan() -> List[ModelRecord]: + """Force re-scan of all model root directories.""" + return await model_store.scan_all_roots_async() + + +@app.delete("/api/v1/models/roots/{root_id}") +async def remove_model_root(root_id: str) -> dict[str, Any]: + """Remove a configured model root directory without deleting files (L10).""" + removed = model_store.remove_root(root_id) + if not removed: + raise HTTPException(status_code=404, detail=f"Model root '{root_id}' not found") + return {"status": "removed", "root_id": root_id} + + + +# --------------------------------------------------------------------------- +# Creative Primary Canvas Endpoints (M3) +# --------------------------------------------------------------------------- + +@app.post("/api/v1/creative/execute", response_model=CreativeActionResult) +async def execute_creative_action(req: CreativeActionRequest) -> CreativeActionResult: + """Execute high-level image creation action (txt2img, img2img, inpaint, upscale).""" + return await creative_runner.execute(req) + + +@app.post("/api/v1/creative/upload", response_model=AssetRecord) +async def upload_creative_asset(file: UploadFile = File(...)) -> AssetRecord: + """Upload user image asset directly onto the creative canvas.""" + data = await file.read() + filename = file.filename or "uploaded_image.png" + content_type = file.content_type or "image/png" + return await asset_store.save_bytes(data, filename=filename, media_type=content_type) + + +# --------------------------------------------------------------------------- +# Cloud Providers & BYOK Credentials Endpoints (M4) +# --------------------------------------------------------------------------- + +@app.get("/api/v1/cloud/providers", response_model=List[CloudProviderInfo]) +async def list_cloud_providers() -> List[CloudProviderInfo]: + """List available BYOK cloud providers with redacted secrets and capabilities.""" + return credentials_manager.list_providers() + + +@app.post("/api/v1/cloud/credentials", response_model=CloudProviderInfo) +async def set_cloud_credential(req: SetCredentialRequest) -> CloudProviderInfo: + """Store or update a BYOK cloud API key locally without exposure.""" + credentials_manager.set_key(req.provider_id, req.api_key) + providers = credentials_manager.list_providers() + target = next((p for p in providers if p.id == req.provider_id), None) + if not target: + raise HTTPException(status_code=404, detail="Provider not found") + return target + + +@app.post("/api/v1/cloud/credentials/test", response_model=TestKeyResult) +async def test_cloud_credential(req: TestKeyRequest) -> TestKeyResult: + """Validate a BYOK API key against the external provider endpoint.""" + return await credentials_manager.test_key(req.provider_id, req.api_key) + + +@app.delete("/api/v1/cloud/credentials/{provider_id}") +async def delete_cloud_credential(provider_id: CloudProviderId) -> dict[str, bool]: + """Delete a stored BYOK API key.""" + credentials_manager.delete_key(provider_id) + return {"success": True} + + @app.get("/api/v1/nodes", response_model=List[NodeDefinition]) async def list_nodes() -> List[NodeDefinition]: """Retrieve all registered node specifications.""" return registry.list_all() +# --------------------------------------------------------------------------- +# Project & Asset REST Endpoints (M1 Persistence) +# --------------------------------------------------------------------------- + +@app.get("/api/v1/projects", response_model=List[Project]) +async def list_projects() -> List[Project]: + return await project_store.list_projects() + + +@app.post("/api/v1/projects", response_model=Project) +async def create_project(data: ProjectCreate) -> Project: + return await project_store.create_project(data) + + +@app.get("/api/v1/projects/{project_id}", response_model=Project) +async def get_project(project_id: str) -> Project: + proj = await project_store.get_project(project_id) + if not proj: + raise HTTPException(status_code=404, detail="Project not found") + return proj + + +@app.put("/api/v1/projects/{project_id}", response_model=Project) +async def update_project(project_id: str, data: ProjectUpdate) -> Project: + proj = await project_store.update_project(project_id, data) + if not proj: + raise HTTPException(status_code=404, detail="Project not found") + return proj + + +@app.delete("/api/v1/projects/{project_id}") +async def delete_project(project_id: str) -> dict[str, bool]: + success = await project_store.delete_project(project_id) + if not success: + raise HTTPException(status_code=404, detail="Project not found") + return {"success": True} + + +@app.get("/api/v1/assets", response_model=List[AssetRecord]) +async def list_assets(project_id: Optional[str] = None) -> List[AssetRecord]: + return await asset_store.list_assets(project_id=project_id) + + +@app.get("/api/v1/assets/{asset_id}", response_model=AssetRecord) +async def get_asset(asset_id: str) -> AssetRecord: + rec = await asset_store.get_asset(asset_id) + if not rec: + raise HTTPException(status_code=404, detail="Asset not found") + return rec + + +@app.get("/api/v1/assets/{asset_id}/content") +async def get_asset_content(asset_id: str) -> FileResponse: + rec = await asset_store.get_asset(asset_id) + if not rec: + raise HTTPException(status_code=404, detail="Asset not found") + abs_path = asset_store.get_absolute_path(rec) + if not abs_path.is_file(): + raise HTTPException(status_code=404, detail="Asset file missing on disk") + return FileResponse(abs_path, media_type="image/png") + + +@app.post("/api/v1/workflow/cancel/{run_id}") +async def cancel_workflow_run(run_id: str) -> dict[str, Any]: + if run_id in active_cancellations: + active_cancellations[run_id].set() + return {"success": True, "run_id": run_id, "status": "cancel-requested"} + return {"success": False, "message": f"Run '{run_id}' not found or already completed"} + + +# --------------------------------------------------------------------------- +# Workflow Planning & Execution +# --------------------------------------------------------------------------- + @app.post("/api/v1/workflow/plan", response_model=ExecutionPlan) async def generate_plan(graph: WorkflowGraph, target_node: str | None = None) -> ExecutionPlan: """Analyze workflow graph, validate DAG cycles, and compute dirty-check cache hashes.""" @@ -116,10 +610,22 @@ async def generate_plan(graph: WorkflowGraph, target_node: str | None = None) -> cached_count = 0 for node in sorted_nodes: + # Build deterministic port bindings from upstream edges + bindings: List[tuple[str, str, str]] = [] parents = resolver.get_parent_ids(node.id) parent_hashes = [node_hashes[pid] for pid in parents if pid in node_hashes] - h = compute_node_hash(node.type, node.params, parent_hashes) + for edge in graph.edges: + if edge.target == node.id and edge.source in node_hashes: + # Use source computation hash as upstream content proxy for planning + bindings.append((edge.target_handle, node_hashes[edge.source], edge.source_handle)) + + # Support both port-aware semantic hash and legacy hash + if bindings: + h = compute_semantic_node_hash(node.type, node.params, bindings) + else: + h = compute_node_hash(node.type, node.params, parent_hashes) + node_hashes[node.id] = h is_cached = cache_store.has(h) if is_cached: @@ -145,102 +651,217 @@ async def generate_plan(graph: WorkflowGraph, target_node: str | None = None) -> @app.websocket("/ws/workflow/run") async def websocket_run_workflow(websocket: WebSocket) -> None: """ - WebSocket endpoint for real-time workflow execution. - - Accepts a WorkflowGraph JSON payload, resolves DAG order, applies cache, - and streams NodeStatusEvent / NodeOutputEvent / GraphFinishedEvent back - to the canvas for live status badge updates. + WebSocket endpoint for real-time workflow execution with: + - Targeted single-node execution + - Port-aware semantic caching + - Strict failure boundaries (aborts dependent child nodes) + - Run cancellation support """ await websocket.accept() try: raw = await websocket.receive_text() - graph = WorkflowGraph.model_validate_json(raw) + parsed = json.loads(raw) + # Parse either WorkflowRunRequest or raw WorkflowGraph + if "graph" in parsed: + run_request = WorkflowRunRequest.model_validate(parsed) + graph = run_request.graph + target_node = run_request.target_node_id + run_id = run_request.run_id or str(uuid.uuid4()) + else: + graph = WorkflowGraph.model_validate(parsed) + target_node = None + run_id = str(uuid.uuid4()) except Exception as e: await websocket.send_text(json.dumps({"type": "ERROR", "message": f"Invalid graph payload: {e}"})) await websocket.close() return + # Register cancellation token + cancel_event = asyncio.Event() + active_cancellations[run_id] = cancel_event + try: resolver = DAGResolver(graph) - sorted_nodes = resolver.topological_sort() + sorted_nodes = resolver.topological_sort(target_node_id=target_node) except CyclicDependencyError as e: - await websocket.send_text(json.dumps({"type": "ERROR", "message": str(e)})) + await websocket.send_text(json.dumps({"type": "ERROR", "message": str(e), "run_id": run_id})) await websocket.close() + active_cancellations.pop(run_id, None) + return + except ValueError as e: + await websocket.send_text(json.dumps({"type": "ERROR", "message": str(e), "run_id": run_id})) + await websocket.close() + active_cancellations.pop(run_id, None) return - # Compute execution plan with cache - node_hashes: Dict[str, str] = {} - node_outputs: Dict[str, Dict[str, Any]] = {} + # Check cache status for active nodes cached_count = 0 - - for node in sorted_nodes: - parents = resolver.get_parent_ids(node.id) - parent_hashes = [node_hashes[pid] for pid in parents if pid in node_hashes] - h = compute_node_hash(node.type, node.params, parent_hashes) - node_hashes[node.id] = h - if cache_store.has(h): - cached_count += 1 + node_outputs: Dict[str, Dict[str, Any]] = {} start_time = time.monotonic() await websocket.send_text( - GraphStartedEvent(total_nodes=len(sorted_nodes), cached_nodes=cached_count).model_dump_json() + GraphStartedEvent(run_id=run_id, total_nodes=len(sorted_nodes), cached_nodes=cached_count).model_dump_json() ) + failed_node_ids: set[str] = set() + for node in sorted_nodes: - node_hash = node_hashes[node.id] - - # Serve from cache if available - if cache_store.has(node_hash): - cached_output = cache_store.get(node_hash) - await websocket.send_text(NodeStatusEvent(node_id=node.id, status="cached").model_dump_json()) - await websocket.send_text(NodeOutputEvent(node_id=node.id, output=cached_output or {}).model_dump_json()) - node_outputs[node.id] = cached_output or {} + # Check cancellation + if cancel_event.is_set(): + await websocket.send_text(RunCancelledEvent(run_id=run_id).model_dump_json()) + elapsed_ms = (time.monotonic() - start_time) * 1000 + await websocket.send_text( + GraphFinishedEvent(run_id=run_id, execution_time_ms=elapsed_ms, status="cancelled").model_dump_json() + ) + break + + # Check if any upstream ancestor failed + ancestors = resolver._get_ancestors(node.id) + if any(anc in failed_node_ids for anc in ancestors): + # Abort this node due to upstream failure + await websocket.send_text( + NodeStatusEvent(node_id=node.id, status="cancelled", run_id=run_id).model_dump_json() + ) continue - # Resolve inputs from upstream node outputs + # Resolve inputs and compute deterministic port bindings inputs: Dict[str, Any] = {} + bindings: List[tuple[str, str, str]] = [] + for edge in graph.edges: if edge.target == node.id and edge.source in node_outputs: source_output = node_outputs[edge.source] if edge.source_handle in source_output: - inputs[edge.target_handle] = source_output[edge.source_handle] + val = source_output[edge.source_handle] + inputs[edge.target_handle] = val + bindings.append((edge.target_handle, compute_content_hash(val), edge.source_handle)) + + # Compute port-aware semantic hash + if bindings: + node_hash = compute_semantic_node_hash(node.type, node.params, bindings) + else: + node_hash = compute_node_hash(node.type, node.params, []) + + # Check cache (memory or SQLite) + cached_output = await cache_store.get_async(node_hash) + if cached_output is not None: + await websocket.send_text(NodeStatusEvent(node_id=node.id, status="cached", run_id=run_id).model_dump_json()) + await websocket.send_text( + NodeOutputEvent(node_id=node.id, output=cached_output, run_id=run_id).model_dump_json() + ) + node_outputs[node.id] = cached_output + continue - # Dispatch to appropriate runner + # Dispatch to runner runner = NODE_RUNNERS.get(node.type) if runner is None: - await websocket.send_text( - json.dumps({"type": "NODE_ERROR", "node_id": node.id, "message": f"No runner for node type: {node.type}"}) - ) + failed_node_ids.add(node.id) + err_msg = f"No runner for node type: {node.type}" + await websocket.send_text(NodeErrorEvent(node_id=node.id, message=err_msg, run_id=run_id).model_dump_json()) + await websocket.send_text(NodeStatusEvent(node_id=node.id, status="error", run_id=run_id).model_dump_json()) continue output: Dict[str, Any] = {} + node_has_error = False + try: - # input.text has a different signature (no inputs dict) if node.type == "input.text": async for event in run_input_text_node(node.id, node.params): + if cancel_event.is_set(): + break + event.run_id = run_id await websocket.send_text(event.model_dump_json()) if isinstance(event, NodeOutputEvent): output = event.output else: async for event in runner(node.id, inputs, node.params): + if cancel_event.is_set(): + break + event.run_id = run_id await websocket.send_text(event.model_dump_json()) - if isinstance(event, NodeOutputEvent): + if isinstance(event, NodeErrorEvent): + node_has_error = True + elif isinstance(event, NodeOutputEvent): output = event.output except Exception as e: + node_has_error = True + await websocket.send_text(NodeErrorEvent(node_id=node.id, message=str(e), run_id=run_id).model_dump_json()) + await websocket.send_text(NodeStatusEvent(node_id=node.id, status="error", run_id=run_id).model_dump_json()) + + if node_has_error: + failed_node_ids.add(node.id) + # Notify cancellation for all unexecuted descendants + descendants = resolver.get_descendants(node.id) + for desc_id in descendants: + await websocket.send_text( + NodeStatusEvent(node_id=desc_id, status="cancelled", run_id=run_id).model_dump_json() + ) + # Conclude run as failed + elapsed_ms = (time.monotonic() - start_time) * 1000 await websocket.send_text( - json.dumps({"type": "NODE_ERROR", "node_id": node.id, "message": str(e)}) + GraphFinishedEvent(run_id=run_id, execution_time_ms=elapsed_ms, status="failed").model_dump_json() ) - continue + break - # Cache the result for future runs + # Cache successful output if output: - cache_store.set(node_hash, output) + await cache_store.set_async(node_hash, output) node_outputs[node.id] = output - elapsed_ms = (time.monotonic() - start_time) * 1000 - await websocket.send_text(GraphFinishedEvent(execution_time_ms=elapsed_ms).model_dump_json()) + else: + # Loop completed without break + elapsed_ms = (time.monotonic() - start_time) * 1000 + await websocket.send_text( + GraphFinishedEvent(run_id=run_id, execution_time_ms=elapsed_ms, status="completed").model_dump_json() + ) + + active_cancellations.pop(run_id, None) try: await websocket.close() except Exception: pass + + +# --------------------------------------------------------------------------- +# Static Frontend Single-App Serving (Production/Release Mode) +# --------------------------------------------------------------------------- + +frontend_dist = Path(__file__).resolve().parent.parent.parent / "frontend" / "dist" +if frontend_dist.is_dir() and (frontend_dist / "index.html").is_file(): + from fastapi.staticfiles import StaticFiles + + app.mount("/", StaticFiles(directory=str(frontend_dist), html=True), name="frontend") +else: + @app.get("/", response_class=HTMLResponse) + async def missing_frontend_page(): + """Visible packaging error when frontend dist is missing (L02).""" + return HTMLResponse( + content=""" + + + + Berry AI Studio - Packaging Error + + + +
+ L02 Diagnostic +

⚠️ Packaging Error: Frontend Build Missing

+

Berry AI Studio backend is running and healthy, but the frontend distribution assets (frontend/dist/index.html) are missing.

+

For Developers: Run the frontend build command before launching:

+

cd frontend && pnpm build

+

Status: Backend core, REST APIs, and background engine supervisors remain accessible at /health and /api/v1/manager/status.

+
+ +""", + status_code=200, + ) + diff --git a/backend/app/runners/api_runner.py b/backend/app/runners/api_runner.py index 7d143c5..e0cfcda 100644 --- a/backend/app/runners/api_runner.py +++ b/backend/app/runners/api_runner.py @@ -1,8 +1,9 @@ """ Cloud API execution driver. -Supports OpenAI-compatible LLM endpoints (DeepSeek, GPT-4o, etc.) and -image generation APIs (Fal.ai / SiliconFlow / OpenAI Images). +Supports OpenAI-compatible LLM endpoints (DeepSeek, GPT-4o, etc.), +image generation APIs (Fal.ai / SiliconFlow / OpenAI Images), +and the output.preview presentation card runner. Reads credentials from environment variables or from node-level params. Never hardcodes secrets. @@ -21,6 +22,7 @@ NodeStatusEvent, WorkflowEvent, ) +from app.storage.asset_store import asset_store class APIRunnerError(Exception): @@ -160,8 +162,22 @@ async def run_image_gen_node( yield NodeStatusEvent(node_id=node_id, status="error") return - yield NodeProgressEvent(node_id=node_id, progress=0.95, message="Image generation completed.") - yield NodeOutputEvent(node_id=node_id, output={"image": image_url}) + yield NodeProgressEvent(node_id=node_id, progress=0.8, message="Persisting cloud image into managed storage...") + output_dict: Dict[str, Any] = {"image": image_url} + try: + asset = await asset_store.save_image_from_url(image_url) + output_dict = { + "image": f"/api/v1/assets/{asset.id}/content", + "asset_id": asset.id, + "content_hash": asset.content_hash, + "remote_url": image_url, + } + except Exception: + # Fall back to remote URL if download fails (e.g. offline mock or local test) + pass + + yield NodeProgressEvent(node_id=node_id, progress=1.0, message="Image generation completed.") + yield NodeOutputEvent(node_id=node_id, output=output_dict) yield NodeStatusEvent(node_id=node_id, status="completed") @@ -229,6 +245,33 @@ async def run_input_text_node( yield NodeStatusEvent(node_id=node_id, status="completed") +# --------------------------------------------------------------------------- +# Preview node: display media input card +# --------------------------------------------------------------------------- + +async def run_preview_node( + node_id: str, + inputs: Dict[str, Any], + params: Dict[str, Any], +) -> AsyncGenerator[WorkflowEvent, None]: + """Presentation card handler for output.preview nodes.""" + yield NodeStatusEvent(node_id=node_id, status="running") + media = inputs.get("media") or inputs.get("image") or inputs.get("text") or inputs.get("result") + output: Dict[str, Any] = {} + if isinstance(media, str): + if media.startswith("http://") or media.startswith("https://") or media.startswith("/api/v1/assets/"): + output["image"] = media + else: + output["result"] = media + elif isinstance(media, dict): + output.update(media) + elif media is not None: + output["result"] = str(media) + + yield NodeOutputEvent(node_id=node_id, output=output) + yield NodeStatusEvent(node_id=node_id, status="completed") + + # --------------------------------------------------------------------------- # Runner dispatch table # --------------------------------------------------------------------------- @@ -237,4 +280,5 @@ async def run_input_text_node( "input.text": run_input_text_node, "text.llm": run_llm_node, "image.generate": run_image_gen_node, + "output.preview": run_preview_node, } diff --git a/backend/app/runners/comfy_runner.py b/backend/app/runners/comfy_runner.py index 898646c..1d6bf13 100644 --- a/backend/app/runners/comfy_runner.py +++ b/backend/app/runners/comfy_runner.py @@ -3,8 +3,10 @@ Provides asynchronous HTTP and WebSocket connectivity to an external or sandboxed ComfyUI instance running on localhost (default: 8188). +Polls /history for real output files and persists them to managed asset storage. """ +import asyncio import json import logging import uuid @@ -19,6 +21,7 @@ NodeStatusEvent, WorkflowEvent, ) +from app.storage.asset_store import asset_store logger = logging.getLogger(__name__) @@ -104,6 +107,47 @@ async def queue_prompt(self, prompt_graph: Dict[str, Any], client_id: Optional[s resp.raise_for_status() return resp.json() + async def get_history(self, prompt_id: str) -> Dict[str, Any]: + """Retrieve execution history for a given prompt_id.""" + async with httpx.AsyncClient(timeout=self.timeout) as client: + resp = await client.get(f"{self.base_url}/history/{prompt_id}") + resp.raise_for_status() + return resp.json() + + async def poll_history_outputs( + self, prompt_id: str, max_wait: float = 120.0, interval: float = 0.5 + ) -> List[Dict[str, Any]]: + """ + Poll /history/{prompt_id} until completion, returning actual image output metadata. + """ + elapsed = 0.0 + while elapsed < max_wait: + try: + history_data = await self.get_history(prompt_id) + if prompt_id in history_data: + prompt_record = history_data[prompt_id] + status_info = prompt_record.get("status", {}) + if status_info.get("status_str") == "error": + messages = status_info.get("messages", []) + raise RuntimeError(f"ComfyUI execution error: {messages}") + + outputs = prompt_record.get("outputs", {}) + images: List[Dict[str, Any]] = [] + for node_output in outputs.values(): + if "images" in node_output: + images.extend(node_output["images"]) + if images: + return images + except Exception as err: + if "ComfyUI execution error" in str(err): + raise + logger.debug("Polling history for %s: %s", prompt_id, err) + + await asyncio.sleep(interval) + elapsed += interval + + raise TimeoutError(f"ComfyUI execution timed out after {max_wait}s for prompt {prompt_id}") + def get_view_url(self, filename: str, subfolder: str = "", folder_type: str = "output") -> str: """Construct the view/download URL for a generated media file.""" return f"{self.base_url}/view?filename={filename}&subfolder={subfolder}&type={folder_type}" @@ -120,7 +164,7 @@ async def run_comfy_txt2img_node( ) -> AsyncGenerator[WorkflowEvent, None]: """ Execute an image.comfy.txt2img node by compiling it to a ComfyUI prompt DAG. - Yields real-time streaming events. + Yields real-time streaming events and retrieves real outputs. """ yield NodeStatusEvent(node_id=node_id, status="running") yield NodeProgressEvent(node_id=node_id, progress=0.05, message="Checking local ComfyUI connectivity...") @@ -158,15 +202,37 @@ async def run_comfy_txt2img_node( yield NodeProgressEvent(node_id=node_id, progress=0.25, message="Submitting workflow to ComfyUI /prompt...") result = await comfy_client.queue_prompt(prompt_graph) prompt_id = result.get("prompt_id") + if not prompt_id: + raise RuntimeError("ComfyUI did not return a valid prompt_id") + + yield NodeProgressEvent(node_id=node_id, progress=0.4, message=f"Queued in ComfyUI (ID: {prompt_id[:8]}...)...") + + # Poll history for real rendered image outputs + yield NodeProgressEvent(node_id=node_id, progress=0.6, message="Waiting for ComfyUI rendering to finish...") + images = await comfy_client.poll_history_outputs(prompt_id) + if not images: + raise RuntimeError("ComfyUI completed but produced no image output records") - yield NodeProgressEvent(node_id=node_id, progress=0.5, message=f"Queued in ComfyUI (ID: {prompt_id[:8]}...)...") + primary_image = images[0] + filename = primary_image.get("filename", "") + subfolder = primary_image.get("subfolder", "") + folder_type = primary_image.get("type", "output") + remote_view_url = comfy_client.get_view_url(filename, subfolder, folder_type) - # In bridge mode, construct expected output view URL or poll /history - mock_filename = f"AI-Workflow_{prompt_id[:8]}_0001.png" - image_url = comfy_client.get_view_url(mock_filename) + yield NodeProgressEvent(node_id=node_id, progress=0.85, message="Adopting output into local asset store...") + asset = await asset_store.save_image_from_url(remote_view_url, filename=filename) - yield NodeProgressEvent(node_id=node_id, progress=0.95, message="Image rendered successfully.") - yield NodeOutputEvent(node_id=node_id, output={"image": image_url}) + local_media_url = f"/api/v1/assets/{asset.id}/content" + yield NodeProgressEvent(node_id=node_id, progress=1.0, message="Image rendered and persisted successfully.") + yield NodeOutputEvent( + node_id=node_id, + output={ + "image": local_media_url, + "asset_id": asset.id, + "filename": filename, + "content_hash": asset.content_hash, + }, + ) yield NodeStatusEvent(node_id=node_id, status="completed") except Exception as err: diff --git a/backend/app/runners/creative_runner.py b/backend/app/runners/creative_runner.py new file mode 100644 index 0000000..dec9172 --- /dev/null +++ b/backend/app/runners/creative_runner.py @@ -0,0 +1,315 @@ +""" +Creative Action Runner. + +High-level dispatcher for primary canvas actions (txt2img, img2img, inpaint, upscale). +Coordinates execution across ComfyUI, Stable Diffusion WebUI, and Cloud APIs, +enforcing deterministic caching, provenance tracking, and content-addressable storage. +""" + +import asyncio +import hashlib +import json +import logging +import random +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Optional + +from app.core.cache import cache_store +from app.runners.api_runner import _call_fal_ai, _call_openai_images, _call_siliconflow +from app.runners.comfy_runner import comfy_client +from app.runners.macro_compiler import ( + ASPECT_RATIO_DIMENSIONS, + build_comfy_img2img_graph, + build_comfy_inpaint_graph, + build_comfy_txt2img_graph, + build_comfy_upscale_graph, +) +from app.runners.webui_runner import WebUIRunner +from app.runtime.credentials import credentials_manager +from app.schemas.cloud import CloudProviderId +from app.schemas.creative import ( + CreativeActionRequest, + CreativeActionResult, + CreativeActionType, + GenerationProvenance, +) +from app.storage.asset_store import asset_store + +logger = logging.getLogger(__name__) + + +def compute_creative_cache_hash(req: CreativeActionRequest, input_hash: str = "", mask_hash: str = "") -> str: + """Compute deterministic semantic cache hash for a creative action.""" + canonical_payload = { + "action": req.action.value, + "prompt": req.prompt.strip(), + "negative_prompt": req.negative_prompt.strip(), + "model": req.model, + "engine_id": req.engine_id, + "width": req.width, + "height": req.height, + "steps": req.steps, + "cfg_scale": req.cfg_scale, + "seed": req.seed, + "denoise": req.denoise, + "input_hash": input_hash, + "mask_hash": mask_hash, + "upscale_factor": req.upscale_factor, + "upscaler_name": req.upscaler_name, + } + raw = json.dumps(canonical_payload, sort_keys=True) + return hashlib.sha256(raw.encode("utf-8")).hexdigest() + + +class CreativeRunner: + """Coordinates canvas-level image generation, inpainting, and upscaling.""" + + async def execute(self, req: CreativeActionRequest) -> CreativeActionResult: + task_id = f"task_{Date_now() if False else int(time.time() * 1000)}" + start_time = time.monotonic() + + # Resolve aspect ratio dimensions if default 512 + if req.aspect_ratio in ASPECT_RATIO_DIMENSIONS: + req.width, req.height = ASPECT_RATIO_DIMENSIONS[req.aspect_ratio] + + # Resolve randomized seed if -1 + actual_seed = req.seed if req.seed >= 0 else random.randint(1, 2147483647) + req.seed = actual_seed + + # Retrieve content hashes of input/mask assets if supplied + input_hash = "" + mask_hash = "" + input_file_path: Optional[Path] = None + mask_file_path: Optional[Path] = None + + if req.input_image_id: + rec = await asset_store.get_asset(req.input_image_id) + if rec: + input_hash = rec.content_hash + input_file_path = asset_store.get_absolute_path(rec) + + if req.mask_image_id: + m_rec = await asset_store.get_asset(req.mask_image_id) + if m_rec: + mask_hash = m_rec.content_hash + mask_file_path = asset_store.get_absolute_path(m_rec) + + # Check deterministic cache + cache_key = compute_creative_cache_hash(req, input_hash, mask_hash) + cached_result = await cache_store.get_async(cache_key) + if cached_result: + return CreativeActionResult( + success=True, + task_id=task_id, + asset_id=cached_result.get("asset_id"), + image_url=cached_result.get("image_url"), + width=cached_result.get("width", req.width), + height=cached_result.get("height", req.height), + provenance=GenerationProvenance.model_validate(cached_result["provenance"]), + is_cached=True, + ) + + # Dispatch based on engine + try: + if "webui" in req.engine_id: + runner = WebUIRunner() + action_data = await runner.execute_action(req) + asset_id = action_data["asset_id"] + image_url = action_data["image_url"] + out_w = action_data.get("width", req.width) + out_h = action_data.get("height", req.height) + elif "comfy" in req.engine_id: + action_data = await self._run_comfy(req, input_file_path, mask_file_path) + asset_id = action_data["asset_id"] + image_url = action_data["image_url"] + out_w = action_data.get("width", req.width) + out_h = action_data.get("height", req.height) + elif "cloud" in req.engine_id: + action_data = await self._run_cloud(req) + asset_id = action_data["asset_id"] + image_url = action_data["image_url"] + out_w = action_data.get("width", req.width) + out_h = action_data.get("height", req.height) + else: + # Default to ComfyUI fallback + action_data = await self._run_comfy(req, input_file_path, mask_file_path) + asset_id = action_data["asset_id"] + image_url = action_data["image_url"] + out_w = action_data.get("width", req.width) + out_h = action_data.get("height", req.height) + + elapsed_ms = round((time.monotonic() - start_time) * 1000, 2) + + provenance = GenerationProvenance( + action=req.action, + prompt=req.prompt, + negative_prompt=req.negative_prompt, + model=req.model, + engine_id=req.engine_id, + seed=req.seed, + steps=req.steps, + cfg_scale=req.cfg_scale, + dimensions=f"{out_w}x{out_h}", + created_at=datetime.now(timezone.utc).isoformat(), + source_asset_id=req.input_image_id, + mask_asset_id=req.mask_image_id, + execution_time_ms=elapsed_ms, + ) + + result = CreativeActionResult( + success=True, + task_id=task_id, + asset_id=asset_id, + image_url=image_url, + width=out_w, + height=out_h, + provenance=provenance, + is_cached=False, + ) + + # Store in cache + cache_payload = { + "asset_id": asset_id, + "image_url": image_url, + "width": out_w, + "height": out_h, + "provenance": provenance.model_dump(), + } + await cache_store.set_async(cache_key, cache_payload) + + return result + + except Exception as e: + logger.error(f"Creative action {req.action} failed: {e}") + return CreativeActionResult( + success=False, + task_id=task_id, + error_message=str(e), + width=req.width, + height=req.height, + ) + + async def _run_comfy( + self, + req: CreativeActionRequest, + input_file: Optional[Path], + mask_file: Optional[Path], + ) -> Dict[str, Any]: + """Compile and submit ComfyUI macro graph.""" + if req.action == CreativeActionType.TXT2IMG: + prompt_graph = build_comfy_txt2img_graph( + prompt=req.prompt, + negative_prompt=req.negative_prompt, + checkpoint=req.model, + steps=req.steps, + cfg=req.cfg_scale, + aspect_ratio=req.aspect_ratio, + seed=req.seed, + ) + elif req.action == CreativeActionType.IMG2IMG: + if not input_file: + raise ValueError("Source image required for img2img") + prompt_graph = build_comfy_img2img_graph( + prompt=req.prompt, + image_filename=input_file.name, + negative_prompt=req.negative_prompt, + checkpoint=req.model, + steps=req.steps, + cfg=req.cfg_scale, + denoise=req.denoise, + seed=req.seed, + ) + elif req.action == CreativeActionType.INPAINT: + if not input_file or not mask_file: + raise ValueError("Source image and mask required for inpaint") + prompt_graph = build_comfy_inpaint_graph( + prompt=req.prompt, + image_filename=input_file.name, + mask_filename=mask_file.name, + negative_prompt=req.negative_prompt, + checkpoint=req.model, + steps=req.steps, + cfg=req.cfg_scale, + denoise=req.denoise, + seed=req.seed, + ) + elif req.action == CreativeActionType.UPSCALE: + if not input_file: + raise ValueError("Source image required for upscale") + prompt_graph = build_comfy_upscale_graph( + image_filename=input_file.name, + upscaler_model=req.upscaler_name or "RealESRGAN_x4plus.pth", + ) + else: + raise ValueError(f"Unsupported action: {req.action}") + + # Queue prompt and poll history + prompt_res = await comfy_client.queue_prompt(prompt_graph) + prompt_id = prompt_res.get("prompt_id") + if not prompt_id: + raise RuntimeError(f"Failed to queue ComfyUI prompt: {prompt_res}") + + outputs = await comfy_client.poll_history_outputs(prompt_id) + if not outputs: + raise RuntimeError("ComfyUI finished with no output images") + + first_img = outputs[0] + # Download and store in local asset store + filename = first_img.get("filename") + subfolder = first_img.get("subfolder", "") + img_type = first_img.get("type", "output") + view_url = f"{comfy_client.base_url}/view?filename={filename}&subfolder={subfolder}&type={img_type}" + + asset = await asset_store.save_image_from_url(view_url, filename=filename or "comfy_output.png") + return { + "asset_id": asset.id, + "image_url": f"/api/v1/assets/{asset.id}/content", + "width": req.width, + "height": req.height, + } + + async def _run_cloud(self, req: CreativeActionRequest) -> Dict[str, Any]: + """Dispatch creative action to cloud API using BYOK key.""" + model = req.model.lower() + if "flux" in model or req.engine_id == "cloud_fal": + provider_id = CloudProviderId.FAL + key = credentials_manager.get_key(provider_id) + if not key: + raise RuntimeError("Fal.ai API key is missing. Configure it in Cloud Providers (BYOK).") + target_model = "flux-schnell" if "schnell" in model else "flux-dev" + remote_url = await _call_fal_ai(target_model, req.prompt, req.width, req.height, key) + elif "dall-e" in model or "openai" in req.engine_id: + provider_id = CloudProviderId.OPENAI + key = credentials_manager.get_key(provider_id) + if not key: + raise RuntimeError("OpenAI API key is missing. Configure it in Cloud Providers (BYOK).") + remote_url = await _call_openai_images(req.prompt, key) + else: + # Default to SiliconFlow / SDXL or Fal.ai + provider_id = CloudProviderId.SILICONFLOW + key = credentials_manager.get_key(provider_id) + if not key: + # Check Fal.ai fallback + key = credentials_manager.get_key(CloudProviderId.FAL) + if key: + remote_url = await _call_fal_ai("flux-schnell", req.prompt, req.width, req.height, key) + else: + raise RuntimeError("No cloud provider API key configured. Please set up a BYOK key in Cloud Settings.") + else: + remote_url = await _call_siliconflow(req.prompt, req.width, req.height, key) + + if not remote_url: + raise RuntimeError("Cloud provider returned no image URL.") + + asset = await asset_store.save_image_from_url(remote_url, filename="cloud_output.png") + return { + "asset_id": asset.id, + "image_url": f"/api/v1/assets/{asset.id}/content", + "width": req.width, + "height": req.height, + } + + +creative_runner = CreativeRunner() diff --git a/backend/app/runners/macro_compiler.py b/backend/app/runners/macro_compiler.py index 84ee955..b821cdf 100644 --- a/backend/app/runners/macro_compiler.py +++ b/backend/app/runners/macro_compiler.py @@ -1,9 +1,9 @@ """ Macro Subgraph Compiler. -Translates high-level, human-friendly node parameters into ComfyUI's internal -low-level DAG (/prompt JSON specification), abstracting away latents, VAEs, -and conditioning linkages from the user. +Translates high-level, human-friendly creative requests (txt2img, img2img, inpaint, upscale) +into ComfyUI's internal low-level DAG (/prompt JSON specification), abstracting away latents, +VAEs, and conditioning linkages from the user. """ import random @@ -32,20 +32,16 @@ def build_comfy_txt2img_graph( lora_name: Optional[str] = None, lora_strength: float = 1.0, ) -> Dict[str, Any]: - """ - Compile a high-level text-to-image request into a complete ComfyUI prompt graph. - """ + """Compile a high-level text-to-image request into a complete ComfyUI prompt graph.""" width, height = ASPECT_RATIO_DIMENSIONS.get(aspect_ratio, (1024, 1024)) - actual_seed = seed if seed is not None else random.randint(1, 1125899906842624) + actual_seed = seed if seed is not None and seed >= 0 else random.randint(1, 1125899906842624) graph: Dict[str, Any] = {} # Node 4: Load Checkpoint graph["4"] = { "class_type": "CheckpointLoaderSimple", - "inputs": { - "ckpt_name": checkpoint, - }, + "inputs": {"ckpt_name": checkpoint}, } model_source = ["4", 0] @@ -69,29 +65,19 @@ def build_comfy_txt2img_graph( # Node 5: Empty Latent Image graph["5"] = { "class_type": "EmptyLatentImage", - "inputs": { - "width": width, - "height": height, - "batch_size": 1, - }, + "inputs": {"width": width, "height": height, "batch_size": 1}, } # Node 6: Positive Prompt CLIP Text Encode graph["6"] = { "class_type": "CLIPTextEncode", - "inputs": { - "text": prompt, - "clip": clip_source, - }, + "inputs": {"text": prompt, "clip": clip_source}, } # Node 7: Negative Prompt CLIP Text Encode graph["7"] = { "class_type": "CLIPTextEncode", - "inputs": { - "text": negative_prompt, - "clip": clip_source, - }, + "inputs": {"text": negative_prompt, "clip": clip_source}, } # Node 3: KSampler @@ -114,19 +100,197 @@ def build_comfy_txt2img_graph( # Node 8: VAE Decode graph["8"] = { "class_type": "VAEDecode", + "inputs": {"samples": ["3", 0], "vae": ["4", 2]}, + } + + # Node 9: Save Image Output + graph["9"] = { + "class_type": "SaveImage", + "inputs": {"filename_prefix": "Berry-Txt2Img", "images": ["8", 0]}, + } + + return graph + + +def build_comfy_img2img_graph( + prompt: str, + image_filename: str, + negative_prompt: str = "", + checkpoint: str = "v1-5-pruned-emaonly.safetensors", + steps: int = 20, + cfg: float = 7.0, + denoise: float = 0.75, + seed: Optional[int] = None, + sampler_name: str = "euler", + scheduler: str = "normal", +) -> Dict[str, Any]: + """Compile an image-to-image request into a complete ComfyUI prompt graph.""" + actual_seed = seed if seed is not None and seed >= 0 else random.randint(1, 1125899906842624) + + graph: Dict[str, Any] = {} + + graph["4"] = { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": checkpoint}, + } + + graph["1"] = { + "class_type": "LoadImage", + "inputs": {"image": image_filename}, + } + + graph["2"] = { + "class_type": "VAEEncode", + "inputs": {"pixels": ["1", 0], "vae": ["4", 2]}, + } + + graph["6"] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": prompt, "clip": ["4", 1]}, + } + + graph["7"] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": negative_prompt, "clip": ["4", 1]}, + } + + graph["3"] = { + "class_type": "KSampler", "inputs": { - "samples": ["3", 0], + "seed": actual_seed, + "steps": steps, + "cfg": cfg, + "sampler_name": sampler_name, + "scheduler": scheduler, + "denoise": denoise, + "model": ["4", 0], + "positive": ["6", 0], + "negative": ["7", 0], + "latent_image": ["2", 0], + }, + } + + graph["8"] = { + "class_type": "VAEDecode", + "inputs": {"samples": ["3", 0], "vae": ["4", 2]}, + } + + graph["9"] = { + "class_type": "SaveImage", + "inputs": {"filename_prefix": "Berry-Img2Img", "images": ["8", 0]}, + } + + return graph + + +def build_comfy_inpaint_graph( + prompt: str, + image_filename: str, + mask_filename: str, + negative_prompt: str = "", + checkpoint: str = "v1-5-pruned-emaonly.safetensors", + steps: int = 20, + cfg: float = 7.0, + denoise: float = 0.85, + seed: Optional[int] = None, +) -> Dict[str, Any]: + """Compile an inpainting request into a complete ComfyUI prompt graph.""" + actual_seed = seed if seed is not None and seed >= 0 else random.randint(1, 1125899906842624) + + graph: Dict[str, Any] = {} + + graph["4"] = { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": checkpoint}, + } + + graph["1"] = { + "class_type": "LoadImage", + "inputs": {"image": image_filename}, + } + + graph["2"] = { + "class_type": "LoadImage", + "inputs": {"image": mask_filename}, + } + + graph["5"] = { + "class_type": "VAEEncodeForInpaint", + "inputs": { + "pixels": ["1", 0], + "mask": ["2", 1], # mask channel "vae": ["4", 2], + "grow_mask_by": 6, }, } - # Node 9: Save Image Output + graph["6"] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": prompt, "clip": ["4", 1]}, + } + + graph["7"] = { + "class_type": "CLIPTextEncode", + "inputs": {"text": negative_prompt, "clip": ["4", 1]}, + } + + graph["3"] = { + "class_type": "KSampler", + "inputs": { + "seed": actual_seed, + "steps": steps, + "cfg": cfg, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": denoise, + "model": ["4", 0], + "positive": ["6", 0], + "negative": ["7", 0], + "latent_image": ["5", 0], + }, + } + + graph["8"] = { + "class_type": "VAEDecode", + "inputs": {"samples": ["3", 0], "vae": ["4", 2]}, + } + graph["9"] = { "class_type": "SaveImage", + "inputs": {"filename_prefix": "Berry-Inpaint", "images": ["8", 0]}, + } + + return graph + + +def build_comfy_upscale_graph( + image_filename: str, + upscaler_model: str = "RealESRGAN_x4plus.pth", +) -> Dict[str, Any]: + """Compile an upscaling request into a ComfyUI prompt graph.""" + graph: Dict[str, Any] = {} + + graph["1"] = { + "class_type": "LoadImage", + "inputs": {"image": image_filename}, + } + + graph["2"] = { + "class_type": "UpscaleModelLoader", + "inputs": {"model_name": upscaler_model}, + } + + graph["3"] = { + "class_type": "ImageUpscaleWithModel", "inputs": { - "filename_prefix": "AI-Workflow", - "images": ["8", 0], + "upscale_model": ["2", 0], + "image": ["1", 0], }, } + graph["4"] = { + "class_type": "SaveImage", + "inputs": {"filename_prefix": "Berry-Upscale", "images": ["3", 0]}, + } + return graph diff --git a/backend/app/runners/webui_runner.py b/backend/app/runners/webui_runner.py new file mode 100644 index 0000000..344b782 --- /dev/null +++ b/backend/app/runners/webui_runner.py @@ -0,0 +1,174 @@ +""" +Stable Diffusion WebUI REST API Runner. + +Executes text-to-image, image-to-image, inpainting, and upscaling directly against +SD WebUI (/sdapi/v1/*), converting base64 image responses into managed assets. +""" + +import base64 +import logging +from typing import Any, Dict, Optional +import httpx + +from app.schemas.creative import CreativeActionRequest, CreativeActionType +from app.storage.asset_store import asset_store + +logger = logging.getLogger(__name__) + + +class WebUIRunner: + """Dispatches creative generation requests to Stable Diffusion WebUI API.""" + + def __init__(self, endpoint_url: str = "http://127.0.0.1:7860") -> None: + self.endpoint_url = endpoint_url.rstrip("/") + + async def execute_action(self, req: CreativeActionRequest) -> Dict[str, Any]: + """Execute creative action via SD WebUI REST API.""" + async with httpx.AsyncClient(timeout=120.0) as client: + if req.action == CreativeActionType.TXT2IMG: + return await self._run_txt2img(client, req) + elif req.action == CreativeActionType.IMG2IMG: + return await self._run_img2img(client, req) + elif req.action == CreativeActionType.INPAINT: + return await self._run_inpaint(client, req) + elif req.action == CreativeActionType.UPSCALE: + return await self._run_upscale(client, req) + else: + raise ValueError(f"Unsupported action for WebUI: {req.action}") + + async def _run_txt2img(self, client: httpx.AsyncClient, req: CreativeActionRequest) -> Dict[str, Any]: + payload = { + "prompt": req.prompt, + "negative_prompt": req.negative_prompt, + "steps": req.steps, + "cfg_scale": req.cfg_scale, + "width": req.width, + "height": req.height, + "seed": req.seed if req.seed >= 0 else -1, + } + resp = await client.post(f"{self.endpoint_url}/sdapi/v1/txt2img", json=payload) + resp.raise_for_status() + data = resp.json() + images = data.get("images", []) + if not images: + raise RuntimeError("WebUI returned no images") + + img_bytes = base64.b64decode(images[0]) + asset = await asset_store.save_bytes(img_bytes, filename="webui_txt2img.png", media_type="image/png") + return { + "asset_id": asset.id, + "image_url": f"/api/v1/assets/{asset.id}/content", + "width": req.width, + "height": req.height, + } + + async def _run_img2img(self, client: httpx.AsyncClient, req: CreativeActionRequest) -> Dict[str, Any]: + if not req.input_image_id: + raise ValueError("Input image ID required for img2img") + + rec = await asset_store.get_asset(req.input_image_id) + if not rec: + raise ValueError(f"Asset '{req.input_image_id}' not found") + abs_path = asset_store.get_absolute_path(rec) + init_b64 = base64.b64encode(abs_path.read_bytes()).decode("utf-8") + + payload = { + "init_images": [init_b64], + "prompt": req.prompt, + "negative_prompt": req.negative_prompt, + "steps": req.steps, + "cfg_scale": req.cfg_scale, + "denoising_strength": req.denoise, + "seed": req.seed if req.seed >= 0 else -1, + } + resp = await client.post(f"{self.endpoint_url}/sdapi/v1/img2img", json=payload) + resp.raise_for_status() + data = resp.json() + images = data.get("images", []) + if not images: + raise RuntimeError("WebUI returned no images") + + img_bytes = base64.b64decode(images[0]) + asset = await asset_store.save_bytes(img_bytes, filename="webui_img2img.png", media_type="image/png") + return { + "asset_id": asset.id, + "image_url": f"/api/v1/assets/{asset.id}/content", + "width": req.width, + "height": req.height, + } + + async def _run_inpaint(self, client: httpx.AsyncClient, req: CreativeActionRequest) -> Dict[str, Any]: + if not req.input_image_id or not req.mask_image_id: + raise ValueError("Both input image and mask image required for inpaint") + + img_rec = await asset_store.get_asset(req.input_image_id) + mask_rec = await asset_store.get_asset(req.mask_image_id) + if not img_rec or not mask_rec: + raise ValueError("Input or mask asset missing on disk") + + img_b64 = base64.b64encode(asset_store.get_absolute_path(img_rec).read_bytes()).decode("utf-8") + mask_b64 = base64.b64encode(asset_store.get_absolute_path(mask_rec).read_bytes()).decode("utf-8") + + payload = { + "init_images": [img_b64], + "mask": mask_b64, + "prompt": req.prompt, + "negative_prompt": req.negative_prompt, + "steps": req.steps, + "cfg_scale": req.cfg_scale, + "denoising_strength": req.denoise, + "inpainting_fill": 1, + "inpaint_full_res": True, + "seed": req.seed if req.seed >= 0 else -1, + } + resp = await client.post(f"{self.endpoint_url}/sdapi/v1/img2img", json=payload) + resp.raise_for_status() + data = resp.json() + images = data.get("images", []) + if not images: + raise RuntimeError("WebUI returned no images") + + img_bytes = base64.b64decode(images[0]) + asset = await asset_store.save_bytes(img_bytes, filename="webui_inpaint.png", media_type="image/png") + return { + "asset_id": asset.id, + "image_url": f"/api/v1/assets/{asset.id}/content", + "width": req.width, + "height": req.height, + } + + async def _run_upscale(self, client: httpx.AsyncClient, req: CreativeActionRequest) -> Dict[str, Any]: + if not req.input_image_id: + raise ValueError("Input image ID required for upscaling") + + rec = await asset_store.get_asset(req.input_image_id) + if not rec: + raise ValueError(f"Asset '{req.input_image_id}' not found") + abs_path = asset_store.get_absolute_path(rec) + img_b64 = base64.b64encode(abs_path.read_bytes()).decode("utf-8") + + payload = { + "image": img_b64, + "upscaling_resize": req.upscale_factor, + "upscaler_1": req.upscaler_name or "R-ESRGAN 4x+", + } + resp = await client.post(f"{self.endpoint_url}/sdapi/v1/extra-single-image", json=payload) + resp.raise_for_status() + data = resp.json() + img_b64_out = data.get("image") + if not img_b64_out: + raise RuntimeError("WebUI returned no upscaled image") + + img_bytes = base64.b64decode(img_b64_out) + asset = await asset_store.save_bytes(img_bytes, filename="webui_upscaled.png", media_type="image/png") + new_width = int(req.width * req.upscale_factor) + new_height = int(req.height * req.upscale_factor) + return { + "asset_id": asset.id, + "image_url": f"/api/v1/assets/{asset.id}/content", + "width": new_width, + "height": new_height, + } + + +webui_runner = WebUIRunner() diff --git a/backend/app/runtime/credentials.py b/backend/app/runtime/credentials.py new file mode 100644 index 0000000..4ca9365 --- /dev/null +++ b/backend/app/runtime/credentials.py @@ -0,0 +1,184 @@ +""" +Secure BYOK Credential Manager. + +Manages API keys for external cloud providers (OpenAI, Fal.ai, SiliconFlow). +Ensures secrets are stored locally, never exposed in workflow files or logs, +and strictly redacted in all diagnostic and status APIs. +""" + +import json +import logging +import os +from pathlib import Path +from typing import Dict, List, Optional +import httpx + +from app.schemas.cloud import ( + CloudProviderId, + CloudProviderInfo, + TestKeyResult, +) +from app.storage.db import get_default_data_dir + +logger = logging.getLogger(__name__) + + +def redact_key(key: str) -> str: + """Mask an API key for safe display, preserving prefix and last 4 characters.""" + if not key: + return "" + if len(key) <= 8: + return "****" + prefix = key[:3] + suffix = key[-4:] + return f"{prefix}...{suffix}" + + +class CredentialManager: + """Manages secure persistence and live validation of BYOK cloud API keys.""" + + def __init__(self, data_dir: Optional[Path] = None) -> None: + self.data_dir = data_dir or get_default_data_dir() + self.creds_file = self.data_dir / "credentials.json" + self._memory_creds: Dict[str, str] = {} + self._load() + + def _load(self) -> None: + if self.creds_file.is_file(): + try: + data = json.loads(self.creds_file.read_text(encoding="utf-8")) + if isinstance(data, dict): + self._memory_creds = data + except Exception as e: + logger.warning(f"Error loading credentials from {self.creds_file}: {e}") + + def _save(self) -> None: + self.data_dir.mkdir(parents=True, exist_ok=True) + try: + self.creds_file.write_text(json.dumps(self._memory_creds, indent=2), encoding="utf-8") + except Exception as e: + logger.error(f"Error saving credentials: {e}") + + def get_key(self, provider_id: CloudProviderId) -> Optional[str]: + """Resolve API key from stored BYOK credentials or environment variable.""" + # 1. Check user-stored BYOK key + stored = self._memory_creds.get(provider_id.value) + if stored: + return stored + + # 2. Check environment variables + if provider_id == CloudProviderId.OPENAI: + return os.environ.get("OPENAI_API_KEY") + elif provider_id == CloudProviderId.FAL: + return os.environ.get("FAL_KEY") or os.environ.get("FAL_API_KEY") + elif provider_id == CloudProviderId.SILICONFLOW: + return os.environ.get("SILICONFLOW_API_KEY") + + return None + + def set_key(self, provider_id: CloudProviderId, api_key: str) -> None: + """Store an API key for a cloud provider.""" + self._memory_creds[provider_id.value] = api_key.strip() + self._save() + + def delete_key(self, provider_id: CloudProviderId) -> None: + """Remove stored API key.""" + if provider_id.value in self._memory_creds: + del self._memory_creds[provider_id.value] + self._save() + + def list_providers(self) -> List[CloudProviderInfo]: + """List cloud providers with capability matrices and redacted keys.""" + providers = [ + CloudProviderInfo( + id=CloudProviderId.OPENAI, + name="OpenAI", + description="High-fidelity image generation via DALL-E 3 and GPT-4o.", + is_configured=bool(self.get_key(CloudProviderId.OPENAI)), + redacted_key=redact_key(self.get_key(CloudProviderId.OPENAI) or ""), + supported_models=["dall-e-3", "gpt-4o"], + capabilities=["txt2img"], + upload_disclosure="Prompts are sent to OpenAI servers. Zero local GPU required.", + website_url="https://platform.openai.com/api-keys", + ), + CloudProviderInfo( + id=CloudProviderId.FAL, + name="Fal.ai", + description="Ultra-fast FLUX.1 [schnell] and FLUX.1 [dev] inference on serverless GPUs.", + is_configured=bool(self.get_key(CloudProviderId.FAL)), + redacted_key=redact_key(self.get_key(CloudProviderId.FAL) or ""), + supported_models=["flux-schnell", "flux-dev"], + capabilities=["txt2img", "img2img"], + upload_disclosure="Prompts and reference images are processed on Fal.ai cloud infrastructure.", + website_url="https://fal.ai/dashboard/keys", + ), + CloudProviderInfo( + id=CloudProviderId.SILICONFLOW, + name="SiliconFlow", + description="Affordable high-throughput SDXL and FLUX inference.", + is_configured=bool(self.get_key(CloudProviderId.SILICONFLOW)), + redacted_key=redact_key(self.get_key(CloudProviderId.SILICONFLOW) or ""), + supported_models=["sdxl-turbo", "sd-3-5-large"], + capabilities=["txt2img"], + upload_disclosure="Requests are dispatched to SiliconFlow cloud API endpoints.", + website_url="https://cloud.siliconflow.cn/account/ak", + ), + ] + return providers + + async def test_key(self, provider_id: CloudProviderId, api_key: Optional[str] = None) -> TestKeyResult: + """Validate key against the provider's authentication endpoint.""" + key = api_key or self.get_key(provider_id) + if not key: + return TestKeyResult( + provider_id=provider_id, + valid=False, + message="No API key supplied or found in configuration.", + ) + + async with httpx.AsyncClient(timeout=10.0) as client: + try: + if provider_id == CloudProviderId.OPENAI: + resp = await client.get( + "https://api.openai.com/v1/models", + headers={"Authorization": f"Bearer {key}"}, + ) + valid = resp.status_code == 200 + msg = "OpenAI API key verified successfully." if valid else f"Invalid key (HTTP {resp.status_code})" + return TestKeyResult(provider_id=provider_id, valid=valid, message=msg, status_code=resp.status_code) + + elif provider_id == CloudProviderId.FAL: + # Fal.ai key check: probe user info or model status + resp = await client.get( + "https://rest.alpha.fal.ai/tokens/current", + headers={"Authorization": f"Key {key}"}, + ) + # Even if 404 on endpoint, 401/403 indicates invalid key, while 200 indicates valid + if resp.status_code == 200: + return TestKeyResult(provider_id=provider_id, valid=True, message="Fal.ai key verified successfully.", status_code=200) + elif resp.status_code in (401, 403): + return TestKeyResult(provider_id=provider_id, valid=False, message="Invalid Fal.ai API key.", status_code=resp.status_code) + else: + # Fallback for alternative Fal auth format + return TestKeyResult(provider_id=provider_id, valid=True, message="Fal.ai key accepted.", status_code=resp.status_code) + + elif provider_id == CloudProviderId.SILICONFLOW: + resp = await client.get( + "https://api.siliconflow.cn/v1/user/info", + headers={"Authorization": f"Bearer {key}"}, + ) + valid = resp.status_code == 200 + msg = "SiliconFlow API key verified successfully." if valid else f"Invalid key (HTTP {resp.status_code})" + return TestKeyResult(provider_id=provider_id, valid=valid, message=msg, status_code=resp.status_code) + + except Exception as e: + return TestKeyResult( + provider_id=provider_id, + valid=False, + message=f"Network error testing key: {e}", + ) + + return TestKeyResult(provider_id=provider_id, valid=False, message="Unsupported provider for testing.") + + +credentials_manager = CredentialManager() diff --git a/backend/app/runtime/engine_manager.py b/backend/app/runtime/engine_manager.py new file mode 100644 index 0000000..6099457 --- /dev/null +++ b/backend/app/runtime/engine_manager.py @@ -0,0 +1,168 @@ +""" +Engine Connection Manager. + +Connects to managed (supervised) and external (user-running) ComfyUI and +Stable Diffusion WebUI instances. For external engines, guarantees ZERO process +ownership: never attempts to kill external processes or mutate user installations. +""" + +import asyncio +import logging +import time +from typing import Dict, List, Optional +import httpx + +from app.runtime.supervisor import supervisor as comfy_supervisor +from app.runtime.webui_supervisor import webui_supervisor +from app.schemas.engine import ( + EngineConnection, + EngineOwnership, + EngineStatus, + EngineType, +) + +logger = logging.getLogger(__name__) + + +class EngineManager: + """Registry and coordinator for local inference engine connections.""" + + def __init__(self) -> None: + self._connections: Dict[str, EngineConnection] = {} + self._init_managed_engines() + + def _init_managed_engines(self) -> None: + """Register the built-in managed engine slots.""" + # Managed ComfyUI + self._connections["managed_comfyui"] = EngineConnection( + id="managed_comfyui", + name="Managed ComfyUI (Isolated)", + engine_type=EngineType.COMFYUI, + ownership=EngineOwnership.MANAGED, + endpoint_url=f"http://127.0.0.1:{comfy_supervisor.port}", + native_ui_url=f"http://127.0.0.1:{comfy_supervisor.port}", + status=EngineStatus.STOPPED if comfy_supervisor.is_installed() else EngineStatus.NOT_INSTALLED, + capabilities=["txt2img", "img2img", "workflows"], + ) + + # Managed WebUI + self._connections["managed_webui"] = EngineConnection( + id="managed_webui", + name="Managed SD WebUI (Isolated)", + engine_type=EngineType.WEBUI, + ownership=EngineOwnership.MANAGED, + endpoint_url=f"http://127.0.0.1:{webui_supervisor.port}", + native_ui_url=f"http://127.0.0.1:{webui_supervisor.port}", + status=EngineStatus.STOPPED if webui_supervisor.is_installed() else EngineStatus.NOT_INSTALLED, + capabilities=["txt2img", "img2img", "inpaint", "extras"], + ) + + def list_engines(self) -> List[EngineConnection]: + """Return list of all registered engines with refreshed managed statuses.""" + # Sync managed ComfyUI status + comfy_conn = self._connections.get("managed_comfyui") + if comfy_conn: + if not comfy_supervisor.is_installed(): + comfy_conn.status = EngineStatus.NOT_INSTALLED + elif comfy_supervisor.is_running(): + comfy_conn.status = EngineStatus.RUNNING + else: + comfy_conn.status = EngineStatus.STOPPED + + # Sync managed WebUI status + webui_conn = self._connections.get("managed_webui") + if webui_conn: + if not webui_supervisor.is_installed(): + webui_conn.status = EngineStatus.NOT_INSTALLED + elif webui_supervisor.is_running(): + webui_conn.status = EngineStatus.RUNNING + else: + webui_conn.status = EngineStatus.STOPPED + + return list(self._connections.values()) + + def get_engine(self, engine_id: str) -> Optional[EngineConnection]: + """Retrieve connection details for a specific engine.""" + self.list_engines() # refresh statuses + return self._connections.get(engine_id) + + async def connect_external_engine( + self, engine_type: EngineType, endpoint_url: str, name: Optional[str] = None + ) -> EngineConnection: + """ + Connect to an existing user-managed engine instance. + Zero process ownership: does not touch process or write PID. + """ + endpoint = endpoint_url.rstrip("/") + engine_id = f"external_{engine_type.value}_{abs(hash(endpoint)) % 10000}" + display_name = name or f"External {engine_type.value.capitalize()} ({endpoint})" + + connection = EngineConnection( + id=engine_id, + name=display_name, + engine_type=engine_type, + ownership=EngineOwnership.EXTERNAL, + endpoint_url=endpoint, + native_ui_url=endpoint, + status=EngineStatus.OFFLINE, + capabilities=["txt2img", "img2img"], + ) + + # Validate external engine health immediately + connection = await self.test_engine_connection(connection) + self._connections[engine_id] = connection + return connection + + async def test_engine_connection(self, connection: EngineConnection) -> EngineConnection: + """Probe engine endpoint and update status and capabilities.""" + endpoint = connection.endpoint_url + async with httpx.AsyncClient(timeout=4.0) as client: + try: + if connection.engine_type == EngineType.COMFYUI: + # Query /system_stats + resp = await client.get(f"{endpoint}/system_stats") + if resp.status_code == 200: + data = resp.json() + connection.status = EngineStatus.READY + devices = data.get("devices", []) + if devices and isinstance(devices, list): + vram_free = devices[0].get("vram_free", 0) + connection.vram_free_mb = int(vram_free / (1024 * 1024)) + connection.error_message = None + else: + connection.status = EngineStatus.DEGRADED + connection.error_message = f"ComfyUI returned HTTP {resp.status_code}" + elif connection.engine_type == EngineType.WEBUI: + # Query /sdapi/v1/options or /sdapi/v1/progress + resp = await client.get(f"{endpoint}/sdapi/v1/options") + if resp.status_code == 200: + connection.status = EngineStatus.READY + connection.error_message = None + else: + # Try /sdapi/v1/progress as fallback + resp_prog = await client.get(f"{endpoint}/sdapi/v1/progress") + if resp_prog.status_code == 200: + connection.status = EngineStatus.READY + connection.error_message = None + else: + connection.status = EngineStatus.DEGRADED + connection.error_message = f"WebUI returned HTTP {resp.status_code}" + + connection.last_heartbeat = time.time() + + except Exception as e: + connection.status = EngineStatus.OFFLINE + connection.error_message = f"Connection failed: {e}" + + return connection + + async def test_connection_by_id(self, engine_id: str) -> Optional[EngineConnection]: + """Test health of a registered engine connection by its ID.""" + conn = self._connections.get(engine_id) + if not conn: + return None + return await self.test_engine_connection(conn) + + +# Global engine manager singleton +engine_manager = EngineManager() diff --git a/backend/app/runtime/hardware.py b/backend/app/runtime/hardware.py new file mode 100644 index 0000000..bcd154c --- /dev/null +++ b/backend/app/runtime/hardware.py @@ -0,0 +1,143 @@ +""" +Hardware and storage readiness diagnostics. + +Queries NVIDIA GPU stats via nvidia-smi and checks filesystem storage thresholds. +Operates with zero heavy dependencies (no PyTorch, no CUDA C-extensions). +""" + +import asyncio +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import List, Optional + +from app.runtime.supervisor import get_default_engine_dir +from app.schemas.hardware import GpuInfo, HardwareReadiness, StorageInfo + + +def detect_gpus() -> List[GpuInfo]: + """Query NVIDIA GPUs using nvidia-smi without importing PyTorch.""" + gpus: List[GpuInfo] = [] + try: + # Query: index, name, memory.total, memory.free, driver_version, temperature.gpu + res = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,name,memory.total,memory.free,driver_version,temperature.gpu", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=5, + check=False, + creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0, + ) + if res.returncode == 0 and res.stdout.strip(): + for line in res.stdout.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) >= 4: + try: + idx = int(parts[0]) + name = parts[1] + vram_total = int(float(parts[2])) + vram_free = int(float(parts[3])) + driver_ver = parts[4] if len(parts) > 4 else None + temp = int(float(parts[5])) if len(parts) > 5 and parts[5].isdigit() else None + gpus.append( + GpuInfo( + index=idx, + name=name, + vram_total_mb=vram_total, + vram_free_mb=vram_free, + driver_version=driver_ver, + temperature_c=temp, + ) + ) + except (ValueError, IndexError): + continue + except (FileNotFoundError, OSError, subprocess.TimeoutExpired): + pass + + return gpus + + +def get_storage_readiness(directory: Path, min_free_gb: float = 15.0) -> StorageInfo: + """Calculate disk storage metrics for the target directory.""" + directory.mkdir(parents=True, exist_ok=True) + try: + usage = shutil.disk_usage(directory) + total_gb = round(usage.total / (1024**3), 2) + free_gb = round(usage.free / (1024**3), 2) + used_gb = round(usage.used / (1024**3), 2) + is_sufficient = free_gb >= min_free_gb + return StorageInfo( + path=str(directory), + total_gb=total_gb, + free_gb=free_gb, + used_gb=used_gb, + is_sufficient=is_sufficient, + ) + except Exception: + return StorageInfo( + path=str(directory), + total_gb=0.0, + free_gb=0.0, + used_gb=0.0, + is_sufficient=False, + ) + + +async def check_hardware_readiness(engine_dir: Optional[Path] = None) -> HardwareReadiness: + """Assess overall system readiness for local AI inference and engines.""" + root_dir = engine_dir or get_default_engine_dir() + models_dir = root_dir / "models" + + # Run detection in thread to avoid blocking event loop + gpus = await asyncio.to_thread(detect_gpus) + engine_storage = await asyncio.to_thread(get_storage_readiness, root_dir, 10.0) + model_storage = await asyncio.to_thread(get_storage_readiness, models_dir, 15.0) + + has_nvidia = len(gpus) > 0 + guidance: List[str] = [] + + if has_nvidia: + primary_gpu = gpus[0] + if primary_gpu.vram_total_mb >= 12000: + ready = True + rec_engine = "comfyui" + summary = f"Ready: {primary_gpu.name} ({primary_gpu.vram_total_mb}MB VRAM) has excellent capacity for local generation." + guidance.append("SDXL, Flux Schnell, and SD 1.5 can run comfortably at high batch sizes.") + elif primary_gpu.vram_total_mb >= 6000: + ready = True + rec_engine = "comfyui" + summary = f"Ready: {primary_gpu.name} ({primary_gpu.vram_total_mb}MB VRAM) is capable of local generation." + guidance.append("SD 1.5 and SDXL are supported; Flux may require low-vram offload.") + else: + ready = True + rec_engine = "comfyui" + summary = f"Limited: {primary_gpu.name} has only {primary_gpu.vram_total_mb}MB VRAM. Low-VRAM flags recommended." + guidance.append("SD 1.5 supported in low-vram mode; cloud API inference recommended for larger models.") + else: + ready = False + rec_engine = "cloud_only" + summary = "No NVIDIA GPU detected. Local inference unavailable or not recommended." + guidance.append("Use Cloud API inference (OpenAI, Fal.ai, SiliconFlow) for high-speed generation without GPU hardware.") + + if not engine_storage.is_sufficient: + guidance.append(f"Low engine disk space: {engine_storage.free_gb}GB available on {engine_storage.path} (10GB recommended).") + + if not model_storage.is_sufficient: + guidance.append(f"Low model storage space: {model_storage.free_gb}GB available on {model_storage.path} (15GB recommended).") + + return HardwareReadiness( + has_nvidia_gpu=has_nvidia, + gpus=gpus, + engine_storage=engine_storage, + model_storage=model_storage, + ready_for_local_inference=ready, + recommended_engine=rec_engine, + summary_message=summary, + guidance_notes=guidance, + ) diff --git a/backend/app/runtime/installer.py b/backend/app/runtime/installer.py new file mode 100644 index 0000000..b1abf3b --- /dev/null +++ b/backend/app/runtime/installer.py @@ -0,0 +1,358 @@ +""" +Sandboxed Engine Installer with Manifest Tracking and Interruption Safety. + +Installs ComfyUI or Stable Diffusion WebUI into dedicated, isolated directories +using hermetic virtual environments. Strictly adheres to zero host environment pollution: +never executes global pip install commands. +""" + +import asyncio +import json +import logging +import os +import shutil +import subprocess +import sys +import venv +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Dict, Optional + +from app.runtime.supervisor import get_default_engine_dir +from app.schemas.engine import ( + EngineInstallManifest, + EngineType, + InstallPhase, + EngineUpdateManifest, + EngineUpdateStatus, +) + +logger = logging.getLogger(__name__) + +COMFYUI_GIT_REPO = "https://github.com/comfyanonymous/ComfyUI.git" +WEBUI_GIT_REPO = "https://github.com/AUTOMATIC1111/stable-diffusion-webui.git" + + +class IsolatedEngineInstaller: + """Manages staged installation of local engines inside hermetic environments.""" + + def __init__(self, engine_dir: Optional[Path] = None) -> None: + self.engine_dir = engine_dir or get_default_engine_dir() + + def get_manifest_path(self, engine_type: EngineType) -> Path: + """Return the path to the installation manifest for the specified engine.""" + return self.engine_dir / f"{engine_type.value}_manifest.json" + + def read_manifest(self, engine_type: EngineType) -> EngineInstallManifest: + """Read existing manifest or return initial state.""" + manifest_path = self.get_manifest_path(engine_type) + engine_target = self.engine_dir / ("comfyui" if engine_type == EngineType.COMFYUI else "webui") + runtime_target = self.engine_dir / ("runtime" if engine_type == EngineType.COMFYUI else "webui_runtime") + + if manifest_path.is_file(): + try: + data = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest = EngineInstallManifest.model_validate(data) + # Check for uncompleted previous run (interruption detection) + if manifest.phase in ( + InstallPhase.CREATING_VENV, + InstallPhase.DOWNLOADING, + InstallPhase.INSTALLING_DEPS, + ): + manifest.phase = InstallPhase.INTERRUPTED + manifest.error_message = "Installation was interrupted before completion." + self.write_manifest(manifest) + return manifest + except Exception as e: + logger.warning(f"Error reading manifest for {engine_type}: {e}") + + # Default clean manifest + return EngineInstallManifest( + engine_type=engine_type, + phase=InstallPhase.IDLE, + engine_dir=str(engine_target), + runtime_dir=str(runtime_target), + ) + + def write_manifest(self, manifest: EngineInstallManifest) -> None: + """Persist installation manifest to disk.""" + self.engine_dir.mkdir(parents=True, exist_ok=True) + manifest_path = self.get_manifest_path(manifest.engine_type) + manifest_path.write_text(manifest.model_dump_json(indent=2), encoding="utf-8") + + def _get_python_bin(self, runtime_dir: Path) -> Path: + if sys.platform == "win32": + return runtime_dir / "Scripts" / "python.exe" + return runtime_dir / "bin" / "python" + + def _get_pip_bin(self, runtime_dir: Path) -> Path: + if sys.platform == "win32": + return runtime_dir / "Scripts" / "pip.exe" + return runtime_dir / "bin" / "pip" + + async def create_isolated_venv( + self, runtime_dir: Path, on_log: Optional[Callable[[str], None]] = None + ) -> bool: + """Create a hermetic virtual environment without polluting host.""" + if on_log: + on_log(f"Creating isolated virtualenv in {runtime_dir}...") + + # If previous broken venv exists, remove it cleanly + python_bin = self._get_python_bin(runtime_dir) + if runtime_dir.exists() and not python_bin.is_file(): + shutil.rmtree(runtime_dir, ignore_errors=True) + + if not python_bin.is_file(): + # Run venv creation in threadpool to avoid blocking event loop + def _create_env(): + builder = venv.EnvBuilder(with_pip=True, clear=True) + builder.create(runtime_dir) + + await asyncio.to_thread(_create_env) + + return python_bin.is_file() + + async def install_engine( + self, + engine_type: EngineType, + on_progress: Optional[Callable[[EngineInstallManifest], None]] = None, + ) -> EngineInstallManifest: + """ + Run complete isolated installation: + 1. Check environment & disk space + 2. Create isolated sandboxed venv + 3. Clone / stage repository code + 4. Install isolated dependencies via sandboxed pip + """ + manifest = self.read_manifest(engine_type) + now_str = datetime.now(timezone.utc).isoformat() + manifest.created_at = now_str + manifest.phase = InstallPhase.CHECKING + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + engine_target = Path(manifest.engine_dir) + runtime_target = Path(manifest.runtime_dir) + + try: + # Phase 1: Virtualenv Creation + manifest.phase = InstallPhase.CREATING_VENV + manifest.last_log_line = "Initializing sandboxed virtual environment..." + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + venv_ok = await self.create_isolated_venv(runtime_target) + if not venv_ok: + raise RuntimeError(f"Failed to create virtual environment at {runtime_target}") + + python_bin = self._get_python_bin(runtime_target) + manifest.python_bin = str(python_bin) + + # Phase 2: Stage / Download code + manifest.phase = InstallPhase.DOWNLOADING + manifest.last_log_line = f"Staging {engine_type.value} repository..." + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + repo_url = COMFYUI_GIT_REPO if engine_type == EngineType.COMFYUI else WEBUI_GIT_REPO + if not engine_target.exists(): + clone_cmd = ["git", "clone", "--depth", "1", repo_url, str(engine_target)] + proc = await asyncio.create_subprocess_exec( + *clone_cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate() + if proc.returncode != 0: + err_msg = stderr.decode(errors="replace").strip() + raise RuntimeError(f"Git clone failed: {err_msg}") + + # Phase 3: Install isolated dependencies + manifest.phase = InstallPhase.INSTALLING_DEPS + manifest.last_log_line = "Installing engine requirements into isolated environment..." + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + req_file = engine_target / "requirements.txt" + pip_bin = self._get_pip_bin(runtime_target) + + if req_file.is_file() and pip_bin.is_file(): + # Strict: execute ONLY pip inside runtime_target! + install_cmd = [str(pip_bin), "install", "--no-warn-script-location", "-r", str(req_file)] + proc = await asyncio.create_subprocess_exec( + *install_cmd, + cwd=str(engine_target), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate() + if proc.returncode != 0: + err_msg = stderr.decode(errors="replace").strip() + raise RuntimeError(f"Pip install failed: {err_msg[:300]}") + + # Completed! + manifest.phase = InstallPhase.COMPLETED + manifest.completed_at = datetime.now(timezone.utc).isoformat() + manifest.last_log_line = f"{engine_type.value} installed successfully." + manifest.error_message = None + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + except Exception as e: + manifest.phase = InstallPhase.FAILED + manifest.error_message = str(e) + manifest.last_log_line = f"Installation error: {e}" + self.write_manifest(manifest) + if on_progress: + on_progress(manifest) + + return manifest + + def get_update_manifest_path(self, engine_type: EngineType) -> Path: + """Return the path to the update manifest for the specified engine.""" + return self.engine_dir / f"{engine_type.value}_update_manifest.json" + + def read_update_manifest(self, engine_type: EngineType) -> EngineUpdateManifest: + """Read update manifest or return initial state.""" + path = self.get_update_manifest_path(engine_type) + if path.is_file(): + try: + data = json.loads(path.read_text(encoding="utf-8")) + return EngineUpdateManifest.model_validate(data) + except Exception as e: + logger.warning(f"Error reading update manifest for {engine_type}: {e}") + return EngineUpdateManifest(engine_type=engine_type) + + def write_update_manifest(self, manifest: EngineUpdateManifest) -> None: + """Persist update manifest to disk.""" + self.engine_dir.mkdir(parents=True, exist_ok=True) + path = self.get_update_manifest_path(manifest.engine_type) + path.write_text(manifest.model_dump_json(indent=2), encoding="utf-8") + + async def update_engine( + self, + engine_type: EngineType, + has_active_tasks_fn: Optional[Callable[[], bool]] = None, + ) -> EngineUpdateManifest: + """ + Safely update managed engine: + 1. Check active jobs (abort if running) + 2. Record current git commit hash as rollback checkpoint + 3. Pull git updates + 4. Install updated requirements into isolated venv + 5. On failure, rollback to previous commit + """ + # Guard: active jobs + if has_active_tasks_fn and has_active_tasks_fn(): + manifest = self.read_update_manifest(engine_type) + manifest.status = EngineUpdateStatus.FAILED + manifest.error_message = "Cannot update engine while active generation tasks are running." + self.write_update_manifest(manifest) + return manifest + + manifest = EngineUpdateManifest( + engine_type=engine_type, + status=EngineUpdateStatus.CHECKING, + updated_at=datetime.now(timezone.utc).isoformat(), + ) + self.write_update_manifest(manifest) + + engine_target = self.engine_dir / ("comfyui" if engine_type == EngineType.COMFYUI else "webui") + runtime_target = self.engine_dir / ("runtime" if engine_type == EngineType.COMFYUI else "webui_runtime") + + if not engine_target.is_dir() or not (engine_target / ".git").is_dir(): + manifest.status = EngineUpdateStatus.FAILED + manifest.error_message = f"Engine directory {engine_target} is not a valid git repository or not installed." + self.write_update_manifest(manifest) + return manifest + + previous_commit: Optional[str] = None + try: + # Step 1: Capture previous commit hash + rev_proc = await asyncio.create_subprocess_exec( + "git", "rev-parse", "HEAD", + cwd=str(engine_target), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + rev_out, _ = await rev_proc.communicate() + if rev_proc.returncode == 0: + previous_commit = rev_out.decode().strip() + manifest.previous_commit = previous_commit + + manifest.status = EngineUpdateStatus.UPDATING + self.write_update_manifest(manifest) + + # Step 2: Fetch and pull latest updates + pull_proc = await asyncio.create_subprocess_exec( + "git", "pull", "--ff-only", + cwd=str(engine_target), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + pull_out, pull_err = await pull_proc.communicate() + if pull_proc.returncode != 0: + raise RuntimeError(f"Git pull failed: {pull_err.decode(errors='replace').strip()}") + + # Get new commit hash + rev_proc2 = await asyncio.create_subprocess_exec( + "git", "rev-parse", "HEAD", + cwd=str(engine_target), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + rev_out2, _ = await rev_proc2.communicate() + if rev_proc2.returncode == 0: + manifest.target_commit = rev_out2.decode().strip() + + # Step 3: Update dependencies in isolated venv + req_file = engine_target / "requirements.txt" + pip_bin = self._get_pip_bin(runtime_target) + if req_file.is_file() and pip_bin.is_file(): + pip_proc = await asyncio.create_subprocess_exec( + str(pip_bin), "install", "--no-warn-script-location", "-r", str(req_file), + cwd=str(engine_target), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _, pip_err = await pip_proc.communicate() + if pip_proc.returncode != 0: + raise RuntimeError(f"Pip dependency update failed: {pip_err.decode(errors='replace')[:200]}") + + manifest.status = EngineUpdateStatus.COMPLETED + manifest.error_message = None + self.write_update_manifest(manifest) + + except Exception as e: + manifest.status = EngineUpdateStatus.FAILED + manifest.error_message = str(e) + + # Perform rollback if previous commit was captured + if previous_commit: + try: + logger.info(f"Rolling back {engine_type.value} to previous commit {previous_commit}...") + rollback_proc = await asyncio.create_subprocess_exec( + "git", "checkout", previous_commit, + cwd=str(engine_target), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await rollback_proc.communicate() + if rollback_proc.returncode == 0: + manifest.status = EngineUpdateStatus.ROLLED_BACK + manifest.rollback_performed = True + except Exception as rb_err: + logger.error(f"Rollback failed for {engine_type.value}: {rb_err}") + + self.write_update_manifest(manifest) + + return manifest + + +installer = IsolatedEngineInstaller() + diff --git a/backend/app/runtime/supervisor.py b/backend/app/runtime/supervisor.py index fdafd7b..48bbd9d 100644 --- a/backend/app/runtime/supervisor.py +++ b/backend/app/runtime/supervisor.py @@ -2,8 +2,8 @@ Sandboxed ComfyUI Runtime Process Supervisor. Manages isolated execution of the local ComfyUI engine in a dedicated, -hermetic application directory (~/.ai-workflow/engine). -Guarantees zero system environment pollution. +hermetic application directory (~/.ai-workflow/engine or %LOCALAPPDATA%/AI-Workflow/engine). +Guarantees zero system environment pollution and strictly disallows host Python fallback. """ import os @@ -52,23 +52,55 @@ def is_installed(self) -> bool: main_py = self.comfy_dir / "main.py" return main_py.is_file() + def get_python_bin(self) -> Path: + """Return the expected path to the isolated virtualenv python binary.""" + if sys.platform == "win32": + return self.runtime_dir / "Scripts" / "python.exe" + return self.runtime_dir / "bin" / "python" + + def has_isolated_env(self) -> bool: + """Check if the sandboxed isolated virtual environment is ready.""" + return self.get_python_bin().is_file() + + def _verify_process_identity(self, pid: int) -> bool: + """Verify the process at pid is actually python/comfy, avoiding recycled PID collision.""" + if sys.platform == "win32": + try: + import ctypes + from ctypes import wintypes + kernel32 = ctypes.windll.kernel32 + handle = kernel32.OpenProcess(0x1000, False, pid) # PROCESS_QUERY_LIMITED_INFORMATION + if not handle: + return False + buf = ctypes.create_unicode_buffer(1024) + size = wintypes.DWORD(1024) + success = kernel32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size)) + kernel32.CloseHandle(handle) + if success: + exe_name = buf.value.lower() + return "python" in exe_name + return False + except Exception: + return False + else: + try: + proc_cmdline = Path(f"/proc/{pid}/cmdline") + if proc_cmdline.is_file(): + content = proc_cmdline.read_text(encoding="latin1", errors="ignore").lower() + return "python" in content + os.kill(pid, 0) + return True + except OSError: + return False + def get_pid(self) -> Optional[int]: - """Read active PID from the supervisor PID file if process is alive.""" + """Read active PID from the supervisor PID file if process is alive and verified.""" if not self.pid_file.is_file(): return None try: pid = int(self.pid_file.read_text(encoding="utf-8").strip()) - # Verify if process is alive - if sys.platform == "win32": - import ctypes - kernel32 = ctypes.windll.kernel32 - handle = kernel32.OpenProcess(0x1000, False, pid) # PROCESS_QUERY_LIMITED_INFORMATION - if handle != 0: - kernel32.CloseHandle(handle) - return pid - else: - os.kill(pid, 0) + if self._verify_process_identity(pid): return pid except (ValueError, OSError): pass @@ -86,22 +118,37 @@ def get_status(self) -> Dict[str, Any]: pid = self.get_pid() return { "installed": self.is_installed(), + "isolated_env_ready": self.has_isolated_env(), "running": pid is not None, "pid": pid, "port": self.port, "engine_dir": str(self.engine_dir), + "runtime_dir": str(self.runtime_dir), "comfy_dir": str(self.comfy_dir), "models_dir": str(self.models_dir), } def start(self) -> Dict[str, Any]: - """Launch the isolated ComfyUI process.""" + """Launch the isolated ComfyUI process strictly inside its hermetic environment.""" if not self.is_installed(): return { "success": False, + "code": "NOT_INSTALLED", "message": f"ComfyUI is not installed in {self.comfy_dir}. Run installer first.", } + python_bin = self.get_python_bin() + if not python_bin.is_file(): + # STRICT REQUIREMENT: Zero host Python fallback! + return { + "success": False, + "code": "ENV_MISSING", + "message": ( + f"Sandboxed virtual environment not found at {python_bin}. " + "Run the isolated installer to set up the engine runtime without host pollution." + ), + } + if self.is_running(): return { "success": True, @@ -109,16 +156,6 @@ def start(self) -> Dict[str, Any]: "pid": self.get_pid(), } - # Resolve isolated python executable - if sys.platform == "win32": - python_bin = self.runtime_dir / "Scripts" / "python.exe" - else: - python_bin = self.runtime_dir / "bin" / "python" - - if not python_bin.is_file(): - # Fallback to current virtualenv python if sandboxed venv is uninitialized - python_bin = Path(sys.executable) - cmd = [ str(python_bin), str(self.comfy_dir / "main.py"), @@ -155,11 +192,18 @@ def start(self) -> Dict[str, Any]: } def stop(self) -> Dict[str, Any]: - """Terminate the managed ComfyUI process.""" + """Terminate the managed ComfyUI process with process identity verification.""" pid = self.get_pid() if pid is None: return {"success": True, "message": "ComfyUI is not running"} + if not self._verify_process_identity(pid): + self._clean_pid_file() + return { + "success": False, + "message": f"Process {pid} is not a verified ComfyUI process; refusing to kill.", + } + try: if sys.platform == "win32": subprocess.run(["taskkill", "/F", "/T", "/PID", str(pid)], capture_output=True, check=False) diff --git a/backend/app/runtime/webui_supervisor.py b/backend/app/runtime/webui_supervisor.py new file mode 100644 index 0000000..7456594 --- /dev/null +++ b/backend/app/runtime/webui_supervisor.py @@ -0,0 +1,241 @@ +""" +Sandboxed Stable Diffusion WebUI Process Supervisor. + +Manages isolated execution of the local SD WebUI engine in a dedicated, +hermetic application directory (~/.ai-workflow/engine/webui or %LOCALAPPDATA%/AI-Workflow/engine/webui). +Guarantees zero system environment pollution, separate virtual environment from ComfyUI, +and strictly disallows host Python fallback. +""" + +import os +import signal +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, Optional + +from app.runtime.supervisor import get_default_engine_dir + + +class WebUISupervisor: + """Supervisor managing the lifecycle of an isolated Stable Diffusion WebUI subprocess.""" + + def __init__(self, engine_dir: Optional[Path] = None, port: int = 7860) -> None: + self.engine_dir = engine_dir or get_default_engine_dir() + self.port = port + self.webui_dir = self.engine_dir / "webui" + self.runtime_dir = self.engine_dir / "webui_runtime" + self.models_dir = self.engine_dir / "models" + self.pid_file = self.engine_dir / "webui.pid" + self._process: Optional[subprocess.Popen] = None + + def ensure_directories(self) -> None: + """Ensure the isolated WebUI engine directory layout exists.""" + self.engine_dir.mkdir(parents=True, exist_ok=True) + self.webui_dir.mkdir(parents=True, exist_ok=True) + self.models_dir.mkdir(parents=True, exist_ok=True) + (self.models_dir / "checkpoints").mkdir(exist_ok=True) + (self.models_dir / "loras").mkdir(exist_ok=True) + (self.models_dir / "vae").mkdir(exist_ok=True) + + def is_installed(self) -> bool: + """Check if the isolated WebUI installation exists and has entrypoint.""" + launch_py = self.webui_dir / "launch.py" + webui_py = self.webui_dir / "webui.py" + return launch_py.is_file() or webui_py.is_file() + + def get_entrypoint(self) -> Optional[Path]: + """Return the WebUI entrypoint script.""" + launch_py = self.webui_dir / "launch.py" + if launch_py.is_file(): + return launch_py + webui_py = self.webui_dir / "webui.py" + if webui_py.is_file(): + return webui_py + return None + + def get_python_bin(self) -> Path: + """Return the expected path to the isolated virtualenv python binary.""" + if sys.platform == "win32": + return self.runtime_dir / "Scripts" / "python.exe" + return self.runtime_dir / "bin" / "python" + + def has_isolated_env(self) -> bool: + """Check if the sandboxed isolated virtual environment is ready.""" + return self.get_python_bin().is_file() + + def _verify_process_identity(self, pid: int) -> bool: + """Verify the process at pid is actually python, avoiding recycled PID collision.""" + if sys.platform == "win32": + try: + import ctypes + from ctypes import wintypes + kernel32 = ctypes.windll.kernel32 + handle = kernel32.OpenProcess(0x1000, False, pid) # PROCESS_QUERY_LIMITED_INFORMATION + if not handle: + return False + buf = ctypes.create_unicode_buffer(1024) + size = wintypes.DWORD(1024) + success = kernel32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size)) + kernel32.CloseHandle(handle) + if success: + exe_name = buf.value.lower() + return "python" in exe_name + return False + except Exception: + return False + else: + try: + proc_cmdline = Path(f"/proc/{pid}/cmdline") + if proc_cmdline.is_file(): + content = proc_cmdline.read_text(encoding="latin1", errors="ignore").lower() + return "python" in content + os.kill(pid, 0) + return True + except OSError: + return False + + def get_pid(self) -> Optional[int]: + """Read active PID from the supervisor PID file if process is alive and verified.""" + if not self.pid_file.is_file(): + return None + + try: + pid = int(self.pid_file.read_text(encoding="utf-8").strip()) + if self._verify_process_identity(pid): + return pid + except (ValueError, OSError): + pass + + self._clean_pid_file() + return None + + def is_running(self) -> bool: + """Check if the managed WebUI process is actively running.""" + return self.get_pid() is not None + + def get_status(self) -> Dict[str, Any]: + """Return comprehensive status of the isolated WebUI runtime.""" + pid = self.get_pid() + return { + "installed": self.is_installed(), + "isolated_env_ready": self.has_isolated_env(), + "running": pid is not None, + "pid": pid, + "port": self.port, + "engine_dir": str(self.engine_dir), + "webui_dir": str(self.webui_dir), + "runtime_dir": str(self.runtime_dir), + "models_dir": str(self.models_dir), + } + + def start(self) -> Dict[str, Any]: + """Launch the isolated WebUI process strictly inside its hermetic environment.""" + if not self.is_installed(): + return { + "success": False, + "code": "NOT_INSTALLED", + "message": f"Stable Diffusion WebUI is not installed in {self.webui_dir}. Run installer first.", + } + + python_bin = self.get_python_bin() + if not python_bin.is_file(): + # STRICT REQUIREMENT: Zero host Python fallback! + return { + "success": False, + "code": "ENV_MISSING", + "message": ( + f"Sandboxed virtual environment not found at {python_bin}. " + "Run the isolated installer to set up the WebUI runtime without host pollution." + ), + } + + if self.is_running(): + return { + "success": True, + "message": f"WebUI is already running (PID: {self.get_pid()})", + "pid": self.get_pid(), + } + + entrypoint = self.get_entrypoint() + if not entrypoint: + return { + "success": False, + "code": "ENTRYPOINT_MISSING", + "message": f"No launch.py or webui.py entrypoint found in {self.webui_dir}", + } + + cmd = [ + str(python_bin), + str(entrypoint), + "--port", + str(self.port), + "--listen", + "127.0.0.1", + "--api", + "--nowebui", + "--ckpt-dir", + str(self.models_dir / "checkpoints"), + "--lora-dir", + str(self.models_dir / "loras"), + "--vae-dir", + str(self.models_dir / "vae"), + ] + + try: + self.ensure_directories() + self._process = subprocess.Popen( + cmd, + cwd=str(self.webui_dir), + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0, + ) + pid = self._process.pid + self.pid_file.write_text(str(pid), encoding="utf-8") + return { + "success": True, + "message": f"WebUI launched successfully (PID: {pid})", + "pid": pid, + } + except Exception as err: + return { + "success": False, + "message": f"Failed to launch WebUI: {err}", + } + + def stop(self) -> Dict[str, Any]: + """Terminate the managed WebUI process with process identity verification.""" + pid = self.get_pid() + if pid is None: + return {"success": True, "message": "WebUI is not running"} + + if not self._verify_process_identity(pid): + self._clean_pid_file() + return { + "success": False, + "message": f"Process {pid} is not a verified WebUI process; refusing to kill.", + } + + try: + if sys.platform == "win32": + subprocess.run(["taskkill", "/F", "/T", "/PID", str(pid)], capture_output=True, check=False) + else: + os.kill(pid, signal.SIGTERM) + except Exception as err: + return {"success": False, "message": f"Error stopping WebUI process: {err}"} + finally: + self._clean_pid_file() + + return {"success": True, "message": f"WebUI process {pid} stopped"} + + def _clean_pid_file(self) -> None: + try: + if self.pid_file.is_file(): + self.pid_file.unlink() + except OSError: + pass + + +# Global default WebUI supervisor singleton +webui_supervisor = WebUISupervisor() diff --git a/backend/app/schemas/cloud.py b/backend/app/schemas/cloud.py new file mode 100644 index 0000000..b0d0a23 --- /dev/null +++ b/backend/app/schemas/cloud.py @@ -0,0 +1,40 @@ +"""Schemas for BYOK cloud providers, credentials, capabilities, and key tests.""" + +from enum import Enum +from typing import List, Optional +from pydantic import BaseModel, Field + + +class CloudProviderId(str, Enum): + OPENAI = "openai" + FAL = "fal" + SILICONFLOW = "siliconflow" + + +class CloudProviderInfo(BaseModel): + id: CloudProviderId + name: str + description: str + is_configured: bool + redacted_key: Optional[str] = None # e.g. "sk-...1234", never raw secret + supported_models: List[str] = Field(default_factory=list) + capabilities: List[str] = Field(default_factory=list) # ["txt2img", "img2img", "inpaint", "upscale"] + upload_disclosure: str + website_url: str + + +class SetCredentialRequest(BaseModel): + provider_id: CloudProviderId + api_key: str + + +class TestKeyRequest(BaseModel): + provider_id: CloudProviderId + api_key: Optional[str] = None # If omitted, test the currently stored key + + +class TestKeyResult(BaseModel): + provider_id: CloudProviderId + valid: bool + message: str + status_code: Optional[int] = None diff --git a/backend/app/schemas/creative.py b/backend/app/schemas/creative.py new file mode 100644 index 0000000..186da0a --- /dev/null +++ b/backend/app/schemas/creative.py @@ -0,0 +1,59 @@ +"""Schemas for high-level creative image actions, provenance, and results.""" + +from enum import Enum +from typing import Any, Dict, List, Optional +from pydantic import BaseModel, Field + + +class CreativeActionType(str, Enum): + TXT2IMG = "txt2img" + IMG2IMG = "img2img" + INPAINT = "inpaint" + UPSCALE = "upscale" + + +class CreativeActionRequest(BaseModel): + action: CreativeActionType + prompt: str = "" + negative_prompt: str = "low quality, blurry, deformed, bad anatomy" + model: str = "v1-5-pruned-emaonly.safetensors" + engine_id: str = "managed_comfyui" # managed_comfyui, managed_webui, or cloud + aspect_ratio: str = "1:1" # 1:1, 16:9, 9:16, 4:3 + width: int = 512 + height: int = 512 + steps: int = 20 + cfg_scale: float = 7.0 + seed: int = -1 # -1 for random + denoise: float = 0.75 # For img2img and inpaint (0.0 to 1.0) + input_image_id: Optional[str] = None # Reference asset ID for img2img / inpaint / upscale + mask_image_id: Optional[str] = None # Binary mask asset ID for inpaint + upscale_factor: float = 2.0 # 2.0 or 4.0 + upscaler_name: str = "R-ESRGAN 4x+" + + +class GenerationProvenance(BaseModel): + action: CreativeActionType + prompt: str + negative_prompt: Optional[str] = None + model: str + engine_id: str + seed: int + steps: int + cfg_scale: float + dimensions: str # e.g. "512x512" + created_at: str + source_asset_id: Optional[str] = None + mask_asset_id: Optional[str] = None + execution_time_ms: Optional[float] = None + + +class CreativeActionResult(BaseModel): + success: bool + task_id: str + asset_id: Optional[str] = None + image_url: Optional[str] = None + width: int = 512 + height: int = 512 + provenance: Optional[GenerationProvenance] = None + is_cached: bool = False + error_message: Optional[str] = None diff --git a/backend/app/schemas/engine.py b/backend/app/schemas/engine.py new file mode 100644 index 0000000..6d2df0c --- /dev/null +++ b/backend/app/schemas/engine.py @@ -0,0 +1,126 @@ +"""Engine connection and installation manifest schemas.""" + +from enum import Enum +from typing import Any, Dict, List, Optional +from pydantic import BaseModel, Field + + +class EngineType(str, Enum): + COMFYUI = "comfyui" + WEBUI = "webui" + + +class EngineOwnership(str, Enum): + MANAGED = "managed" # Controlled by Berry supervisor (can start/stop) + EXTERNAL = "external" # Pre-existing user installation (connect-only, no process control) + + +class EngineStatus(str, Enum): + READY = "ready" + RUNNING = "running" + STOPPED = "stopped" + OFFLINE = "offline" + DEGRADED = "degraded" + NOT_INSTALLED = "not_installed" + + +class InstallPhase(str, Enum): + IDLE = "idle" + CHECKING = "checking" + CREATING_VENV = "creating_venv" + DOWNLOADING = "downloading" + INSTALLING_DEPS = "installing_deps" + COMPLETED = "completed" + FAILED = "failed" + INTERRUPTED = "interrupted" + + +class EngineInstallManifest(BaseModel): + engine_type: EngineType + phase: InstallPhase = InstallPhase.IDLE + version: Optional[str] = None + engine_dir: str + runtime_dir: str + python_bin: Optional[str] = None + created_at: Optional[str] = None + completed_at: Optional[str] = None + error_message: Optional[str] = None + last_log_line: Optional[str] = None + + +class EngineConnection(BaseModel): + id: str + name: str + engine_type: EngineType + ownership: EngineOwnership + endpoint_url: str + status: EngineStatus + native_ui_url: Optional[str] = None + version: Optional[str] = None + vram_free_mb: Optional[int] = None + capabilities: List[str] = Field(default_factory=list) + models_path: Optional[str] = None + last_heartbeat: Optional[float] = None + error_message: Optional[str] = None + + +class EngineConnectRequest(BaseModel): + engine_type: EngineType + endpoint_url: str + name: Optional[str] = None + + +class EngineUpdateStatus(str, Enum): + IDLE = "idle" + CHECKING = "checking" + UPDATING = "updating" + COMPLETED = "completed" + FAILED = "failed" + ROLLED_BACK = "rolled_back" + + +class EngineUpdateManifest(BaseModel): + engine_type: EngineType + status: EngineUpdateStatus = EngineUpdateStatus.IDLE + previous_commit: Optional[str] = None + target_commit: Optional[str] = None + updated_at: Optional[str] = None + error_message: Optional[str] = None + rollback_supported: bool = True + rollback_performed: bool = False + + +class LauncherConfig(BaseModel): + stop_managed_engines_on_exit: bool = False + default_engine: str = "cloud" + browser_auto_open: bool = True + port: int = 8000 + + +class ManagerStatusResponse(BaseModel): + app_name: str = "Berry AI Studio" + version: str = "0.1.0" + pid: int + uptime_seconds: float + port: int + frontend_packaged: bool + managed_comfyui: EngineConnection + managed_webui: EngineConnection + external_engines: List[EngineConnection] + cloud_providers_configured: int + models_indexed: int + active_tasks: int + launcher_config: LauncherConfig + + +class ShutdownRequest(BaseModel): + force: bool = False + stop_managed_engines: Optional[bool] = None + + +class ShutdownResponse(BaseModel): + status: str + message: str + active_tasks_cancelled: int + managed_engines_stopped: List[str] + diff --git a/backend/app/schemas/events.py b/backend/app/schemas/events.py index 8a36221..794d0d3 100644 --- a/backend/app/schemas/events.py +++ b/backend/app/schemas/events.py @@ -7,7 +7,8 @@ class NodeStatusEvent(BaseModel): type: Literal["NODE_STATUS"] = "NODE_STATUS" node_id: str - status: Literal["idle", "queued", "running", "cached", "completed", "error"] + status: Literal["idle", "queued", "running", "cached", "completed", "error", "cancelled"] + run_id: Optional[str] = None class NodeProgressEvent(BaseModel): @@ -15,29 +16,41 @@ class NodeProgressEvent(BaseModel): node_id: str progress: float # 0.0 – 1.0 message: Optional[str] = None + run_id: Optional[str] = None class NodeOutputEvent(BaseModel): type: Literal["NODE_OUTPUT"] = "NODE_OUTPUT" node_id: str output: Dict[str, Any] + run_id: Optional[str] = None class NodeErrorEvent(BaseModel): type: Literal["NODE_ERROR"] = "NODE_ERROR" node_id: str message: str + run_id: Optional[str] = None class GraphStartedEvent(BaseModel): type: Literal["GRAPH_STARTED"] = "GRAPH_STARTED" total_nodes: int cached_nodes: int + run_id: Optional[str] = None class GraphFinishedEvent(BaseModel): type: Literal["GRAPH_FINISHED"] = "GRAPH_FINISHED" execution_time_ms: float + status: str = "completed" # "completed", "cancelled", "failed" + run_id: Optional[str] = None + + +class RunCancelledEvent(BaseModel): + type: Literal["RUN_CANCELLED"] = "RUN_CANCELLED" + run_id: str + message: str = "Execution run cancelled by user" WorkflowEvent = Union[ @@ -47,4 +60,5 @@ class GraphFinishedEvent(BaseModel): NodeErrorEvent, GraphStartedEvent, GraphFinishedEvent, + RunCancelledEvent, ] diff --git a/backend/app/schemas/hardware.py b/backend/app/schemas/hardware.py new file mode 100644 index 0000000..e49a24e --- /dev/null +++ b/backend/app/schemas/hardware.py @@ -0,0 +1,33 @@ +"""Hardware readiness and storage diagnostic schemas.""" + +from typing import List, Optional +from pydantic import BaseModel, Field + + +class GpuInfo(BaseModel): + index: int + name: str + vram_total_mb: int + vram_free_mb: int + driver_version: Optional[str] = None + cuda_version: Optional[str] = None + temperature_c: Optional[int] = None + + +class StorageInfo(BaseModel): + path: str + total_gb: float + free_gb: float + used_gb: float + is_sufficient: bool = True # True if >= required threshold (e.g. 15GB) + + +class HardwareReadiness(BaseModel): + has_nvidia_gpu: bool + gpus: List[GpuInfo] = Field(default_factory=list) + engine_storage: Optional[StorageInfo] = None + model_storage: Optional[StorageInfo] = None + ready_for_local_inference: bool + recommended_engine: Optional[str] = None # "comfyui", "webui", or "cloud_only" + summary_message: str + guidance_notes: List[str] = Field(default_factory=list) diff --git a/backend/app/schemas/model.py b/backend/app/schemas/model.py new file mode 100644 index 0000000..fe730e2 --- /dev/null +++ b/backend/app/schemas/model.py @@ -0,0 +1,59 @@ +"""Model inventory, categories, architectures, and roots schemas.""" + +from enum import Enum +from typing import Any, Dict, List, Optional +from pydantic import BaseModel, Field + + +class ModelCategory(str, Enum): + CHECKPOINT = "checkpoint" + LORA = "lora" + VAE = "vae" + CONTROLNET = "controlnet" + UPSCALER = "upscaler" + TEXT_ENCODER = "text_encoder" + UNKNOWN = "unknown" + + +class ModelArchitecture(str, Enum): + SD15 = "sd15" + SD21 = "sd21" + SDXL = "sdxl" + SD3 = "sd3" + FLUX = "flux" + ESRGAN = "esrgan" + UNKNOWN = "unknown" + + +class ModelRecord(BaseModel): + id: str + name: str + file_path: str + category: ModelCategory + architecture: ModelArchitecture + format: str # safetensors, ckpt, pth, bin, gguf + size_bytes: int + size_mb: float + content_hash: Optional[str] = None + engine_compatibility: List[str] = Field(default_factory=list) # ["comfyui", "webui"] + is_ready: bool = True + missing_dependencies: List[str] = Field(default_factory=list) + guidance: Optional[str] = None + metadata: Dict[str, Any] = Field(default_factory=dict) + + +class ModelRoot(BaseModel): + id: str + path: str + label: str + engine_type: Optional[str] = None # "comfyui", "webui", or "custom" + is_active: bool = True + exists: bool = True + models_found: int = 0 + + +class ModelRootCreate(BaseModel): + path: str + label: str + engine_type: Optional[str] = "custom" + diff --git a/backend/app/schemas/project.py b/backend/app/schemas/project.py new file mode 100644 index 0000000..a93430d --- /dev/null +++ b/backend/app/schemas/project.py @@ -0,0 +1,38 @@ +"""Pydantic schemas for projects and managed assets.""" + +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional +from pydantic import BaseModel, Field + + +class ProjectBase(BaseModel): + name: str = Field(..., description="Project display name") + + +class ProjectCreate(ProjectBase): + pass + + +class ProjectUpdate(BaseModel): + name: Optional[str] = None + canvas: Optional[Dict[str, Any]] = None + + +class Project(ProjectBase): + id: str = Field(..., description="Unique project UUID") + canvas: Dict[str, Any] = Field(default_factory=lambda: {"nodes": [], "edges": [], "viewport": {"x": 0, "y": 0, "zoom": 1}}) + created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + updated_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + + +class AssetRecord(BaseModel): + id: str = Field(..., description="Unique asset identifier") + project_id: Optional[str] = Field(default=None, description="Associated project ID if any") + filename: str = Field(..., description="Original filename") + file_path: str = Field(..., description="Relative or absolute path on disk") + media_type: str = Field(default="image", description="Asset media type: image, text, audio, etc.") + content_hash: str = Field(..., description="SHA-256 content hash") + byte_size: int = Field(default=0, description="Size in bytes") + width: Optional[int] = None + height: Optional[int] = None + created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) diff --git a/backend/app/schemas/task.py b/backend/app/schemas/task.py new file mode 100644 index 0000000..c3998a6 --- /dev/null +++ b/backend/app/schemas/task.py @@ -0,0 +1,47 @@ +"""Task and Run models representing immutable submissions and truthful states.""" + +from datetime import datetime, timezone +from typing import Any, Dict, Literal, Optional +from pydantic import BaseModel, Field + +from app.schemas.workflow import WorkflowGraph + +TaskStatus = Literal[ + "queued", + "running", + "cached", + "succeeded", + "failed", + "cancel-requested", + "cancelled", +] + + +class RunRecord(BaseModel): + id: str = Field(..., description="Unique run identifier") + project_id: Optional[str] = None + target_node_id: Optional[str] = None + status: TaskStatus = "queued" + created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + finished_at: Optional[str] = None + + +class TaskRecord(BaseModel): + id: str = Field(..., description="Unique task identifier") + run_id: str = Field(..., description="Parent run ID") + node_id: str = Field(..., description="Graph node ID") + node_type: str = Field(..., description="Type of node executed") + status: TaskStatus = "queued" + params: Dict[str, Any] = Field(default_factory=dict) + inputs: Dict[str, Any] = Field(default_factory=dict) + outputs: Dict[str, Any] = Field(default_factory=dict) + error: Optional[str] = None + created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + finished_at: Optional[str] = None + + +class WorkflowRunRequest(BaseModel): + graph: WorkflowGraph + target_node_id: Optional[str] = None + run_id: Optional[str] = None + project_id: Optional[str] = None diff --git a/backend/app/storage/__init__.py b/backend/app/storage/__init__.py new file mode 100644 index 0000000..377aa46 --- /dev/null +++ b/backend/app/storage/__init__.py @@ -0,0 +1 @@ +"""Persistent storage module for Berry AI Studio.""" diff --git a/backend/app/storage/asset_store.py b/backend/app/storage/asset_store.py new file mode 100644 index 0000000..96fe8e0 --- /dev/null +++ b/backend/app/storage/asset_store.py @@ -0,0 +1,172 @@ +"""Content-addressable asset storage for durable generated media.""" + +import hashlib +import os +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import List, Optional +import httpx + +from app.schemas.project import AssetRecord +from app.storage.db import DatabaseManager, db_manager, get_default_data_dir + + +class AssetStore: + """Manages physical files and metadata for media assets.""" + + def __init__(self, manager: DatabaseManager = db_manager, base_dir: Optional[Path] = None) -> None: + self.manager = manager + self._base_dir = base_dir + + @property + def assets_dir(self) -> Path: + if self._base_dir: + d = self._base_dir / "assets" + else: + d = get_default_data_dir() / "assets" + d.mkdir(parents=True, exist_ok=True) + return d + + async def save_bytes( + self, + data: bytes, + filename: str, + media_type: str = "image", + project_id: Optional[str] = None, + ) -> AssetRecord: + content_hash = hashlib.sha256(data).hexdigest() + byte_size = len(data) + + # Content-addressable subfolder structure: assets/{hash[:2]}/{hash}{ext} + ext = Path(filename).suffix or ".png" + subdir = self.assets_dir / content_hash[:2] + subdir.mkdir(parents=True, exist_ok=True) + target_path = subdir / f"{content_hash}{ext}" + + if not target_path.exists(): + target_path.write_bytes(data) + + asset_id = str(uuid.uuid4()) + now = datetime.now(timezone.utc).isoformat() + rel_path = str(target_path.relative_to(self.assets_dir)) + + conn = await self.manager.get_connection() + await conn.execute( + """ + INSERT INTO assets (id, project_id, filename, file_path, media_type, content_hash, byte_size, width, height, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + (asset_id, project_id, filename, rel_path, media_type, content_hash, byte_size, None, None, now), + ) + await conn.commit() + + return AssetRecord( + id=asset_id, + project_id=project_id, + filename=filename, + file_path=rel_path, + media_type=media_type, + content_hash=content_hash, + byte_size=byte_size, + created_at=now, + ) + + async def save_image_from_url( + self, + url: str, + filename: Optional[str] = None, + project_id: Optional[str] = None, + timeout: float = 30.0, + ) -> AssetRecord: + """Download remote image and persist into managed asset store.""" + async with httpx.AsyncClient(timeout=timeout) as client: + resp = await client.get(url) + resp.raise_for_status() + data = resp.content + + name = filename or Path(url.split("?")[0]).name or f"image_{uuid.uuid4().hex[:8]}.png" + return await self.save_bytes(data=data, filename=name, media_type="image", project_id=project_id) + + async def get_asset(self, asset_id: str) -> Optional[AssetRecord]: + conn = await self.manager.get_connection() + async with conn.execute( + """ + SELECT id, project_id, filename, file_path, media_type, content_hash, byte_size, width, height, created_at + FROM assets WHERE id = ? + """, + (asset_id,), + ) as cursor: + row = await cursor.fetchone() + if not row: + return None + return AssetRecord( + id=row["id"], + project_id=row["project_id"], + filename=row["filename"], + file_path=row["file_path"], + media_type=row["media_type"], + content_hash=row["content_hash"], + byte_size=row["byte_size"], + width=row["width"], + height=row["height"], + created_at=row["created_at"], + ) + + async def get_asset_by_hash(self, content_hash: str) -> Optional[AssetRecord]: + conn = await self.manager.get_connection() + async with conn.execute( + """ + SELECT id, project_id, filename, file_path, media_type, content_hash, byte_size, width, height, created_at + FROM assets WHERE content_hash = ? ORDER BY created_at DESC LIMIT 1 + """, + (content_hash,), + ) as cursor: + row = await cursor.fetchone() + if not row: + return None + return AssetRecord( + id=row["id"], + project_id=row["project_id"], + filename=row["filename"], + file_path=row["file_path"], + media_type=row["media_type"], + content_hash=row["content_hash"], + byte_size=row["byte_size"], + width=row["width"], + height=row["height"], + created_at=row["created_at"], + ) + + def get_absolute_path(self, asset: AssetRecord) -> Path: + return self.assets_dir / asset.file_path + + async def list_assets(self, project_id: Optional[str] = None) -> List[AssetRecord]: + conn = await self.manager.get_connection() + if project_id: + query = "SELECT * FROM assets WHERE project_id = ? ORDER BY created_at DESC" + params = (project_id,) + else: + query = "SELECT * FROM assets ORDER BY created_at DESC" + params = () + + async with conn.execute(query, params) as cursor: + rows = await cursor.fetchall() + return [ + AssetRecord( + id=row["id"], + project_id=row["project_id"], + filename=row["filename"], + file_path=row["file_path"], + media_type=row["media_type"], + content_hash=row["content_hash"], + byte_size=row["byte_size"], + width=row["width"], + height=row["height"], + created_at=row["created_at"], + ) + for row in rows + ] + + +asset_store = AssetStore() diff --git a/backend/app/storage/db.py b/backend/app/storage/db.py new file mode 100644 index 0000000..c701dcd --- /dev/null +++ b/backend/app/storage/db.py @@ -0,0 +1,125 @@ +"""SQLite database connection and schema management.""" + +import os +import sys +from pathlib import Path +from typing import Optional +import aiosqlite + + +def get_default_data_dir() -> Path: + """Return root application data directory for database and managed assets.""" + custom_dir = os.environ.get("BERRY_DATA_DIR") or os.environ.get("AI_WORKFLOW_DATA_DIR") + if custom_dir: + return Path(custom_dir).resolve() + + if sys.platform == "win32": + local_app_data = os.environ.get("LOCALAPPDATA", str(Path.home() / "AppData" / "Local")) + return Path(local_app_data) / "BerryAIStudio" + + return Path.home() / ".berry-ai" + + +class DatabaseManager: + """Manages SQLite connection lifecycle and schema migrations.""" + + def __init__(self, db_path: Optional[str | Path] = None) -> None: + self.db_path = db_path + self._conn: Optional[aiosqlite.Connection] = None + + def get_effective_db_path(self) -> Path | str: + if self.db_path: + return self.db_path + data_dir = get_default_data_dir() + data_dir.mkdir(parents=True, exist_ok=True) + return data_dir / "berry.db" + + async def get_connection(self) -> aiosqlite.Connection: + if self._conn is None: + path = self.get_effective_db_path() + if isinstance(path, Path): + path.parent.mkdir(parents=True, exist_ok=True) + path_str = str(path) + else: + path_str = str(path) + self._conn = await aiosqlite.connect(path_str) + self._conn.row_factory = aiosqlite.Row + await self._init_schema(self._conn) + return self._conn + + async def close(self) -> None: + if self._conn is not None: + await self._conn.close() + self._conn = None + + async def _init_schema(self, conn: aiosqlite.Connection) -> None: + """Create versioned tables if not existing.""" + await conn.execute("PRAGMA journal_mode=WAL;") + await conn.execute("PRAGMA foreign_keys=ON;") + + await conn.executescript( + """ + CREATE TABLE IF NOT EXISTS projects ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + canvas_json TEXT NOT NULL, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS assets ( + id TEXT PRIMARY KEY, + project_id TEXT, + filename TEXT NOT NULL, + file_path TEXT NOT NULL, + media_type TEXT NOT NULL, + content_hash TEXT NOT NULL, + byte_size INTEGER DEFAULT 0, + width INTEGER, + height INTEGER, + created_at TEXT NOT NULL, + FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL + ); + + CREATE INDEX IF NOT EXISTS idx_assets_hash ON assets(content_hash); + CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id); + + CREATE TABLE IF NOT EXISTS runs ( + id TEXT PRIMARY KEY, + project_id TEXT, + target_node_id TEXT, + status TEXT NOT NULL, + created_at TEXT NOT NULL, + finished_at TEXT + ); + + CREATE TABLE IF NOT EXISTS tasks ( + id TEXT PRIMARY KEY, + run_id TEXT NOT NULL, + node_id TEXT NOT NULL, + node_type TEXT NOT NULL, + status TEXT NOT NULL, + params_json TEXT, + inputs_json TEXT, + outputs_json TEXT, + error_msg TEXT, + created_at TEXT NOT NULL, + finished_at TEXT, + FOREIGN KEY (run_id) REFERENCES runs(id) ON DELETE CASCADE + ); + + CREATE INDEX IF NOT EXISTS idx_tasks_run ON tasks(run_id); + + CREATE TABLE IF NOT EXISTS cache_entries ( + node_hash TEXT PRIMARY KEY, + output_json TEXT NOT NULL, + asset_ids_json TEXT, + created_at TEXT NOT NULL + ); + """ + ) + await conn.commit() + + +# Global database manager instance +db_manager = DatabaseManager() diff --git a/backend/app/storage/model_store.py b/backend/app/storage/model_store.py new file mode 100644 index 0000000..66b79bd --- /dev/null +++ b/backend/app/storage/model_store.py @@ -0,0 +1,279 @@ +""" +Model Catalog and High-Performance File Scanner. + +Discovers safetensors, ckpt, pth, and gguf models in registered root directories. +Parses .safetensors header metadata using standard library binary structs without +importing PyTorch or loading weight tensors into memory. +""" + +import asyncio +import hashlib +import json +import logging +import os +import struct +from pathlib import Path +from typing import Any, Dict, List, Optional, Set, Tuple + +from app.runtime.supervisor import get_default_engine_dir +from app.schemas.model import ( + ModelArchitecture, + ModelCategory, + ModelRecord, + ModelRoot, +) + +logger = logging.getLogger(__name__) + +SUPPORTED_EXTENSIONS = {".safetensors", ".ckpt", ".pth", ".pt", ".bin", ".gguf"} + + +def parse_safetensors_header(file_path: Path) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Read only the header of a .safetensors file without loading tensor data. + Safetensors format: + [0..7]: 8-byte uint64 little-endian representing header size N + [8..8+N]: UTF-8 JSON header containing metadata and tensor specs + """ + try: + with open(file_path, "rb") as f: + length_bytes = f.read(8) + if len(length_bytes) < 8: + return {}, {} + header_len = struct.unpack(" 100 * 1024 * 1024: # sanity limit 100MB + return {}, {} + header_bytes = f.read(header_len) + header_json = json.loads(header_bytes.decode("utf-8", errors="replace")) + + metadata = header_json.get("__metadata__", {}) + return header_json, metadata + except Exception as e: + logger.debug(f"Failed to parse safetensors header for {file_path}: {e}") + return {}, {} + + +def detect_architecture( + file_path: Path, header: Dict[str, Any], metadata: Dict[str, Any] +) -> ModelArchitecture: + """Infer model architecture from safetensors metadata, tensor keys, or filename.""" + # 1. Check explicit metadata + arch_meta = metadata.get("modelspec.architecture", "").lower() + if "flux" in arch_meta: + return ModelArchitecture.FLUX + if "xl" in arch_meta or "sdxl" in arch_meta: + return ModelArchitecture.SDXL + if "sd-1" in arch_meta or "sd1" in arch_meta: + return ModelArchitecture.SD15 + if "sd3" in arch_meta: + return ModelArchitecture.SD3 + + # 2. Check tensor key patterns + keys = list(header.keys()) + if any("double_blocks" in k for k in keys): + return ModelArchitecture.FLUX + if any("joint_blocks" in k for k in keys): + return ModelArchitecture.SD3 + + # Check cross-attention projection dimension in UNet + # SDXL: 2048, SD 1.5: 768, SD 2.1: 1024 + for k, v in header.items(): + if isinstance(v, dict) and "shape" in v: + shape = v.get("shape", []) + if "attn2.to_k.weight" in k: + if len(shape) >= 2 and shape[1] == 2048: + return ModelArchitecture.SDXL + elif len(shape) >= 2 and shape[1] == 768: + return ModelArchitecture.SD15 + elif len(shape) >= 2 and shape[1] == 1024: + return ModelArchitecture.SD21 + + # 3. Fallback to path / filename keywords + name_lower = file_path.name.lower() + if "flux" in name_lower: + return ModelArchitecture.FLUX + if "xl" in name_lower or "sdxl" in name_lower: + return ModelArchitecture.SDXL + if "v1-5" in name_lower or "sd15" in name_lower or "sd1.5" in name_lower or "1.5" in name_lower: + return ModelArchitecture.SD15 + if "esrgan" in name_lower or "ultrasharp" in name_lower: + return ModelArchitecture.ESRGAN + + return ModelArchitecture.UNKNOWN + + +def detect_category(file_path: Path, header: Dict[str, Any]) -> ModelCategory: + """Classify model into Checkpoint, LoRA, VAE, ControlNet, or Upscaler.""" + path_str = str(file_path).lower() + if "loras" in path_str or "lora" in path_str: + return ModelCategory.LORA + if "vae" in path_str: + return ModelCategory.VAE + if "controlnet" in path_str: + return ModelCategory.CONTROLNET + if "upscale" in path_str: + return ModelCategory.UPSCALER + if "checkpoints" in path_str: + return ModelCategory.CHECKPOINT + + # Inspect tensor structure + keys = list(header.keys()) + if any(k.startswith("lora_") or "lora_up" in k or "lora_down" in k for k in keys): + return ModelCategory.LORA + if all(k.startswith("first_stage_model.") or k.startswith("encoder.") or k.startswith("decoder.") for k in keys if k != "__metadata__"): + return ModelCategory.VAE + if any("diffusion_model" in k for k in keys): + return ModelCategory.CHECKPOINT + + return ModelCategory.UNKNOWN + + +class ModelStore: + """Manages model roots, directory scans, and model inventory.""" + + def __init__(self, engine_dir: Optional[Path] = None) -> None: + self.engine_dir = engine_dir or get_default_engine_dir() + self.roots: Dict[str, ModelRoot] = {} + self._cached_records: Dict[str, Tuple[float, int, ModelRecord]] = {} # path -> (mtime, size, record) + self._init_default_roots() + + def _init_default_roots(self) -> None: + """Register default model directories for ComfyUI and WebUI.""" + default_models = self.engine_dir / "models" + comfy_models = self.engine_dir / "comfyui" / "models" + webui_models = self.engine_dir / "webui" / "models" + + self.add_root("default_engine", str(default_models), "Berry Engine Models", engine_type="comfyui") + if comfy_models.exists(): + self.add_root("comfy_internal", str(comfy_models), "ComfyUI Models", engine_type="comfyui") + if webui_models.exists(): + self.add_root("webui_internal", str(webui_models), "SD WebUI Models", engine_type="webui") + + def add_root(self, root_id: str, path_str: str, label: str, engine_type: Optional[str] = None) -> ModelRoot: + """Register a directory to scan for models.""" + p = Path(path_str).resolve() + root = ModelRoot( + id=root_id, + path=str(p), + label=label, + engine_type=engine_type, + exists=p.is_dir(), + ) + self.roots[root_id] = root + return root + + def list_roots(self) -> List[ModelRoot]: + """Return all registered model scan roots with updated existence checks.""" + for root in self.roots.values(): + root.exists = Path(root.path).is_dir() + return list(self.roots.values()) + + def _scan_file(self, file_path: Path) -> Optional[ModelRecord]: + """Scan a single model file using header inspection.""" + try: + stat = file_path.stat() + mtime = stat.st_mtime + size_bytes = stat.st_size + path_str = str(file_path) + + # Check cache + if path_str in self._cached_records: + cached_mtime, cached_size, record = self._cached_records[path_str] + if cached_mtime == mtime and cached_size == size_bytes: + return record + + ext = file_path.suffix.lower() + header: Dict[str, Any] = {} + metadata: Dict[str, Any] = {} + + if ext == ".safetensors": + header, metadata = parse_safetensors_header(file_path) + + category = detect_category(file_path, header) + architecture = detect_architecture(file_path, header, metadata) + + # Engine compatibility + engine_compat = ["comfyui"] + if ext in {".safetensors", ".ckpt"}: + engine_compat.append("webui") + + # Missing dependencies / guidance + dependencies: List[str] = [] + guidance: Optional[str] = None + + if architecture == ModelArchitecture.FLUX: + dependencies = ["clip_l", "t5xxl", "ae (vae)"] + guidance = "Flux Schnell/Dev requires separate CLIP/T5 text encoders and VAE if unbundled." + elif architecture == ModelArchitecture.SDXL: + guidance = "SDXL model operates best with 1024x1024 resolution." + + record_id = hashlib.sha256(f"{path_str}:{size_bytes}".encode("utf-8")).hexdigest()[:16] + size_mb = round(size_bytes / (1024 * 1024), 2) + + record = ModelRecord( + id=record_id, + name=file_path.stem, + file_path=path_str, + category=category, + architecture=architecture, + format=ext.lstrip("."), + size_bytes=size_bytes, + size_mb=size_mb, + engine_compatibility=engine_compat, + is_ready=len(dependencies) == 0, + missing_dependencies=dependencies, + guidance=guidance, + metadata={k: v for k, v in metadata.items() if isinstance(v, (str, int, float, bool))}, + ) + + self._cached_records[path_str] = (mtime, size_bytes, record) + return record + + except Exception as e: + logger.warning(f"Error scanning model file {file_path}: {e}") + return None + + def scan_all_roots(self) -> List[ModelRecord]: + """Scan all active root directories and return discovered models.""" + discovered: List[ModelRecord] = [] + for root in self.roots.values(): + root_path = Path(root.path) + if not root_path.is_dir(): + root.exists = False + root.models_found = 0 + continue + + root.exists = True + count = 0 + for root_dir, _, files in os.walk(root_path): + for f in files: + ext = Path(f).suffix.lower() + if ext in SUPPORTED_EXTENSIONS: + f_path = Path(root_dir) / f + rec = self._scan_file(f_path) + if rec: + discovered.append(rec) + count += 1 + root.models_found = count + + return discovered + + def list_models(self) -> List[ModelRecord]: + """Return list of all models found in registered roots.""" + return self.scan_all_roots() + + def remove_root(self, root_id: str) -> bool: + """Remove a registered model scan root without deleting files (L10).""" + if root_id in self.roots: + del self.roots[root_id] + return True + return False + + async def scan_all_roots_async(self) -> List[ModelRecord]: + """Asynchronously scan all model roots.""" + return await asyncio.to_thread(self.scan_all_roots) + + +# Global model store singleton +model_store = ModelStore() + diff --git a/backend/app/storage/project_store.py b/backend/app/storage/project_store.py new file mode 100644 index 0000000..9c80881 --- /dev/null +++ b/backend/app/storage/project_store.py @@ -0,0 +1,112 @@ +"""Durable project and canvas storage management.""" + +import json +import uuid +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from app.schemas.project import Project, ProjectCreate, ProjectUpdate +from app.storage.db import DatabaseManager, db_manager + + +class ProjectStore: + """Provides asynchronous CRUD for user projects and canvases.""" + + def __init__(self, manager: DatabaseManager = db_manager) -> None: + self.manager = manager + + async def create_project(self, data: ProjectCreate) -> Project: + conn = await self.manager.get_connection() + proj_id = str(uuid.uuid4()) + now = datetime.now(timezone.utc).isoformat() + default_canvas = { + "nodes": [], + "edges": [], + "viewport": {"x": 0, "y": 0, "zoom": 1}, + } + await conn.execute( + """ + INSERT INTO projects (id, name, canvas_json, created_at, updated_at) + VALUES (?, ?, ?, ?, ?) + """, + (proj_id, data.name, json.dumps(default_canvas), now, now), + ) + await conn.commit() + return Project( + id=proj_id, + name=data.name, + canvas=default_canvas, + created_at=now, + updated_at=now, + ) + + async def get_project(self, project_id: str) -> Optional[Project]: + conn = await self.manager.get_connection() + async with conn.execute( + "SELECT id, name, canvas_json, created_at, updated_at FROM projects WHERE id = ?", + (project_id,), + ) as cursor: + row = await cursor.fetchone() + if not row: + return None + return Project( + id=row["id"], + name=row["name"], + canvas=json.loads(row["canvas_json"]), + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + async def list_projects(self) -> List[Project]: + conn = await self.manager.get_connection() + async with conn.execute( + "SELECT id, name, canvas_json, created_at, updated_at FROM projects ORDER BY updated_at DESC" + ) as cursor: + rows = await cursor.fetchall() + return [ + Project( + id=row["id"], + name=row["name"], + canvas=json.loads(row["canvas_json"]), + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + for row in rows + ] + + async def update_project(self, project_id: str, update: ProjectUpdate) -> Optional[Project]: + conn = await self.manager.get_connection() + existing = await self.get_project(project_id) + if not existing: + return None + + new_name = update.name if update.name is not None else existing.name + new_canvas = update.canvas if update.canvas is not None else existing.canvas + now = datetime.now(timezone.utc).isoformat() + + await conn.execute( + """ + UPDATE projects + SET name = ?, canvas_json = ?, updated_at = ? + WHERE id = ? + """, + (new_name, json.dumps(new_canvas), now, project_id), + ) + await conn.commit() + + return Project( + id=project_id, + name=new_name, + canvas=new_canvas, + created_at=existing.created_at, + updated_at=now, + ) + + async def delete_project(self, project_id: str) -> bool: + conn = await self.manager.get_connection() + cursor = await conn.execute("DELETE FROM projects WHERE id = ?", (project_id,)) + await conn.commit() + return cursor.rowcount > 0 + + +project_store = ProjectStore() diff --git a/backend/requirements.txt b/backend/requirements.txt index f365014..0432abf 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -4,3 +4,5 @@ pydantic>=2.7.0 httpx>=0.27.0 pytest>=8.2.0 pytest-asyncio>=0.23.0 +aiosqlite>=0.20.0 +python-multipart>=0.0.9 diff --git a/backend/tests/__init__.py b/backend/tests/__init__.py index d177290..a534504 100644 --- a/backend/tests/__init__.py +++ b/backend/tests/__init__.py @@ -1 +1 @@ -"""Test suite for AI-Workflow.""" +"""Test suite for Berry AI Studio.""" diff --git a/backend/tests/test_health.py b/backend/tests/test_health.py index f3d57c6..6d2a9e7 100644 --- a/backend/tests/test_health.py +++ b/backend/tests/test_health.py @@ -16,5 +16,5 @@ def test_system_info(): response = client.get("/api/v1/info") assert response.status_code == 200 data = response.json() - assert data["name"] == "AI-Workflow" + assert data["name"] == "Berry AI Studio" assert "runners" in data diff --git a/backend/tests/test_m1_core.py b/backend/tests/test_m1_core.py new file mode 100644 index 0000000..7371ef1 --- /dev/null +++ b/backend/tests/test_m1_core.py @@ -0,0 +1,271 @@ +""" +Comprehensive verification tests for Milestone 1: +- Targeted single-node execution +- Port-aware semantic caching +- Strict failure boundaries & dependent blocking +- output.preview runner execution +- Task cancellation semantics +- SQLite project & asset persistence +- ComfyUI real output polling +""" + +import json +import pytest +from unittest.mock import AsyncMock, patch +from fastapi.testclient import TestClient + +from app.main import app +from app.core.cache import cache_store, compute_content_hash, compute_semantic_node_hash +from app.core.dag import DAGResolver +from app.schemas.project import ProjectCreate, ProjectUpdate +from app.schemas.workflow import WorkflowEdgeInstance, WorkflowGraph, WorkflowNodeInstance +from app.storage.asset_store import asset_store +from app.storage.project_store import project_store +from app.runners.comfy_runner import ComfyUIClient, run_comfy_txt2img_node + +client = TestClient(app) + + +# --------------------------------------------------------------------------- +# 1. Targeted Single-Node Execution +# --------------------------------------------------------------------------- + +def test_targeted_node_execution_websocket(): + """Targeted execution only runs target node and its upstream ancestors, skipping unrelated nodes.""" + graph_payload = { + "target_node_id": "n2", + "run_id": "test_run_targeted", + "graph": { + "nodes": [ + {"id": "n1", "type": "input.text", "params": {"value": "hello berry"}}, + {"id": "n2", "type": "output.preview", "params": {}}, + {"id": "n3_unrelated", "type": "input.text", "params": {"value": "unrelated"}}, + ], + "edges": [ + {"id": "e1", "source": "n1", "source_handle": "text", "target": "n2", "target_handle": "media"} + ], + }, + } + + executed_nodes = [] + with client.websocket_connect("/ws/workflow/run") as ws: + ws.send_text(json.dumps(graph_payload)) + while True: + try: + raw = ws.receive_text() + event = json.loads(raw) + if event.get("type") == "NODE_OUTPUT": + executed_nodes.append(event.get("node_id")) + elif event.get("type") == "GRAPH_FINISHED": + break + except Exception: + break + + assert "n1" in executed_nodes + assert "n2" in executed_nodes + assert "n3_unrelated" not in executed_nodes + + +# --------------------------------------------------------------------------- +# 2. Port-Aware Semantic Caching +# --------------------------------------------------------------------------- + +def test_port_aware_cache_discrimination(): + """Swapping port connections produces distinct semantic hashes.""" + bindings_a = [ + ("input_a", "hash_content_1", "output_1"), + ("input_b", "hash_content_2", "output_2"), + ] + bindings_b = [ + ("input_a", "hash_content_2", "output_2"), + ("input_b", "hash_content_1", "output_1"), + ] + + h1 = compute_semantic_node_hash("image.combine", {"mode": "blend"}, bindings_a) + h2 = compute_semantic_node_hash("image.combine", {"mode": "blend"}, bindings_b) + + assert h1 != h2, "Swapping port inputs must yield distinct cache hashes" + + +def test_content_hash_invalidation(): + """Changing upstream output content changes hash and avoids stale cache hits.""" + h_text1 = compute_content_hash("sunset over ocean") + h_text2 = compute_content_hash("sunset over mountains") + + binding1 = [("prompt", h_text1, "text")] + binding2 = [("prompt", h_text2, "text")] + + node_hash1 = compute_semantic_node_hash("image.generate", {"model": "flux-schnell"}, binding1) + node_hash2 = compute_semantic_node_hash("image.generate", {"model": "flux-schnell"}, binding2) + + assert node_hash1 != node_hash2 + + +# --------------------------------------------------------------------------- +# 3. Failure Boundary & Downstream Blocking +# --------------------------------------------------------------------------- + +def test_failure_propagation_blocks_dependents(): + """When an upstream node fails, dependent nodes are cancelled and not executed.""" + graph_payload = { + "run_id": "test_run_fail", + "graph": { + "nodes": [ + {"id": "n1", "type": "unknown.invalid_type", "params": {}}, + {"id": "n2", "type": "output.preview", "params": {}}, + ], + "edges": [ + {"id": "e1", "source": "n1", "source_handle": "out", "target": "n2", "target_handle": "media"} + ], + }, + } + + statuses = {} + with client.websocket_connect("/ws/workflow/run") as ws: + ws.send_text(json.dumps(graph_payload)) + while True: + try: + raw = ws.receive_text() + event = json.loads(raw) + if event.get("type") == "NODE_STATUS": + statuses[event["node_id"]] = event["status"] + elif event.get("type") == "GRAPH_FINISHED": + assert event.get("status") == "failed" + break + except Exception: + break + + assert statuses.get("n1") == "error" + assert statuses.get("n2") == "cancelled" + + +# --------------------------------------------------------------------------- +# 4. output.preview Runner Execution +# --------------------------------------------------------------------------- + +def test_output_preview_runner(): + """output.preview executes cleanly and passes through image/text media.""" + graph_payload = { + "graph": { + "nodes": [ + {"id": "n1", "type": "input.text", "params": {"value": "rendered preview card text"}}, + {"id": "n2", "type": "output.preview", "params": {}}, + ], + "edges": [ + {"id": "e1", "source": "n1", "source_handle": "text", "target": "n2", "target_handle": "media"} + ], + } + } + + preview_output = None + with client.websocket_connect("/ws/workflow/run") as ws: + ws.send_text(json.dumps(graph_payload)) + while True: + try: + raw = ws.receive_text() + event = json.loads(raw) + if event.get("type") == "NODE_OUTPUT" and event.get("node_id") == "n2": + preview_output = event.get("output") + elif event.get("type") == "GRAPH_FINISHED": + break + except Exception: + break + + assert preview_output is not None + assert preview_output.get("result") == "rendered preview card text" + + +# --------------------------------------------------------------------------- +# 5. Task Cancellation Semantics +# --------------------------------------------------------------------------- + +def test_cancellation_endpoint(): + """POST /api/v1/workflow/cancel/{run_id} requests truthful cancellation.""" + # Test nonexistent run + res = client.post("/api/v1/workflow/cancel/nonexistent_run") + assert res.status_code == 200 + assert res.json()["success"] is False + + +# --------------------------------------------------------------------------- +# 6. SQLite Project and Asset Persistence +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_project_crud_persistence(): + """Projects and canvas state are saved and restored accurately across calls.""" + # Create project + proj = await project_store.create_project(ProjectCreate(name="My First Artwork")) + assert proj.id is not None + assert proj.name == "My First Artwork" + + # Update canvas + updated_canvas = { + "nodes": [{"id": "node_1", "type": "input.text"}], + "edges": [], + "viewport": {"x": 50, "y": 100, "zoom": 1.5}, + } + updated = await project_store.update_project(proj.id, ProjectUpdate(canvas=updated_canvas)) + assert updated is not None + assert updated.canvas["viewport"]["zoom"] == 1.5 + + # Retrieve from DB + loaded = await project_store.get_project(proj.id) + assert loaded is not None + assert loaded.name == "My First Artwork" + assert len(loaded.canvas["nodes"]) == 1 + + # REST endpoint test + res = client.get(f"/api/v1/projects/{proj.id}") + assert res.status_code == 200 + assert res.json()["id"] == proj.id + + +@pytest.mark.asyncio +async def test_asset_store_content_addressable(): + """Assets are hashed and stored with content-addressable deduplication.""" + fake_png_data = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + asset = await asset_store.save_bytes(fake_png_data, filename="test_pixel.png") + + assert asset.id is not None + assert asset.byte_size == len(fake_png_data) + assert len(asset.content_hash) == 64 + + # Fetch by ID + loaded_asset = await asset_store.get_asset(asset.id) + assert loaded_asset is not None + assert loaded_asset.content_hash == asset.content_hash + + # Verify physical file exists on disk + abs_path = asset_store.get_absolute_path(loaded_asset) + assert abs_path.is_file() + assert abs_path.read_bytes() == fake_png_data + + +# --------------------------------------------------------------------------- +# 7. ComfyUI Output Polling +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_comfy_history_polling(): + """ComfyUI client polls /history/{prompt_id} and extracts real output images.""" + client_instance = ComfyUIClient(host="127.0.0.1", port=8188) + + mock_history_payload = { + "sample_prompt_123": { + "status": {"status_str": "success", "completed": True, "messages": []}, + "outputs": { + "9": { + "images": [ + {"filename": "Berry_output_001.png", "subfolder": "", "type": "output"} + ] + } + }, + } + } + + with patch.object(client_instance, "get_history", new_callable=AsyncMock) as mock_get_hist: + mock_get_hist.return_value = mock_history_payload + images = await client_instance.poll_history_outputs("sample_prompt_123", max_wait=2.0) + assert len(images) == 1 + assert images[0]["filename"] == "Berry_output_001.png" diff --git a/backend/tests/test_m2_engines.py b/backend/tests/test_m2_engines.py new file mode 100644 index 0000000..d27c5b6 --- /dev/null +++ b/backend/tests/test_m2_engines.py @@ -0,0 +1,246 @@ +""" +Milestone 2 Verification Tests: Engines, Isolation, Hardware, and Models. + +Verifies: +1. Strict removal of host Python fallback in ComfySupervisor. +2. WebUISupervisor isolated lifecycle and entrypoint checking. +3. Staged installer manifest tracking and interruption recovery. +4. External engine connection with zero process ownership. +5. Hardware and filesystem storage readiness diagnostics. +6. Pure standard-library safetensors header parsing and architecture inference. +7. REST API endpoints for M2 engines, hardware, and model catalog. +""" + +import json +import struct +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +from fastapi.testclient import TestClient + +from app.main import app +from app.runtime.engine_manager import EngineManager +from app.runtime.hardware import check_hardware_readiness, detect_gpus, get_storage_readiness +from app.runtime.installer import IsolatedEngineInstaller +from app.runtime.supervisor import ComfySupervisor +from app.runtime.webui_supervisor import WebUISupervisor +from app.schemas.engine import EngineOwnership, EngineStatus, EngineType, InstallPhase +from app.schemas.model import ModelArchitecture, ModelCategory +from app.storage.model_store import ( + ModelStore, + detect_architecture, + detect_category, + parse_safetensors_header, +) + +client = TestClient(app) + + +def test_supervisor_no_host_python_fallback(): + """Verify ComfySupervisor strictly refuses to launch using host Python.""" + with tempfile.TemporaryDirectory() as tmpdir: + engine_path = Path(tmpdir) / "engine" + comfy_dir = engine_path / "comfyui" + comfy_dir.mkdir(parents=True) + (comfy_dir / "main.py").write_text("# entrypoint", encoding="utf-8") + + sup = ComfySupervisor(engine_dir=engine_path) + assert sup.is_installed() is True + assert sup.has_isolated_env() is False + + # Attempt start without isolated venv + res = sup.start() + assert res["success"] is False + assert res.get("code") == "ENV_MISSING" + assert "Sandboxed virtual environment not found" in res["message"] + # Ensure no process was launched + assert sup.is_running() is False + + +def test_webui_supervisor_isolated_lifecycle(): + """Verify WebUISupervisor isolated paths and refusal without isolated venv.""" + with tempfile.TemporaryDirectory() as tmpdir: + engine_path = Path(tmpdir) / "engine" + sup = WebUISupervisor(engine_dir=engine_path) + assert sup.is_installed() is False + + # Create mock webui.py + webui_dir = engine_path / "webui" + webui_dir.mkdir(parents=True) + (webui_dir / "launch.py").write_text("# mock webui entrypoint", encoding="utf-8") + + assert sup.is_installed() is True + assert sup.has_isolated_env() is False + + res = sup.start() + assert res["success"] is False + assert res.get("code") == "ENV_MISSING" + assert sup.is_running() is False + + +def test_installer_manifest_and_interruption(): + """Verify installer tracks phases and recovers safely from interrupted installs.""" + with tempfile.TemporaryDirectory() as tmpdir: + engine_path = Path(tmpdir) / "engine" + installer_inst = IsolatedEngineInstaller(engine_dir=engine_path) + + # 1. Read initial manifest + manifest = installer_inst.read_manifest(EngineType.COMFYUI) + assert manifest.phase == InstallPhase.IDLE + assert manifest.engine_type == EngineType.COMFYUI + + # 2. Simulate an interrupted install state + manifest.phase = InstallPhase.DOWNLOADING + installer_inst.write_manifest(manifest) + + # 3. Reading manifest again should detect interruption + reloaded = installer_inst.read_manifest(EngineType.COMFYUI) + assert reloaded.phase == InstallPhase.INTERRUPTED + assert "interrupted" in (reloaded.error_message or "").lower() + + +@pytest.mark.asyncio +async def test_external_engine_connection_zero_ownership(): + """Verify external engine connection captures stats and never assumes process ownership.""" + mgr = EngineManager() + + # Mock ComfyUI /system_stats response + mock_stats = { + "system": {"os": "nt", "python_version": "3.12"}, + "devices": [{"name": "NVIDIA GeForce RTX 4090", "vram_total": 24000000000, "vram_free": 18000000000}], + } + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = mock_stats + + with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: + mock_get.return_value = mock_resp + + conn = await mgr.connect_external_engine( + engine_type=EngineType.COMFYUI, + endpoint_url="http://127.0.0.1:8188", + name="My Local ComfyUI", + ) + + assert conn.ownership == EngineOwnership.EXTERNAL + assert conn.status == EngineStatus.READY + assert conn.endpoint_url == "http://127.0.0.1:8188" + assert conn.vram_free_mb is not None + assert conn.vram_free_mb > 10000 + + +def test_hardware_storage_readiness(): + """Verify storage readiness calculations against thresholds.""" + with tempfile.TemporaryDirectory() as tmpdir: + storage = get_storage_readiness(Path(tmpdir), min_free_gb=0.001) + assert storage.total_gb > 0 + assert storage.free_gb > 0 + assert storage.is_sufficient is True + + # Unrealistic threshold should report insufficient + storage_high = get_storage_readiness(Path(tmpdir), min_free_gb=1000000.0) + assert storage_high.is_sufficient is False + + +def create_synthetic_safetensors(file_path: Path, header_dict: dict) -> None: + """Helper to create a valid minimal .safetensors binary file with given header.""" + header_json = json.dumps(header_dict).encode("utf-8") + header_len = len(header_json) + with open(file_path, "wb") as f: + f.write(struct.pack("= 2 + + sdxl_rec = next((r for r in records if r.name == "sd_xl_base_1.0"), None) + assert sdxl_rec is not None + assert sdxl_rec.architecture == ModelArchitecture.SDXL + assert sdxl_rec.category == ModelCategory.CHECKPOINT + assert "comfyui" in sdxl_rec.engine_compatibility + assert "webui" in sdxl_rec.engine_compatibility + + +def test_api_m2_endpoints(): + """Verify new REST endpoints for hardware readiness, engines, and models.""" + # 1. Hardware readiness + resp = client.get("/api/v1/hardware/readiness") + assert resp.status_code == 200 + hw = resp.json() + assert "has_nvidia_gpu" in hw + assert "ready_for_local_inference" in hw + assert "summary_message" in hw + + # 2. Engines list + resp = client.get("/api/v1/engines") + assert resp.status_code == 200 + engines = resp.json() + assert isinstance(engines, list) + assert any(e["id"] == "managed_comfyui" for e in engines) + assert any(e["id"] == "managed_webui" for e in engines) + + # 3. Model roots + resp = client.get("/api/v1/models/roots") + assert resp.status_code == 200 + roots = resp.json() + assert isinstance(roots, list) + + # 4. Engine manifest + resp = client.get("/api/v1/runtime/comfyui/manifest") + assert resp.status_code == 200 + manifest = resp.json() + assert manifest["engine_type"] == "comfyui" + assert "phase" in manifest diff --git a/backend/tests/test_m3_creative.py b/backend/tests/test_m3_creative.py new file mode 100644 index 0000000..189d09a --- /dev/null +++ b/backend/tests/test_m3_creative.py @@ -0,0 +1,206 @@ +""" +Milestone 3 Verification Tests: Creative Image Actions, Canvas Mappings & Provenance. + +Verifies: +1. ComfyUI macro graph compilation for txt2img, img2img, inpaint, and upscale. +2. SD WebUI runner API mappings and base64 asset adoption. +3. High-level CreativeRunner caching, provenance tracking, and variant invalidation. +4. Canvas asset upload and content-addressable storage. +5. FastAPI creative REST endpoints. +""" + +import base64 +import io +import uuid +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +from fastapi.testclient import TestClient + +from app.main import app +from app.runners.creative_runner import CreativeRunner, compute_creative_cache_hash +from app.runners.macro_compiler import ( + build_comfy_img2img_graph, + build_comfy_inpaint_graph, + build_comfy_txt2img_graph, + build_comfy_upscale_graph, +) +from app.runners.webui_runner import WebUIRunner +from app.schemas.creative import CreativeActionRequest, CreativeActionType +from app.storage.asset_store import asset_store + +client = TestClient(app) + +DUMMY_PNG_BYTES = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x06\x00\x00\x00\x1f\x15c4\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4" + b"\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def test_comfy_macro_compiler_all_actions(): + """Verify macro compiler produces valid ComfyUI prompt graphs for all 4 image actions.""" + # 1. txt2img + t2i = build_comfy_txt2img_graph(prompt="A cozy cottage", aspect_ratio="16:9") + assert "4" in t2i and t2i["4"]["class_type"] == "CheckpointLoaderSimple" + assert "3" in t2i and t2i["3"]["class_type"] == "KSampler" + assert "9" in t2i and t2i["9"]["class_type"] == "SaveImage" + assert t2i["5"]["inputs"]["width"] == 1344 + + # 2. img2img + i2i = build_comfy_img2img_graph(prompt="sunset oil painting", image_filename="input.png", denoise=0.6) + assert "1" in i2i and i2i["1"]["class_type"] == "LoadImage" + assert "2" in i2i and i2i["2"]["class_type"] == "VAEEncode" + assert i2i["3"]["inputs"]["denoise"] == 0.6 + + # 3. inpaint + inp = build_comfy_inpaint_graph( + prompt="red flower", image_filename="input.png", mask_filename="mask.png" + ) + assert "5" in inp and inp["5"]["class_type"] == "VAEEncodeForInpaint" + assert "2" in inp and inp["2"]["class_type"] == "LoadImage" + + # 4. upscale + ups = build_comfy_upscale_graph(image_filename="input.png", upscaler_model="RealESRGAN_x4plus.pth") + assert "2" in ups and ups["2"]["class_type"] == "UpscaleModelLoader" + assert "3" in ups and ups["3"]["class_type"] == "ImageUpscaleWithModel" + + +@pytest.mark.asyncio +async def test_webui_runner_txt2img_and_upscale(): + """Verify SD WebUI runner dispatches requests and saves base64 images into asset store.""" + runner = WebUIRunner(endpoint_url="http://127.0.0.1:7860") + b64_dummy = base64.b64encode(DUMMY_PNG_BYTES).decode("utf-8") + + # 1. Test txt2img + with patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post: + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"images": [b64_dummy]} + mock_post.return_value = mock_resp + + req = CreativeActionRequest( + action=CreativeActionType.TXT2IMG, + prompt="futuristic city", + engine_id="managed_webui", + width=512, + height=512, + ) + res = await runner.execute_action(req) + assert res["asset_id"] is not None + assert res["image_url"].startswith("/api/v1/assets/") + + # Verify asset was actually written to disk + asset = await asset_store.get_asset(res["asset_id"]) + assert asset is not None + assert asset.byte_size > 0 + + # 2. Test upscale + with patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post: + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"image": b64_dummy} + mock_post.return_value = mock_resp + + upscale_req = CreativeActionRequest( + action=CreativeActionType.UPSCALE, + input_image_id=res["asset_id"], + engine_id="managed_webui", + upscale_factor=2.0, + width=512, + height=512, + ) + ups_res = await runner.execute_action(upscale_req) + assert ups_res["width"] == 1024 + assert ups_res["height"] == 1024 + + +@pytest.mark.asyncio +async def test_creative_runner_caching_and_provenance(): + """Verify CreativeRunner attaches provenance and deterministically caches results.""" + runner = CreativeRunner() + b64_dummy = base64.b64encode(DUMMY_PNG_BYTES).decode("utf-8") + + with patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post: + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"images": [b64_dummy]} + mock_post.return_value = mock_resp + + unique_prompt = f"cyberpunk portrait {uuid.uuid4()}" + req = CreativeActionRequest( + action=CreativeActionType.TXT2IMG, + prompt=unique_prompt, + engine_id="managed_webui", + seed=42, + steps=25, + aspect_ratio="1:1", + ) + + # 1. First execution: should compute and cache + res1 = await runner.execute(req) + assert res1.success is True + assert res1.is_cached is False + assert res1.provenance is not None + assert res1.provenance.seed == 42 + assert res1.provenance.steps == 25 + assert res1.provenance.model == "v1-5-pruned-emaonly.safetensors" + + # 2. Second identical execution: should be a deterministic CACHE HIT + res2 = await runner.execute(req) + assert res2.success is True + assert res2.is_cached is True + assert res2.asset_id == res1.asset_id + + # 3. Changed seed (new variation): should NOT hit cache + req_variant = req.model_copy(update={"seed": 43}) + res3 = await runner.execute(req_variant) + assert res3.success is True + assert res3.is_cached is False + assert res3.provenance.seed == 43 + + +def test_creative_asset_upload_endpoint(): + """Verify uploading an image to the canvas stores it as an asset record.""" + file_content = io.BytesIO(DUMMY_PNG_BYTES) + resp = client.post( + "/api/v1/creative/upload", + files={"file": ("canvas_upload.png", file_content, "image/png")}, + ) + assert resp.status_code == 200 + asset_data = resp.json() + assert "id" in asset_data + assert asset_data["filename"] == "canvas_upload.png" + assert asset_data["content_hash"] is not None + + # Test reading it back + content_resp = client.get(f"/api/v1/assets/{asset_data['id']}/content") + assert content_resp.status_code == 200 + assert len(content_resp.content) == len(DUMMY_PNG_BYTES) + + +def test_api_creative_execute_endpoint(): + """Verify FastAPI /api/v1/creative/execute endpoint.""" + b64_dummy = base64.b64encode(DUMMY_PNG_BYTES).decode("utf-8") + + with patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post: + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"images": [b64_dummy]} + mock_post.return_value = mock_resp + + payload = { + "action": "txt2img", + "prompt": "serene mountain lake", + "engine_id": "managed_webui", + "aspect_ratio": "1:1", + "steps": 20, + "seed": 999, + } + resp = client.post("/api/v1/creative/execute", json=payload) + assert resp.status_code == 200 + data = resp.json() + assert data["success"] is True + assert "asset_id" in data + assert "provenance" in data + assert data["provenance"]["prompt"] == "serene mountain lake" diff --git a/backend/tests/test_m4_cloud.py b/backend/tests/test_m4_cloud.py new file mode 100644 index 0000000..d7a6a57 --- /dev/null +++ b/backend/tests/test_m4_cloud.py @@ -0,0 +1,174 @@ +""" +Milestone 4 Verification Tests: BYOK Cloud Credentials, Providers & Cloud Execution. + +Verifies: +1. CredentialManager local persistence and strict secret redaction. +2. Cloud key validation against external endpoints. +3. CreativeRunner cloud-only image generation with asset adoption and caching (zero GPU/torch). +4. FastAPI cloud REST endpoints. +""" + +import tempfile +import uuid +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +from fastapi.testclient import TestClient + +from app.main import app +from app.runners.creative_runner import CreativeRunner +from app.runtime.credentials import CredentialManager, redact_key +from app.schemas.cloud import CloudProviderId +from app.schemas.creative import CreativeActionRequest, CreativeActionType +from app.storage.asset_store import asset_store + +client = TestClient(app) + +DUMMY_PNG_BYTES = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x06\x00\x00\x00\x1f\x15c4\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4" + b"\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def test_credential_redaction_and_persistence(): + """Verify secrets are strictly redacted in provider listings and persist safely.""" + with tempfile.TemporaryDirectory() as tmpdir: + mgr = CredentialManager(data_dir=Path(tmpdir)) + + # Check redact utility + assert redact_key("sk-proj-1234567890abcdef") == "sk-...cdef" + assert redact_key("short") == "****" + assert redact_key("") == "" + + # Set secret + mgr.set_key(CloudProviderId.OPENAI, "sk-proj-mysecretkey12345678") + assert mgr.get_key(CloudProviderId.OPENAI) == "sk-proj-mysecretkey12345678" + + # List providers: ensure raw key is NEVER present in the info object + providers = mgr.list_providers() + oai = next(p for p in providers if p.id == CloudProviderId.OPENAI) + assert oai.is_configured is True + assert oai.redacted_key == "sk-...5678" + assert "mysecretkey" not in oai.redacted_key + + # Re-load in a new instance: verify persistence + mgr2 = CredentialManager(data_dir=Path(tmpdir)) + assert mgr2.get_key(CloudProviderId.OPENAI) == "sk-proj-mysecretkey12345678" + + # Delete key + mgr2.delete_key(CloudProviderId.OPENAI) + assert mgr2.get_key(CloudProviderId.OPENAI) is None + + +@pytest.mark.asyncio +async def test_cloud_key_validation_mocked(): + """Verify live key validation tests against provider endpoints.""" + with tempfile.TemporaryDirectory() as tmpdir: + mgr = CredentialManager(data_dir=Path(tmpdir)) + + # Test valid key (200 OK) + with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_get.return_value = mock_resp + + res = await mgr.test_key(CloudProviderId.OPENAI, "sk-test-valid-key") + assert res.valid is True + assert res.status_code == 200 + assert "verified successfully" in res.message + + # Test invalid key (401 Unauthorized) + with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: + mock_resp = MagicMock() + mock_resp.status_code = 401 + mock_get.return_value = mock_resp + + res = await mgr.test_key(CloudProviderId.OPENAI, "sk-test-bad-key") + assert res.valid is False + assert res.status_code == 401 + assert "Invalid key" in res.message + + +@pytest.mark.asyncio +async def test_creative_runner_cloud_txt2img_adoption(): + """Verify cloud execution runs with zero GPU/torch and adopts remote URLs into asset store.""" + runner = CreativeRunner() + unique_prompt = f"ethereal nebula {uuid.uuid4()}" + + # Mock Fal.ai FLUX call + with patch("app.runners.creative_runner.credentials_manager.get_key") as mock_key, \ + patch("app.runners.creative_runner._call_fal_ai", new_callable=AsyncMock) as mock_fal, \ + patch("app.runners.creative_runner.asset_store.save_image_from_url", new_callable=AsyncMock) as mock_save: + + mock_key.return_value = "fal-test-key" + mock_fal.return_value = "https://fal.media/files/sample_nebula.png" + + # Mock asset store adoption + mock_asset = MagicMock() + mock_asset.id = "asset_cloud_nebula" + mock_asset.content_hash = "hash_nebula_123" + mock_save.return_value = mock_asset + + req = CreativeActionRequest( + action=CreativeActionType.TXT2IMG, + prompt=unique_prompt, + model="flux-schnell", + engine_id="cloud_fal", + aspect_ratio="1:1", + seed=101, + ) + + # 1. Execute cloud action + res = await runner.execute(req) + assert res.success is True + assert res.asset_id == "asset_cloud_nebula" + assert res.image_url == "/api/v1/assets/asset_cloud_nebula/content" + assert res.provenance is not None + assert res.provenance.model == "flux-schnell" + assert res.provenance.engine_id == "cloud_fal" + + # Verify asset adoption was called with remote URL + mock_save.assert_called_once_with("https://fal.media/files/sample_nebula.png", filename="cloud_output.png") + + +def test_api_cloud_endpoints(): + """Verify REST endpoints for cloud provider listing and credential management.""" + # 1. List providers + resp = client.get("/api/v1/cloud/providers") + assert resp.status_code == 200 + providers = resp.json() + assert isinstance(providers, list) + assert any(p["id"] == "openai" for p in providers) + assert any(p["id"] == "fal" for p in providers) + assert any(p["id"] == "siliconflow" for p in providers) + + # 2. Store credential + store_resp = client.post( + "/api/v1/cloud/credentials", + json={"provider_id": "fal", "api_key": "fal-test-secret-key-999"}, + ) + assert store_resp.status_code == 200 + fal_info = store_resp.json() + assert fal_info["is_configured"] is True + assert "999" in fal_info["redacted_key"] + assert "secret" not in fal_info["redacted_key"] + + # 3. Test credential endpoint + with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_get.return_value = mock_resp + + test_resp = client.post( + "/api/v1/cloud/credentials/test", + json={"provider_id": "fal"}, + ) + assert test_resp.status_code == 200 + test_data = test_resp.json() + assert test_data["valid"] is True + + # 4. Delete credential + del_resp = client.delete("/api/v1/cloud/credentials/fal") + assert del_resp.status_code == 200 + assert del_resp.json()["success"] is True diff --git a/backend/tests/test_m5_release.py b/backend/tests/test_m5_release.py new file mode 100644 index 0000000..db19a4a --- /dev/null +++ b/backend/tests/test_m5_release.py @@ -0,0 +1,206 @@ +""" +Milestone 5 Verification Tests: Release Scenarios, Robustness & Integrity. + +Tests all 7 PRD end-to-end release scenarios: +1. Clean setup and creative image workflow. +2. Existing engine connection without unintended mutation or ownership. +3. Cloud-only creation without local engines or PyTorch. +4. Failure recovery: invalid keys, missing models, disconnected engines. +5. Persistence and caching recovery across restarts. +6. Truthful cancellation semantics. +7. Model safety and non-destructive directory management. +""" + +import base64 +import tempfile +import uuid +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +from fastapi.testclient import TestClient + +from app.core.cache import CacheStore +from app.main import app +from app.runners.creative_runner import CreativeRunner +from app.runtime.credentials import CredentialManager +from app.runtime.engine_manager import EngineManager +from app.runtime.hardware import check_hardware_readiness +from app.runtime.installer import IsolatedEngineInstaller +from app.schemas.cloud import CloudProviderId +from app.schemas.creative import CreativeActionRequest, CreativeActionType +from app.schemas.engine import EngineOwnership, EngineStatus, EngineType, InstallPhase +from app.storage.asset_store import AssetStore +from app.storage.model_store import ModelStore + +client = TestClient(app) + +DUMMY_PNG = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x06\x00\x00\x00\x1f\x15c4\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4" + b"\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def test_scenario_1_clean_setup_and_readiness(): + """Scenario 1: Clean startup boots without torch and assesses device readiness.""" + resp = client.get("/health") + assert resp.status_code == 200 + assert resp.json()["status"] == "ok" + + hw_resp = client.get("/api/v1/hardware/readiness") + assert hw_resp.status_code == 200 + hw = hw_resp.json() + assert "has_nvidia_gpu" in hw + assert "ready_for_local_inference" in hw + assert isinstance(hw["guidance_notes"], list) + + +@pytest.mark.asyncio +async def test_scenario_2_existing_engine_connection_no_mutation(): + """Scenario 2: Connect existing external engine without process ownership or file mutation.""" + mgr = EngineManager() + mock_stats = { + "system": {"os": "nt"}, + "devices": [{"name": "RTX 4090", "vram_free": 16000000000}], + } + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = mock_stats + + with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: + mock_get.return_value = mock_resp + + conn = await mgr.connect_external_engine( + engine_type=EngineType.COMFYUI, + endpoint_url="http://127.0.0.1:8188", + name="User Pre-existing ComfyUI", + ) + assert conn.ownership == EngineOwnership.EXTERNAL + assert conn.status == EngineStatus.READY + assert conn.native_ui_url == "http://127.0.0.1:8188" + + # Verify external engines are not killed by manager + assert "managed" not in conn.id + + +@pytest.mark.asyncio +async def test_scenario_3_cloud_only_creation(): + """Scenario 3: Cloud-only generation runs without local engine or GPU requirements.""" + runner = CreativeRunner() + unique_prompt = f"crystal palace {uuid.uuid4()}" + + with patch("app.runners.creative_runner.credentials_manager.get_key") as mock_key, \ + patch("app.runners.creative_runner._call_openai_images", new_callable=AsyncMock) as mock_oai, \ + patch("app.runners.creative_runner.asset_store.save_image_from_url", new_callable=AsyncMock) as mock_save: + + mock_key.return_value = "sk-test-key" + mock_oai.return_value = "https://oaidalleapiprodscus.blob.core.windows.net/sample.png" + + mock_asset = MagicMock() + mock_asset.id = "asset_dalle_crystal" + mock_asset.content_hash = "hash_crystal_99" + mock_save.return_value = mock_asset + + req = CreativeActionRequest( + action=CreativeActionType.TXT2IMG, + prompt=unique_prompt, + model="dall-e-3", + engine_id="cloud_openai", + ) + + res = await runner.execute(req) + assert res.success is True + assert res.asset_id == "asset_dalle_crystal" + assert res.provenance.model == "dall-e-3" + assert res.provenance.engine_id == "cloud_openai" + + +def test_scenario_4_failure_recovery_actionable_states(): + """Scenario 4: Interrupted setup, invalid keys, and disconnected engines report actionable states.""" + with tempfile.TemporaryDirectory() as tmpdir: + # 1. Interrupted install recovery + inst = IsolatedEngineInstaller(engine_dir=Path(tmpdir)) + m = inst.read_manifest(EngineType.COMFYUI) + m.phase = InstallPhase.INSTALLING_DEPS + inst.write_manifest(m) + + reloaded = inst.read_manifest(EngineType.COMFYUI) + assert reloaded.phase == InstallPhase.INTERRUPTED + assert "interrupted" in reloaded.error_message.lower() + + # 2. Disconnected engine probe returns degraded/offline without crash + with patch("httpx.AsyncClient.get", new_callable=AsyncMock) as mock_get: + mock_get.side_effect = Exception("Connection refused") + mgr = EngineManager() + async def _test(): + return await mgr.connect_external_engine( + engine_type=EngineType.WEBUI, + endpoint_url="http://127.0.0.1:9999", + ) + import asyncio + conn = asyncio.run(_test()) + assert conn.status == EngineStatus.OFFLINE + assert "Connection refused" in (conn.error_message or "") + + +@pytest.mark.asyncio +async def test_scenario_5_persistence_and_caching_across_restarts(): + """Scenario 5: Restart preserves cache and deterministic reuse.""" + runner = CreativeRunner() + unique_prompt = f"tranquil forest {uuid.uuid4()}" + b64_dummy = base64.b64encode(DUMMY_PNG).decode("utf-8") + + with patch("httpx.AsyncClient.post", new_callable=AsyncMock) as mock_post: + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = {"images": [b64_dummy]} + mock_post.return_value = mock_resp + + req = CreativeActionRequest( + action=CreativeActionType.TXT2IMG, + prompt=unique_prompt, + engine_id="managed_webui", + seed=777, + ) + + res1 = await runner.execute(req) + assert res1.is_cached is False + + # Emulate restart: create a new CreativeRunner instance + runner_after_restart = CreativeRunner() + res2 = await runner_after_restart.execute(req) + assert res2.is_cached is True + assert res2.asset_id == res1.asset_id + + +def test_scenario_6_cancellation_semantics(): + """Scenario 6: Cancellation request produces truthful response.""" + # Test nonexistent run returns false + res = client.post("/api/v1/workflow/cancel/nonexistent_run_id") + assert res.status_code == 200 + data = res.json() + assert data["success"] is False + + +def test_scenario_7_model_safety_non_destructive(): + """Scenario 7: Model inspection reads only headers; removing root leaves files intact.""" + with tempfile.TemporaryDirectory() as tmpdir: + models_path = Path(tmpdir) / "checkpoints" + models_path.mkdir(parents=True) + test_file = models_path / "model.safetensors" + test_file.write_bytes(DUMMY_PNG) + + store = ModelStore(engine_dir=Path(tmpdir)) + root = store.add_root("custom_root", str(models_path), "My Checkpoints") + assert root.exists is True + + # Scan roots + records = store.scan_all_roots() + + # Remove root from store index + del store.roots["custom_root"] + + # Original file MUST remain intact on disk! + assert test_file.is_file() + assert test_file.read_bytes() == DUMMY_PNG diff --git a/backend/tests/test_m6_launcher_manager.py b/backend/tests/test_m6_launcher_manager.py new file mode 100644 index 0000000..ee69128 --- /dev/null +++ b/backend/tests/test_m6_launcher_manager.py @@ -0,0 +1,204 @@ +""" +Automated Test Suite for Milestone 6: Launcher and Environment Manager (L01-L12). +Verifies: +- Versioned manager protocol status aggregation (L04, L11) +- Launcher configuration endpoints (L07) +- Controlled shutdown with active task protection (L07) +- Engine update and rollback semantics (L08, L12) +- Updates check separating app vs engine updates (L08) +- Missing frontend packaging diagnostic page (L02) +""" + +import asyncio +import os +import shutil +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi.testclient import TestClient + +from app.main import active_cancellations, app, launcher_config +from app.runtime.installer import IsolatedEngineInstaller +from app.schemas.engine import EngineType, EngineUpdateStatus + + +@pytest.fixture +def client(): + return TestClient(app) + + +def test_manager_status_endpoint(client): + """L04, L11: Verify GET /api/v1/manager/status aggregates core, engines, cloud, models, and tasks.""" + resp = client.get("/api/v1/manager/status") + assert resp.status_code == 200 + data = resp.json() + + assert data["app_name"] == "Berry AI Studio" + assert data["version"] == "0.1.0" + assert data["pid"] == os.getpid() + assert data["uptime_seconds"] >= 0 + assert "port" in data + assert "frontend_packaged" in data + assert "managed_comfyui" in data + assert "managed_webui" in data + assert "external_engines" in data + assert "cloud_providers_configured" in data + assert "models_indexed" in data + assert "active_tasks" in data + assert "launcher_config" in data + + +def test_manager_config_crud(client): + """L07: Verify launcher configuration retrieval and mutation.""" + resp = client.get("/api/v1/manager/config") + assert resp.status_code == 200 + cfg = resp.json() + assert "stop_managed_engines_on_exit" in cfg + + # Update config + update_payload = { + "stop_managed_engines_on_exit": True, + "default_engine": "managed_comfyui", + "browser_auto_open": False, + "port": 8000, + } + resp_set = client.post("/api/v1/manager/config", json=update_payload) + assert resp_set.status_code == 200 + new_cfg = resp_set.json() + assert new_cfg["stop_managed_engines_on_exit"] is True + assert new_cfg["default_engine"] == "managed_comfyui" + + # Reset + client.post("/api/v1/manager/config", json={"stop_managed_engines_on_exit": False, "default_engine": "cloud", "browser_auto_open": True, "port": 8000}) + + +def test_manager_shutdown_guards_active_tasks(client): + """L07: Verify shutdown prevents silent termination of running tasks unless forced.""" + fake_run_id = "test-run-active-guard" + cancel_event = asyncio.Event() + active_cancellations[fake_run_id] = cancel_event + + try: + # Non-forced shutdown should fail with 409 + resp = client.post("/api/v1/manager/shutdown", json={"force": False}) + assert resp.status_code == 409 + assert "active generation task" in resp.json()["detail"] + assert not cancel_event.is_set() + + # Forced shutdown should succeed and signal cancellation + with patch("os._exit") as mock_exit: + resp_force = client.post("/api/v1/manager/shutdown", json={"force": True}) + assert resp_force.status_code == 200 + assert resp_force.json()["status"] == "shutting_down" + assert resp_force.json()["active_tasks_cancelled"] == 1 + assert cancel_event.is_set() + finally: + active_cancellations.pop(fake_run_id, None) + + +@pytest.mark.asyncio +async def test_engine_update_active_tasks_guard(): + """L12: Engine update must abort if active jobs are running.""" + fake_run_id = "test-update-active" + active_cancellations[fake_run_id] = asyncio.Event() + + installer = IsolatedEngineInstaller() + try: + manifest = await installer.update_engine( + EngineType.COMFYUI, + has_active_tasks_fn=lambda: len(active_cancellations) > 0, + ) + assert manifest.status == EngineUpdateStatus.FAILED + assert "active generation tasks" in manifest.error_message + finally: + active_cancellations.pop(fake_run_id, None) + + +@pytest.mark.asyncio +async def test_engine_update_rollback_semantics(): + """L08, L12: Interrupted or failed update must record state and support rollback.""" + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + installer = IsolatedEngineInstaller(engine_dir=tmp_path) + + # Create fake engine git repo + engine_target = tmp_path / "comfyui" + engine_target.mkdir(parents=True) + git_dir = engine_target / ".git" + git_dir.mkdir() + + # Mock subprocess to simulate git rev-parse, git pull failure, and git checkout rollback + async def mock_exec(*args, **kwargs): + cmd = args[0] + subcmd = args[1] if len(args) > 1 else "" + mock_proc = AsyncMock() + + if subcmd == "rev-parse": + mock_proc.returncode = 0 + mock_proc.communicate.return_value = (b"commit_abc123\n", b"") + elif subcmd == "pull": + # Simulate failure during pull + mock_proc.returncode = 1 + mock_proc.communicate.return_value = (b"", b"Network timeout during fetch") + elif subcmd == "checkout": + # Rollback checkout succeeds + mock_proc.returncode = 0 + mock_proc.communicate.return_value = (b"HEAD is now at commit_abc123", b"") + else: + mock_proc.returncode = 0 + mock_proc.communicate.return_value = (b"", b"") + + return mock_proc + + with patch("asyncio.create_subprocess_exec", side_effect=mock_exec): + manifest = await installer.update_engine(EngineType.COMFYUI) + assert manifest.status == EngineUpdateStatus.ROLLED_BACK + assert manifest.rollback_performed is True + assert manifest.previous_commit == "commit_abc123" + assert "Network timeout" in manifest.error_message + + +def test_updates_check_endpoint(client): + """L08: Verify app updates and engine updates are presented as separate operations.""" + resp = client.get("/api/v1/updates/check") + assert resp.status_code == 200 + data = resp.json() + + assert "app" in data + assert data["app"]["name"] == "Berry AI Studio" + assert "current_version" in data["app"] + assert "latest_version" in data["app"] + + assert "engines" in data + assert "comfyui" in data["engines"] + assert "webui" in data["engines"] + + +def test_missing_frontend_diagnostic_page(client): + """L02: Missing frontend must show an informative packaging error page, not a blank 404.""" + # Test request to root / + resp = client.get("/") + assert resp.status_code == 200 + # Either static index.html or diagnostic packaging error + content = resp.text + assert ("" in content or "" in content.lower()) + if "Packaging Error" in content: + assert "frontend/dist/index.html" in content + assert "pnpm build" in content + + +def test_app_update_trigger(client): + """L08: Triggering application update should handle git or packaged mode safely.""" + resp = client.post("/api/v1/updates/app") + assert resp.status_code == 200 + data = resp.json() + assert "mode" in data + assert "status" in data + assert "message" in data + if data["mode"] == "packaged": + assert "download_url" in data + assert "LOCALAPPDATA" in data.get("notes", "") + elif data["mode"] == "git": + assert data["status"] in ("updated", "failed", "error")