From 494bb3d78b955312b6fb03aea4a5146ce1a3a94e Mon Sep 17 00:00:00 2001
From: ashim-hq <ashim.hq@gmail.com>
Date: Sat, 18 Apr 2026 02:42:34 +0800
Subject: [PATCH] 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.
---
 packages/ai/python/dispatcher.py | 43 ++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/packages/ai/python/dispatcher.py b/packages/ai/python/dispatcher.py
index 8b58bbdc..303955fb 100644
--- a/packages/ai/python/dispatcher.py
+++ b/packages/ai/python/dispatcher.py
@@ -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