feat: add feature gating to Python sidecar dispatcher

Check installed.json before exec()-ing AI scripts so that requests
for uninstalled feature bundles return a structured error instead
of crashing with an ImportError. Also sets U2NET_HOME to the
bundled model directory when present.
This commit is contained in:
ashim-hq
2026-04-18 02:42:34 +08:00
parent 3c3aa74e98
commit 494bb3d78b
+43
View File
@@ -17,6 +17,33 @@ import os
import traceback
INSTALLED_PATH = os.path.join(os.environ.get("DATA_DIR", "/data"), "ai", "installed.json")
MODELS_DIR = os.path.join(os.environ.get("DATA_DIR", "/data"), "ai", "models")
TOOL_BUNDLE_MAP = {
"remove_bg": "background-removal",
"detect_faces": "face-detection",
"face_landmarks": "face-detection",
"red_eye_removal": "face-detection",
"inpaint": "object-eraser-colorize",
"colorize": "object-eraser-colorize",
"upscale": "upscale-enhance",
"enhance_faces": "upscale-enhance",
"noise_removal": "upscale-enhance",
"restore": "photo-restoration",
"ocr": "ocr",
}
def _get_installed_bundles():
try:
with open(INSTALLED_PATH) as f:
data = json.load(f)
return set(data.get("bundles", {}).keys())
except (FileNotFoundError, json.JSONDecodeError):
return set()
def emit_progress(percent, stage):
"""Emit structured progress to stderr."""
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
@@ -44,6 +71,10 @@ _try_import("gpu", lambda: __import__("gpu"))
# Heavy ML libraries - import but don't fail if unavailable
_try_import("rembg", lambda: __import__("rembg"))
# Point rembg at the bundled model directory if it exists
if os.path.isdir(MODELS_DIR):
os.environ.setdefault("U2NET_HOME", os.path.join(MODELS_DIR, "rembg"))
# ── Script handlers ─────────────────────────────────────────────────
# Each handler sets sys.argv and calls the script's main() function,
@@ -59,6 +90,18 @@ def _run_script_main(script_name, args):
"""
script_dir = os.path.dirname(os.path.abspath(__file__))
# ── Feature gate: reject scripts whose bundle is not installed ──
bundle_id = TOOL_BUNDLE_MAP.get(script_name)
if bundle_id:
installed = _get_installed_bundles()
if bundle_id not in installed:
return (json.dumps({
"success": False,
"error": "feature_not_installed",
"feature": bundle_id,
"message": f"Feature bundle '{bundle_id}' is not installed"
}), 1)
# Save original state
old_argv = sys.argv