mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix: reliable, self-healing AI feature-bundle installs (#472)
Make on-demand AI feature-bundle installs reliable and self-healing, closing the failure modes behind most "some tool doesn't work" reports. Multi-bundle installs: tools needing more than one bundle (Passport Photo, Enhance Faces) install every required bundle from one action and stay not-installed until all are present. Verified across all 19 AI tools. Downloads: self-heal the accelerated Hugging Face (Xet) client so an upgraded venv no longer silently falls back to slow urllib; restart instead of corrupting a resumed partial when a proxy ignores Range and returns 200; verify the completed size; fail fast on disk-full and HTTP 4xx; retry transient errors five times; add hf_transfer fallback and document Xet egress. Install integrity: crash-atomic venv writes so a killed or out-of-space install can no longer tear the shared venv and break other tools; a boot breadcrumb reseeds a torn venv to a clean state automatically; a post-install smoke import test refuses to record a bundle whose libraries cannot load; an install watchdog stops a wedged installer that would otherwise hold the venv writer lock forever. Adds unit and end-to-end tests for every failure mode above.
This commit is contained in:
SnapOtter
@@ -108,6 +108,12 @@ const envSchema = z
|
||||
SYNC_WAIT_MS: z.coerce.number().default(8000),
|
||||
JOB_TIMEOUT_FAST_S: z.coerce.number().default(120),
|
||||
JOB_TIMEOUT_LONG_S: z.coerce.number().default(7200),
|
||||
// AI bundle install watchdog. A wedged installer (dead download socket,
|
||||
// hung pip) otherwise holds the venv writer lock forever and blocks every
|
||||
// other install. STALL = max time with no progress frame before a kill;
|
||||
// MAX = absolute wall-clock ceiling. 0 disables that check.
|
||||
INSTALL_STALL_MS: z.coerce.number().default(1_200_000),
|
||||
INSTALL_MAX_MS: z.coerce.number().default(7_200_000),
|
||||
JOBS_RETENTION_DAYS: z.coerce.number().default(30),
|
||||
AUDIT_RETENTION_DAYS: z.coerce.number().default(0),
|
||||
LOG_DIR: z.string().default("./data/logs"),
|
||||
|
||||
@@ -27,12 +27,16 @@ import { getQueuedBundleIds } from "./feature-install-queue.js";
|
||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
const PROJECT_ROOT = resolve(__dirname, "../../../..");
|
||||
|
||||
const DATA_DIR = process.env.DATA_DIR || "/data";
|
||||
const DATA_DIR = process.env.DATA_DIR || "./data";
|
||||
const AI_DIR = join(DATA_DIR, "ai");
|
||||
const MODELS_DIR = join(AI_DIR, "models");
|
||||
const INSTALLED_PATH = join(AI_DIR, "installed.json");
|
||||
const INSTALLED_TMP_PATH = `${INSTALLED_PATH}.tmp`;
|
||||
const LOCK_PATH = join(AI_DIR, "install.lock");
|
||||
// Breadcrumb the installer drops right before it writes into the shared venv
|
||||
// site-packages and clears the instant that write completes. A survivor on boot
|
||||
// means the process died mid-write, so the venv may be torn (see move_tree).
|
||||
const VENV_WRITING_MARKER = join(AI_DIR, "venv.writing");
|
||||
const MANIFEST_PATH =
|
||||
process.env.FEATURE_MANIFEST_PATH || join(PROJECT_ROOT, "docker/feature-manifest.json");
|
||||
|
||||
@@ -61,10 +65,8 @@ export function ensureAiDirs(): void {
|
||||
mkdirSync(MODELS_DIR, { recursive: true });
|
||||
mkdirSync(join(AI_DIR, "pip-cache"), { recursive: true });
|
||||
} catch (err: unknown) {
|
||||
// Never refuse to boot over the AI data dir. On native checkouts the
|
||||
// default DATA_DIR (/data) is often uncreatable (ENOENT/EROFS on a
|
||||
// sealed macOS root, EACCES on restrictive volumes); AI tools simply
|
||||
// report as not installed until DATA_DIR points somewhere writable.
|
||||
// Never refuse to boot over the AI data dir. AI tools simply report as not
|
||||
// installed until DATA_DIR points somewhere writable.
|
||||
const code = (err as NodeJS.ErrnoException).code;
|
||||
console.error(
|
||||
`WARNING: Cannot create AI directories under "${AI_DIR}" (${code}). AI features will be unavailable. Set DATA_DIR to a writable path (or check volume permissions / PUID / PGID in Docker).`,
|
||||
@@ -440,6 +442,37 @@ export function recoverInterruptedInstalls(): void {
|
||||
}
|
||||
}
|
||||
|
||||
// 4b. Heal a torn shared venv. If the venv-writing breadcrumb survived, an
|
||||
// install died while rewriting the shared site-packages, which can leave a
|
||||
// package half-replaced and break EVERY AI tool (not just the one installing).
|
||||
// Model weight files under MODELS_DIR are unaffected, so the safe, automatic
|
||||
// recovery is to reseed the venv from the image base and reset the install
|
||||
// ledger; the user reinstalls bundles from a known-good state. This is the
|
||||
// same repair as "Reset AI Environment", done automatically on next boot so a
|
||||
// crash-broken install self-heals instead of leaving mysteriously dead tools.
|
||||
if (existsSync(VENV_WRITING_MARKER)) {
|
||||
if (isDockerEnvironment() && existsSync("/opt/venv")) {
|
||||
console.warn(
|
||||
"[feature-status] An install was interrupted mid venv-write; reseeding the AI venv to a clean state and clearing installed bundles (reinstall to restore).",
|
||||
);
|
||||
try {
|
||||
execFileSync("/usr/local/bin/reseed-ai-venv.sh", { stdio: "ignore", timeout: 120_000 });
|
||||
writeInstalled({ bundles: {} });
|
||||
} catch (err) {
|
||||
console.error("[feature-status] Failed to reseed AI venv after interrupted install:", err);
|
||||
}
|
||||
} else {
|
||||
console.warn(
|
||||
"[feature-status] An install was interrupted mid venv-write (non-Docker); the AI venv may be inconsistent. Reinstall the affected bundle.",
|
||||
);
|
||||
}
|
||||
try {
|
||||
unlinkSync(VENV_WRITING_MARKER);
|
||||
} catch {
|
||||
// best-effort
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Verify installed bundles still have their model files
|
||||
const manifest = readManifest();
|
||||
if (manifest) {
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
/**
|
||||
* Pure decision for the AI-bundle install watchdog.
|
||||
*
|
||||
* A wedged installer (dead download socket, hung pip) holds the venv writer lock
|
||||
* forever and blocks every other install with no exit but a server restart. The
|
||||
* watchdog kills such a child so its close handler frees the locks and the user
|
||||
* gets a retryable failure. This function holds the (timer-free, side-effect
|
||||
* free) rule so it can be unit-tested exhaustively; features.ts owns the timers.
|
||||
*
|
||||
* @param now current time (ms)
|
||||
* @param lastProgressAt time of the last progress frame (ms)
|
||||
* @param startedAt time the install started (ms)
|
||||
* @param stallMs max time with no progress before a kill (0 disables)
|
||||
* @param maxMs absolute wall-clock ceiling (0 disables)
|
||||
*/
|
||||
export function evaluateInstallWatchdog(
|
||||
now: number,
|
||||
lastProgressAt: number,
|
||||
startedAt: number,
|
||||
stallMs: number,
|
||||
maxMs: number,
|
||||
): { kill: boolean; reason: string | null } {
|
||||
const overMax = maxMs > 0 && now - startedAt > maxMs;
|
||||
if (overMax) {
|
||||
return {
|
||||
kill: true,
|
||||
reason: `Installation exceeded the ${Math.round(
|
||||
maxMs / 60_000,
|
||||
)} minute time limit and was stopped. Please retry.`,
|
||||
};
|
||||
}
|
||||
const stalled = stallMs > 0 && now - lastProgressAt > stallMs;
|
||||
if (stalled) {
|
||||
return {
|
||||
kill: true,
|
||||
reason: `Installation made no progress for ${Math.round(
|
||||
stallMs / 60_000,
|
||||
)} minutes and was stopped. Check your connection and retry.`,
|
||||
};
|
||||
}
|
||||
return { kill: false, reason: null };
|
||||
}
|
||||
+145
-28
@@ -2,11 +2,12 @@
|
||||
* Feature bundle management routes.
|
||||
*
|
||||
* GET /api/v1/features - List feature bundles and their statuses
|
||||
* POST /api/v1/admin/features/:bundleId/install - Install a feature bundle (async)
|
||||
* POST /api/v1/admin/features/:bundleId/uninstall - Uninstall a feature bundle
|
||||
* POST /api/v1/admin/features/reset - Wipe the AI venv/models, reset all bundles
|
||||
* GET /api/v1/admin/features/disk-usage - Get AI model disk usage
|
||||
* POST /api/v1/admin/features/import - Import an offline bundle archive
|
||||
* POST /api/v1/admin/features/:bundleId/install - Install a feature bundle (async)
|
||||
* POST /api/v1/admin/tools/:toolId/features/install - Install every bundle a tool requires
|
||||
* POST /api/v1/admin/features/:bundleId/uninstall - Uninstall a feature bundle
|
||||
* POST /api/v1/admin/features/reset - Wipe the AI venv/models, reset all bundles
|
||||
* GET /api/v1/admin/features/disk-usage - Get AI model disk usage
|
||||
* POST /api/v1/admin/features/import - Import an offline bundle archive
|
||||
*/
|
||||
|
||||
import { spawn } from "node:child_process";
|
||||
@@ -22,8 +23,9 @@ import {
|
||||
} from "node:fs";
|
||||
import { join } from "node:path";
|
||||
import { acquireVenvLock, shutdownDispatcher } from "@snapotter/ai";
|
||||
import { ANALYTICS_EVENTS, FEATURE_BUNDLES } from "@snapotter/shared";
|
||||
import { ANALYTICS_EVENTS, FEATURE_BUNDLES, getRequiredBundlesForTool } from "@snapotter/shared";
|
||||
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
|
||||
import { env } from "../config.js";
|
||||
import { trackEvent } from "../lib/analytics.js";
|
||||
import {
|
||||
clearActive,
|
||||
@@ -52,6 +54,7 @@ import {
|
||||
setInstallProgress,
|
||||
verifyBundleModels,
|
||||
} from "../lib/feature-status.js";
|
||||
import { evaluateInstallWatchdog } from "../lib/install-watchdog.js";
|
||||
import { requirePermission } from "../permissions.js";
|
||||
import { requireAuth } from "../plugins/auth.js";
|
||||
import { updateSingleFileProgress } from "./progress.js";
|
||||
@@ -99,10 +102,32 @@ function startInstall(bundleId: string, jobId: string): void {
|
||||
// would release the file lock and active slot that pump() just handed to
|
||||
// the next queued bundle, letting two installers run into the same venv
|
||||
// at once (the corruption the lock exists to prevent).
|
||||
// Install watchdog: a wedged installer (dead download socket, hung pip)
|
||||
// otherwise holds the venv writer lock forever and blocks every other
|
||||
// install with no way out but a server restart. Track the last progress
|
||||
// frame; if the child goes silent past the stall budget, or blows the
|
||||
// absolute ceiling, kill it so its close handler frees the locks and the
|
||||
// user sees a retryable failure.
|
||||
let lastProgressAt = Date.now();
|
||||
let watchdog: ReturnType<typeof setInterval> | null = null;
|
||||
let killGrace: ReturnType<typeof setTimeout> | null = null;
|
||||
let watchdogError: string | null = null;
|
||||
const clearWatchdog = () => {
|
||||
if (watchdog) {
|
||||
clearInterval(watchdog);
|
||||
watchdog = null;
|
||||
}
|
||||
if (killGrace) {
|
||||
clearTimeout(killGrace);
|
||||
killGrace = null;
|
||||
}
|
||||
};
|
||||
|
||||
let finalized = false;
|
||||
const finalizeOnce = (): boolean => {
|
||||
if (finalized) return false;
|
||||
finalized = true;
|
||||
clearWatchdog();
|
||||
releaseVenvOnce();
|
||||
releaseInstallLock();
|
||||
clearActive();
|
||||
@@ -118,6 +143,42 @@ function startInstall(bundleId: string, jobId: string): void {
|
||||
},
|
||||
});
|
||||
|
||||
const stallMs = env.INSTALL_STALL_MS;
|
||||
const maxMs = env.INSTALL_MAX_MS;
|
||||
if (stallMs > 0 || maxMs > 0) {
|
||||
watchdog = setInterval(() => {
|
||||
const verdict = evaluateInstallWatchdog(
|
||||
Date.now(),
|
||||
lastProgressAt,
|
||||
installStartTime,
|
||||
stallMs,
|
||||
maxMs,
|
||||
);
|
||||
if (!verdict.kill) return;
|
||||
watchdogError = verdict.reason;
|
||||
setInstallProgress(bundleId, null, watchdogError);
|
||||
if (watchdog) {
|
||||
clearInterval(watchdog);
|
||||
watchdog = null;
|
||||
}
|
||||
try {
|
||||
child.kill("SIGTERM");
|
||||
} catch {
|
||||
// child already gone
|
||||
}
|
||||
// Escalate to SIGKILL if SIGTERM does not land (e.g. a C extension
|
||||
// ignoring the signal); the close handler clears this grace timer.
|
||||
killGrace = setTimeout(() => {
|
||||
try {
|
||||
child.kill("SIGKILL");
|
||||
} catch {
|
||||
// already gone
|
||||
}
|
||||
}, 10_000);
|
||||
}, 30_000);
|
||||
watchdog.unref?.();
|
||||
}
|
||||
|
||||
let stderrBuffer = "";
|
||||
let stdoutBuffer = "";
|
||||
const lastStderrLines: string[] = [];
|
||||
@@ -142,6 +203,7 @@ function startInstall(bundleId: string, jobId: string): void {
|
||||
try {
|
||||
const parsed = JSON.parse(trimmed) as { progress?: number; stage?: string };
|
||||
if (typeof parsed.progress === "number") {
|
||||
lastProgressAt = Date.now();
|
||||
setInstallProgress(
|
||||
bundleId,
|
||||
{ percent: parsed.progress, stage: parsed.stage ?? "" },
|
||||
@@ -174,10 +236,12 @@ function startInstall(bundleId: string, jobId: string): void {
|
||||
duration_ms: Date.now() - installStartTime,
|
||||
});
|
||||
} else {
|
||||
// A watchdog kill wins: the child's exit code/stderr would otherwise
|
||||
// read as a generic signal death and bury why it was stopped.
|
||||
let errorMsg: string | undefined = watchdogError ?? undefined;
|
||||
// Extract the structured error from Python's fail() function first.
|
||||
// fail() writes {"error": "..."} to stderr - prefer this over raw lines.
|
||||
let errorMsg: string | undefined;
|
||||
for (let i = lastStderrLines.length - 1; i >= 0; i--) {
|
||||
for (let i = lastStderrLines.length - 1; !errorMsg && i >= 0; i--) {
|
||||
const line = lastStderrLines[i];
|
||||
if (line.startsWith("{")) {
|
||||
try {
|
||||
@@ -250,6 +314,23 @@ interface BundleIdParams {
|
||||
bundleId: string;
|
||||
}
|
||||
|
||||
interface ToolIdParams {
|
||||
toolId: string;
|
||||
}
|
||||
|
||||
interface EnqueuedBundleInstall {
|
||||
bundleId: string;
|
||||
jobId: string;
|
||||
queued: boolean;
|
||||
}
|
||||
|
||||
interface ToolBundleInstallResult {
|
||||
bundleId: string;
|
||||
jobId?: string;
|
||||
queued: boolean;
|
||||
skipped?: boolean;
|
||||
}
|
||||
|
||||
interface ManifestModel {
|
||||
id: string;
|
||||
path?: string;
|
||||
@@ -275,6 +356,32 @@ function readManifest(): Manifest | null {
|
||||
}
|
||||
}
|
||||
|
||||
function queueBundleInstallIfNeeded(bundleId: string): EnqueuedBundleInstall | null {
|
||||
if (isFeatureInstalled(bundleId)) {
|
||||
const modelError = verifyBundleModels(bundleId);
|
||||
if (!modelError) return null;
|
||||
markUninstalled(bundleId);
|
||||
}
|
||||
|
||||
// Queue the install on the server so the POST is durable: enqueue()
|
||||
// dedups an already-active/queued bundle and returns the effective jobId,
|
||||
// then pump() starts it immediately if nothing else is installing. A bundle
|
||||
// that lands behind another install stays queued server-side and starts
|
||||
// automatically when the running install finishes.
|
||||
const jobId = crypto.randomUUID();
|
||||
const effectiveJobId = enqueue({ bundleId, jobId });
|
||||
pump();
|
||||
|
||||
// queued === true means it did NOT start right now (another install is
|
||||
// active, or an offline import holds the lock). The client polls queued
|
||||
// bundles and opens SSE only for the active install.
|
||||
return {
|
||||
bundleId,
|
||||
jobId: effectiveJobId,
|
||||
queued: getActiveBundleId() !== bundleId,
|
||||
};
|
||||
}
|
||||
|
||||
/** Recursively calculate total size of a directory in bytes. */
|
||||
function getDirSize(dirPath: string): number {
|
||||
if (!existsSync(dirPath)) return 0;
|
||||
@@ -348,30 +455,40 @@ export async function registerFeatureRoutes(app: FastifyInstance): Promise<void>
|
||||
return reply.status(404).send({ error: `Unknown bundle: ${bundleId}` });
|
||||
}
|
||||
|
||||
if (isFeatureInstalled(bundleId)) {
|
||||
const modelError = verifyBundleModels(bundleId);
|
||||
if (!modelError) {
|
||||
return reply.status(409).send({ error: `Bundle "${bundleId}" is already installed` });
|
||||
}
|
||||
markUninstalled(bundleId);
|
||||
const result = queueBundleInstallIfNeeded(bundleId);
|
||||
if (!result) {
|
||||
return reply.status(409).send({ error: `Bundle "${bundleId}" is already installed` });
|
||||
}
|
||||
|
||||
// Queue the install on the server so the POST is durable: enqueue()
|
||||
// dedups an already-active/queued bundle and returns the effective jobId,
|
||||
// then pump() starts it immediately if nothing else is installing. A
|
||||
// bundle that lands behind another install stays queued server-side and
|
||||
// starts automatically when the running install finishes.
|
||||
const jobId = crypto.randomUUID();
|
||||
const effectiveJobId = enqueue({ bundleId, jobId });
|
||||
pump();
|
||||
return reply.status(202).send({ jobId: result.jobId, queued: result.queued });
|
||||
},
|
||||
);
|
||||
|
||||
// queued === true means it did NOT start right now (another install is
|
||||
// active, or an offline import holds the lock). The client still opens
|
||||
// the SSE stream for the returned jobId; the progress route tolerates a
|
||||
// not-yet-started jobId and streams once the installer begins.
|
||||
const queued = getActiveBundleId() !== bundleId;
|
||||
// POST /api/v1/admin/tools/:toolId/features/install - Install all bundles a tool needs
|
||||
app.post(
|
||||
"/api/v1/admin/tools/:toolId/features/install",
|
||||
{ config: { rateLimit: { max: 300, timeWindow: "1 minute" } } },
|
||||
async (request: FastifyRequest<{ Params: ToolIdParams }>, reply: FastifyReply) => {
|
||||
const admin = await requirePermission("features:manage")(request, reply);
|
||||
if (!admin) return;
|
||||
|
||||
return reply.status(202).send({ jobId: effectiveJobId, queued });
|
||||
const { toolId } = request.params;
|
||||
const requiredBundles = getRequiredBundlesForTool(toolId);
|
||||
if (requiredBundles.length === 0) {
|
||||
return reply.status(404).send({ error: `No feature bundles required for tool: ${toolId}` });
|
||||
}
|
||||
|
||||
const bundles: ToolBundleInstallResult[] = [];
|
||||
for (const bundleId of requiredBundles) {
|
||||
if (!FEATURE_BUNDLES[bundleId]) {
|
||||
return reply.status(404).send({ error: `Unknown bundle: ${bundleId}` });
|
||||
}
|
||||
|
||||
const result = queueBundleInstallIfNeeded(bundleId);
|
||||
bundles.push(result ?? { bundleId, queued: false, skipped: true });
|
||||
}
|
||||
|
||||
return reply.status(202).send({ bundles });
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ import { tmpdir } from "node:os";
|
||||
import { extname, join } from "node:path";
|
||||
import {
|
||||
apiToolPath,
|
||||
getBundleForTool,
|
||||
FEATURE_BUNDLES,
|
||||
type Section,
|
||||
TOOL_BUNDLE_MAP,
|
||||
TOOLS,
|
||||
@@ -16,7 +16,7 @@ import { env } from "../config.js";
|
||||
import { db, schema } from "../db/index.js";
|
||||
import { enqueueToolJob, waitForJob } from "../jobs/enqueue.js";
|
||||
import { formatZodErrors, friendlyError, stripInternalPaths } from "../lib/errors.js";
|
||||
import { isToolInstalled } from "../lib/feature-status.js";
|
||||
import { getFirstMissingBundleForTool, isToolInstalled } from "../lib/feature-status.js";
|
||||
import { getObjectBuffer, putObject } from "../lib/object-storage.js";
|
||||
import { resolveToolPool, shouldSkipSyncWindow } from "../lib/pool.js";
|
||||
import { getSettingNumber } from "../lib/settings-helpers.js";
|
||||
@@ -470,13 +470,14 @@ export function createToolRoute<T>(app: FastifyInstance, config: ToolRouteConfig
|
||||
// Guard: check if the tool's AI feature bundle is installed
|
||||
const bundleId = TOOL_BUNDLE_MAP[config.toolId];
|
||||
if (bundleId && !isToolInstalled(config.toolId)) {
|
||||
const bundle = getBundleForTool(config.toolId);
|
||||
const missingBundleId = getFirstMissingBundleForTool(config.toolId) ?? bundleId;
|
||||
const bundle = FEATURE_BUNDLES[missingBundleId];
|
||||
// Orphaned uploads/<jobId>/ dir will be cleaned by T10 TTL sweeper
|
||||
return reply.status(501).send({
|
||||
error: "Feature not installed",
|
||||
code: "FEATURE_NOT_INSTALLED",
|
||||
feature: bundleId,
|
||||
featureName: bundle?.name ?? bundleId,
|
||||
feature: missingBundleId,
|
||||
featureName: bundle?.name ?? missingBundleId,
|
||||
estimatedSize: bundle?.estimatedSize ?? "unknown",
|
||||
});
|
||||
}
|
||||
|
||||
@@ -3,14 +3,14 @@ import { mkdir, rm } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { enhanceFaces } from "@snapotter/ai";
|
||||
import { getBundleForTool, TOOL_BUNDLE_MAP } from "@snapotter/shared";
|
||||
import { FEATURE_BUNDLES, TOOL_BUNDLE_MAP } from "@snapotter/shared";
|
||||
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
|
||||
import { z } from "zod";
|
||||
import { registerAiJobHandler } from "../../jobs/ai-handlers.js";
|
||||
import { enqueueToolJob } from "../../jobs/enqueue.js";
|
||||
import { autoOrient } from "../../lib/auto-orient.js";
|
||||
import { formatZodErrors, stripInternalPaths } from "../../lib/errors.js";
|
||||
import { isToolInstalled } from "../../lib/feature-status.js";
|
||||
import { getFirstMissingBundleForTool, isToolInstalled } from "../../lib/feature-status.js";
|
||||
import { validateImageBuffer } from "../../lib/file-validation.js";
|
||||
import { decodeToSharpCompat, needsCliDecode } from "../../lib/format-decoders.js";
|
||||
import { decodeHeic } from "../../lib/heic-converter.js";
|
||||
@@ -64,12 +64,13 @@ export function registerEnhanceFaces(app: FastifyInstance) {
|
||||
async (request: FastifyRequest, reply: FastifyReply) => {
|
||||
const toolId = "enhance-faces";
|
||||
if (!isToolInstalled(toolId)) {
|
||||
const bundle = getBundleForTool(toolId);
|
||||
const missingBundleId = getFirstMissingBundleForTool(toolId) ?? TOOL_BUNDLE_MAP[toolId];
|
||||
const bundle = FEATURE_BUNDLES[missingBundleId];
|
||||
return reply.status(501).send({
|
||||
error: "Feature not installed",
|
||||
code: "FEATURE_NOT_INSTALLED",
|
||||
feature: TOOL_BUNDLE_MAP[toolId],
|
||||
featureName: bundle?.name ?? toolId,
|
||||
feature: missingBundleId,
|
||||
featureName: bundle?.name ?? missingBundleId,
|
||||
estimatedSize: bundle?.estimatedSize ?? "unknown",
|
||||
});
|
||||
}
|
||||
|
||||
+23
-10
@@ -45,17 +45,30 @@ A separate "docs" dispatcher profile replaces the AI allowlist with document-pro
|
||||
|
||||
## Feature Bundles
|
||||
|
||||
Each AI tool requires a model bundle to be installed before use. Bundles are installed on demand via the admin UI or `install_feature.py`.
|
||||
AI models are packaged by shared dependency stack, not one archive per tool. A feature bundle can enable several tools when they use the same model family, Python wheels, or native libraries. This keeps the release Docker image smaller and avoids storing duplicate copies of the same background matting, face detection, OCR, restoration, and speech models.
|
||||
|
||||
| Bundle | Size | Tools |
|
||||
|--------|------|-------|
|
||||
| `background-removal` | 4-5 GB | remove-background, passport-photo, transparency-fixer, background-replace, blur-background |
|
||||
| `face-detection` | 200-300 MB | blur-faces, red-eye-removal, smart-crop |
|
||||
| `object-eraser-colorize` | 1-2 GB | erase-object, colorize, ai-canvas-expand |
|
||||
| `upscale-enhance` | 5-6 GB | upscale, enhance-faces, noise-removal |
|
||||
| `photo-restoration` | 4-5 GB | restore-photo |
|
||||
| `ocr` | 5-6 GB | ocr, ocr-pdf |
|
||||
| `transcription` | ~600 MB | transcribe-audio, auto-subtitles |
|
||||
The Docker image ships the application plus the common runtime. Large model archives are downloaded on demand into the persistent `/data/ai` volume, then reused by every tool that needs them. If a bundle is already installed because another tool needed it, enabling a new dependent tool does not download that bundle again.
|
||||
|
||||
Each AI tool requires one or more feature bundles before it can run. The admin UI installs by tool through `POST /api/v1/admin/tools/:toolId/features/install`, which resolves the full bundle list, skips bundles that are already installed, and queues only the missing downloads. For example, enabling Passport Photo on a fresh instance queues `background-removal` and `face-detection`; enabling it after Background Removal is already installed queues only `face-detection`.
|
||||
|
||||
| Bundle | Size | Shared dependency group | Tools that use it |
|
||||
|--------|------|-------------------------|-------------------|
|
||||
| `background-removal` | 4-5 GB | rembg / BiRefNet background matting | remove-background, passport-photo, transparency-fixer, background-replace, blur-background |
|
||||
| `face-detection` | 200-300 MB | MediaPipe face detection and landmarks | blur-faces, red-eye-removal, smart-crop |
|
||||
| `object-eraser-colorize` | 1-2 GB | LaMa inpainting/outpainting and DDColor | erase-object, colorize, ai-canvas-expand |
|
||||
| `upscale-enhance` | 5-6 GB | RealESRGAN, GFPGAN / CodeFormer, denoising | upscale, enhance-faces, noise-removal |
|
||||
| `photo-restoration` | 4-5 GB | scratch repair and restoration pipeline | restore-photo |
|
||||
| `ocr` | 5-6 GB | PaddleOCR / Tesseract OCR stack | ocr, ocr-pdf |
|
||||
| `transcription` | ~600 MB | faster-whisper speech-to-text models | transcribe-audio, auto-subtitles |
|
||||
|
||||
Tools with cross-bundle dependencies:
|
||||
|
||||
| Tool | Required bundles | Why |
|
||||
|------|------------------|-----|
|
||||
| `passport-photo` | `background-removal`, `face-detection` | Removes the background, then uses face landmarks to frame the crop to passport and ID photo rules. |
|
||||
| `enhance-faces` | `upscale-enhance`, `face-detection` | Detects faces before running GFPGAN or CodeFormer enhancement on the selected face regions. |
|
||||
|
||||
A tool is available only when all of its required bundles are installed. Partial installs are valid and are handled incrementally: installed bundles are reused, missing bundles are shown as downloads, and queued installs run one at a time so the shared Python environment is not modified concurrently.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -589,12 +589,13 @@ Query parameters:
|
||||
|
||||
## Features / AI Bundles
|
||||
|
||||
Manage AI feature bundles (install/uninstall AI model packages in the Docker environment).
|
||||
Manage AI feature bundles (install/uninstall AI model packages in the Docker environment). Prefer the tool-level install endpoint when enabling a tool from custom automation: some AI tools need more than one shared bundle, and this endpoint skips already-installed bundles while queuing only the missing ones.
|
||||
|
||||
| Method | Path | Access | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `GET` | `/api/v1/features` | Auth | List all feature bundles and their install status |
|
||||
| `POST` | `/api/v1/admin/features/:bundleId/install` | Admin (`features:manage`) | Install a feature bundle (async, returns `jobId` for progress tracking) |
|
||||
| `POST` | `/api/v1/admin/tools/:toolId/features/install` | Admin (`features:manage`) | Install every bundle a tool requires; returns per-bundle queued/skipped status |
|
||||
| `POST` | `/api/v1/admin/features/:bundleId/uninstall` | Admin (`features:manage`) | Uninstall a feature bundle and clean up model files |
|
||||
| `GET` | `/api/v1/admin/features/disk-usage` | Admin (`features:manage`) | Get total disk usage of AI models |
|
||||
| `POST` | `/api/v1/admin/features/import` | Admin (`features:manage`) | Import an offline AI bundle archive |
|
||||
|
||||
@@ -264,6 +264,10 @@ Most AI tools are perfectly usable on CPU; a couple really want a GPU. Measured
|
||||
| AI upscale (RealESRGAN) | ~33 s small; minutes on large images | Marginal — GPU strongly recommended |
|
||||
| Photo restoration (full pipeline) | several minutes | No — needs a GPU or a fast many-core CPU |
|
||||
|
||||
SnapOtter intentionally does not bake these model downloads into the Docker image. AI bundles are pulled only when an admin enables the related tool, stored in the persistent `/data/ai` volume, and shared by every tool that depends on the same model stack. This keeps the final container image small while still letting a full AI installation reach the larger storage numbers below.
|
||||
|
||||
Some tools depend on more than one shared bundle. For example, Passport Photo needs both `background-removal` and `face-detection`; if `background-removal` is already installed, enabling Passport Photo only downloads the missing `face-detection` bundle. The same reuse applies across all AI tools.
|
||||
|
||||
AI model download sizes:
|
||||
|
||||
| Bundle | Disk Size |
|
||||
|
||||
@@ -122,16 +122,18 @@ During normal operation, the container makes **zero outbound network connections
|
||||
Browser --> Reverse Proxy (TLS) --> SnapOtter container --> (nothing)
|
||||
```
|
||||
|
||||
The only exception is **AI model downloads**: when a user installs an AI feature bundle through the UI, the container downloads model files from GitHub Releases and PyPI. These downloads happen once per bundle and are stored in the `/data` volume.
|
||||
The only exception is **AI model downloads**: when a user installs an AI feature bundle through the UI, the container downloads the pre-built bundle archive from Hugging Face, plus a few individual model files from GitHub Releases, Google Storage, and PyPI. These downloads happen once per bundle and are stored in the `/data` volume.
|
||||
|
||||
**Firewall recommendations:**
|
||||
|
||||
| Scenario | Outbound rule |
|
||||
|---|---|
|
||||
| Air-gapped (no AI) | Block all outbound traffic from the container |
|
||||
| AI bundles needed | Allow HTTPS to `github.com`, `objects.githubusercontent.com`, `pypi.org`, `files.pythonhosted.org` during install, then block |
|
||||
| AI bundles needed | Allow HTTPS to `huggingface.co`, `*.xethub.hf.co`, `cdn-lfs.huggingface.co`, `github.com`, `objects.githubusercontent.com`, `storage.googleapis.com`, `pypi.org`, `files.pythonhosted.org` during install, then block |
|
||||
| After AI install | Block all outbound traffic - models are cached locally |
|
||||
|
||||
Bundle archives are served from Hugging Face's Xet storage, which transfers over the `*.xethub.hf.co` endpoints in parallel and is what makes multi-GB bundle downloads fast. If your firewall allows `huggingface.co` but blocks `*.xethub.hf.co`, installs still succeed but fall back to a slower single-stream download, so allowlist the Xet hosts to stay on the fast path. Fully offline installs can skip all of this and use [Offline Bundle Import](/guide/deployment) instead.
|
||||
|
||||
For reverse proxy configuration (Nginx, Traefik, Caddy, Cloudflare Tunnels), see the [Deployment guide](/guide/deployment#reverse-proxy).
|
||||
|
||||
## Docker Secrets
|
||||
|
||||
@@ -1,5 +1,10 @@
|
||||
import type { Tool } from "@snapotter/shared";
|
||||
import { PYTHON_SIDECAR_TOOLS, SECTIONS, TOOL_BUNDLE_MAP, toolSection } from "@snapotter/shared";
|
||||
import {
|
||||
getRequiredBundlesForTool,
|
||||
PYTHON_SIDECAR_TOOLS,
|
||||
SECTIONS,
|
||||
toolSection,
|
||||
} from "@snapotter/shared";
|
||||
import { Clock, Download, FileImage, Loader2, Pin } from "lucide-react";
|
||||
import { useMemo } from "react";
|
||||
import { Link } from "react-router-dom";
|
||||
@@ -63,12 +68,15 @@ export function ToolCard({ tool, variant = "compact", showModalityBadge, showPin
|
||||
const queued = useFeaturesStore((s) => s.queued);
|
||||
const aiStatus = useMemo(() => {
|
||||
if (!isAiTool) return "installed";
|
||||
const bundleId = TOOL_BUNDLE_MAP[tool.id];
|
||||
if (!bundleId) return "installed";
|
||||
if (queued.includes(bundleId)) return "queued";
|
||||
if (installing[bundleId]) return "installing";
|
||||
const bundle = bundles.find((b) => b.id === bundleId);
|
||||
return bundle?.status === "installed" ? "installed" : "not_installed";
|
||||
const requiredBundleIds = getRequiredBundlesForTool(tool.id);
|
||||
if (requiredBundleIds.length === 0) return "installed";
|
||||
if (requiredBundleIds.some((bundleId) => queued.includes(bundleId))) return "queued";
|
||||
if (requiredBundleIds.some((bundleId) => installing[bundleId])) return "installing";
|
||||
return requiredBundleIds.every(
|
||||
(bundleId) => bundles.find((bundle) => bundle.id === bundleId)?.status === "installed",
|
||||
)
|
||||
? "installed"
|
||||
: "not_installed";
|
||||
}, [isAiTool, tool.id, bundles, installing, queued]);
|
||||
|
||||
const section = toolSection(tool);
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
import type { FeatureBundleState } from "@snapotter/shared";
|
||||
import {
|
||||
FEATURE_BUNDLES,
|
||||
type FeatureBundleState,
|
||||
getRequiredBundlesForTool,
|
||||
} from "@snapotter/shared";
|
||||
import { AlertCircle, Clock, Download, Loader2, RotateCcw } from "lucide-react";
|
||||
import { useEffect, useState } from "react";
|
||||
import { useTranslation } from "@/contexts/i18n-context";
|
||||
@@ -48,18 +52,45 @@ function formatTimeRemaining(ms: number): string {
|
||||
interface FeatureInstallPromptProps {
|
||||
bundle: FeatureBundleState;
|
||||
isAdmin: boolean;
|
||||
toolId?: string;
|
||||
toolName?: string;
|
||||
toolDescription?: string;
|
||||
}
|
||||
|
||||
function fallbackBundleState(bundleId: string): FeatureBundleState | null {
|
||||
const info = FEATURE_BUNDLES[bundleId];
|
||||
if (!info) return null;
|
||||
return {
|
||||
id: info.id,
|
||||
name: info.name,
|
||||
description: info.description,
|
||||
status: "not_installed",
|
||||
installedVersion: null,
|
||||
estimatedSize: info.estimatedSize,
|
||||
enablesTools: info.enablesTools,
|
||||
progress: null,
|
||||
error: null,
|
||||
};
|
||||
}
|
||||
|
||||
export function FeatureInstallPrompt({
|
||||
bundle,
|
||||
isAdmin,
|
||||
toolId,
|
||||
toolName,
|
||||
toolDescription,
|
||||
}: FeatureInstallPromptProps) {
|
||||
const { t } = useTranslation();
|
||||
const { installBundle, clearError, installing, errors, startTimes, queued } = useFeaturesStore();
|
||||
const {
|
||||
bundles,
|
||||
installBundle,
|
||||
installTool,
|
||||
clearError,
|
||||
installing,
|
||||
errors,
|
||||
startTimes,
|
||||
queued,
|
||||
} = useFeaturesStore();
|
||||
const progress = installing[bundle.id] ?? null;
|
||||
const error = errors[bundle.id] ?? null;
|
||||
const isInstalling = !!progress;
|
||||
@@ -68,6 +99,37 @@ export function FeatureInstallPrompt({
|
||||
const displayName = toolName || bundle.name;
|
||||
const displayDescription = toolDescription || bundle.description;
|
||||
const isRepair = bundle.status === "error";
|
||||
const requiredBundleIds = toolId ? getRequiredBundlesForTool(toolId) : [bundle.id];
|
||||
const requiredBundles = requiredBundleIds
|
||||
.map(
|
||||
(bundleId) =>
|
||||
bundles.find((candidate) => candidate.id === bundleId) ??
|
||||
(bundle.id === bundleId ? bundle : fallbackBundleState(bundleId)),
|
||||
)
|
||||
.filter((candidate): candidate is FeatureBundleState => candidate !== null);
|
||||
const bundlesNeedingDownload = requiredBundles.filter(
|
||||
(candidate) => candidate.status !== "installed",
|
||||
);
|
||||
const downloadSizeLabel =
|
||||
bundlesNeedingDownload
|
||||
.map((candidate) =>
|
||||
candidate.downloadBytes ? formatFileSize(candidate.downloadBytes) : candidate.estimatedSize,
|
||||
)
|
||||
.join(" + ") ||
|
||||
(bundle.downloadBytes ? formatFileSize(bundle.downloadBytes) : bundle.estimatedSize);
|
||||
// Show the per-bundle breakdown for any multi-bundle tool, including the
|
||||
// repair state: when one bundle of a multi-bundle tool (e.g. passport-photo)
|
||||
// errors, the user still needs to see that the sibling bundle is installing,
|
||||
// queued, or already done. Hiding it during repair is exactly when it hurts.
|
||||
const showBundleBreakdown = toolId !== undefined && requiredBundles.length > 1;
|
||||
|
||||
function bundleStatusLabel(candidate: FeatureBundleState): string {
|
||||
if (installing[candidate.id]) return t.settings.aiFeatures.installing;
|
||||
if (queued.includes(candidate.id)) return t.settings.aiFeatures.queued;
|
||||
if (candidate.status === "installed") return t.settings.aiFeatures.installed;
|
||||
if (candidate.status === "error") return t.settings.aiFeatures.repair;
|
||||
return t.settings.aiFeatures.notInstalled;
|
||||
}
|
||||
|
||||
const [messageIndex, setMessageIndex] = useState(() =>
|
||||
Math.floor(Math.random() * PROGRESS_MESSAGES.length),
|
||||
@@ -94,7 +156,11 @@ export function FeatureInstallPrompt({
|
||||
|
||||
function handleInstall() {
|
||||
clearError(bundle.id);
|
||||
installBundle(bundle.id);
|
||||
if (toolId) {
|
||||
installTool(toolId);
|
||||
} else {
|
||||
installBundle(bundle.id);
|
||||
}
|
||||
}
|
||||
|
||||
// Defensive guard: if the bundle is already installed (status may have
|
||||
@@ -126,14 +192,49 @@ export function FeatureInstallPrompt({
|
||||
{!isRepair && (
|
||||
<p className="text-sm text-muted-foreground">
|
||||
{format(t.features.requiresDownload, {
|
||||
size: bundle.downloadBytes
|
||||
? formatFileSize(bundle.downloadBytes)
|
||||
: bundle.estimatedSize,
|
||||
size: downloadSizeLabel,
|
||||
})}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{showBundleBreakdown && (
|
||||
<div className="w-full max-w-md rounded-lg border border-border bg-background text-start overflow-hidden">
|
||||
{requiredBundles.map((candidate) => {
|
||||
const isCandidateInstalling = !!installing[candidate.id];
|
||||
const isCandidateQueued = queued.includes(candidate.id);
|
||||
const isCandidateInstalled = candidate.status === "installed";
|
||||
const statusClass = isCandidateInstalled
|
||||
? "bg-emerald-500/10 text-emerald-700 dark:text-emerald-300"
|
||||
: isCandidateInstalling || isCandidateQueued
|
||||
? "bg-primary/10 text-primary"
|
||||
: candidate.status === "error"
|
||||
? "bg-destructive/10 text-destructive"
|
||||
: "bg-muted text-muted-foreground";
|
||||
return (
|
||||
<div
|
||||
key={candidate.id}
|
||||
className="flex items-center justify-between gap-3 px-4 py-3 border-b border-border last:border-b-0"
|
||||
>
|
||||
<div className="min-w-0">
|
||||
<p className="text-sm font-medium text-foreground truncate">{candidate.name}</p>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{candidate.downloadBytes
|
||||
? formatFileSize(candidate.downloadBytes)
|
||||
: candidate.estimatedSize}
|
||||
</p>
|
||||
</div>
|
||||
<span
|
||||
className={`shrink-0 rounded-full px-2.5 py-1 text-xs font-medium ${statusClass}`}
|
||||
>
|
||||
{bundleStatusLabel(candidate)}
|
||||
</span>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{(error || (isRepair && bundle.error)) && (
|
||||
<div className="flex items-center gap-2 bg-destructive/10 text-destructive rounded-lg px-4 py-3 max-w-md w-full">
|
||||
<AlertCircle className="h-4 w-4 shrink-0" />
|
||||
|
||||
@@ -600,6 +600,7 @@ export function ToolPage() {
|
||||
<FeatureInstallPrompt
|
||||
bundle={featureBundle}
|
||||
isAdmin={isAdmin}
|
||||
toolId={toolId}
|
||||
toolName={tool?.name}
|
||||
toolDescription={tool?.description}
|
||||
/>
|
||||
|
||||
@@ -20,6 +20,13 @@ interface BundleProgress {
|
||||
stage: string;
|
||||
}
|
||||
|
||||
interface ToolBundleInstallResult {
|
||||
bundleId: string;
|
||||
jobId?: string;
|
||||
queued?: boolean;
|
||||
skipped?: boolean;
|
||||
}
|
||||
|
||||
interface FeaturesState {
|
||||
bundles: FeatureBundleState[];
|
||||
loaded: boolean;
|
||||
@@ -35,6 +42,7 @@ interface FeaturesState {
|
||||
isToolInstalled: (toolId: string) => boolean;
|
||||
getBundleForTool: (toolId: string) => FeatureBundleState | null;
|
||||
installBundle: (bundleId: string) => Promise<void>;
|
||||
installTool: (toolId: string) => Promise<void>;
|
||||
uninstallBundle: (bundleId: string) => Promise<void>;
|
||||
reinstallBundle: (bundleId: string) => Promise<void>;
|
||||
installAll: () => Promise<void>;
|
||||
@@ -346,6 +354,74 @@ export const useFeaturesStore = create<FeaturesState>((set, get) => {
|
||||
}
|
||||
},
|
||||
|
||||
installTool: async (toolId: string) => {
|
||||
const required = requiredBundlesForTool(toolId);
|
||||
const targets = required.filter(
|
||||
(bundleId) => get().bundles.find((b) => b.id === bundleId)?.status !== "installed",
|
||||
);
|
||||
if (targets.length === 0) return;
|
||||
|
||||
const errors = { ...get().errors };
|
||||
const installing = { ...get().installing };
|
||||
const startTimes = { ...get().startTimes };
|
||||
const now = Date.now();
|
||||
|
||||
for (const bundleId of targets) {
|
||||
delete errors[bundleId];
|
||||
if (!get().queued.includes(bundleId)) {
|
||||
installing[bundleId] = installing[bundleId] ?? { percent: 5, stage: "Starting..." };
|
||||
}
|
||||
startTimes[bundleId] = startTimes[bundleId] ?? now;
|
||||
}
|
||||
set({ errors, installing, startTimes });
|
||||
|
||||
try {
|
||||
const result = await apiPost<{ bundles: ToolBundleInstallResult[] }>(
|
||||
`/v1/admin/tools/${toolId}/features/install`,
|
||||
{},
|
||||
);
|
||||
|
||||
for (const item of result.bundles) {
|
||||
if (item.skipped) {
|
||||
stopTracking(item.bundleId);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (item.queued) {
|
||||
const nextInstalling = { ...get().installing };
|
||||
delete nextInstalling[item.bundleId];
|
||||
set({
|
||||
installing: nextInstalling,
|
||||
queued: get().queued.includes(item.bundleId)
|
||||
? get().queued
|
||||
: [...get().queued, item.bundleId],
|
||||
});
|
||||
startPolling(item.bundleId);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (item.jobId) {
|
||||
listenToProgress(item.bundleId, item.jobId);
|
||||
}
|
||||
}
|
||||
|
||||
if (result.bundles.every((item) => item.skipped)) {
|
||||
await refreshBundles();
|
||||
}
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : "Failed to start installation";
|
||||
const nextErrors = { ...get().errors };
|
||||
|
||||
for (const bundleId of targets) {
|
||||
stopTracking(bundleId);
|
||||
nextErrors[bundleId] = message;
|
||||
}
|
||||
|
||||
set({ errors: nextErrors });
|
||||
maybeFinishInstallAll();
|
||||
}
|
||||
},
|
||||
|
||||
uninstallBundle: async (bundleId: string) => {
|
||||
try {
|
||||
await apiPost(`/v1/admin/features/${bundleId}/uninstall`, {});
|
||||
|
||||
@@ -366,6 +366,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
Pillow==12.2.0 \
|
||||
numpy==1.26.4 \
|
||||
opencv-python-headless==4.10.0.84 \
|
||||
"huggingface-hub[hf_xet,hf_transfer]==0.36.2" \
|
||||
pikepdf==10.8.0 \
|
||||
PyMuPDF==1.27.2.3 \
|
||||
weasyprint==69.0 \
|
||||
|
||||
@@ -2,7 +2,12 @@
|
||||
"manifestVersion": 2,
|
||||
"imageVersion": "2.0.0",
|
||||
"pythonVersion": { "amd64": "3.12", "arm64": "3.11" },
|
||||
"basePackages": ["numpy==1.26.4", "Pillow==12.2.0", "opencv-python-headless==4.10.0.84"],
|
||||
"basePackages": [
|
||||
"numpy==1.26.4",
|
||||
"Pillow==12.2.0",
|
||||
"opencv-python-headless==4.10.0.84",
|
||||
"huggingface-hub[hf_xet,hf_transfer]==0.36.2"
|
||||
],
|
||||
"constraints": [
|
||||
"numpy==1.26.4",
|
||||
"scipy==1.12.0",
|
||||
@@ -79,6 +84,7 @@
|
||||
"args": ["birefnet-hr-matting"]
|
||||
}
|
||||
],
|
||||
"smokeImports": ["rembg", "onnxruntime"],
|
||||
"enablesTools": [
|
||||
"remove-background",
|
||||
"passport-photo",
|
||||
@@ -126,6 +132,7 @@
|
||||
"minSize": 1000000
|
||||
}
|
||||
],
|
||||
"smokeImports": ["mediapipe"],
|
||||
"enablesTools": ["blur-faces", "red-eye-removal", "smart-crop"]
|
||||
},
|
||||
"object-eraser-colorize": {
|
||||
@@ -147,7 +154,7 @@
|
||||
}
|
||||
},
|
||||
"packages": {
|
||||
"common": ["huggingface-hub"],
|
||||
"common": ["huggingface-hub[hf_xet]==0.36.2"],
|
||||
"amd64": ["onnxruntime-gpu==1.20.1"],
|
||||
"arm64": ["onnxruntime==1.20.1", "protobuf>=4.25.3,<5"]
|
||||
},
|
||||
@@ -192,6 +199,7 @@
|
||||
"minSize": 0
|
||||
}
|
||||
],
|
||||
"smokeImports": ["onnxruntime"],
|
||||
"enablesTools": ["erase-object", "colorize", "ai-canvas-expand"]
|
||||
},
|
||||
"upscale-enhance": {
|
||||
@@ -213,7 +221,12 @@
|
||||
}
|
||||
},
|
||||
"packages": {
|
||||
"common": ["codeformer-pip==0.0.4", "huggingface-hub", "einops", "setuptools<75"],
|
||||
"common": [
|
||||
"codeformer-pip==0.0.4",
|
||||
"huggingface-hub[hf_xet]==0.36.2",
|
||||
"einops",
|
||||
"setuptools<75"
|
||||
],
|
||||
"amd64": [
|
||||
"torch==2.7.0+cu126 torchvision==0.22.0+cu126 --index-url https://download.pytorch.org/whl/cu126",
|
||||
"lpips",
|
||||
@@ -301,6 +314,7 @@
|
||||
"replace": "from torchvision.transforms.functional import rgb_to_grayscale"
|
||||
}
|
||||
],
|
||||
"smokeImports": ["torch", "torchvision"],
|
||||
"enablesTools": ["upscale", "enhance-faces", "noise-removal"]
|
||||
},
|
||||
"photo-restoration": {
|
||||
@@ -322,7 +336,7 @@
|
||||
}
|
||||
},
|
||||
"packages": {
|
||||
"common": ["codeformer-pip==0.0.4", "huggingface-hub", "setuptools<75"],
|
||||
"common": ["codeformer-pip==0.0.4", "huggingface-hub[hf_xet]==0.36.2", "setuptools<75"],
|
||||
"amd64": [
|
||||
"onnxruntime-gpu==1.20.1",
|
||||
"mediapipe>=0.10.21",
|
||||
@@ -406,6 +420,7 @@
|
||||
"replace": "from torchvision.transforms.functional import rgb_to_grayscale"
|
||||
}
|
||||
],
|
||||
"smokeImports": ["onnxruntime", "torch", "mediapipe"],
|
||||
"enablesTools": ["restore-photo"]
|
||||
},
|
||||
"ocr": {
|
||||
@@ -427,7 +442,7 @@
|
||||
}
|
||||
},
|
||||
"packages": {
|
||||
"common": ["huggingface-hub"],
|
||||
"common": ["huggingface-hub[hf_xet]==0.36.2"],
|
||||
"amd64": [
|
||||
"paddlepaddle-gpu>=3.2.1 --index-url https://www.paddlepaddle.org.cn/packages/stable/cu126/ --extra-index-url https://pypi.org/simple/",
|
||||
"paddleocr[doc-parser]>=3.4.0,<3.5.0"
|
||||
@@ -482,6 +497,7 @@
|
||||
"args": ["PaddlePaddle/PaddleOCR-VL-1.5", "PaddleOCR-VL-1.5"]
|
||||
}
|
||||
],
|
||||
"smokeImports": ["paddle", "paddleocr"],
|
||||
"enablesTools": ["ocr", "ocr-pdf"]
|
||||
},
|
||||
"transcription": {
|
||||
@@ -516,6 +532,7 @@
|
||||
"args": ["Systran/faster-whisper-small", "faster-whisper-small"]
|
||||
}
|
||||
],
|
||||
"smokeImports": ["faster_whisper", "ctranslate2"],
|
||||
"enablesTools": ["transcribe-audio", "auto-subtitles"]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ Final result is a JSON object on stdout.
|
||||
import errno
|
||||
import glob
|
||||
import hashlib
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
@@ -27,6 +28,9 @@ import urllib.error
|
||||
import urllib.request
|
||||
from datetime import datetime, timezone
|
||||
|
||||
DOWNLOAD_CHUNK_BYTES = 4 * 1024 * 1024
|
||||
DOWNLOAD_META_BYTES = 64 * 1024 * 1024
|
||||
|
||||
|
||||
# -- Helpers --
|
||||
|
||||
@@ -149,6 +153,166 @@ def verify_sha256(filepath: str, expected: str) -> bool:
|
||||
|
||||
# -- Download with resume --
|
||||
|
||||
def _set_env_temporarily(key: str, value: str):
|
||||
previous = os.environ.get(key)
|
||||
os.environ[key] = value
|
||||
return previous
|
||||
|
||||
|
||||
def _restore_env(key: str, previous) -> None:
|
||||
if previous is None:
|
||||
os.environ.pop(key, None)
|
||||
else:
|
||||
os.environ[key] = previous
|
||||
|
||||
|
||||
def _cleanup_hf_local_dir(local_dir: str, archive_file: str) -> None:
|
||||
if "/" in archive_file:
|
||||
top_level = archive_file.split("/", 1)[0]
|
||||
shutil.rmtree(os.path.join(local_dir, top_level), ignore_errors=True)
|
||||
shutil.rmtree(os.path.join(local_dir, ".cache"), ignore_errors=True)
|
||||
|
||||
|
||||
def ensure_hf_hub(venv_path: str) -> None:
|
||||
"""Guarantee the accelerated Hugging Face client is importable before the
|
||||
download so the multi-GB bundle transfer takes the fast Xet path.
|
||||
|
||||
The installer runs under the on-disk venv (PYTHON_VENV_PATH, i.e.
|
||||
/data/ai/venv in Docker). That venv is normally seeded from the image's
|
||||
/opt/venv, which bakes huggingface-hub[hf_xet]. But an install whose venv
|
||||
predates the base package (an upgrade where the reseed stamp didn't move, a
|
||||
hand-copied or offline-imported venv) would import-fail in
|
||||
download_with_hf_hub and silently fall back to the slow single-stream urllib
|
||||
downloader. Self-heal by pip-installing the client into this same venv.
|
||||
|
||||
A bundle install already requires network and lifts the offline guard (see
|
||||
main()), so this adds no new offline dependency; if the pip install fails we
|
||||
fall through to the resumable urllib downloader, the correct degraded path.
|
||||
"""
|
||||
try:
|
||||
import huggingface_hub # noqa: F401
|
||||
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
python_path = os.path.join(venv_path, "bin", "python3")
|
||||
if not os.path.exists(python_path):
|
||||
return
|
||||
|
||||
emit_progress(1, "Preparing accelerated download client...")
|
||||
try:
|
||||
subprocess.run(
|
||||
[
|
||||
python_path, "-m", "pip", "install", "--quiet",
|
||||
"huggingface-hub[hf_xet,hf_transfer]==0.36.2",
|
||||
],
|
||||
capture_output=True, text=True, timeout=300, check=True,
|
||||
)
|
||||
# The finder caches the venv's site-packages listing; drop it so the
|
||||
# just-installed package is visible to the import in download_with_hf_hub.
|
||||
importlib.invalidate_caches()
|
||||
except Exception as e:
|
||||
emit_progress(1, f"Accelerated client unavailable ({e}); using resumable download.")
|
||||
|
||||
|
||||
def download_with_hf_hub(
|
||||
bundle_repo: str,
|
||||
archive_file: str,
|
||||
dest: str,
|
||||
expected_size: int,
|
||||
progress_start: int,
|
||||
progress_end: int,
|
||||
force_download: bool = False,
|
||||
) -> bool:
|
||||
"""Download through huggingface_hub when available.
|
||||
|
||||
huggingface_hub 0.32+ can use hf_xet for faster large-file transfers and
|
||||
manages retries/resume internally. Return False when the client is missing
|
||||
or fails so callers can fall back to the manual urllib downloader.
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Enable hf_transfer (Rust multi-connection downloader) ONLY when the
|
||||
# package is actually importable. For a Xet-backed repo hf_xet takes
|
||||
# precedence and this is a no-op, but if the Xet CAS endpoint is unreachable
|
||||
# (e.g. a firewall that allows huggingface.co but blocks transfer.xethub.hf.co)
|
||||
# hf_hub_download falls back to plain HTTP, and hf_transfer makes that
|
||||
# fallback multi-connection instead of single-stream. Gating on the import
|
||||
# avoids the "HF_HUB_ENABLE_HF_TRANSFER set but package missing" hard error
|
||||
# on a venv that only has hf_xet.
|
||||
try:
|
||||
import hf_transfer # noqa: F401
|
||||
|
||||
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
local_dir = os.path.dirname(dest)
|
||||
os.makedirs(local_dir, exist_ok=True)
|
||||
emit_progress(progress_start, "Downloading with accelerated Hugging Face client...")
|
||||
|
||||
previous_progress = _set_env_temporarily("HF_HUB_DISABLE_PROGRESS_BARS", "1")
|
||||
try:
|
||||
downloaded_path = hf_hub_download(
|
||||
repo_id=bundle_repo,
|
||||
filename=archive_file,
|
||||
repo_type="model",
|
||||
local_dir=local_dir,
|
||||
force_download=force_download,
|
||||
)
|
||||
except Exception as e:
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Accelerated download unavailable, using resumable fallback: {e}",
|
||||
)
|
||||
# Reclaim any partial blob/metadata hf_hub_download staged under
|
||||
# local_dir/.cache so the urllib fallback starts clean and disk is freed.
|
||||
_cleanup_hf_local_dir(local_dir, archive_file)
|
||||
return False
|
||||
finally:
|
||||
_restore_env("HF_HUB_DISABLE_PROGRESS_BARS", previous_progress)
|
||||
|
||||
try:
|
||||
if not os.path.exists(downloaded_path):
|
||||
emit_progress(
|
||||
progress_start,
|
||||
"Accelerated download did not produce an archive, using resumable fallback...",
|
||||
)
|
||||
return False
|
||||
|
||||
if os.path.abspath(downloaded_path) != os.path.abspath(dest):
|
||||
if os.path.exists(dest):
|
||||
os.unlink(dest)
|
||||
os.replace(downloaded_path, dest)
|
||||
|
||||
size = os.path.getsize(dest)
|
||||
if expected_size > 0:
|
||||
pct = min(size / expected_size, 1.0)
|
||||
progress = int(progress_start + pct * (progress_end - progress_start))
|
||||
else:
|
||||
progress = progress_end
|
||||
emit_progress(
|
||||
min(progress, progress_end),
|
||||
f"Downloaded with accelerated client ({size / (1024**3):.1f} GB)",
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Accelerated download post-processing failed, using resumable fallback: {e}",
|
||||
)
|
||||
return False
|
||||
finally:
|
||||
# Always drop the hf staging tree (local_dir/<top>, local_dir/.cache).
|
||||
# On success the archive is already moved to dest; on any failure this
|
||||
# stops the transient hf cache copy from leaking across the fallback.
|
||||
_cleanup_hf_local_dir(local_dir, archive_file)
|
||||
|
||||
|
||||
def download_with_resume(
|
||||
url: str,
|
||||
dest: str,
|
||||
@@ -180,7 +344,15 @@ def download_with_resume(
|
||||
if bytes_downloaded == 0 and os.path.exists(partial_path):
|
||||
os.unlink(partial_path)
|
||||
|
||||
max_retries = 3
|
||||
def _cleanup_partial() -> None:
|
||||
for p in (partial_path, meta_path):
|
||||
if os.path.exists(p):
|
||||
try:
|
||||
os.unlink(p)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
max_retries = 5
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
headers = {"User-Agent": "snapotter-installer/2.0"}
|
||||
@@ -193,10 +365,19 @@ def download_with_resume(
|
||||
|
||||
req = urllib.request.Request(url, headers=headers)
|
||||
with urllib.request.urlopen(req, timeout=300) as resp:
|
||||
status = getattr(resp, "status", None) or resp.getcode()
|
||||
# If we asked to resume (sent a Range) but the server sent the
|
||||
# whole file back (200 instead of 206 Partial Content -- a proxy
|
||||
# or CDN that ignores Range), restart from byte 0. Appending a
|
||||
# full body onto the existing partial would corrupt the archive
|
||||
# and fail the checksum on every retry.
|
||||
if bytes_downloaded > 0 and status != 206:
|
||||
bytes_downloaded = 0
|
||||
mode = "ab" if bytes_downloaded > 0 else "wb"
|
||||
next_meta_at = bytes_downloaded + DOWNLOAD_META_BYTES
|
||||
with open(partial_path, mode) as f:
|
||||
while True:
|
||||
chunk = resp.read(65536)
|
||||
chunk = resp.read(DOWNLOAD_CHUNK_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
f.write(chunk)
|
||||
@@ -212,10 +393,20 @@ def download_with_resume(
|
||||
stage = f"Downloading... {bytes_downloaded / (1024**3):.1f} GB"
|
||||
emit_progress(progress, stage)
|
||||
|
||||
# Write meta periodically (every 10 MB)
|
||||
if bytes_downloaded % (10 * 1024 * 1024) < 65536:
|
||||
# Write meta periodically so a crash can resume.
|
||||
if bytes_downloaded >= next_meta_at:
|
||||
with open(meta_path, "w") as mf:
|
||||
json.dump({"bytesDownloaded": bytes_downloaded}, mf)
|
||||
next_meta_at = bytes_downloaded + DOWNLOAD_META_BYTES
|
||||
|
||||
# Guard against a truncated body or an error page served as the
|
||||
# archive: the completed size must match what the manifest expects.
|
||||
# A mismatch is retryable (transient truncation / a stale CDN edge).
|
||||
if expected_size > 0 and bytes_downloaded != expected_size:
|
||||
raise RuntimeError(
|
||||
f"incomplete download: got {bytes_downloaded} bytes, "
|
||||
f"expected {expected_size} (truncated response or error page)"
|
||||
)
|
||||
|
||||
# Download complete
|
||||
os.rename(partial_path, dest)
|
||||
@@ -223,27 +414,58 @@ def download_with_resume(
|
||||
os.unlink(meta_path)
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
# Write meta for resume on next attempt
|
||||
with open(meta_path, "w") as mf:
|
||||
json.dump({"bytesDownloaded": bytes_downloaded}, mf)
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
delay = 10 * (2 ** attempt)
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Download failed (attempt {attempt + 1}/{max_retries}), "
|
||||
f"retrying in {delay}s: {e}",
|
||||
)
|
||||
time.sleep(delay)
|
||||
else:
|
||||
# Clean up on final failure
|
||||
for p in (partial_path, meta_path):
|
||||
if os.path.exists(p):
|
||||
os.unlink(p)
|
||||
except urllib.error.HTTPError as e:
|
||||
# HTTPError subclasses OSError, so it MUST be caught before the
|
||||
# OSError clause below. 4xx (except 408 Timeout / 429 Too Many
|
||||
# Requests) won't fix on retry -- a wrong URL, a private repo, or a
|
||||
# removed archive -- so fail fast with the manual-download hint.
|
||||
if 400 <= e.code < 500 and e.code not in (408, 429):
|
||||
_cleanup_partial()
|
||||
raise RuntimeError(
|
||||
f"Failed to download after {max_retries} attempts: {e}"
|
||||
f"Download failed with HTTP {e.code} ({e.reason}). The archive "
|
||||
f"URL may be wrong or access-restricted."
|
||||
)
|
||||
_retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, _cleanup_partial)
|
||||
except OSError as e:
|
||||
# Disk full is not transient: retrying can't create space. Fail fast
|
||||
# with an actionable message instead of burning the backoff budget.
|
||||
# (URLError/connection errors also land here; their errno is None, so
|
||||
# they fall through to the retry path.)
|
||||
if getattr(e, "errno", None) == errno.ENOSPC:
|
||||
_cleanup_partial()
|
||||
raise RuntimeError(
|
||||
"Ran out of disk space while downloading the bundle. "
|
||||
"Free up space and retry."
|
||||
)
|
||||
_retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, _cleanup_partial)
|
||||
except Exception as e:
|
||||
_retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, _cleanup_partial)
|
||||
|
||||
|
||||
def _retry_or_raise(err, attempt, max_retries, bytes_downloaded, meta_path,
|
||||
progress_start, cleanup) -> None:
|
||||
"""Shared transient-failure handler for download_with_resume: persist resume
|
||||
metadata and back off, or clean up and raise on the final attempt."""
|
||||
try:
|
||||
with open(meta_path, "w") as mf:
|
||||
json.dump({"bytesDownloaded": bytes_downloaded}, mf)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
delay = min(60, 5 * (2 ** attempt))
|
||||
emit_progress(
|
||||
progress_start,
|
||||
f"Download failed (attempt {attempt + 1}/{max_retries}), "
|
||||
f"retrying in {delay}s: {err}",
|
||||
)
|
||||
time.sleep(delay)
|
||||
else:
|
||||
cleanup()
|
||||
raise RuntimeError(f"Failed to download after {max_retries} attempts: {err}")
|
||||
|
||||
|
||||
# -- Safe tar extraction --
|
||||
@@ -270,11 +492,15 @@ def safe_extract(tar_path: str, staging_dir: str) -> None:
|
||||
# -- File move --
|
||||
|
||||
def move_tree(src: str, dst: str) -> None:
|
||||
"""Merge src into dst, overwriting existing files. Renames entries where
|
||||
possible so that on the same filesystem no copy (and thus no transient
|
||||
doubling of the payload on disk) occurs; falls back to a copy only across
|
||||
filesystems. The old copytree+rmtree approach duplicated the whole tree on
|
||||
disk during the move, which could exhaust the host on a tight-disk node."""
|
||||
"""Merge src into dst, replacing entries crash-atomically where possible.
|
||||
|
||||
This writes into the SHARED /data/ai/venv site-packages, so a crash mid-move
|
||||
must never leave a package in a half-replaced state (that tears the venv and
|
||||
breaks every other AI tool). For a file replacing a file, os.replace swaps in
|
||||
place with NO delete-then-write window, so an interruption leaves either the
|
||||
old or the new file intact, never a missing one. Cross-filesystem copies go
|
||||
through a temp sibling then an atomic rename for the same reason. Renames
|
||||
(vs copytree) also avoid transiently doubling the payload on disk."""
|
||||
if not os.path.isdir(src):
|
||||
return
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
@@ -285,22 +511,31 @@ def move_tree(src: str, dst: str) -> None:
|
||||
# Both dirs exist: merge recursively rather than replace.
|
||||
move_tree(s, d)
|
||||
continue
|
||||
if os.path.exists(d):
|
||||
if os.path.isdir(d):
|
||||
shutil.rmtree(d)
|
||||
else:
|
||||
os.remove(d)
|
||||
try:
|
||||
os.rename(s, d)
|
||||
except OSError as e:
|
||||
if getattr(e, "errno", None) == errno.EXDEV:
|
||||
# Cross-filesystem: rename isn't allowed, fall back to a copy.
|
||||
if os.path.isdir(s):
|
||||
shutil.copytree(s, d)
|
||||
# A type mismatch (dir<->file) can't be atomically swapped by rename,
|
||||
# so clear the destination first. A file-over-file or new entry needs
|
||||
# no pre-delete: os.replace is atomic and leaves no torn window.
|
||||
if os.path.exists(d) and os.path.isdir(d) != os.path.isdir(s):
|
||||
if os.path.isdir(d):
|
||||
shutil.rmtree(d)
|
||||
else:
|
||||
shutil.copy2(s, d)
|
||||
else:
|
||||
os.remove(d)
|
||||
os.replace(s, d)
|
||||
continue
|
||||
except OSError as e:
|
||||
if getattr(e, "errno", None) != errno.EXDEV:
|
||||
raise
|
||||
# Cross-filesystem: rename isn't allowed. Copy to a temp sibling and then
|
||||
# atomically replace, so a mid-copy ENOSPC never leaves a truncated file
|
||||
# where a working one used to be.
|
||||
if os.path.isdir(s):
|
||||
if os.path.exists(d):
|
||||
shutil.rmtree(d)
|
||||
shutil.copytree(s, d)
|
||||
else:
|
||||
tmp = d + ".part"
|
||||
shutil.copy2(s, tmp)
|
||||
os.replace(tmp, d)
|
||||
# Remove whatever remains of src (emptied by renames, or copied originals).
|
||||
shutil.rmtree(src, ignore_errors=True)
|
||||
|
||||
@@ -461,8 +696,20 @@ def _install() -> None:
|
||||
|
||||
emit_progress(2, f"Downloading {bundle.get('name', bundle_id)} bundle...")
|
||||
|
||||
# Make sure the accelerated Xet client is importable in this venv, so a
|
||||
# drifted/upgraded venv doesn't silently fall back to slow urllib.
|
||||
ensure_hf_hub(venv_path)
|
||||
|
||||
try:
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
if not download_with_hf_hub(
|
||||
bundle_repo,
|
||||
archive_file,
|
||||
tar_path,
|
||||
compressed_size,
|
||||
2,
|
||||
85,
|
||||
):
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
except RuntimeError as e:
|
||||
fail(
|
||||
f"{e}\n\n"
|
||||
@@ -478,7 +725,16 @@ def _install() -> None:
|
||||
os.unlink(tar_path)
|
||||
emit_progress(86, "Checksum mismatch, retrying download...")
|
||||
try:
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
if not download_with_hf_hub(
|
||||
bundle_repo,
|
||||
archive_file,
|
||||
tar_path,
|
||||
compressed_size,
|
||||
2,
|
||||
85,
|
||||
force_download=True,
|
||||
):
|
||||
download_with_resume(url, tar_path, compressed_size, 2, 85)
|
||||
except RuntimeError as e:
|
||||
fail(str(e))
|
||||
|
||||
@@ -501,6 +757,11 @@ def _install() -> None:
|
||||
except Exception as e:
|
||||
if os.path.exists(staging_dir):
|
||||
shutil.rmtree(staging_dir, ignore_errors=True)
|
||||
if isinstance(e, OSError) and getattr(e, "errno", None) == errno.ENOSPC:
|
||||
fail(
|
||||
"Ran out of disk space while extracting the bundle. "
|
||||
"Free up space and retry."
|
||||
)
|
||||
fail(f"Failed to extract archive: {e}")
|
||||
|
||||
# -- Read bundle.json from tar --
|
||||
@@ -539,10 +800,28 @@ def _install() -> None:
|
||||
# -- Move site-packages --
|
||||
emit_progress(92, "Installing packages...")
|
||||
site_packages_dir = get_site_packages_dir(venv_path)
|
||||
venv_writing_marker = os.path.join(ai_dir, "venv.writing")
|
||||
|
||||
try:
|
||||
if os.path.isdir(staging_sp) and site_packages_dir:
|
||||
# Breadcrumb the destructive shared-venv write. If the process is
|
||||
# killed mid-move (OOM/SIGKILL/power loss), move_tree can leave the
|
||||
# venv torn, which breaks OTHER installed tools. The marker survives
|
||||
# the crash; on next boot recoverInterruptedInstalls sees it and
|
||||
# reseeds the venv back to a known-good base. We clear it the instant
|
||||
# the site-packages move completes, since the venv is consistent
|
||||
# again then (a later models-move failure can't tear the venv).
|
||||
with open(venv_writing_marker, "w") as mf:
|
||||
json.dump(
|
||||
{
|
||||
"bundleId": bundle_id,
|
||||
"startedAt": datetime.now(timezone.utc).isoformat(),
|
||||
},
|
||||
mf,
|
||||
)
|
||||
move_tree(staging_sp, site_packages_dir)
|
||||
if os.path.exists(venv_writing_marker):
|
||||
os.unlink(venv_writing_marker)
|
||||
|
||||
# -- Move models --
|
||||
emit_progress(95, "Installing models...")
|
||||
@@ -560,6 +839,39 @@ def _install() -> None:
|
||||
emit_progress(97, "Finalizing...")
|
||||
apply_fixups(staging_dir, venv_path)
|
||||
|
||||
# -- Verify the bundle actually imports --
|
||||
# File-copy completion does NOT prove the bundle works: an incomplete
|
||||
# extraction or an ABI mismatch (e.g. a numpy/torch/protobuf skew) can leave
|
||||
# every file present yet the module unimportable, so the tool "installs" but
|
||||
# fails at first use. Import the bundle's key native libraries in the venv
|
||||
# now; if that fails, refuse to mark the bundle installed so the user gets a
|
||||
# clear retry instead of a silently broken tool.
|
||||
smoke_imports = bundle.get("smokeImports") or []
|
||||
if smoke_imports and os.environ.get("SNAPOTTER_SKIP_INSTALL_SMOKE") != "1":
|
||||
emit_progress(99, "Verifying installation...")
|
||||
venv_python = os.path.join(venv_path, "bin", "python3")
|
||||
if os.path.exists(venv_python):
|
||||
import_stmt = "\n".join(f"import {mod}" for mod in smoke_imports)
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[venv_python, "-c", import_stmt],
|
||||
capture_output=True, text=True, timeout=300,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
shutil.rmtree(staging_dir, ignore_errors=True)
|
||||
fail("Installation verification timed out. Please retry the install.")
|
||||
if proc.returncode != 0:
|
||||
shutil.rmtree(staging_dir, ignore_errors=True)
|
||||
tail = "\n".join((proc.stderr or "").strip().splitlines()[-6:])
|
||||
fail(
|
||||
"Installation verification failed: the bundle installed but its "
|
||||
"libraries could not be loaded, so the tool would not work.\n"
|
||||
f"{tail}\n\n"
|
||||
"This usually means an interrupted or corrupted install. Retry the "
|
||||
"install; if it keeps failing, use Settings > AI Features > Reset AI "
|
||||
"Environment, then reinstall."
|
||||
)
|
||||
|
||||
# -- Write installed.json --
|
||||
emit_progress(98, "Recording installation...")
|
||||
installed = read_installed(ai_dir)
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
|
||||
def load_installer():
|
||||
script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
|
||||
spec = importlib.util.spec_from_file_location("install_feature_under_test", script_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_download_with_hf_hub_uses_accelerated_client(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
downloaded = tmp_path / "hf-cache" / "bundle.tar.gz"
|
||||
downloaded.parent.mkdir()
|
||||
|
||||
calls = {}
|
||||
|
||||
def fake_hf_hub_download(**kwargs):
|
||||
calls.update(kwargs)
|
||||
downloaded.write_bytes(b"archive")
|
||||
return str(downloaded)
|
||||
|
||||
fake_module = types.ModuleType("huggingface_hub")
|
||||
fake_module.hf_hub_download = fake_hf_hub_download
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
|
||||
|
||||
progress = []
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: progress.append((p, s)))
|
||||
|
||||
dest = tmp_path / "staging" / "object-eraser-colorize-amd64-gpu.tar.gz"
|
||||
dest.parent.mkdir()
|
||||
|
||||
assert (
|
||||
installer.download_with_hf_hub(
|
||||
"snapotter/feature-bundles",
|
||||
"v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
|
||||
str(dest),
|
||||
100,
|
||||
2,
|
||||
85,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
assert dest.read_bytes() == b"archive"
|
||||
assert calls["repo_id"] == "snapotter/feature-bundles"
|
||||
assert calls["repo_type"] == "model"
|
||||
assert calls["filename"] == "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz"
|
||||
assert any("accelerated" in stage.lower() for _, stage in progress)
|
||||
|
||||
|
||||
def test_download_with_hf_hub_cleans_cache_when_download_raises(monkeypatch, tmp_path):
|
||||
"""A failed accelerated download must not leak its .cache staging tree onto
|
||||
disk before the urllib fallback runs."""
|
||||
installer = load_installer()
|
||||
staging = tmp_path / "staging"
|
||||
staging.mkdir()
|
||||
# Simulate a partial hf cache tree left behind by a failed transfer.
|
||||
leaked_cache = staging / ".cache" / "huggingface" / "download"
|
||||
leaked_cache.mkdir(parents=True)
|
||||
(leaked_cache / "blob.incomplete").write_bytes(b"partial")
|
||||
leaked_nested = staging / "v2.0.0"
|
||||
leaked_nested.mkdir()
|
||||
|
||||
def fake_hf_hub_download(**_kwargs):
|
||||
raise RuntimeError("xet CAS unreachable")
|
||||
|
||||
fake_module = types.ModuleType("huggingface_hub")
|
||||
fake_module.hf_hub_download = fake_hf_hub_download
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
|
||||
dest = staging / "object-eraser-colorize-amd64-gpu.tar.gz"
|
||||
|
||||
assert (
|
||||
installer.download_with_hf_hub(
|
||||
"deepsafe/feature-bundles",
|
||||
"v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
|
||||
str(dest),
|
||||
100,
|
||||
2,
|
||||
85,
|
||||
)
|
||||
is False
|
||||
)
|
||||
# Both the .cache tree and the nested archive dir are reclaimed.
|
||||
assert not (staging / ".cache").exists()
|
||||
assert not (staging / "v2.0.0").exists()
|
||||
|
||||
|
||||
def test_ensure_hf_hub_noops_when_client_already_importable(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
fake_module = types.ModuleType("huggingface_hub")
|
||||
monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
|
||||
|
||||
ran = {"pip": False}
|
||||
monkeypatch.setattr(
|
||||
installer.subprocess, "run", lambda *a, **k: ran.__setitem__("pip", True)
|
||||
)
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
|
||||
installer.ensure_hf_hub(str(tmp_path))
|
||||
assert ran["pip"] is False
|
||||
|
||||
|
||||
def test_ensure_hf_hub_self_heals_missing_client(monkeypatch, tmp_path):
|
||||
"""On a drifted venv where huggingface_hub is missing, ensure_hf_hub must
|
||||
pip-install it into that venv rather than let the caller fall back to the
|
||||
slow single-stream urllib downloader silently."""
|
||||
installer = load_installer()
|
||||
monkeypatch.delitem(sys.modules, "huggingface_hub", raising=False)
|
||||
|
||||
# Make the huggingface_hub import fail deterministically so ensure_hf_hub
|
||||
# takes its self-heal branch.
|
||||
import builtins
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def blocked_import(name, *a, **k):
|
||||
if name == "huggingface_hub":
|
||||
raise ImportError("No module named 'huggingface_hub'")
|
||||
return real_import(name, *a, **k)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", blocked_import)
|
||||
|
||||
venv = tmp_path / "venv"
|
||||
(venv / "bin").mkdir(parents=True)
|
||||
(venv / "bin" / "python3").write_text("")
|
||||
|
||||
pip_calls = []
|
||||
monkeypatch.setattr(
|
||||
installer.subprocess,
|
||||
"run",
|
||||
lambda cmd, **k: pip_calls.append(cmd) or types.SimpleNamespace(returncode=0),
|
||||
)
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
|
||||
installer.ensure_hf_hub(str(venv))
|
||||
|
||||
assert len(pip_calls) == 1
|
||||
cmd = pip_calls[0]
|
||||
assert cmd[0] == str(venv / "bin" / "python3")
|
||||
assert "install" in cmd
|
||||
spec = next(part for part in cmd if part.startswith("huggingface-hub["))
|
||||
assert "hf_xet" in spec
|
||||
assert "hf_transfer" in spec
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Crash-atomicity tests for move_tree, which writes into the SHARED venv.
|
||||
|
||||
The invariant under test: a crash (or ENOSPC) partway through move_tree must
|
||||
never leave a destination entry missing. Because the venv is shared by every AI
|
||||
tool, a half-replaced package is what tears the venv and breaks unrelated tools.
|
||||
"""
|
||||
|
||||
import errno
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def load_installer():
|
||||
script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
|
||||
spec = importlib.util.spec_from_file_location("install_feature_move_under_test", script_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_file_over_file_is_replaced_without_a_delete_window(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
(dst / "pkg.py").write_text("OLD")
|
||||
(src / "pkg.py").write_text("NEW")
|
||||
|
||||
removed = []
|
||||
real_remove = os.remove
|
||||
monkeypatch.setattr(installer.os, "remove", lambda p: removed.append(p) or real_remove(p))
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "pkg.py").read_text() == "NEW"
|
||||
# The old file was atomically replaced, never deleted-then-rewritten.
|
||||
assert str(dst / "pkg.py") not in removed
|
||||
|
||||
|
||||
def test_crash_mid_move_leaves_every_dest_old_or_new_never_missing(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
for name in ("a.py", "b.py", "c.py"):
|
||||
(dst / name).write_text(f"OLD_{name}")
|
||||
(src / name).write_text(f"NEW_{name}")
|
||||
|
||||
real_replace = os.replace
|
||||
state = {"n": 0}
|
||||
|
||||
def flaky_replace(s, d):
|
||||
state["n"] += 1
|
||||
if state["n"] == 2:
|
||||
raise OSError("simulated crash mid-move")
|
||||
return real_replace(s, d)
|
||||
|
||||
monkeypatch.setattr(installer.os, "replace", flaky_replace)
|
||||
|
||||
with pytest.raises(OSError):
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
# Regardless of listdir order, every destination file must still exist and
|
||||
# hold either its old or its new content -- never a torn/missing entry.
|
||||
for name in ("a.py", "b.py", "c.py"):
|
||||
assert (dst / name).exists()
|
||||
assert (dst / name).read_text() in (f"OLD_{name}", f"NEW_{name}")
|
||||
|
||||
|
||||
def test_merges_directories_and_overwrites_files(tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
(src / "pkg").mkdir(parents=True)
|
||||
(dst / "pkg").mkdir(parents=True)
|
||||
(dst / "pkg" / "keep.py").write_text("KEEP")
|
||||
(src / "pkg" / "keep.py").write_text("UPDATED")
|
||||
(src / "pkg" / "new.py").write_text("NEW")
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "pkg" / "keep.py").read_text() == "UPDATED"
|
||||
assert (dst / "pkg" / "new.py").read_text() == "NEW"
|
||||
|
||||
|
||||
def test_type_mismatch_dir_replaces_file(tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
(dst / "x").write_text("i-am-a-file")
|
||||
(src / "x").mkdir()
|
||||
(src / "x" / "inner.py").write_text("dir-content")
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "x").is_dir()
|
||||
assert (dst / "x" / "inner.py").read_text() == "dir-content"
|
||||
|
||||
|
||||
def test_exdev_file_copies_via_temp_then_atomic_replace(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
src = tmp_path / "src"
|
||||
dst = tmp_path / "dst"
|
||||
src.mkdir()
|
||||
dst.mkdir()
|
||||
(dst / "f.py").write_text("OLD")
|
||||
(src / "f.py").write_text("NEW")
|
||||
|
||||
real_replace = os.replace
|
||||
seen = {"exdev": False, "part": False}
|
||||
|
||||
def exdev_for_direct_move(s, d):
|
||||
if s == str(src / "f.py"):
|
||||
seen["exdev"] = True
|
||||
raise OSError(errno.EXDEV, "cross-device link")
|
||||
if s.endswith(".part"):
|
||||
seen["part"] = True
|
||||
return real_replace(s, d)
|
||||
|
||||
monkeypatch.setattr(installer.os, "replace", exdev_for_direct_move)
|
||||
|
||||
installer.move_tree(str(src), str(dst))
|
||||
|
||||
assert (dst / "f.py").read_text() == "NEW"
|
||||
assert seen["exdev"] and seen["part"]
|
||||
# No leftover temp file.
|
||||
assert not (dst / "f.py.part").exists()
|
||||
@@ -0,0 +1,238 @@
|
||||
"""Robustness tests for download_with_resume: the urllib fallback downloader.
|
||||
|
||||
These simulate the network/disk failure modes users hit (flaky connections,
|
||||
Range-ignoring proxies, truncated bodies, disk full, dead URLs) with no real
|
||||
network, and assert the downloader either self-recovers or fails fast with a
|
||||
clear, actionable error instead of corrupting the archive or hanging.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
import errno
|
||||
import importlib.util
|
||||
import os
|
||||
import urllib.error
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def load_installer():
|
||||
script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
|
||||
spec = importlib.util.spec_from_file_location("install_feature_resume_under_test", script_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class FakeResp:
|
||||
"""Minimal stand-in for the urlopen response context manager."""
|
||||
|
||||
def __init__(self, data: bytes, status: int = 200):
|
||||
self._data = data
|
||||
self._pos = 0
|
||||
self.status = status
|
||||
|
||||
def read(self, n: int) -> bytes:
|
||||
chunk = self._data[self._pos : self._pos + n]
|
||||
self._pos += len(chunk)
|
||||
return chunk
|
||||
|
||||
def getcode(self) -> int:
|
||||
return self.status
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_):
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_sleep(monkeypatch):
|
||||
"""Never actually sleep during backoff; record calls instead."""
|
||||
installer = load_installer()
|
||||
calls = []
|
||||
monkeypatch.setattr(installer.time, "sleep", lambda s: calls.append(s))
|
||||
return calls
|
||||
|
||||
|
||||
def _patch_urlopen(monkeypatch, installer, handler):
|
||||
"""handler(req, call_index) -> FakeResp or raises."""
|
||||
state = {"n": 0}
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
idx = state["n"]
|
||||
state["n"] += 1
|
||||
return handler(req, idx)
|
||||
|
||||
monkeypatch.setattr(installer.urllib.request, "urlopen", fake_urlopen)
|
||||
return state
|
||||
|
||||
|
||||
def test_fresh_download_succeeds(monkeypatch, tmp_path):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"x" * 4096
|
||||
_patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
|
||||
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
assert dest.read_bytes() == data
|
||||
assert not (tmp_path / "bundle.tar.gz.partial").exists()
|
||||
assert not (tmp_path / "bundle.tar.gz.meta").exists()
|
||||
|
||||
|
||||
def test_range_ignored_200_restarts_instead_of_corrupting(monkeypatch, tmp_path, no_sleep):
|
||||
"""A resumable partial exists, but the server ignores Range and returns 200
|
||||
with the full body. The downloader must restart (truncate) rather than
|
||||
append the full body onto the partial and corrupt the archive."""
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"GOOD" * 1024
|
||||
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
partial = tmp_path / "bundle.tar.gz.partial"
|
||||
meta = tmp_path / "bundle.tar.gz.meta"
|
||||
# A stale partial that would corrupt the file if appended to.
|
||||
partial.write_bytes(b"STALEPARTIAL")
|
||||
meta.write_text('{"bytesDownloaded": 12}')
|
||||
|
||||
# Server ignores the Range header and always sends the full body as 200.
|
||||
_patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
|
||||
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
# Exactly the good bytes, not STALEPARTIAL + data.
|
||||
assert dest.read_bytes() == data
|
||||
assert no_sleep == [] # no retry needed; it restarted cleanly in one pass
|
||||
|
||||
|
||||
def test_proper_206_resume_appends(monkeypatch, tmp_path):
|
||||
"""When the server honors Range with 206, the partial is kept and only the
|
||||
remaining bytes are appended."""
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
full = b"HEADER" + b"TAILDATA" * 512
|
||||
head_len = 6
|
||||
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
partial = tmp_path / "bundle.tar.gz.partial"
|
||||
meta = tmp_path / "bundle.tar.gz.meta"
|
||||
partial.write_bytes(full[:head_len])
|
||||
meta.write_text(f'{{"bytesDownloaded": {head_len}}}')
|
||||
|
||||
def handler(req, i):
|
||||
assert req.get_header("Range") == f"bytes={head_len}-"
|
||||
return FakeResp(full[head_len:], 206)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
installer.download_with_resume("https://h/f", str(dest), len(full), 2, 85)
|
||||
assert dest.read_bytes() == full
|
||||
|
||||
|
||||
def test_disk_full_fails_fast_without_retry(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"y" * 4096
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
partial_path = str(dest) + ".partial"
|
||||
|
||||
real_open = builtins.open
|
||||
|
||||
def fake_open(path, mode="r", *a, **k):
|
||||
if str(path) == partial_path and ("w" in mode or "a" in mode):
|
||||
real = real_open(path, mode, *a, **k)
|
||||
|
||||
class NoSpace:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_):
|
||||
real.close()
|
||||
return False
|
||||
|
||||
def write(self, _data):
|
||||
raise OSError(errno.ENOSPC, "No space left on device")
|
||||
|
||||
return NoSpace()
|
||||
return real_open(path, mode, *a, **k)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", fake_open)
|
||||
_patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
|
||||
|
||||
with pytest.raises(RuntimeError, match="disk space"):
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
assert no_sleep == [] # disk-full is not retried
|
||||
assert not os.path.exists(partial_path)
|
||||
|
||||
|
||||
def test_http_404_fails_fast_without_retry(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
raise urllib.error.HTTPError("https://h/f", 404, "Not Found", {}, None)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
|
||||
with pytest.raises(RuntimeError, match="HTTP 404"):
|
||||
installer.download_with_resume("https://h/f", str(dest), 4096, 2, 85)
|
||||
|
||||
assert no_sleep == [] # a 404 won't fix on retry
|
||||
|
||||
|
||||
def test_429_is_retried(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
data = b"z" * 2048
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
if i == 0:
|
||||
raise urllib.error.HTTPError("https://h/f", 429, "Too Many Requests", {}, None)
|
||||
return FakeResp(data, 200)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
|
||||
|
||||
assert dest.read_bytes() == data
|
||||
assert len(no_sleep) == 1 # backed off once, then succeeded
|
||||
|
||||
|
||||
def test_truncated_body_is_retried_then_succeeds(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
full = b"c" * 4096
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
if i == 0:
|
||||
return FakeResp(full[:-100], 200) # truncated / error page
|
||||
return FakeResp(full, 200)
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
installer.download_with_resume("https://h/f", str(dest), len(full), 2, 85)
|
||||
|
||||
assert dest.read_bytes() == full
|
||||
assert len(no_sleep) == 1
|
||||
|
||||
|
||||
def test_connection_error_retried_then_raises_after_max(monkeypatch, tmp_path, no_sleep):
|
||||
installer = load_installer()
|
||||
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
|
||||
dest = tmp_path / "bundle.tar.gz"
|
||||
|
||||
def handler(req, i):
|
||||
raise urllib.error.URLError("connection reset by peer")
|
||||
|
||||
_patch_urlopen(monkeypatch, installer, handler)
|
||||
|
||||
with pytest.raises(RuntimeError, match="after 5 attempts"):
|
||||
installer.download_with_resume("https://h/f", str(dest), 4096, 2, 85)
|
||||
|
||||
assert len(no_sleep) == 4 # 5 attempts => 4 backoffs
|
||||
assert not os.path.exists(str(dest) + ".partial")
|
||||
assert not os.path.exists(str(dest) + ".meta")
|
||||
@@ -9,6 +9,11 @@ import { acquireVenvRead, tryAcquireVenvRead } from "./venv-lock.js";
|
||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
const PYTHON_DIR = resolve(__dirname, "../python");
|
||||
|
||||
function appendEnvPath(base: string, suffix: string): string {
|
||||
const normalizedBase = base.replace(/\/+$/, "");
|
||||
return `${normalizedBase || "/"}${normalizedBase === "" ? "" : "/"}${suffix}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build a minimal environment for spawned Python processes.
|
||||
* Only passes through variables needed for venv, CUDA, model cache,
|
||||
@@ -44,6 +49,10 @@ function buildMinimalEnv(): Record<string, string> {
|
||||
env[key] = process.env[key] as string;
|
||||
}
|
||||
}
|
||||
|
||||
env.DATA_DIR ??= "./data";
|
||||
env.MODELS_PATH ??= appendEnvPath(env.DATA_DIR, "ai/models");
|
||||
|
||||
// Runtime model downloads are allowed by default (public model weights
|
||||
// only, never user data). SNAPOTTER_ALLOW_MODEL_DOWNLOAD=0 enables strict
|
||||
// offline mode for airgapped deployments: the sidecar then gets the
|
||||
|
||||
@@ -29,9 +29,9 @@ export const SCRIPT_BUNDLE_MAP: Record<string, string> = {
|
||||
|
||||
/** Bundle ids currently recorded as installed in DATA_DIR/ai/installed.json. */
|
||||
function installedBundles(): Set<string> {
|
||||
// Resolve DATA_DIR the same way the Python dispatcher does
|
||||
// (os.environ.get("DATA_DIR", "/data")).
|
||||
const installedPath = join(process.env.DATA_DIR || "/data", "ai", "installed.json");
|
||||
// Resolve DATA_DIR the same way the API config does for native checkouts.
|
||||
// Docker images set DATA_DIR=/data explicitly.
|
||||
const installedPath = join(process.env.DATA_DIR || "./data", "ai", "installed.json");
|
||||
try {
|
||||
const data = JSON.parse(readFileSync(installedPath, "utf-8")) as {
|
||||
bundles?: Record<string, unknown>;
|
||||
|
||||
@@ -175,6 +175,14 @@ describe("POST /api/v1/admin/features/:bundleId/install queue", () => {
|
||||
});
|
||||
}
|
||||
|
||||
async function postToolInstall(toolId: string) {
|
||||
return app.inject({
|
||||
method: "POST",
|
||||
url: `/api/v1/admin/tools/${toolId}/features/install`,
|
||||
headers: auth(),
|
||||
});
|
||||
}
|
||||
|
||||
async function getFeatures() {
|
||||
const res = await app.inject({ method: "GET", url: "/api/v1/features", headers: auth() });
|
||||
return JSON.parse(res.body).bundles as Array<{ id: string; status: string }>;
|
||||
@@ -209,6 +217,29 @@ describe("POST /api/v1/admin/features/:bundleId/install queue", () => {
|
||||
expect(bundles.find((b) => b.id === "face-detection")?.status).toBe("queued");
|
||||
});
|
||||
|
||||
it("tool install enqueues every missing hard dependency in one request", async () => {
|
||||
const res = await postToolInstall("passport-photo");
|
||||
expect(res.statusCode).toBe(202);
|
||||
const body = JSON.parse(res.body) as {
|
||||
bundles: Array<{ bundleId: string; jobId: string; queued: boolean }>;
|
||||
};
|
||||
|
||||
expect(body.bundles.map((b) => b.bundleId)).toEqual(["background-removal", "face-detection"]);
|
||||
expect(body.bundles[0].queued).toBe(false);
|
||||
expect(body.bundles[1].queued).toBe(true);
|
||||
|
||||
await waitFor(() => hoisted.spawnCalls.length === 1);
|
||||
expect(hoisted.spawnCalls[0].bundleId).toBe("background-removal");
|
||||
|
||||
const queuedBundles = await getFeatures();
|
||||
expect(queuedBundles.find((b) => b.id === "background-removal")?.status).toBe("installing");
|
||||
expect(queuedBundles.find((b) => b.id === "face-detection")?.status).toBe("queued");
|
||||
|
||||
hoisted.spawnCalls[0].emit("close", 0);
|
||||
await waitFor(() => hoisted.spawnCalls.length === 2);
|
||||
expect(hoisted.spawnCalls[1].bundleId).toBe("face-detection");
|
||||
});
|
||||
|
||||
it("dedups a duplicate install POST of the active bundle (no second entry, same job)", async () => {
|
||||
const r1 = await postInstall("ocr");
|
||||
const jobId1 = JSON.parse(r1.body).jobId;
|
||||
|
||||
@@ -1309,6 +1309,7 @@ describe("bridge - env passthrough for sidecar-specific vars", () => {
|
||||
vi.restoreAllMocks();
|
||||
delete process.env.U2NET_HOME;
|
||||
delete process.env.DATA_DIR;
|
||||
delete process.env.MODELS_PATH;
|
||||
delete process.env.DISPATCHER_MAX_REQUESTS;
|
||||
});
|
||||
|
||||
@@ -1360,9 +1361,10 @@ describe("bridge - env passthrough for sidecar-specific vars", () => {
|
||||
expect(getLastSpawnEnv()?.DISPATCHER_MAX_REQUESTS).toBe("100");
|
||||
});
|
||||
|
||||
it("does not include unset vars in subprocess env", async () => {
|
||||
it("defaults data/model paths but does not include other unset vars in subprocess env", async () => {
|
||||
delete process.env.U2NET_HOME;
|
||||
delete process.env.DATA_DIR;
|
||||
delete process.env.MODELS_PATH;
|
||||
delete process.env.DISPATCHER_MAX_REQUESTS;
|
||||
|
||||
const mock = createMockProcess();
|
||||
@@ -1375,7 +1377,8 @@ describe("bridge - env passthrough for sidecar-specific vars", () => {
|
||||
|
||||
const env = getLastSpawnEnv();
|
||||
expect(env?.U2NET_HOME).toBeUndefined();
|
||||
expect(env?.DATA_DIR).toBeUndefined();
|
||||
expect(env?.DATA_DIR).toBe("./data");
|
||||
expect(env?.MODELS_PATH).toBe("./data/ai/models");
|
||||
expect(env?.DISPATCHER_MAX_REQUESTS).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -9,15 +9,18 @@ import {
|
||||
|
||||
let tempDir: string;
|
||||
let savedDataDir: string | undefined;
|
||||
let savedCwd: string;
|
||||
|
||||
beforeEach(() => {
|
||||
savedDataDir = process.env.DATA_DIR;
|
||||
savedCwd = process.cwd();
|
||||
tempDir = mkdtempSync(join(tmpdir(), "snapotter-gate-"));
|
||||
mkdirSync(join(tempDir, "ai"), { recursive: true });
|
||||
process.env.DATA_DIR = tempDir;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
process.chdir(savedCwd);
|
||||
if (savedDataDir === undefined) delete process.env.DATA_DIR;
|
||||
else process.env.DATA_DIR = savedDataDir;
|
||||
rmSync(tempDir, { recursive: true, force: true });
|
||||
@@ -73,6 +76,28 @@ describe("missingBundleForScript", () => {
|
||||
expect(missingBundleForScript("face_landmarks")).toBe("face-detection");
|
||||
expect(missingBundleForScript("remove_bg")).toBeNull();
|
||||
});
|
||||
|
||||
it("uses the native ./data fallback when DATA_DIR is unset", () => {
|
||||
delete process.env.DATA_DIR;
|
||||
process.chdir(tempDir);
|
||||
const defaultAiDir = join(tempDir, "data", "ai");
|
||||
mkdirSync(defaultAiDir, { recursive: true });
|
||||
writeFileSync(
|
||||
join(defaultAiDir, "installed.json"),
|
||||
JSON.stringify({
|
||||
bundles: {
|
||||
"object-eraser-colorize": {
|
||||
version: "1.0.0-test",
|
||||
installedAt: "2026-01-01T00:00:00.000Z",
|
||||
models: [],
|
||||
},
|
||||
},
|
||||
}),
|
||||
"utf-8",
|
||||
);
|
||||
|
||||
expect(missingBundleForScript("inpaint.py")).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("SCRIPT_BUNDLE_MAP drift vs dispatcher.py", () => {
|
||||
|
||||
@@ -67,6 +67,7 @@ describe("buildMinimalEnv - env passthrough", () => {
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
delete process.env.DATA_DIR;
|
||||
delete process.env.SNAPOTTER_GPU;
|
||||
delete process.env.MODELS_PATH;
|
||||
delete process.env.MODELS_DIR;
|
||||
@@ -110,6 +111,24 @@ describe("buildMinimalEnv - env passthrough", () => {
|
||||
expect(env?.MODELS_PATH).toBe("/data/ai/models");
|
||||
});
|
||||
|
||||
it("sets matching native DATA_DIR and MODELS_PATH defaults", async () => {
|
||||
delete process.env.DATA_DIR;
|
||||
delete process.env.MODELS_PATH;
|
||||
|
||||
const mock = createMockProcess();
|
||||
vi.mocked(spawn).mockReturnValue(mock.process);
|
||||
|
||||
const promise = runPythonWithProgress("test.py", []);
|
||||
mock.stdout.emit("data", Buffer.from('{"ok": true}\n'));
|
||||
mock.emitEvent("close", 0, null);
|
||||
await promise;
|
||||
|
||||
const env = getSpawnEnv();
|
||||
expect(env).toBeDefined();
|
||||
expect(env?.DATA_DIR).toBe("./data");
|
||||
expect(env?.MODELS_PATH).toBe("./data/ai/models");
|
||||
});
|
||||
|
||||
it("does not pass MODELS_DIR (removed dead entry)", async () => {
|
||||
process.env.MODELS_DIR = "/some/path";
|
||||
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { evaluateInstallWatchdog } from "../../../apps/api/src/lib/install-watchdog.js";
|
||||
|
||||
const STALL = 20 * 60_000; // 20 min
|
||||
const MAX = 120 * 60_000; // 2 h
|
||||
|
||||
describe("evaluateInstallWatchdog", () => {
|
||||
it("does not kill an install making steady progress", () => {
|
||||
const now = 1_000_000;
|
||||
const v = evaluateInstallWatchdog(now, now - 5_000, now - 60_000, STALL, MAX);
|
||||
expect(v.kill).toBe(false);
|
||||
expect(v.reason).toBeNull();
|
||||
});
|
||||
|
||||
it("kills when no progress frame has arrived within the stall budget", () => {
|
||||
const now = 1_000_000;
|
||||
const v = evaluateInstallWatchdog(now, now - (STALL + 1), now - (STALL + 1), STALL, MAX);
|
||||
expect(v.kill).toBe(true);
|
||||
expect(v.reason).toContain("no progress");
|
||||
});
|
||||
|
||||
it("does not kill exactly at the stall boundary (strictly greater than)", () => {
|
||||
const now = 1_000_000;
|
||||
const v = evaluateInstallWatchdog(now, now - STALL, now - STALL, STALL, MAX);
|
||||
expect(v.kill).toBe(false);
|
||||
});
|
||||
|
||||
it("kills when the absolute time ceiling is exceeded even if progress is recent", () => {
|
||||
const now = 1_000_000;
|
||||
// Progress 1s ago (not stalled) but the install started > MAX ago.
|
||||
const v = evaluateInstallWatchdog(now, now - 1_000, now - (MAX + 1), STALL, MAX);
|
||||
expect(v.kill).toBe(true);
|
||||
expect(v.reason).toContain("time limit");
|
||||
});
|
||||
|
||||
it("prefers the absolute-ceiling reason when both conditions hold", () => {
|
||||
const now = 1_000_000;
|
||||
const v = evaluateInstallWatchdog(now, now - (STALL + 1), now - (MAX + 1), STALL, MAX);
|
||||
expect(v.kill).toBe(true);
|
||||
expect(v.reason).toContain("time limit");
|
||||
});
|
||||
|
||||
it("treats 0 as disabled for each check independently", () => {
|
||||
const now = 1_000_000;
|
||||
// Stall disabled, max active: a long stall alone must not kill.
|
||||
expect(evaluateInstallWatchdog(now, now - 10 * STALL, now - 1_000, 0, MAX).kill).toBe(false);
|
||||
// Max disabled, stall active: an old start alone must not kill.
|
||||
expect(evaluateInstallWatchdog(now, now - 1_000, now - 10 * MAX, STALL, 0).kill).toBe(false);
|
||||
// Both disabled: never kills.
|
||||
expect(evaluateInstallWatchdog(now, now - 10 * STALL, now - 10 * MAX, 0, 0).kill).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -100,6 +100,7 @@ vi.mock("../../../apps/api/src/lib/svg-sanitize.js", () => ({
|
||||
}));
|
||||
|
||||
vi.mock("../../../apps/api/src/lib/feature-status.js", () => ({
|
||||
getFirstMissingBundleForTool: vi.fn(() => null),
|
||||
isToolInstalled: vi.fn(() => true),
|
||||
}));
|
||||
|
||||
@@ -139,7 +140,10 @@ vi.mock("sharp", () => ({
|
||||
|
||||
import { apiToolPath } from "@snapotter/shared";
|
||||
import { enqueueToolJob, waitForJob } from "../../../apps/api/src/jobs/enqueue.js";
|
||||
import { isToolInstalled } from "../../../apps/api/src/lib/feature-status.js";
|
||||
import {
|
||||
getFirstMissingBundleForTool,
|
||||
isToolInstalled,
|
||||
} from "../../../apps/api/src/lib/feature-status.js";
|
||||
import { validateImageBuffer } from "../../../apps/api/src/lib/file-validation.js";
|
||||
import type { AnyToolRouteConfig } from "../../../apps/api/src/routes/tool-factory.js";
|
||||
import {
|
||||
@@ -281,6 +285,8 @@ function createMockRequest(opts: {
|
||||
describe("createToolRoute", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(getFirstMissingBundleForTool).mockReset();
|
||||
vi.mocked(getFirstMissingBundleForTool).mockReturnValue(null);
|
||||
vi.mocked(isToolInstalled).mockReset();
|
||||
vi.mocked(isToolInstalled).mockReturnValue(true);
|
||||
});
|
||||
@@ -424,6 +430,31 @@ describe("createToolRoute", () => {
|
||||
// This validates the guard code path exists without false positives.
|
||||
});
|
||||
|
||||
it("returns 501 naming the first missing extra bundle for multi-bundle AI tools", async () => {
|
||||
vi.mocked(isToolInstalled).mockReturnValueOnce(false);
|
||||
vi.mocked(getFirstMissingBundleForTool).mockReturnValueOnce("face-detection");
|
||||
const app = createMockApp();
|
||||
const id = "enhance-faces";
|
||||
createToolRoute(app as never, makeMockConfig(id));
|
||||
const handler = app.routes[apiToolPath(id)];
|
||||
const reply = createMockReply();
|
||||
const req = createMockRequest({
|
||||
fileBuffer: Buffer.from("png-data"),
|
||||
settings: JSON.stringify({}),
|
||||
});
|
||||
|
||||
await handler(req, reply);
|
||||
|
||||
expect(reply.status).toHaveBeenCalledWith(501);
|
||||
expect(reply.send).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
code: "FEATURE_NOT_INSTALLED",
|
||||
feature: "face-detection",
|
||||
featureName: "Face Detection",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("returns 200 success envelope when waitForJob resolves", async () => {
|
||||
const app = createMockApp();
|
||||
const id = "resize";
|
||||
|
||||
@@ -388,6 +388,16 @@ describe("Feature status queries", () => {
|
||||
expect(mod.isToolInstalled("passport-photo")).toBe(true);
|
||||
});
|
||||
|
||||
it("isToolInstalled is false for enhance-faces when only upscale-enhance is installed", () => {
|
||||
mod.markInstalled("upscale-enhance", "1.0.0", []);
|
||||
expect(mod.isToolInstalled("enhance-faces")).toBe(false);
|
||||
});
|
||||
|
||||
it("getFirstMissingBundleForTool names face-detection for enhance-faces when only upscale-enhance is installed", () => {
|
||||
mod.markInstalled("upscale-enhance", "1.0.0", []);
|
||||
expect(mod.getFirstMissingBundleForTool("enhance-faces")).toBe("face-detection");
|
||||
});
|
||||
|
||||
it("getFirstMissingBundleForTool names face-detection when only background-removal is installed", () => {
|
||||
mod.markInstalled("background-removal", "1.0.0", []);
|
||||
expect(mod.getFirstMissingBundleForTool("passport-photo")).toBe("face-detection");
|
||||
@@ -559,6 +569,23 @@ describe("Crash recovery - recoverInterruptedInstalls", () => {
|
||||
expect(existsSync(lockPath)).toBe(false);
|
||||
});
|
||||
|
||||
it("consumes a surviving venv.writing breadcrumb (interrupted venv write)", () => {
|
||||
const marker = join(aiDir, "venv.writing");
|
||||
writeFileSync(
|
||||
marker,
|
||||
JSON.stringify({ bundleId: "ocr", startedAt: "2026-01-01T00:00:00.000Z" }),
|
||||
);
|
||||
const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
|
||||
|
||||
expect(() => mod.recoverInterruptedInstalls()).not.toThrow();
|
||||
|
||||
// The breadcrumb is always consumed so it can't retrigger recovery forever.
|
||||
expect(existsSync(marker)).toBe(false);
|
||||
// And the interruption is surfaced in the logs.
|
||||
expect(warn.mock.calls.flat().join(" ")).toMatch(/interrupted|venv/i);
|
||||
warn.mockRestore();
|
||||
});
|
||||
|
||||
it("handles missing directories gracefully", async () => {
|
||||
vi.resetModules();
|
||||
const emptyTemp = mkdtempSync(join(tmpdir(), "snapotter-empty-"));
|
||||
|
||||
@@ -1,6 +1,14 @@
|
||||
import { execFileSync, spawnSync } from "node:child_process";
|
||||
import { createHash } from "node:crypto";
|
||||
import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
|
||||
import {
|
||||
existsSync,
|
||||
mkdirSync,
|
||||
mkdtempSync,
|
||||
readFileSync,
|
||||
rmSync,
|
||||
symlinkSync,
|
||||
writeFileSync,
|
||||
} from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
@@ -58,7 +66,12 @@ function createTestTar(bundleId: string): { tarPath: string; sha256: string } {
|
||||
return { tarPath, sha256: hash };
|
||||
}
|
||||
|
||||
function writeManifest(bundleId: string, tarPath: string, sha256: string) {
|
||||
function writeManifest(
|
||||
bundleId: string,
|
||||
tarPath: string,
|
||||
sha256: string,
|
||||
extra: Record<string, unknown> = {},
|
||||
) {
|
||||
const size = readFileSync(tarPath).length;
|
||||
const manifest = {
|
||||
manifestVersion: 2,
|
||||
@@ -75,12 +88,20 @@ function writeManifest(bundleId: string, tarPath: string, sha256: string) {
|
||||
},
|
||||
models: [{ id: "testmodel", path: "testmodel/weights.bin", minSize: 0 }],
|
||||
enablesTools: [],
|
||||
...extra,
|
||||
},
|
||||
},
|
||||
};
|
||||
writeFileSync(manifestPath, JSON.stringify(manifest));
|
||||
}
|
||||
|
||||
/** Put a real python3 at venv/bin/python3 so the post-install smoke check runs. */
|
||||
function linkVenvPython() {
|
||||
const py = execFileSync("python3", ["-c", "import sys; print(sys.executable)"]).toString().trim();
|
||||
mkdirSync(join(venvDir, "bin"), { recursive: true });
|
||||
symlinkSync(py, join(venvDir, "bin", "python3"));
|
||||
}
|
||||
|
||||
describe("install_feature.py prebuilt mode", () => {
|
||||
it("extracts models and site-packages from a local tar", () => {
|
||||
const { tarPath, sha256 } = createTestTar("face-detection");
|
||||
@@ -150,4 +171,73 @@ describe("install_feature.py prebuilt mode", () => {
|
||||
const last = JSON.parse(progressLines[progressLines.length - 1]);
|
||||
expect(last.progress).toBe(100);
|
||||
});
|
||||
|
||||
it("passes the post-install smoke import check and clears the venv-writing breadcrumb", () => {
|
||||
linkVenvPython();
|
||||
const { tarPath, sha256 } = createTestTar("face-detection");
|
||||
writeManifest("face-detection", tarPath, sha256, { smokeImports: ["json", "sys"] });
|
||||
|
||||
const result = spawnSync("python3", [scriptPath, "face-detection", manifestPath, modelsDir], {
|
||||
env: {
|
||||
...process.env,
|
||||
DATA_DIR: tempDir,
|
||||
PYTHON_VENV_PATH: venvDir,
|
||||
SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath,
|
||||
},
|
||||
timeout: 30_000,
|
||||
});
|
||||
|
||||
expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0);
|
||||
const installed = JSON.parse(readFileSync(join(aiDir, "installed.json"), "utf-8"));
|
||||
expect(installed.bundles["face-detection"]).toBeDefined();
|
||||
// The breadcrumb is cleared once the venv write completes cleanly.
|
||||
expect(existsSync(join(aiDir, "venv.writing"))).toBe(false);
|
||||
});
|
||||
|
||||
it("fails the install (and does NOT record it) when the smoke import cannot load", () => {
|
||||
linkVenvPython();
|
||||
const { tarPath, sha256 } = createTestTar("face-detection");
|
||||
writeManifest("face-detection", tarPath, sha256, {
|
||||
smokeImports: ["snapotter_not_a_real_module_zzz"],
|
||||
});
|
||||
|
||||
const result = spawnSync("python3", [scriptPath, "face-detection", manifestPath, modelsDir], {
|
||||
env: {
|
||||
...process.env,
|
||||
DATA_DIR: tempDir,
|
||||
PYTHON_VENV_PATH: venvDir,
|
||||
SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath,
|
||||
},
|
||||
timeout: 30_000,
|
||||
});
|
||||
|
||||
expect(result.status).not.toBe(0);
|
||||
expect(result.stderr?.toString()).toContain("verification failed");
|
||||
// Not marked installed, so the tool shows as needing install and a retry is clean.
|
||||
const installed = JSON.parse(readFileSync(join(aiDir, "installed.json"), "utf-8"));
|
||||
expect(installed.bundles["face-detection"]).toBeUndefined();
|
||||
});
|
||||
|
||||
it("honors SNAPOTTER_SKIP_INSTALL_SMOKE=1 as a safety valve for false positives", () => {
|
||||
linkVenvPython();
|
||||
const { tarPath, sha256 } = createTestTar("face-detection");
|
||||
writeManifest("face-detection", tarPath, sha256, {
|
||||
smokeImports: ["snapotter_not_a_real_module_zzz"],
|
||||
});
|
||||
|
||||
const result = spawnSync("python3", [scriptPath, "face-detection", manifestPath, modelsDir], {
|
||||
env: {
|
||||
...process.env,
|
||||
DATA_DIR: tempDir,
|
||||
PYTHON_VENV_PATH: venvDir,
|
||||
SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath,
|
||||
SNAPOTTER_SKIP_INSTALL_SMOKE: "1",
|
||||
},
|
||||
timeout: 30_000,
|
||||
});
|
||||
|
||||
expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0);
|
||||
const installed = JSON.parse(readFileSync(join(aiDir, "installed.json"), "utf-8"));
|
||||
expect(installed.bundles["face-detection"]).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
// @vitest-environment jsdom
|
||||
|
||||
import type { FeatureBundleState } from "@snapotter/shared";
|
||||
import { cleanup, fireEvent, render, screen } from "@testing-library/react";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { FeatureInstallPrompt } from "@/components/features/feature-install-prompt";
|
||||
import { useFeaturesStore } from "@/stores/features-store";
|
||||
|
||||
function makeBundleState(overrides: Partial<FeatureBundleState> = {}): FeatureBundleState {
|
||||
return {
|
||||
id: "background-removal",
|
||||
name: "Background Removal",
|
||||
description: "Remove backgrounds",
|
||||
status: "not_installed",
|
||||
installedVersion: null,
|
||||
estimatedSize: "4-5 GB",
|
||||
enablesTools: ["remove-background", "passport-photo"],
|
||||
progress: null,
|
||||
error: null,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe("FeatureInstallPrompt", () => {
|
||||
beforeEach(() => {
|
||||
useFeaturesStore.setState({
|
||||
bundles: [],
|
||||
loaded: true,
|
||||
loadError: false,
|
||||
installing: {},
|
||||
errors: {},
|
||||
queued: [],
|
||||
installAllActive: false,
|
||||
startTimes: {},
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
cleanup();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it("uses the tool-aware install action when a tool id is provided", () => {
|
||||
const installTool = vi.fn();
|
||||
const installBundle = vi.fn();
|
||||
useFeaturesStore.setState({ installTool, installBundle });
|
||||
|
||||
render(
|
||||
<FeatureInstallPrompt
|
||||
bundle={makeBundleState()}
|
||||
isAdmin
|
||||
toolId="passport-photo"
|
||||
toolName="Passport Photo"
|
||||
/>,
|
||||
);
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: "Enable Passport Photo" }));
|
||||
|
||||
expect(installTool).toHaveBeenCalledWith("passport-photo");
|
||||
expect(installBundle).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("shows every required bundle and keeps installed dependencies clear", () => {
|
||||
const installTool = vi.fn();
|
||||
const installBundle = vi.fn();
|
||||
const backgroundRemoval = makeBundleState({ status: "installed" });
|
||||
const faceDetection = makeBundleState({
|
||||
id: "face-detection",
|
||||
name: "Face Detection",
|
||||
description: "Detect faces",
|
||||
status: "not_installed",
|
||||
estimatedSize: "200-300 MB",
|
||||
enablesTools: ["blur-faces", "red-eye-removal", "smart-crop"],
|
||||
});
|
||||
useFeaturesStore.setState({
|
||||
bundles: [backgroundRemoval, faceDetection],
|
||||
installTool,
|
||||
installBundle,
|
||||
});
|
||||
|
||||
render(
|
||||
<FeatureInstallPrompt
|
||||
bundle={faceDetection}
|
||||
isAdmin
|
||||
toolId="passport-photo"
|
||||
toolName="Passport Photo"
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Background Removal")).toBeTruthy();
|
||||
expect(screen.getByText("Face Detection")).toBeTruthy();
|
||||
expect(screen.getByText("Installed")).toBeTruthy();
|
||||
expect(screen.getByText("Not installed")).toBeTruthy();
|
||||
expect(screen.getByText("200-300 MB")).toBeTruthy();
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: "Enable Passport Photo" }));
|
||||
|
||||
expect(installTool).toHaveBeenCalledWith("passport-photo");
|
||||
expect(installBundle).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("keeps the multi-bundle breakdown visible when one bundle is in the error/repair state", () => {
|
||||
const backgroundRemoval = makeBundleState({
|
||||
status: "error",
|
||||
error: "Checksum mismatch",
|
||||
});
|
||||
const faceDetection = makeBundleState({
|
||||
id: "face-detection",
|
||||
name: "Face Detection",
|
||||
description: "Detect faces",
|
||||
status: "installed",
|
||||
estimatedSize: "200-300 MB",
|
||||
enablesTools: ["blur-faces", "red-eye-removal", "smart-crop"],
|
||||
});
|
||||
useFeaturesStore.setState({
|
||||
bundles: [backgroundRemoval, faceDetection],
|
||||
installTool: vi.fn(),
|
||||
installBundle: vi.fn(),
|
||||
});
|
||||
|
||||
render(
|
||||
<FeatureInstallPrompt
|
||||
bundle={backgroundRemoval}
|
||||
isAdmin
|
||||
toolId="passport-photo"
|
||||
toolName="Passport Photo"
|
||||
/>,
|
||||
);
|
||||
|
||||
// The breakdown must still render during repair so the user can see the
|
||||
// sibling bundle's state, not just the single failed one.
|
||||
expect(screen.getByText("Background Removal")).toBeTruthy();
|
||||
expect(screen.getByText("Face Detection")).toBeTruthy();
|
||||
expect(screen.getByText("Installed")).toBeTruthy();
|
||||
});
|
||||
});
|
||||
@@ -341,6 +341,58 @@ describe("useFeaturesStore", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("installTool()", () => {
|
||||
it("starts every missing hard dependency for a multi-bundle AI tool from one server request", async () => {
|
||||
const bundles = [
|
||||
makeBundleState({ id: "background-removal", status: "not_installed" }),
|
||||
makeBundleState({ id: "face-detection", status: "not_installed" }),
|
||||
];
|
||||
useFeaturesStore.setState({ bundles, loaded: true });
|
||||
apiPostMock.mockResolvedValueOnce({
|
||||
bundles: [
|
||||
{ bundleId: "background-removal", jobId: "job-bg", queued: false },
|
||||
{ bundleId: "face-detection", jobId: "job-face", queued: true },
|
||||
],
|
||||
});
|
||||
|
||||
await useFeaturesStore.getState().installTool("passport-photo");
|
||||
|
||||
expect(apiPostMock).toHaveBeenCalledWith(
|
||||
"/v1/admin/tools/passport-photo/features/install",
|
||||
{},
|
||||
);
|
||||
expect(FakeEventSource.instances).toHaveLength(1);
|
||||
expect(FakeEventSource.instances[0].url).toBe("/api/v1/jobs/job-bg/progress");
|
||||
expect(useFeaturesStore.getState().installing["background-removal"]).toBeDefined();
|
||||
expect(useFeaturesStore.getState().installing["face-detection"]).toBeUndefined();
|
||||
expect(useFeaturesStore.getState().queued).toContain("face-detection");
|
||||
});
|
||||
|
||||
it("only tracks missing dependencies when a multi-bundle AI tool is partially installed", async () => {
|
||||
const bundles = [
|
||||
makeBundleState({ id: "background-removal", status: "installed" }),
|
||||
makeBundleState({ id: "face-detection", status: "not_installed" }),
|
||||
];
|
||||
useFeaturesStore.setState({ bundles, loaded: true });
|
||||
apiPostMock.mockResolvedValueOnce({
|
||||
bundles: [
|
||||
{ bundleId: "background-removal", skipped: true },
|
||||
{ bundleId: "face-detection", jobId: "job-face", queued: false },
|
||||
],
|
||||
});
|
||||
|
||||
await useFeaturesStore.getState().installTool("passport-photo");
|
||||
|
||||
expect(apiPostMock).toHaveBeenCalledWith(
|
||||
"/v1/admin/tools/passport-photo/features/install",
|
||||
{},
|
||||
);
|
||||
expect(useFeaturesStore.getState().installing["background-removal"]).toBeUndefined();
|
||||
expect(useFeaturesStore.getState().installing["face-detection"]).toBeDefined();
|
||||
expect(FakeEventSource.instances[0].url).toBe("/api/v1/jobs/job-face/progress");
|
||||
});
|
||||
});
|
||||
|
||||
describe("uninstallBundle()", () => {
|
||||
it("calls API and refreshes", async () => {
|
||||
apiPostMock.mockResolvedValueOnce({});
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
// @vitest-environment jsdom
|
||||
import type { FeatureBundleState } from "@snapotter/shared";
|
||||
import { TOOLS } from "@snapotter/shared";
|
||||
import { cleanup, fireEvent, render, screen } from "@testing-library/react";
|
||||
import { MemoryRouter } from "react-router-dom";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { ToolCard } from "@/components/common/tool-card";
|
||||
import { useFeaturesStore } from "@/stores/features-store";
|
||||
import { usePinnedToolsStore } from "@/stores/pinned-tools-store";
|
||||
|
||||
// Make the store's optimistic persistence a no-op so the test stays server-free.
|
||||
@@ -14,6 +16,23 @@ vi.mock("@/lib/api", () => ({
|
||||
|
||||
const resize = TOOLS.find((tool) => tool.id === "resize");
|
||||
if (!resize) throw new Error("resize tool missing from TOOLS");
|
||||
const passportPhoto = TOOLS.find((tool) => tool.id === "passport-photo");
|
||||
if (!passportPhoto) throw new Error("passport-photo tool missing from TOOLS");
|
||||
|
||||
function makeBundleState(overrides: Partial<FeatureBundleState> = {}): FeatureBundleState {
|
||||
return {
|
||||
id: "background-removal",
|
||||
name: "Background Removal",
|
||||
description: "Remove backgrounds",
|
||||
status: "not_installed",
|
||||
installedVersion: null,
|
||||
estimatedSize: "4-5 GB",
|
||||
enablesTools: ["remove-background", "passport-photo"],
|
||||
progress: null,
|
||||
error: null,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
afterEach(cleanup);
|
||||
|
||||
@@ -24,6 +43,16 @@ beforeEach(() => {
|
||||
loaded: true,
|
||||
loadError: false,
|
||||
});
|
||||
useFeaturesStore.setState({
|
||||
bundles: [],
|
||||
loaded: true,
|
||||
loadError: false,
|
||||
installing: {},
|
||||
errors: {},
|
||||
queued: [],
|
||||
installAllActive: false,
|
||||
startTimes: {},
|
||||
});
|
||||
});
|
||||
|
||||
function renderCard(showPin: boolean) {
|
||||
@@ -56,3 +85,28 @@ describe("ToolCard pin button", () => {
|
||||
expect(usePinnedToolsStore.getState().pinnedTools).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("ToolCard AI bundle status", () => {
|
||||
it("treats multi-bundle tools as not installed when an extra bundle is missing", () => {
|
||||
useFeaturesStore.setState({
|
||||
bundles: [
|
||||
makeBundleState({ id: "background-removal", status: "installed" }),
|
||||
makeBundleState({
|
||||
id: "face-detection",
|
||||
name: "Face Detection",
|
||||
status: "not_installed",
|
||||
enablesTools: ["blur-faces", "red-eye-removal", "smart-crop"],
|
||||
}),
|
||||
],
|
||||
});
|
||||
|
||||
const { container } = render(
|
||||
<MemoryRouter>
|
||||
<ToolCard tool={passportPhoto} variant="descriptive" />
|
||||
</MemoryRouter>,
|
||||
);
|
||||
|
||||
// One SVG is the tool icon, the second is the missing-AI-bundle indicator.
|
||||
expect(container.querySelectorAll("svg")).toHaveLength(2);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user