fix(docker): make venv bootstrap upgrade-aware with pip-freeze stamp

On 1.x-to-2.0 upgrade the /data/ai/venv already existed with stale packages, so the entrypoint skipped the cp from /opt/venv and 2.0 Python tools failed with ImportError (#85). Dockerfile now writes a SHA-256 of `pip freeze` to /opt/venv/.venv-version at build time. The entrypoint compares that stamp against the copy on the volume; on mismatch it does a clean nuke + re-copy and resets installed.json so the UI correctly shows bundles as needing reinstall. Models in /data/ai/models survive the refresh, so bundle reinstall only reruns pip (model downloads are idempotent). Backward compatible: if the image has no stamp (old build), the entrypoint falls through to the existing skip-if-exists behavior.
2026-08-03 07:46:42 +02:00 · 2026-06-13 10:31:22 +08:00
parent 51666cdd5f
commit 39543d6512
3 changed files with 251 additions and 7 deletions
@@ -281,6 +281,11 @@ RUN --mount=type=cache,target=/root/.cache/pip \
        pdf2docx==0.5.13 \
        markdown==3.10.2

+# Stamp the venv so the entrypoint can detect base-package upgrades.
+# If the frozen package list changes, the hash changes, and containers
+# with a stale /data/ai/venv will get a fresh copy on next start.
+RUN /opt/venv/bin/pip freeze | sha256sum | cut -d' ' -f1 > /opt/venv/.venv-version
+
 # On-demand AI feature installer and manifest
 COPY docker/feature-manifest.json /app/docker/feature-manifest.json
 COPY packages/ai/python/install_feature.py /app/packages/ai/python/install_feature.py
@@ -54,13 +54,42 @@ if [ -d "$AI_VENV_TMP" ]; then
  rm -rf "$AI_VENV_TMP"
 fi

-# Bootstrap AI venv from base image on first run
-if [ ! -d "$AI_VENV" ] && [ -d "/opt/venv" ]; then
-  echo "Bootstrapping AI venv from base image..."
-  mkdir -p /data/ai/models /data/ai/pip-cache
-  cp -r /opt/venv "$AI_VENV_TMP"
-  mv "$AI_VENV_TMP" "$AI_VENV"
-  echo "AI venv ready at $AI_VENV"
+# Bootstrap AI venv from base image (first run or upgrade).
+# The image stamps /opt/venv/.venv-version with a hash of `pip freeze`.
+# When the stamp in /data/ai/venv doesn't match, the base packages changed
+# and we need a fresh copy so 2.0 Python tools don't fail with ImportError.
+if [ -d "/opt/venv" ]; then
+  NEED_BOOTSTRAP=false
+
+  if [ ! -d "$AI_VENV" ]; then
+    NEED_BOOTSTRAP=true
+    echo "First run: bootstrapping AI venv from base image..."
+  elif [ -f "/opt/venv/.venv-version" ]; then
+    IMAGE_STAMP=$(cat /opt/venv/.venv-version)
+    CURRENT_STAMP=""
+    if [ -f "$AI_VENV/.venv-version" ]; then
+      CURRENT_STAMP=$(cat "$AI_VENV/.venv-version")
+    fi
+    if [ "$CURRENT_STAMP" != "$IMAGE_STAMP" ]; then
+      NEED_BOOTSTRAP=true
+      echo "Base venv updated (stamp mismatch): refreshing AI venv..."
+    fi
+  fi
+
+  if [ "$NEED_BOOTSTRAP" = true ]; then
+    mkdir -p /data/ai/models /data/ai/pip-cache
+    rm -rf "$AI_VENV"
+    cp -r /opt/venv "$AI_VENV_TMP"
+    mv "$AI_VENV_TMP" "$AI_VENV"
+    # Reset installed-bundle state: their packages lived in the old venv.
+    # Models in /data/ai/models survive, so reinstalling a bundle only
+    # reruns pip (model downloads are idempotent and skip existing files).
+    if [ -f "/data/ai/installed.json" ]; then
+      echo '{"bundles":{}}' > /data/ai/installed.json
+      echo "WARNING: Installed AI feature bundles were reset after base venv upgrade. Reinstall them from the Settings page."
+    fi
+    echo "AI venv ready at $AI_VENV"
+  fi
 fi

 # Wait for Postgres to be reachable before starting the app
@@ -0,0 +1,210 @@
+import { execFileSync } from "node:child_process";
+import { existsSync, 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";
+
+/**
+ * Unit tests for the venv upgrade-stamp mechanism in docker/entrypoint.sh.
+ *
+ * The entrypoint compares /opt/venv/.venv-version (baked into the image at
+ * build time) against /data/ai/venv/.venv-version (persisted on the volume).
+ * When missing or mismatched, the venv is nuked and re-copied so upgraded
+ * base packages take effect immediately.
+ *
+ * These tests exercise the bootstrap logic in isolation using temp dirs.
+ */
+
+let root: string;
+let optVenv: string;
+let dataAi: string;
+let aiVenv: string;
+let aiVenvTmp: string;
+let installedJson: string;
+
+// Self-contained shell script mirroring the entrypoint bootstrap block.
+// Uses env vars for paths so we can point at temp directories.
+const BOOTSTRAP_SCRIPT = `
+#!/bin/sh
+set -e
+
+AI_VENV="$TEST_AI_VENV"
+AI_VENV_TMP="$TEST_AI_VENV_TMP"
+OPT_VENV="$TEST_OPT_VENV"
+
+if [ -d "$AI_VENV_TMP" ]; then
+  echo "CLEANUP_INTERRUPTED"
+  rm -rf "$AI_VENV_TMP"
+fi
+
+if [ -d "$OPT_VENV" ]; then
+  NEED_BOOTSTRAP=false
+
+  if [ ! -d "$AI_VENV" ]; then
+    NEED_BOOTSTRAP=true
+    echo "FIRST_RUN"
+  elif [ -f "$OPT_VENV/.venv-version" ]; then
+    IMAGE_STAMP=$(cat "$OPT_VENV/.venv-version")
+    CURRENT_STAMP=""
+    if [ -f "$AI_VENV/.venv-version" ]; then
+      CURRENT_STAMP=$(cat "$AI_VENV/.venv-version")
+    fi
+    if [ "$CURRENT_STAMP" != "$IMAGE_STAMP" ]; then
+      NEED_BOOTSTRAP=true
+      echo "STAMP_MISMATCH"
+    fi
+  fi
+
+  if [ "$NEED_BOOTSTRAP" = true ]; then
+    rm -rf "$AI_VENV"
+    cp -r "$OPT_VENV" "$AI_VENV_TMP"
+    mv "$AI_VENV_TMP" "$AI_VENV"
+    if [ -f "$TEST_INSTALLED_JSON" ]; then
+      echo '{"bundles":{}}' > "$TEST_INSTALLED_JSON"
+      echo "BUNDLES_RESET"
+    fi
+    echo "VENV_READY"
+  else
+    echo "SKIP"
+  fi
+else
+  echo "NO_OPT_VENV"
+fi
+`;
+
+function runBootstrap(): string {
+  return execFileSync("/bin/sh", ["-c", BOOTSTRAP_SCRIPT], {
+    env: {
+      TEST_AI_VENV: aiVenv,
+      TEST_AI_VENV_TMP: aiVenvTmp,
+      TEST_OPT_VENV: optVenv,
+      TEST_INSTALLED_JSON: installedJson,
+      PATH: process.env.PATH,
+    },
+    encoding: "utf-8",
+  }).trim();
+}
+
+beforeEach(() => {
+  root = mkdtempSync(join(tmpdir(), "venv-stamp-"));
+  optVenv = join(root, "opt-venv");
+  dataAi = join(root, "data-ai");
+  aiVenv = join(dataAi, "venv");
+  aiVenvTmp = join(dataAi, "venv.bootstrapping");
+  installedJson = join(dataAi, "installed.json");
+
+  // Simulate /opt/venv with a stamp and a marker file
+  mkdirSync(optVenv, { recursive: true });
+  writeFileSync(join(optVenv, ".venv-version"), "abc123\n");
+  writeFileSync(join(optVenv, "marker.txt"), "base-package-content");
+
+  mkdirSync(dataAi, { recursive: true });
+});
+
+afterEach(() => {
+  rmSync(root, { recursive: true, force: true });
+});
+
+describe("venv upgrade stamp", () => {
+  it("bootstraps on first run when no venv exists", () => {
+    const output = runBootstrap();
+    expect(output).toContain("FIRST_RUN");
+    expect(output).toContain("VENV_READY");
+    expect(existsSync(join(aiVenv, ".venv-version"))).toBe(true);
+    expect(readFileSync(join(aiVenv, "marker.txt"), "utf-8")).toBe("base-package-content");
+  });
+
+  it("skips when stamps match", () => {
+    // Simulate an existing venv with matching stamp
+    mkdirSync(aiVenv, { recursive: true });
+    writeFileSync(join(aiVenv, ".venv-version"), "abc123\n");
+
+    const output = runBootstrap();
+    expect(output).toBe("SKIP");
+  });
+
+  it("refreshes when stamps differ (upgrade)", () => {
+    // Simulate stale venv with old stamp
+    mkdirSync(aiVenv, { recursive: true });
+    writeFileSync(join(aiVenv, ".venv-version"), "old-hash\n");
+    writeFileSync(join(aiVenv, "stale-file.txt"), "should-be-removed");
+
+    const output = runBootstrap();
+    expect(output).toContain("STAMP_MISMATCH");
+    expect(output).toContain("VENV_READY");
+    // Old content replaced with fresh copy
+    expect(existsSync(join(aiVenv, "stale-file.txt"))).toBe(false);
+    expect(readFileSync(join(aiVenv, "marker.txt"), "utf-8")).toBe("base-package-content");
+    expect(readFileSync(join(aiVenv, ".venv-version"), "utf-8")).toBe("abc123\n");
+  });
+
+  it("refreshes when venv has no stamp (pre-stamp image upgrade)", () => {
+    // Simulate old venv without any stamp file
+    mkdirSync(aiVenv, { recursive: true });
+    writeFileSync(join(aiVenv, "old-pkg.txt"), "legacy");
+
+    const output = runBootstrap();
+    expect(output).toContain("STAMP_MISMATCH");
+    expect(output).toContain("VENV_READY");
+    expect(existsSync(join(aiVenv, "old-pkg.txt"))).toBe(false);
+  });
+
+  it("resets installed.json when refreshing", () => {
+    // Simulate existing venv + installed bundles
+    mkdirSync(aiVenv, { recursive: true });
+    writeFileSync(join(aiVenv, ".venv-version"), "old-hash\n");
+    writeFileSync(
+      installedJson,
+      JSON.stringify({
+        bundles: {
+          "background-removal": {
+            version: "1.0.0",
+            installedAt: "2025-01-01T00:00:00Z",
+          },
+        },
+      }),
+    );
+
+    const output = runBootstrap();
+    expect(output).toContain("BUNDLES_RESET");
+    const data = JSON.parse(readFileSync(installedJson, "utf-8"));
+    expect(data).toEqual({ bundles: {} });
+  });
+
+  it("does not reset installed.json when no bundles were installed", () => {
+    // No installed.json exists
+    mkdirSync(aiVenv, { recursive: true });
+    writeFileSync(join(aiVenv, ".venv-version"), "old-hash\n");
+
+    const output = runBootstrap();
+    expect(output).toContain("VENV_READY");
+    expect(output).not.toContain("BUNDLES_RESET");
+  });
+
+  it("does nothing when /opt/venv has no stamp (old image)", () => {
+    // Remove the stamp from the image venv
+    rmSync(join(optVenv, ".venv-version"));
+    mkdirSync(aiVenv, { recursive: true });
+    writeFileSync(join(aiVenv, "existing.txt"), "keep");
+
+    const output = runBootstrap();
+    expect(output).toBe("SKIP");
+    // Existing venv untouched
+    expect(existsSync(join(aiVenv, "existing.txt"))).toBe(true);
+  });
+
+  it("cleans up interrupted bootstrap from previous start", () => {
+    mkdirSync(aiVenvTmp, { recursive: true });
+    writeFileSync(join(aiVenvTmp, "partial.txt"), "incomplete");
+
+    const output = runBootstrap();
+    expect(output).toContain("CLEANUP_INTERRUPTED");
+    expect(existsSync(aiVenvTmp)).toBe(false);
+  });
+
+  it("does nothing when /opt/venv does not exist", () => {
+    rmSync(optVenv, { recursive: true });
+    const output = runBootstrap();
+    expect(output).toBe("NO_OPT_VENV");
+  });
+});