From 5d5240e48131653e3a64917fae312a044326992c Mon Sep 17 00:00:00 2001
From: SnapOtter <snapotter.hq@gmail.com>
Date: Sat, 13 Jun 2026 16:32:45 +0800
Subject: [PATCH] feat: extend importBundleArchive for site-packages and fixups

---
 apps/api/src/lib/feature-status.ts       | 39 ++++++++++++++++
 tests/integration/feature-import.test.ts | 59 ++++++++++++++++++++++++
 2 files changed, 98 insertions(+)

diff --git a/apps/api/src/lib/feature-status.ts b/apps/api/src/lib/feature-status.ts
index 54dac5c7..0cd36eee 100644
--- a/apps/api/src/lib/feature-status.ts
+++ b/apps/api/src/lib/feature-status.ts
@@ -1,3 +1,4 @@
+import { execFileSync } from "node:child_process";
 import { randomUUID } from "node:crypto";
 import {
   constants,
@@ -632,6 +633,44 @@ export async function importBundleArchive(
       moveTreeRecursive(stagingModels, MODELS_DIR);
     }
 
+    // Move site-packages/* into venv site-packages
+    const stagingSitePackages = join(stagingDir, "site-packages");
+    if (existsSync(stagingSitePackages)) {
+      const venvPath = process.env.PYTHON_VENV_PATH || join(AI_DIR, "venv");
+      let sitePackagesDir = "";
+      const libDir = join(venvPath, "lib");
+      if (existsSync(libDir)) {
+        const pyDirs = readdirSync(libDir).filter((d) => d.startsWith("python"));
+        if (pyDirs.length > 0) {
+          sitePackagesDir = join(libDir, pyDirs[0], "site-packages");
+        }
+      }
+      if (sitePackagesDir && existsSync(sitePackagesDir)) {
+        moveTreeRecursive(stagingSitePackages, sitePackagesDir);
+      }
+    }
+
+    // Apply fixups (NCCL wheel) if present
+    const stagingFixups = join(stagingDir, "fixups");
+    if (existsSync(stagingFixups)) {
+      const wheels = readdirSync(stagingFixups).filter((f) => f.endsWith(".whl"));
+      if (wheels.length > 0) {
+        const venvPython =
+          (process.env.PYTHON_VENV_PATH || join(AI_DIR, "venv")) + "/bin/python3";
+        for (const wheel of wheels) {
+          try {
+            execFileSync(venvPython, [
+              "-m", "pip", "install", "--no-index",
+              `--find-links=${stagingFixups}`,
+              wheel.split("-")[0],
+            ], { stdio: "ignore", timeout: 30_000 });
+          } catch {
+            // Non-fatal
+          }
+        }
+      }
+    }
+
     markInstalled(descriptor.bundleId, descriptor.version, descriptor.models);
 
     return {
diff --git a/tests/integration/feature-import.test.ts b/tests/integration/feature-import.test.ts
index 2dd5e606..71c25116 100644
--- a/tests/integration/feature-import.test.ts
+++ b/tests/integration/feature-import.test.ts
@@ -143,6 +143,41 @@ async function buildSymlinkArchive(): Promise<string> {
   return archivePath;
 }
 
+async function createBundleTarWithSitePackages(
+  bundleId: string,
+  version: string,
+  modelFiles: Record<string, string>,
+  sitePackageFiles: Record<string, string>,
+): Promise<string> {
+  const tarDir = join(testRoot, `tar-src-${randomUUID()}`);
+  mkdirSync(tarDir, { recursive: true });
+
+  writeFileSync(
+    join(tarDir, "bundle.json"),
+    JSON.stringify({ bundleId, version, models: Object.keys(modelFiles) }),
+  );
+
+  const modelsSubdir = join(tarDir, "models");
+  mkdirSync(modelsSubdir, { recursive: true });
+  for (const [name, content] of Object.entries(modelFiles)) {
+    const modelPath = join(modelsSubdir, name);
+    mkdirSync(join(modelPath, ".."), { recursive: true });
+    writeFileSync(modelPath, content);
+  }
+
+  const spSubdir = join(tarDir, "site-packages");
+  mkdirSync(spSubdir, { recursive: true });
+  for (const [name, content] of Object.entries(sitePackageFiles)) {
+    const spPath = join(spSubdir, name);
+    mkdirSync(join(spPath, ".."), { recursive: true });
+    writeFileSync(spPath, content);
+  }
+
+  const tarPath = join(testRoot, `bundle-${randomUUID()}.tar.gz`);
+  await tar.create({ gzip: true, file: tarPath, cwd: tarDir }, ["."]);
+  return tarPath;
+}
+
 function resetState(): void {
   writeFileSync(installedPath, JSON.stringify({ bundles: {} }), "utf-8");
   invalidateCache();
@@ -265,6 +300,30 @@ describe("importBundleArchive", () => {
   });
 });
 
+describe("site-packages import", () => {
+  beforeEach(resetState);
+
+  it("extracts site-packages into venv site-packages directory", async () => {
+    const venvSitePackages = join(aiDir, "venv", "lib", "python3.12", "site-packages");
+    mkdirSync(venvSitePackages, { recursive: true });
+    process.env.PYTHON_VENV_PATH = join(aiDir, "venv");
+
+    const tarPath = await createBundleTarWithSitePackages(
+      testBundleId,
+      testVersion,
+      { "mediapipe/face.tflite": "model-data" },
+      { "fakepkg/__init__.py": "# fake package" },
+    );
+
+    invalidateCache();
+    const result = await importBundleArchive(createReadStream(tarPath));
+    expect(result.bundleId).toBe(testBundleId);
+    expect(existsSync(join(venvSitePackages, "fakepkg", "__init__.py"))).toBe(true);
+
+    delete process.env.PYTHON_VENV_PATH;
+  });
+});
+
 describe("POST /api/v1/admin/features/import", () => {
   let app: Awaited<ReturnType<typeof import("fastify")>>["default"] extends (
     ...args: infer _A