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")