mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix(ai): enforce the feature gate on the per-request fallback path (#331)
The persistent Python dispatcher rejects scripts whose feature bundle is not installed, but the per-request fallback (used when the dispatcher is down, e.g. restarting right after a model repair) spawned scripts directly and bypassed that gate. Behavior was therefore inconsistent: a gated script would fail under the dispatcher but run under the fallback -- the "works once after a repair" symptom from the original report. - add packages/ai/src/feature-gate.ts: SCRIPT_BUNDLE_MAP + missingBundleForScript, mirroring TOOL_BUNDLE_MAP in dispatcher.py, reading the same installed.json and failing closed exactly like dispatcher._get_installed_bundles() - runPerRequest now rejects with "feature_not_installed" (the same message the dispatcher path surfaces) when a gated script's bundle is not installed - unit tests for the gate, plus a drift test pinning the TS map to dispatcher.py Closes #327
This commit is contained in:
@@ -3,6 +3,7 @@ import { randomUUID } from "node:crypto";
|
|||||||
import { dirname, resolve } from "node:path";
|
import { dirname, resolve } from "node:path";
|
||||||
import { fileURLToPath } from "node:url";
|
import { fileURLToPath } from "node:url";
|
||||||
import { context, propagation, SpanStatusCode, trace } from "@opentelemetry/api";
|
import { context, propagation, SpanStatusCode, trace } from "@opentelemetry/api";
|
||||||
|
import { missingBundleForScript } from "./feature-gate.js";
|
||||||
|
|
||||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||||
const PYTHON_DIR = resolve(__dirname, "../python");
|
const PYTHON_DIR = resolve(__dirname, "../python");
|
||||||
@@ -398,6 +399,14 @@ export class PythonDispatcher {
|
|||||||
timeout?: number;
|
timeout?: number;
|
||||||
} = {},
|
} = {},
|
||||||
): Promise<{ stdout: string; stderr: string }> {
|
): Promise<{ stdout: string; stderr: string }> {
|
||||||
|
// Mirror the dispatcher's feature gate. The persistent dispatcher rejects
|
||||||
|
// scripts whose bundle is not installed (in Python); this fallback spawns
|
||||||
|
// scripts directly, so without the same check it would run them ungated
|
||||||
|
// when the dispatcher is down (e.g. right after a model repair). Reject
|
||||||
|
// with the same message the dispatcher path surfaces.
|
||||||
|
if (missingBundleForScript(scriptName)) {
|
||||||
|
return Promise.reject(new Error("feature_not_installed"));
|
||||||
|
}
|
||||||
const scriptPath = resolve(PYTHON_DIR, scriptName);
|
const scriptPath = resolve(PYTHON_DIR, scriptName);
|
||||||
const timeout =
|
const timeout =
|
||||||
options.timeout ??
|
options.timeout ??
|
||||||
|
|||||||
@@ -0,0 +1,59 @@
|
|||||||
|
import { readFileSync } from "node:fs";
|
||||||
|
import { join } from "node:path";
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Python script name -> required feature bundle.
|
||||||
|
*
|
||||||
|
* This MUST mirror TOOL_BUNDLE_MAP in packages/ai/python/dispatcher.py. The
|
||||||
|
* persistent dispatcher enforces this gate in Python; the per-request fallback
|
||||||
|
* path (PythonDispatcher.runPerRequest) spawns scripts directly and would
|
||||||
|
* otherwise bypass it, so it enforces the same gate here. A drift test keeps
|
||||||
|
* the two maps in sync (tests/unit/ai/feature-gate.test.ts).
|
||||||
|
*/
|
||||||
|
export const SCRIPT_BUNDLE_MAP: Record<string, string> = {
|
||||||
|
remove_bg: "background-removal",
|
||||||
|
detect_faces: "face-detection",
|
||||||
|
face_landmarks: "face-detection",
|
||||||
|
red_eye_removal: "face-detection",
|
||||||
|
inpaint: "object-eraser-colorize",
|
||||||
|
outpaint: "object-eraser-colorize",
|
||||||
|
colorize: "object-eraser-colorize",
|
||||||
|
upscale: "upscale-enhance",
|
||||||
|
enhance_faces: "upscale-enhance",
|
||||||
|
noise_removal: "upscale-enhance",
|
||||||
|
restore: "photo-restoration",
|
||||||
|
ocr: "ocr",
|
||||||
|
ocr_pdf: "ocr",
|
||||||
|
transcribe: "transcription",
|
||||||
|
};
|
||||||
|
|
||||||
|
/** 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");
|
||||||
|
try {
|
||||||
|
const data = JSON.parse(readFileSync(installedPath, "utf-8")) as {
|
||||||
|
bundles?: Record<string, unknown>;
|
||||||
|
};
|
||||||
|
return new Set(Object.keys(data.bundles ?? {}));
|
||||||
|
} catch {
|
||||||
|
// Fail closed, exactly like dispatcher._get_installed_bundles(): a missing
|
||||||
|
// or unreadable installed.json reads as "nothing installed".
|
||||||
|
return new Set();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The feature bundle a script requires if that bundle is NOT installed, else
|
||||||
|
* null. Returns null for ungated scripts (e.g. doc-profile scripts) and for
|
||||||
|
* gated scripts whose bundle is installed. Mirrors the dispatcher's gate so the
|
||||||
|
* fallback path behaves identically whether or not the dispatcher is running.
|
||||||
|
*
|
||||||
|
* Accepts the script name with or without a ".py" suffix.
|
||||||
|
*/
|
||||||
|
export function missingBundleForScript(scriptName: string): string | null {
|
||||||
|
const bundle = SCRIPT_BUNDLE_MAP[scriptName.replace(/\.py$/, "")];
|
||||||
|
if (!bundle) return null;
|
||||||
|
return installedBundles().has(bundle) ? null : bundle;
|
||||||
|
}
|
||||||
@@ -16,6 +16,7 @@ export { blurFaces, detectFaces } from "./face-detection.js";
|
|||||||
export { enhanceFaces } from "./face-enhancement.js";
|
export { enhanceFaces } from "./face-enhancement.js";
|
||||||
export type { FaceLandmarkPoint, FaceLandmarks, FaceLandmarksResult } from "./face-landmarks.js";
|
export type { FaceLandmarkPoint, FaceLandmarks, FaceLandmarksResult } from "./face-landmarks.js";
|
||||||
export { detectFaceLandmarks } from "./face-landmarks.js";
|
export { detectFaceLandmarks } from "./face-landmarks.js";
|
||||||
|
export { missingBundleForScript, SCRIPT_BUNDLE_MAP } from "./feature-gate.js";
|
||||||
export { inpaint } from "./inpainting.js";
|
export { inpaint } from "./inpainting.js";
|
||||||
export { noiseRemoval } from "./noise-removal.js";
|
export { noiseRemoval } from "./noise-removal.js";
|
||||||
export type { PdfOcrOptions, PdfOcrResult } from "./ocr.js";
|
export type { PdfOcrOptions, PdfOcrResult } from "./ocr.js";
|
||||||
|
|||||||
@@ -383,7 +383,10 @@ describe("bridge - runPythonWithProgress (per-request fallback)", () => {
|
|||||||
const mock = createMockProcess();
|
const mock = createMockProcess();
|
||||||
vi.mocked(spawn).mockReturnValue(mock.process);
|
vi.mocked(spawn).mockReturnValue(mock.process);
|
||||||
|
|
||||||
const promise = runPythonWithProgress("remove_bg.py", ["/tmp/in.png", "/tmp/out.png"]);
|
// Ungated script name: this test exercises generic spawn-args plumbing, not
|
||||||
|
// a specific bundle. The feature gate on the fallback is covered separately
|
||||||
|
// in tests/unit/ai/feature-gate.test.ts.
|
||||||
|
const promise = runPythonWithProgress("test_script.py", ["/tmp/in.png", "/tmp/out.png"]);
|
||||||
|
|
||||||
mock.stdout.emit("data", Buffer.from('{"success": true}\n'));
|
mock.stdout.emit("data", Buffer.from('{"success": true}\n'));
|
||||||
mock.emitEvent("close", 0, null);
|
mock.emitEvent("close", 0, null);
|
||||||
@@ -402,7 +405,7 @@ describe("bridge - runPythonWithProgress (per-request fallback)", () => {
|
|||||||
expect(perRequestCall).toBeDefined();
|
expect(perRequestCall).toBeDefined();
|
||||||
expect(perRequestCall?.[1]).toEqual(
|
expect(perRequestCall?.[1]).toEqual(
|
||||||
expect.arrayContaining([
|
expect.arrayContaining([
|
||||||
expect.stringContaining("remove_bg.py"),
|
expect.stringContaining("test_script.py"),
|
||||||
"/tmp/in.png",
|
"/tmp/in.png",
|
||||||
"/tmp/out.png",
|
"/tmp/out.png",
|
||||||
]),
|
]),
|
||||||
|
|||||||
@@ -0,0 +1,92 @@
|
|||||||
|
import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
|
||||||
|
import { tmpdir } from "node:os";
|
||||||
|
import { join } from "node:path";
|
||||||
|
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||||
|
import {
|
||||||
|
missingBundleForScript,
|
||||||
|
SCRIPT_BUNDLE_MAP,
|
||||||
|
} from "../../../packages/ai/src/feature-gate.js";
|
||||||
|
|
||||||
|
let tempDir: string;
|
||||||
|
let savedDataDir: string | undefined;
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
savedDataDir = process.env.DATA_DIR;
|
||||||
|
tempDir = mkdtempSync(join(tmpdir(), "snapotter-gate-"));
|
||||||
|
mkdirSync(join(tempDir, "ai"), { recursive: true });
|
||||||
|
process.env.DATA_DIR = tempDir;
|
||||||
|
});
|
||||||
|
|
||||||
|
afterEach(() => {
|
||||||
|
if (savedDataDir === undefined) delete process.env.DATA_DIR;
|
||||||
|
else process.env.DATA_DIR = savedDataDir;
|
||||||
|
rmSync(tempDir, { recursive: true, force: true });
|
||||||
|
});
|
||||||
|
|
||||||
|
function setInstalled(bundleIds: string[]): void {
|
||||||
|
const bundles: Record<string, { version: string; installedAt: string; models: string[] }> = {};
|
||||||
|
for (const id of bundleIds) {
|
||||||
|
bundles[id] = { version: "1.0.0-test", installedAt: "2026-01-01T00:00:00.000Z", models: [] };
|
||||||
|
}
|
||||||
|
writeFileSync(join(tempDir, "ai", "installed.json"), JSON.stringify({ bundles }), "utf-8");
|
||||||
|
}
|
||||||
|
|
||||||
|
describe("missingBundleForScript", () => {
|
||||||
|
it("returns the bundle when a gated script's bundle is not installed", () => {
|
||||||
|
setInstalled([]);
|
||||||
|
expect(missingBundleForScript("face_landmarks")).toBe("face-detection");
|
||||||
|
expect(missingBundleForScript("remove_bg")).toBe("background-removal");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("returns null when the gated script's bundle is installed", () => {
|
||||||
|
setInstalled(["face-detection"]);
|
||||||
|
expect(missingBundleForScript("face_landmarks")).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("handles a .py suffix", () => {
|
||||||
|
setInstalled(["background-removal"]);
|
||||||
|
expect(missingBundleForScript("remove_bg.py")).toBeNull();
|
||||||
|
setInstalled([]);
|
||||||
|
expect(missingBundleForScript("remove_bg.py")).toBe("background-removal");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("returns null for ungated scripts", () => {
|
||||||
|
setInstalled([]);
|
||||||
|
expect(missingBundleForScript("doc_text")).toBeNull();
|
||||||
|
expect(missingBundleForScript("unknown_script")).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("fails closed when installed.json is missing", () => {
|
||||||
|
// No setInstalled(): the file does not exist. Like the dispatcher, an
|
||||||
|
// unreadable installed.json reads as "nothing installed", so a gated
|
||||||
|
// script is blocked.
|
||||||
|
expect(missingBundleForScript("face_landmarks")).toBe("face-detection");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("fails closed when installed.json is corrupt", () => {
|
||||||
|
writeFileSync(join(tempDir, "ai", "installed.json"), "{{{not json", "utf-8");
|
||||||
|
expect(missingBundleForScript("face_landmarks")).toBe("face-detection");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("gates both of passport-photo's scripts (issue #327)", () => {
|
||||||
|
setInstalled(["background-removal"]); // face-detection still missing
|
||||||
|
expect(missingBundleForScript("face_landmarks")).toBe("face-detection");
|
||||||
|
expect(missingBundleForScript("remove_bg")).toBeNull();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe("SCRIPT_BUNDLE_MAP drift vs dispatcher.py", () => {
|
||||||
|
it("matches TOOL_BUNDLE_MAP in dispatcher.py exactly", () => {
|
||||||
|
const src = readFileSync(join(process.cwd(), "packages/ai/python/dispatcher.py"), "utf-8");
|
||||||
|
const block = src.match(/TOOL_BUNDLE_MAP\s*=\s*\{([\s\S]*?)\}/);
|
||||||
|
if (!block) throw new Error("TOOL_BUNDLE_MAP not found in dispatcher.py");
|
||||||
|
|
||||||
|
const pythonMap: Record<string, string> = {};
|
||||||
|
for (const entry of block[1].matchAll(/"(\w+)"\s*:\s*"([a-z0-9-]+)"/g)) {
|
||||||
|
pythonMap[entry[1]] = entry[2];
|
||||||
|
}
|
||||||
|
|
||||||
|
// Both directions: the TS mirror must equal the Python source of truth.
|
||||||
|
expect(pythonMap).toEqual(SCRIPT_BUNDLE_MAP);
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user