mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
docs: add crash recovery and robustness to spec and plan
Atomic model downloads (.downloading suffix + rename), file-based install lock (survives container restart), atomic JSON writes, startup recovery sequence, frontend double-click prevention, SSE fallback polling, disk space pre-checks.
This commit is contained in:
ashim-hq
@@ -0,0 +1,712 @@
|
||||
# On-Demand AI Feature Downloads Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Reduce Docker image from ~30 GB to ~5-6 GB by making AI features downloadable post-install via a UI-driven bundle system.
|
||||
|
||||
**Architecture:** Six feature bundles (Background Removal, Face Detection, Object Eraser & Colorize, Upscale & Enhance, Photo Restoration, OCR) are defined in a JSON manifest baked into the image. A Python install script handles pip + model downloads to a persistent volume. The backend exposes install/uninstall APIs with SSE progress. The frontend shows download badges on uninstalled tools and an install prompt on tool pages.
|
||||
|
||||
**Tech Stack:** Fastify (API), Zustand (frontend state), Python (install script), Docker (image restructuring), SSE (progress streaming)
|
||||
|
||||
**Spec:** `docs/superpowers/specs/2026-04-17-on-demand-ai-features-design.md`
|
||||
|
||||
---
|
||||
|
||||
## File Map
|
||||
|
||||
```
|
||||
NEW FILES:
|
||||
packages/shared/src/features.ts # Bundle definitions, tool-to-bundle map, types
|
||||
docker/feature-manifest.json # Authoritative manifest baked into image
|
||||
apps/api/src/lib/feature-status.ts # Reads manifest + installed.json, provides status
|
||||
apps/api/src/routes/features.ts # GET /features, POST install/uninstall, GET disk-usage
|
||||
packages/ai/python/install_feature.py # Python install script (pip + model downloads)
|
||||
apps/web/src/stores/features-store.ts # Zustand store for bundle statuses
|
||||
apps/web/src/components/features/feature-install-prompt.tsx # Install prompt card for tool pages
|
||||
apps/web/src/components/settings/ai-features-section.tsx # Settings panel section
|
||||
tests/unit/features.test.ts # Unit tests for feature logic
|
||||
|
||||
MODIFIED FILES:
|
||||
packages/ai/src/bridge.ts # restartDispatcher(), FEATURE_NOT_INSTALLED handling
|
||||
packages/ai/src/index.ts # Export restartDispatcher
|
||||
packages/ai/python/dispatcher.py # Read installed.json, gate scripts by feature
|
||||
packages/ai/python/colorize.py # Hard imports to lazy imports
|
||||
packages/ai/python/restore.py # Hard imports to lazy imports
|
||||
apps/api/src/index.ts # Register feature routes, startup venv check
|
||||
apps/api/src/routes/tool-factory.ts # Feature-installed guard before process()
|
||||
apps/api/src/routes/batch.ts # Feature-installed check at gating point
|
||||
apps/api/src/routes/pipeline.ts # Feature-installed check in pre-validation
|
||||
apps/api/src/routes/tools/restore-photo.ts # Feature-installed guard
|
||||
apps/web/src/lib/api.ts # Extend parseApiError for FEATURE_NOT_INSTALLED
|
||||
apps/web/src/components/common/tool-card.tsx # Download badge on uninstalled AI tools
|
||||
apps/web/src/pages/tool-page.tsx # Feature check then install prompt or "not enabled"
|
||||
apps/web/src/components/layout/tool-panel.tsx # Fetch features on mount
|
||||
apps/web/src/pages/fullscreen-grid-page.tsx # Fetch features on mount
|
||||
apps/web/src/components/settings/settings-dialog.tsx # Add AI Features nav item + section
|
||||
docker/Dockerfile # Remove ML packages/models, keep base
|
||||
docker/entrypoint.sh # Venv bootstrap, /data/ai/ setup
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Shared Feature Types and Bundle Definitions
|
||||
|
||||
**Files:**
|
||||
- Create: `packages/shared/src/features.ts`
|
||||
- Modify: `packages/shared/src/index.ts`
|
||||
- Test: `tests/unit/features.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write the failing test for bundle definitions**
|
||||
|
||||
Create `tests/unit/features.test.ts`:
|
||||
|
||||
```ts
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
FEATURE_BUNDLES,
|
||||
getBundleForTool,
|
||||
getToolsForBundle,
|
||||
TOOL_BUNDLE_MAP,
|
||||
} from "@ashim/shared/features";
|
||||
import { PYTHON_SIDECAR_TOOLS } from "@ashim/shared";
|
||||
|
||||
describe("Feature bundles", () => {
|
||||
it("every PYTHON_SIDECAR_TOOL maps to exactly one bundle", () => {
|
||||
for (const toolId of PYTHON_SIDECAR_TOOLS) {
|
||||
const bundle = getBundleForTool(toolId);
|
||||
expect(bundle, `${toolId} has no bundle`).toBeDefined();
|
||||
}
|
||||
});
|
||||
|
||||
it("getBundleForTool returns null for non-AI tools", () => {
|
||||
expect(getBundleForTool("resize")).toBeNull();
|
||||
expect(getBundleForTool("crop")).toBeNull();
|
||||
});
|
||||
|
||||
it("getToolsForBundle returns correct tools", () => {
|
||||
const tools = getToolsForBundle("background-removal");
|
||||
expect(tools).toContain("remove-background");
|
||||
expect(tools).toContain("passport-photo");
|
||||
expect(tools).not.toContain("upscale");
|
||||
});
|
||||
|
||||
it("all 6 bundles are defined", () => {
|
||||
expect(Object.keys(FEATURE_BUNDLES)).toHaveLength(6);
|
||||
expect(FEATURE_BUNDLES["background-removal"]).toBeDefined();
|
||||
expect(FEATURE_BUNDLES["face-detection"]).toBeDefined();
|
||||
expect(FEATURE_BUNDLES["object-eraser-colorize"]).toBeDefined();
|
||||
expect(FEATURE_BUNDLES["upscale-enhance"]).toBeDefined();
|
||||
expect(FEATURE_BUNDLES["photo-restoration"]).toBeDefined();
|
||||
expect(FEATURE_BUNDLES["ocr"]).toBeDefined();
|
||||
});
|
||||
|
||||
it("TOOL_BUNDLE_MAP covers all sidecar tools", () => {
|
||||
const mappedTools = Object.keys(TOOL_BUNDLE_MAP);
|
||||
for (const toolId of PYTHON_SIDECAR_TOOLS) {
|
||||
expect(mappedTools, `${toolId} missing from TOOL_BUNDLE_MAP`).toContain(toolId);
|
||||
}
|
||||
});
|
||||
});
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run test to verify it fails**
|
||||
|
||||
Run: `pnpm test:unit -- tests/unit/features.test.ts`
|
||||
Expected: FAIL with module not found error.
|
||||
|
||||
- [ ] **Step 3: Create the feature definitions module**
|
||||
|
||||
Create `packages/shared/src/features.ts`:
|
||||
|
||||
```ts
|
||||
export interface FeatureBundleInfo {
|
||||
id: string;
|
||||
name: string;
|
||||
description: string;
|
||||
estimatedSize: string;
|
||||
enablesTools: string[];
|
||||
}
|
||||
|
||||
export type FeatureStatus = "not_installed" | "installing" | "installed" | "error";
|
||||
|
||||
export interface FeatureBundleState {
|
||||
id: string;
|
||||
name: string;
|
||||
description: string;
|
||||
status: FeatureStatus;
|
||||
installedVersion: string | null;
|
||||
estimatedSize: string;
|
||||
enablesTools: string[];
|
||||
progress: { percent: number; stage: string } | null;
|
||||
error: string | null;
|
||||
}
|
||||
|
||||
export const FEATURE_BUNDLES: Record<string, FeatureBundleInfo> = {
|
||||
"background-removal": {
|
||||
id: "background-removal",
|
||||
name: "Background Removal",
|
||||
description: "Remove image backgrounds with AI",
|
||||
estimatedSize: "700 MB - 1 GB",
|
||||
enablesTools: ["remove-background", "passport-photo"],
|
||||
},
|
||||
"face-detection": {
|
||||
id: "face-detection",
|
||||
name: "Face Detection",
|
||||
description: "Detect and blur faces, fix red-eye, smart crop",
|
||||
estimatedSize: "200-300 MB",
|
||||
enablesTools: ["blur-faces", "red-eye-removal", "smart-crop"],
|
||||
},
|
||||
"object-eraser-colorize": {
|
||||
id: "object-eraser-colorize",
|
||||
name: "Object Eraser & Colorize",
|
||||
description: "Erase objects from photos and colorize B&W images",
|
||||
estimatedSize: "600-800 MB",
|
||||
enablesTools: ["erase-object", "colorize"],
|
||||
},
|
||||
"upscale-enhance": {
|
||||
id: "upscale-enhance",
|
||||
name: "Upscale & Enhance",
|
||||
description: "AI upscaling, face enhancement, and noise removal",
|
||||
estimatedSize: "4-5 GB",
|
||||
enablesTools: ["upscale", "enhance-faces", "noise-removal"],
|
||||
},
|
||||
"photo-restoration": {
|
||||
id: "photo-restoration",
|
||||
name: "Photo Restoration",
|
||||
description: "Restore old or damaged photos",
|
||||
estimatedSize: "800 MB - 1 GB",
|
||||
enablesTools: ["restore-photo"],
|
||||
},
|
||||
ocr: {
|
||||
id: "ocr",
|
||||
name: "OCR",
|
||||
description: "Extract text from images",
|
||||
estimatedSize: "3-4 GB",
|
||||
enablesTools: ["ocr"],
|
||||
},
|
||||
};
|
||||
|
||||
export const TOOL_BUNDLE_MAP: Record<string, string> = {};
|
||||
for (const [bundleId, bundle] of Object.entries(FEATURE_BUNDLES)) {
|
||||
for (const toolId of bundle.enablesTools) {
|
||||
TOOL_BUNDLE_MAP[toolId] = bundleId;
|
||||
}
|
||||
}
|
||||
|
||||
export function getBundleForTool(toolId: string): FeatureBundleInfo | null {
|
||||
const bundleId = TOOL_BUNDLE_MAP[toolId];
|
||||
return bundleId ? FEATURE_BUNDLES[bundleId] : null;
|
||||
}
|
||||
|
||||
export function getToolsForBundle(bundleId: string): string[] {
|
||||
return FEATURE_BUNDLES[bundleId]?.enablesTools ?? [];
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Export from shared package**
|
||||
|
||||
Add to the end of `packages/shared/src/index.ts`:
|
||||
|
||||
```ts
|
||||
export * from "./features.js";
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Run test to verify it passes**
|
||||
|
||||
Run: `pnpm test:unit -- tests/unit/features.test.ts`
|
||||
Expected: PASS, all 5 tests green.
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add packages/shared/src/features.ts packages/shared/src/index.ts tests/unit/features.test.ts
|
||||
git commit -m "feat: add shared feature bundle definitions and tool-to-bundle mapping"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Feature Manifest File
|
||||
|
||||
**Files:**
|
||||
- Create: `docker/feature-manifest.json`
|
||||
|
||||
- [ ] **Step 1: Create the feature manifest**
|
||||
|
||||
Create `docker/feature-manifest.json` containing the full bundle definitions with exact package versions, pip flags, platform-specific packages, and model download URLs. Source exact versions from the current Dockerfile (lines 167-206) and model URLs from `docker/download_models.py`.
|
||||
|
||||
Key details: amd64 uses `--extra-index-url https://download.pytorch.org/whl/cu126` for torch/realesrgan; amd64 uses `paddlepaddle-gpu>=3.2.1` from `https://www.paddlepaddle.org.cn/packages/stable/cu126/`; arm64 uses `mediapipe==0.10.18`; `codeformer-pip==0.0.4` needs `--no-deps`; `postInstall` re-pins `numpy==1.26.4`.
|
||||
|
||||
The file should contain a top-level `manifestVersion`, `imageVersion`, `pythonVersion`, `basePackages` array, and `bundles` object with all 6 bundles. Each bundle has `name`, `description`, `estimatedSize`, `packages` (with `common`/`amd64`/`arm64` arrays), `pipFlags`, `postInstall`, `models` array, and `enablesTools` array.
|
||||
|
||||
Model entries use either: `{ "id", "url", "path", "minSize" }` for direct downloads, `{ "id", "downloadFn": "rembg_session", "args": [...] }` for rembg models, or `{ "id", "downloadFn": "hf_snapshot", "args": [repo_id, local_subpath] }` for HuggingFace snapshots.
|
||||
|
||||
- [ ] **Step 2: Commit**
|
||||
|
||||
```bash
|
||||
git add docker/feature-manifest.json
|
||||
git commit -m "feat: add feature manifest with all 6 bundle definitions"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Backend Feature Status Service
|
||||
|
||||
**Files:**
|
||||
- Create: `apps/api/src/lib/feature-status.ts`
|
||||
|
||||
- [ ] **Step 1: Create the feature status service**
|
||||
|
||||
Create `apps/api/src/lib/feature-status.ts`. This module reads/writes `/data/ai/installed.json`, provides `isFeatureInstalled(bundleId)`, `isToolInstalled(toolId)`, `getFeatureStates()`, `markInstalled()`, `markUninstalled()`, `setInstallProgress()`, and `ensureAiDirs()`.
|
||||
|
||||
Uses `FEATURE_BUNDLES` and `TOOL_BUNDLE_MAP` from `@ashim/shared`. Caches `installed.json` in memory with `invalidateCache()` for refresh after install/uninstall. Detects Docker environment via `existsSync("/.dockerenv")`.
|
||||
|
||||
See spec section "Persistent Storage" for directory structure: `/data/ai/venv/`, `/data/ai/models/`, `/data/ai/pip-cache/`, `/data/ai/installed.json`.
|
||||
|
||||
**Robustness requirements for this module:**
|
||||
|
||||
- **Atomic JSON writes:** `markInstalled()` and `markUninstalled()` must write to `installed.json.tmp` first, then `renameSync()` to `installed.json`. Never write directly to `installed.json`.
|
||||
- **Corrupt JSON recovery:** `readInstalled()` wraps `JSON.parse` in try/catch. If the file is corrupt, treat as empty `{ bundles: {} }` and log a warning.
|
||||
- **File-based install lock:** Instead of just in-memory `installInProgress`, use `/data/ai/install.lock` file containing `{ bundleId, startedAt, pid }`. Create lock before install, delete on completion/failure. `getInstallingBundle()` reads from the lock file, not memory.
|
||||
- **`recoverInterruptedInstalls()`** function called on startup:
|
||||
1. Delete any `*.downloading` files in `/data/ai/models/` (recursive glob)
|
||||
2. Delete `installed.json.tmp` if it exists
|
||||
3. Delete `/data/ai/venv.bootstrapping/` if it exists
|
||||
4. If `install.lock` exists: check if PID is alive (via `process.kill(pid, 0)` in try/catch). If dead, delete the lock and log a warning. If alive, leave it (install is still running from a previous container lifecycle — unlikely but possible with shared volumes).
|
||||
5. For each bundle in `installed.json`, verify model files exist and meet `minSize` from the feature manifest. If any model is missing/undersized, set the bundle's error field to "Some model files are missing. Reinstall this feature." but do NOT remove from installed.json.
|
||||
- **`acquireInstallLock(bundleId)`** and **`releaseInstallLock()`** functions that create/delete the lock file atomically.
|
||||
|
||||
- [ ] **Step 2: Commit**
|
||||
|
||||
```bash
|
||||
git add apps/api/src/lib/feature-status.ts
|
||||
git commit -m "feat: add backend feature status service for tracking installed bundles"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Feature API Routes
|
||||
|
||||
**Files:**
|
||||
- Create: `apps/api/src/routes/features.ts`
|
||||
- Modify: `apps/api/src/index.ts`
|
||||
|
||||
- [ ] **Step 1: Create the features route file**
|
||||
|
||||
Create `apps/api/src/routes/features.ts` with 4 endpoints:
|
||||
|
||||
1. `GET /api/v1/features` (any authenticated user) — returns `{ bundles: FeatureBundleState[] }`. In non-Docker environments, returns all features as installed.
|
||||
2. `POST /api/v1/admin/features/:bundleId/install` (admin only) — validates bundle exists, checks not already installed, checks no other install in progress (409). Spawns `install_feature.py` as child process via `spawn()`. Parses stderr JSON progress lines, updates progress via `updateSingleFileProgress()` from `progress.ts`. On success, calls `invalidateCache()` and `shutdownDispatcher()` (from `@ashim/ai`). Returns `{ jobId }`.
|
||||
3. `POST /api/v1/admin/features/:bundleId/uninstall` (admin only) — removes model files listed in the manifest, calls `markUninstalled()`, calls `shutdownDispatcher()`. Returns `{ ok: true }`.
|
||||
4. `GET /api/v1/admin/features/disk-usage` (admin only) — returns `{ totalBytes }` by recursively sizing `/data/ai/`.
|
||||
|
||||
Note: Use `spawn()` from `node:child_process` (not `exec()`) for the install script to avoid shell injection. Pass arguments as array elements.
|
||||
|
||||
**Robustness requirements for install endpoint:**
|
||||
- Call `acquireInstallLock(bundleId)` before spawning the child process. If lock acquisition fails (lock file already exists with a live PID), return 409.
|
||||
- Check available disk space before starting: `const { availableParallelism } = require("node:os"); const stats = statfsSync("/data"); const freeBytes = stats.bfree * stats.bsize;`. Compare against a rough estimate for the bundle. If insufficient, return 400 with disk space info.
|
||||
- On child process `close` event with code 0: call `releaseInstallLock()`, `invalidateCache()`, `shutdownDispatcher()`.
|
||||
- On child process `close` event with non-zero code: call `releaseInstallLock()`, set error state. Do NOT leave the lock file behind.
|
||||
- On child process `error` event (spawn failure): call `releaseInstallLock()`, return error.
|
||||
- The install endpoint returns `{ jobId }` immediately. The child process runs asynchronously. The HTTP response does not block on completion.
|
||||
|
||||
- [ ] **Step 2: Register feature routes in index.ts**
|
||||
|
||||
In `apps/api/src/index.ts`: import `registerFeatureRoutes`, call it after the settings routes registration. Also import and call `ensureAiDirs()` and `recoverInterruptedInstalls()` near the top of the startup sequence after `runMigrations()`.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add apps/api/src/routes/features.ts apps/api/src/index.ts
|
||||
git commit -m "feat: add feature install/uninstall API routes with SSE progress"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Python Install Script
|
||||
|
||||
**Files:**
|
||||
- Create: `packages/ai/python/install_feature.py`
|
||||
|
||||
- [ ] **Step 1: Create the install script**
|
||||
|
||||
Create `packages/ai/python/install_feature.py`. Takes 3 CLI args: `bundleId`, `manifestPath`, `modelsDir`. Reads manifest JSON, detects architecture via `platform.machine()`, runs pip install for each package using `subprocess.run([sys.executable, "-m", "pip", "install", ...])`, downloads models with retry logic (exponential backoff, 3 retries, file size assertions).
|
||||
|
||||
Progress reported via stderr JSON lines: `{"progress": N, "stage": "..."}`. Result written to stdout JSON: `{"success": true, "bundleId": "...", "version": "...", "models": [...]}`.
|
||||
|
||||
Port the retry pattern from `docker/download_models.py` `_urlretrieve()` (lines 18-35). Handle rembg models via `rembg.new_session()` and HuggingFace models via `huggingface_hub.snapshot_download()`. Must be idempotent.
|
||||
|
||||
Writes to `/data/ai/installed.json` on success (matching the structure read by `feature-status.ts`).
|
||||
|
||||
**Robustness requirements for the install script:**
|
||||
|
||||
- **Atomic model downloads:** For each URL-based model:
|
||||
1. Check if final path already exists and meets `minSize` — skip if so (idempotent)
|
||||
2. Delete any existing `<path>.downloading` file (orphan from a previous failed attempt)
|
||||
3. Download to `<path>.downloading`
|
||||
4. Verify file size against `minSize`. If too small, delete and raise error.
|
||||
5. `os.rename(<path>.downloading, <path>)` — atomic on same filesystem
|
||||
6. Never leave a `.downloading` file behind on success
|
||||
- **Atomic JSON writes:** When writing `installed.json`:
|
||||
1. Write to `installed.json.tmp`
|
||||
2. `os.rename()` to `installed.json`
|
||||
- **Disk space pre-check:** Before starting, check available disk space via `shutil.disk_usage()`. If free space is less than estimated bundle size, exit with a clear error message.
|
||||
- **pip failure recovery:** If `pip install` fails for one package, emit the error and exit. The packages that were already installed remain (pip is idempotent — re-running skips them). The admin can retry.
|
||||
- **Model failure isolation:** If one model fails to download after retries, continue downloading other models. At the end, report which models failed. Exit with non-zero code so the bundle is NOT marked as installed. On retry, only the failed models need downloading (others pass the exists+size check).
|
||||
|
||||
- [ ] **Step 2: Commit**
|
||||
|
||||
```bash
|
||||
git add packages/ai/python/install_feature.py
|
||||
git commit -m "feat: add Python install script for feature bundles"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Tool Route Guards
|
||||
|
||||
**Files:**
|
||||
- Modify: `apps/api/src/routes/tool-factory.ts`
|
||||
- Modify: `apps/api/src/routes/batch.ts`
|
||||
- Modify: `apps/api/src/routes/pipeline.ts`
|
||||
- Modify: `apps/api/src/routes/tools/restore-photo.ts`
|
||||
|
||||
- [ ] **Step 1: Add feature guard to tool-factory.ts**
|
||||
|
||||
Import `isToolInstalled` from `../lib/feature-status.js` and `TOOL_BUNDLE_MAP`, `getBundleForTool` from `@ashim/shared`. Inside `createToolRoute`, after settings validation and before `config.process()`, add:
|
||||
|
||||
```ts
|
||||
const bundleId = TOOL_BUNDLE_MAP[config.toolId];
|
||||
if (bundleId && !isToolInstalled(config.toolId)) {
|
||||
const bundle = getBundleForTool(config.toolId);
|
||||
return reply.status(501).send({
|
||||
error: "Feature not installed",
|
||||
code: "FEATURE_NOT_INSTALLED",
|
||||
feature: bundleId,
|
||||
featureName: bundle?.name ?? bundleId,
|
||||
estimatedSize: bundle?.estimatedSize ?? "unknown",
|
||||
});
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add feature guard to batch.ts**
|
||||
|
||||
Same imports. After `getToolConfig(toolId)` returns (around line 35-37), add the same guard returning 501 with `FEATURE_NOT_INSTALLED` code.
|
||||
|
||||
- [ ] **Step 3: Add feature guard to pipeline.ts**
|
||||
|
||||
Same imports. In both pre-validation loops (execute at lines 143-172, batch at lines 441-462), after successful `getToolConfig(resolvedToolId)`, add the guard. Return 501 with step number in the error message.
|
||||
|
||||
- [ ] **Step 4: Add feature guard to restore-photo.ts**
|
||||
|
||||
This tool uses its own route handler, not the factory. Import `isToolInstalled` and add the guard before `restorePhoto()` is called.
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add apps/api/src/routes/tool-factory.ts apps/api/src/routes/batch.ts apps/api/src/routes/pipeline.ts apps/api/src/routes/tools/restore-photo.ts
|
||||
git commit -m "feat: add feature-installed guards to tool routes, batch, and pipeline"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Bridge and Python Sidecar Changes
|
||||
|
||||
**Files:**
|
||||
- Modify: `packages/ai/python/dispatcher.py`
|
||||
- Modify: `packages/ai/python/colorize.py`
|
||||
- Modify: `packages/ai/python/restore.py`
|
||||
|
||||
- [ ] **Step 1: Add feature gating to dispatcher.py**
|
||||
|
||||
Add a `TOOL_BUNDLE_MAP` dict mapping Python script names (without `.py`) to bundle IDs: `remove_bg` -> `background-removal`, `detect_faces` -> `face-detection`, `face_landmarks` -> `face-detection`, `red_eye_removal` -> `face-detection`, `inpaint` -> `object-eraser-colorize`, `colorize` -> `object-eraser-colorize`, `upscale` -> `upscale-enhance`, `enhance_faces` -> `upscale-enhance`, `noise_removal` -> `upscale-enhance`, `restore` -> `photo-restoration`, `ocr` -> `ocr`.
|
||||
|
||||
Add `_get_installed_bundles()` that reads `/data/ai/installed.json` and returns a set of installed bundle IDs.
|
||||
|
||||
In `_run_script_main()`, before the `exec()` call, check if the script's bundle is installed. If not, return a JSON error: `{"success": false, "error": "feature_not_installed", "feature": bundle_id, "message": "..."}`.
|
||||
|
||||
Also set `U2NET_HOME` to `/data/ai/models/rembg` on startup if `/data/ai/models` exists.
|
||||
|
||||
- [ ] **Step 2: Convert hard imports in colorize.py**
|
||||
|
||||
Move module-level `import numpy as np`, `import cv2`, `from PIL import Image` (lines 10-12) inside each function that uses them (`colorize_ddcolor`, `colorize_opencv`, `main`).
|
||||
|
||||
- [ ] **Step 3: Convert hard imports in restore.py**
|
||||
|
||||
Move module-level `import numpy as np`, `import cv2`, `from PIL import Image` (lines 13-15) inside each function that uses them.
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add packages/ai/python/dispatcher.py packages/ai/python/colorize.py packages/ai/python/restore.py
|
||||
git commit -m "feat: add feature gating to Python dispatcher, convert hard imports to lazy"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: Frontend Features Store and API Error Extension
|
||||
|
||||
**Files:**
|
||||
- Create: `apps/web/src/stores/features-store.ts`
|
||||
- Modify: `apps/web/src/lib/api.ts`
|
||||
- Modify: `apps/web/src/hooks/use-tool-processor.ts`
|
||||
- Modify: `apps/web/src/hooks/use-pipeline-processor.ts`
|
||||
|
||||
- [ ] **Step 1: Create the features store**
|
||||
|
||||
Create `apps/web/src/stores/features-store.ts` following the `settings-store.ts` pattern. Zustand store with `bundles: FeatureBundleState[]`, `loaded: boolean`, `fetch()` (one-shot), `refresh()` (force re-fetch), `isToolInstalled(toolId)`, `getBundleForTool(toolId)`. Fetches from `GET /api/v1/features`.
|
||||
|
||||
- [ ] **Step 2: Extend parseApiError for FEATURE_NOT_INSTALLED**
|
||||
|
||||
In `apps/web/src/lib/api.ts`, add a `FeatureNotInstalledError` interface export: `{ type: "feature_not_installed"; feature: string; featureName: string; estimatedSize: string }`.
|
||||
|
||||
Modify `parseApiError` return type to `string | FeatureNotInstalledError`. Add early return when `body.code === "FEATURE_NOT_INSTALLED"`.
|
||||
|
||||
- [ ] **Step 3: Update use-tool-processor.ts and use-pipeline-processor.ts**
|
||||
|
||||
In both hooks, where `parseApiError` is called and passed to `setError()`, add a type check:
|
||||
|
||||
```ts
|
||||
const parsed = parseApiError(body, xhr.status);
|
||||
if (typeof parsed === "object" && parsed.type === "feature_not_installed") {
|
||||
setError(`Feature "${parsed.featureName}" is not installed. Enable it in Settings.`);
|
||||
} else {
|
||||
setError(parsed);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add apps/web/src/stores/features-store.ts apps/web/src/lib/api.ts apps/web/src/hooks/use-tool-processor.ts apps/web/src/hooks/use-pipeline-processor.ts
|
||||
git commit -m "feat: add frontend features store and FEATURE_NOT_INSTALLED error handling"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 9: Frontend Tool Grid Badge
|
||||
|
||||
**Files:**
|
||||
- Modify: `apps/web/src/components/common/tool-card.tsx`
|
||||
- Modify: `apps/web/src/components/layout/tool-panel.tsx`
|
||||
- Modify: `apps/web/src/pages/fullscreen-grid-page.tsx`
|
||||
|
||||
- [ ] **Step 1: Add download badge to ToolCard**
|
||||
|
||||
Import `useFeaturesStore`, `PYTHON_SIDECAR_TOOLS`, and `Download` icon from lucide-react. Compute `showDownloadBadge` when the tool is an AI tool and not installed. Render a `<Download className="h-3.5 w-3.5 text-muted-foreground" />` icon after the experimental badge.
|
||||
|
||||
- [ ] **Step 2: Fetch features on app load**
|
||||
|
||||
In `tool-panel.tsx`, add `useFeaturesStore().fetch()` in a useEffect alongside the existing settings fetch. Do the same in `fullscreen-grid-page.tsx`.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add apps/web/src/components/common/tool-card.tsx apps/web/src/components/layout/tool-panel.tsx apps/web/src/pages/fullscreen-grid-page.tsx
|
||||
git commit -m "feat: add download badge to uninstalled AI tools in tool grid"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 10: Frontend Tool Page Install Prompt
|
||||
|
||||
**Files:**
|
||||
- Create: `apps/web/src/components/features/feature-install-prompt.tsx`
|
||||
- Modify: `apps/web/src/pages/tool-page.tsx`
|
||||
|
||||
- [ ] **Step 1: Create the FeatureInstallPrompt component**
|
||||
|
||||
Props: `{ bundle: FeatureBundleState; isAdmin: boolean }`.
|
||||
|
||||
For non-admins: show centered Download icon + "Feature Not Enabled" heading + "Ask your administrator" text.
|
||||
|
||||
For admins: show Download icon + bundle name/description + "requires additional download (~{estimatedSize})" + [Enable Feature] button. On click: POST to install endpoint, open EventSource for SSE progress, show progress bar with stage text and percent. On completion: call `useFeaturesStore().refresh()` to trigger re-render. On error: show error message with retry option.
|
||||
|
||||
Use same Tailwind patterns as existing components: `bg-primary text-primary-foreground` for buttons, `Loader2 animate-spin` for loading, `text-destructive` for errors.
|
||||
|
||||
**Robustness requirements for the frontend:**
|
||||
|
||||
- **Double-click prevention:** Set `installing = true` immediately on first click (before the API call). The button must be `disabled={installing || bundle.status === "installing"}`. This prevents any re-click.
|
||||
- **Browser close / navigate away:** The server-side install continues regardless. On component mount, check `bundle.status` from the features store. If it's `"installing"`, immediately show the progress bar and open EventSource for the in-progress job (fetch `jobId` from the features endpoint or use the bundle's progress data).
|
||||
- **SSE connection loss fallback:** If EventSource fires `onerror`, close it and fall back to polling `GET /api/v1/features` every 3 seconds via `setInterval`. When status changes from `"installing"` to `"installed"` or `"error"`, stop polling and update UI.
|
||||
- **Page refresh during install:** The features store's `fetch()` returns current status. If a bundle is `"installing"`, the component renders progress state immediately — no need for the user to click anything.
|
||||
- **Multiple admin sessions:** All sessions see the same `"installing"` status from the shared `GET /api/v1/features` endpoint. The server's install lock prevents concurrent installs. Any session trying to install gets a 409.
|
||||
- **Retry after error:** Show a "Retry" button when status is `"error"`. On retry, call the install endpoint again (the lock is released on failure, so this works). pip cache means previously-downloaded wheels aren't re-downloaded. Idempotent model downloads skip already-complete files.
|
||||
|
||||
- [ ] **Step 2: Integrate into ToolPage**
|
||||
|
||||
In `tool-page.tsx`: import `useFeaturesStore`, `PYTHON_SIDECAR_TOOLS`, `useAuth`, and `FeatureInstallPrompt`. After the tool/registryEntry lookup, compute `isAiTool`, `toolInstalled`, `featureBundle`, `isAdmin`. After the "Tool not found" guard, add a guard that renders `<FeatureInstallPrompt>` wrapped in `<AppLayout>` when the tool is AI and not installed.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add apps/web/src/components/features/feature-install-prompt.tsx apps/web/src/pages/tool-page.tsx
|
||||
git commit -m "feat: add feature install prompt on uninstalled AI tool pages"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 11: Settings AI Features Section
|
||||
|
||||
**Files:**
|
||||
- Create: `apps/web/src/components/settings/ai-features-section.tsx`
|
||||
- Modify: `apps/web/src/components/settings/settings-dialog.tsx`
|
||||
|
||||
- [ ] **Step 1: Create AiFeaturesSection component**
|
||||
|
||||
Follow the card-based layout of existing sections in `settings-dialog.tsx`. Use `useFeaturesStore()`. Render each bundle as a bordered card (`rounded-lg border border-border`) with: name, description, status indicator (green dot = installed, gray = not installed, spinning = installing), estimated size, Install/Uninstall button. Add "Install All" button at top. Show total disk usage at bottom (fetch from `GET /api/v1/admin/features/disk-usage`). Reuse the toggle/button patterns from `ToolsSection`.
|
||||
|
||||
- [ ] **Step 2: Add section to settings-dialog.tsx**
|
||||
|
||||
Add `"ai-features"` to the `Section` type union. Add to `NAV_ITEMS` between `"api-keys"` and `"tools"`: `{ id: "ai-features", label: "AI Features", icon: Sparkles, requiredPermission: "settings:write" }`. Import `Sparkles` from lucide-react. Add `{section === "ai-features" && <AiFeaturesSection />}` to the conditional render block. Lazy-import `AiFeaturesSection` from `"./ai-features-section"`.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add apps/web/src/components/settings/ai-features-section.tsx apps/web/src/components/settings/settings-dialog.tsx
|
||||
git commit -m "feat: add AI Features settings panel for managing feature bundles"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 12: Dockerfile Restructuring
|
||||
|
||||
**Files:**
|
||||
- Modify: `docker/Dockerfile`
|
||||
- Modify: `docker/entrypoint.sh`
|
||||
|
||||
- [ ] **Step 1: Modify the Dockerfile**
|
||||
|
||||
In `docker/Dockerfile` production stage:
|
||||
|
||||
1. **Keep**: base image selection, Node.js install, pnpm setup, system packages, Python venv creation with base packages (numpy, Pillow, opencv)
|
||||
2. **Remove**: all ML pip install commands (lines 175-206: onnxruntime, rembg, realesrgan, paddlepaddle, mediapipe, codeformer)
|
||||
3. **Remove**: download_models.py COPY and RUN (lines 219-231)
|
||||
4. **Remove**: the `apt-get purge build-essential python3-dev` line (line 251) so build-essential stays for runtime pip installs
|
||||
5. **Add**: `COPY docker/feature-manifest.json /app/docker/feature-manifest.json`
|
||||
6. **Add**: `COPY packages/ai/python/install_feature.py /app/packages/ai/python/install_feature.py`
|
||||
7. **Update** env vars: `PYTHON_VENV_PATH=/data/ai/venv`, add `MODELS_PATH=/data/ai/models`, add `DATA_DIR=/data`
|
||||
|
||||
- [ ] **Step 2: Update entrypoint.sh for venv bootstrap**
|
||||
|
||||
Add venv bootstrap after auth defaults and before volume permission fix. Use atomic directory rename to prevent corrupt venv from partial copy:
|
||||
|
||||
```sh
|
||||
AI_VENV="/data/ai/venv"
|
||||
AI_VENV_TMP="/data/ai/venv.bootstrapping"
|
||||
|
||||
# Clean up any interrupted bootstrap from a previous start
|
||||
if [ -d "$AI_VENV_TMP" ]; then
|
||||
echo "Cleaning up interrupted venv bootstrap..."
|
||||
rm -rf "$AI_VENV_TMP"
|
||||
fi
|
||||
|
||||
# Bootstrap AI venv from base image on first run
|
||||
if [ ! -d "$AI_VENV" ] && [ -d "/opt/venv" ]; then
|
||||
echo "Bootstrapping AI venv from base image..."
|
||||
mkdir -p /data/ai/models /data/ai/pip-cache
|
||||
cp -r /opt/venv "$AI_VENV_TMP"
|
||||
mv "$AI_VENV_TMP" "$AI_VENV"
|
||||
echo "AI venv ready at $AI_VENV"
|
||||
fi
|
||||
```
|
||||
|
||||
The `cp -r` + `mv` pattern ensures `/data/ai/venv` is either fully present or absent — never half-copied. If the container is killed during `cp -r`, the `.bootstrapping` directory is cleaned up on next start.
|
||||
|
||||
- [ ] **Step 3: Build and verify**
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile -t ashim:dev .
|
||||
docker images ashim:dev --format "{{.Size}}"
|
||||
```
|
||||
Expected: Image size ~5-6 GB (amd64) instead of ~30 GB.
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add docker/Dockerfile docker/entrypoint.sh
|
||||
git commit -m "feat: restructure Dockerfile to remove ML packages and models
|
||||
|
||||
Base image now includes only Node.js + Sharp + Python with base deps.
|
||||
AI features are downloaded on-demand via the feature install system.
|
||||
Image reduced from ~30GB to ~5-6GB (amd64) / ~2-3GB (arm64)."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 13: Integration Testing
|
||||
|
||||
**Files:**
|
||||
- Create: `tests/e2e-docker/features.spec.ts`
|
||||
|
||||
- [ ] **Step 1: Create Docker e2e tests for feature system**
|
||||
|
||||
Create `tests/e2e-docker/features.spec.ts` using the existing `playwright.docker.config.ts` infrastructure:
|
||||
|
||||
```ts
|
||||
import { expect, test } from "@playwright/test";
|
||||
|
||||
test.describe("On-demand AI features", () => {
|
||||
test("GET /api/v1/features returns all 6 bundles", async ({ request }) => {
|
||||
const response = await request.get("/api/v1/features");
|
||||
expect(response.ok()).toBeTruthy();
|
||||
const data = await response.json();
|
||||
expect(data.bundles).toHaveLength(6);
|
||||
for (const bundle of data.bundles) {
|
||||
expect(bundle).toHaveProperty("id");
|
||||
expect(bundle).toHaveProperty("name");
|
||||
expect(bundle).toHaveProperty("status");
|
||||
expect(bundle).toHaveProperty("enablesTools");
|
||||
}
|
||||
});
|
||||
|
||||
test("AI tool returns 501 FEATURE_NOT_INSTALLED when bundle not installed", async ({ request }) => {
|
||||
const pngBuffer = Buffer.from(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
|
||||
"base64",
|
||||
);
|
||||
const response = await request.post("/api/v1/tools/remove-background", {
|
||||
multipart: {
|
||||
file: { name: "test.png", mimeType: "image/png", buffer: pngBuffer },
|
||||
settings: JSON.stringify({}),
|
||||
},
|
||||
});
|
||||
expect(response.status()).toBe(501);
|
||||
const body = await response.json();
|
||||
expect(body.code).toBe("FEATURE_NOT_INSTALLED");
|
||||
expect(body.feature).toBe("background-removal");
|
||||
});
|
||||
|
||||
test("uninstalled AI tool page shows install prompt for admin", async ({ page }) => {
|
||||
await page.goto("/remove-background");
|
||||
await expect(page.getByText("Enable")).toBeVisible({ timeout: 10000 });
|
||||
await expect(page.getByText("additional download")).toBeVisible();
|
||||
});
|
||||
});
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Commit**
|
||||
|
||||
```bash
|
||||
git add tests/e2e-docker/features.spec.ts
|
||||
git commit -m "test: add e2e tests for on-demand AI feature system"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task Summary
|
||||
|
||||
| Task | Description | Key Files |
|
||||
|------|------------|-----------|
|
||||
| 1 | Shared types and bundle definitions | `packages/shared/src/features.ts` |
|
||||
| 2 | Feature manifest JSON | `docker/feature-manifest.json` |
|
||||
| 3 | Backend feature status service | `apps/api/src/lib/feature-status.ts` |
|
||||
| 4 | Feature API routes | `apps/api/src/routes/features.ts` |
|
||||
| 5 | Python install script | `packages/ai/python/install_feature.py` |
|
||||
| 6 | Tool route guards | `tool-factory.ts`, `batch.ts`, `pipeline.ts` |
|
||||
| 7 | Bridge + Python sidecar changes | `dispatcher.py`, `colorize.py`, `restore.py` |
|
||||
| 8 | Frontend features store + error handling | `features-store.ts`, `api.ts` |
|
||||
| 9 | Frontend tool grid badge | `tool-card.tsx`, `tool-panel.tsx` |
|
||||
| 10 | Frontend tool page install prompt | `feature-install-prompt.tsx`, `tool-page.tsx` |
|
||||
| 11 | Settings AI Features section | `ai-features-section.tsx`, `settings-dialog.tsx` |
|
||||
| 12 | Dockerfile restructuring | `Dockerfile`, `entrypoint.sh` |
|
||||
| 13 | Integration testing | `tests/e2e-docker/features.spec.ts` |
|
||||
Reference in New Issue
Block a user