@vellumai/cli 0.5.6 → 0.5.7

package/src/commands/rollback.ts

@@ -1,4 +1,5 @@
 import { randomBytes } from "crypto";
+import { join } from "path";
 
 import {
   findAssistantByName,
@@ -9,7 +10,6 @@ import {
 import type { AssistantEntry } from "../lib/assistant-config";
 import {
   captureImageRefs,
-  clearSigningKeyBootstrapLock,
   GATEWAY_INTERNAL_PORT,
   dockerResourceNames,
   migrateCesSecurityFiles,
@@ -18,12 +18,25 @@ import {
   stopContainers,
 } from "../lib/docker";
 import type { ServiceName } from "../lib/docker";
-import { loadBootstrapSecret } from "../lib/guardian-token";
+import {
+  loadBootstrapSecret,
+  saveBootstrapSecret,
+} from "../lib/guardian-token";
+import { restoreBackup } from "../lib/backup-ops.js";
+import { emitCliError, categorizeUpgradeError } from "../lib/cli-error.js";
 import {
   broadcastUpgradeEvent,
+  buildCompleteEvent,
+  buildProgressEvent,
+  buildStartingEvent,
+  buildUpgradeCommitMessage,
   captureContainerEnv,
+  CONTAINER_ENV_EXCLUDE_KEYS,
+  rollbackMigrations,
+  UPGRADE_PROGRESS,
   waitForReady,
-} from "./upgrade";
+} from "../lib/upgrade-lifecycle.js";
+import { commitWorkspaceState } from "../lib/workspace-git.js";
 
 function parseArgs(): { name: string | null } {
   const args = process.argv.slice(3);
@@ -53,6 +66,7 @@ function parseArgs(): { name: string | null } {
       name = arg;
     } else {
       console.error(`Error: Unknown option '${arg}'.`);
+      emitCliError("UNKNOWN", `Unknown option '${arg}'`);
       process.exit(1);
     }
   }
@@ -84,6 +98,10 @@ function resolveTargetAssistant(nameArg: string | null): AssistantEntry {
     const entry = findAssistantByName(nameArg);
     if (!entry) {
       console.error(`No assistant found with name '${nameArg}'.`);
+      emitCliError(
+        "ASSISTANT_NOT_FOUND",
+        `No assistant found with name '${nameArg}'.`,
+      );
       process.exit(1);
     }
     return entry;
@@ -99,11 +117,14 @@ function resolveTargetAssistant(nameArg: string | null): AssistantEntry {
   if (all.length === 1) return all[0];
 
   if (all.length === 0) {
-    console.error("No assistants found. Run 'vellum hatch' first.");
+    const msg = "No assistants found. Run 'vellum hatch' first.";
+    console.error(msg);
+    emitCliError("ASSISTANT_NOT_FOUND", msg);
   } else {
-    console.error(
-      "Multiple assistants found. Specify a name or set an active assistant with 'vellum use <name>'.",
-    );
+    const msg =
+      "Multiple assistants found. Specify a name or set an active assistant with 'vellum use <name>'.";
+    console.error(msg);
+    emitCliError("ASSISTANT_NOT_FOUND", msg);
   }
   process.exit(1);
 }
@@ -115,26 +136,29 @@ export async function rollback(): Promise<void> {
 
   // Only Docker assistants support rollback
   if (cloud !== "docker") {
-    console.error(
-      "Rollback is only supported for Docker assistants. For managed assistants, use the version picker to upgrade to the previous version.",
-    );
+    const msg =
+      "Rollback is only supported for Docker assistants. For managed assistants, use the version picker to upgrade to the previous version.";
+    console.error(msg);
+    emitCliError("UNSUPPORTED_TOPOLOGY", msg);
     process.exit(1);
   }
 
   // Verify rollback state exists
   if (!entry.previousServiceGroupVersion || !entry.previousContainerInfo) {
-    console.error(
-      "No rollback state available. Run `vellum upgrade` first to create a rollback point.",
-    );
+    const msg =
+      "No rollback state available. Run `vellum upgrade` first to create a rollback point.";
+    console.error(msg);
+    emitCliError("ROLLBACK_NO_STATE", msg);
     process.exit(1);
   }
 
   // Verify all three digest fields are present
   const prev = entry.previousContainerInfo;
   if (!prev.assistantDigest || !prev.gatewayDigest || !prev.cesDigest) {
-    console.error(
-      "Incomplete rollback state. Previous container digests are missing.",
-    );
+    const msg =
+      "Incomplete rollback state. Previous container digests are missing.";
+    console.error(msg);
+    emitCliError("ROLLBACK_NO_STATE", msg);
     process.exit(1);
   }
 
@@ -148,133 +172,272 @@ export async function rollback(): Promise<void> {
   const instanceName = entry.assistantId;
   const res = dockerResourceNames(instanceName);
 
-  console.log(
-    `🔄 Rolling back Docker assistant '${instanceName}' to ${entry.previousServiceGroupVersion}...\n`,
-  );
-
-  // Capture current container env
-  console.log("💾 Capturing existing container environment...");
-  const capturedEnv = await captureContainerEnv(res.assistantContainer);
-  console.log(
-    `   Captured ${Object.keys(capturedEnv).length} env var(s) from ${res.assistantContainer}\n`,
-  );
-
-  // Extract CES_SERVICE_TOKEN from captured env, or generate fresh one
-  const cesServiceToken =
-    capturedEnv["CES_SERVICE_TOKEN"] || randomBytes(32).toString("hex");
-
-  // Retrieve or generate a bootstrap secret for the gateway.
-  const bootstrapSecret =
-    loadBootstrapSecret(instanceName) || randomBytes(32).toString("hex");
-
-  // Build extra env vars, excluding keys managed by serviceDockerRunArgs
-  const envKeysSetByRunArgs = new Set([
-    "CES_SERVICE_TOKEN",
-    "VELLUM_ASSISTANT_NAME",
-    "RUNTIME_HTTP_HOST",
-    "PATH",
-  ]);
-  for (const envVar of ["ANTHROPIC_API_KEY", "VELLUM_PLATFORM_URL"]) {
-    if (process.env[envVar]) {
-      envKeysSetByRunArgs.add(envVar);
+  try {
+    const workspaceDir = entry.resources
+      ? join(entry.resources.instanceDir, ".vellum", "workspace")
+      : undefined;
+
+    // Record rollback start in workspace git history
+    if (workspaceDir) {
+      try {
+        await commitWorkspaceState(
+          workspaceDir,
+          buildUpgradeCommitMessage({
+            action: "rollback",
+            phase: "starting",
+            from: entry.serviceGroupVersion ?? "unknown",
+            to: entry.previousServiceGroupVersion ?? "unknown",
+            topology: "docker",
+            assistantId: entry.assistantId,
+          }),
+        );
+      } catch (err) {
+        console.warn(
+          `⚠️  Failed to create pre-rollback workspace commit: ${err instanceof Error ? err.message : String(err)}`,
+        );
+      }
     }
-  }
-  const extraAssistantEnv: Record<string, string> = {};
-  for (const [key, value] of Object.entries(capturedEnv)) {
-    if (!envKeysSetByRunArgs.has(key)) {
-      extraAssistantEnv[key] = value;
+
+    console.log(
+      `🔄 Rolling back Docker assistant '${instanceName}' to ${entry.previousServiceGroupVersion}...\n`,
+    );
+
+    // Capture current container env
+    console.log("💾 Capturing existing container environment...");
+    const capturedEnv = await captureContainerEnv(res.assistantContainer);
+    console.log(
+      `   Captured ${Object.keys(capturedEnv).length} env var(s) from ${res.assistantContainer}\n`,
+    );
+
+    // Extract CES_SERVICE_TOKEN from captured env, or generate fresh one
+    const cesServiceToken =
+      capturedEnv["CES_SERVICE_TOKEN"] || randomBytes(32).toString("hex");
+
+    // Retrieve or generate a bootstrap secret for the gateway.
+    const loadedSecret = loadBootstrapSecret(instanceName);
+    const bootstrapSecret = loadedSecret || randomBytes(32).toString("hex");
+    if (!loadedSecret) {
+      saveBootstrapSecret(instanceName, bootstrapSecret);
     }
-  }
 
-  // Parse gateway port from entry's runtimeUrl, fall back to default
-  let gatewayPort = GATEWAY_INTERNAL_PORT;
-  try {
-    const parsed = new URL(entry.runtimeUrl);
-    const port = parseInt(parsed.port, 10);
-    if (!isNaN(port)) {
-      gatewayPort = port;
+    // Extract or generate the shared JWT signing key.
+    const signingKey =
+      capturedEnv["ACTOR_TOKEN_SIGNING_KEY"] || randomBytes(32).toString("hex");
+
+    // Build extra env vars, excluding keys managed by serviceDockerRunArgs
+    const envKeysSetByRunArgs = new Set(CONTAINER_ENV_EXCLUDE_KEYS);
+    for (const envVar of ["ANTHROPIC_API_KEY", "VELLUM_PLATFORM_URL"]) {
+      if (process.env[envVar]) {
+        envKeysSetByRunArgs.add(envVar);
+      }
+    }
+    const extraAssistantEnv: Record<string, string> = {};
+    for (const [key, value] of Object.entries(capturedEnv)) {
+      if (!envKeysSetByRunArgs.has(key)) {
+        extraAssistantEnv[key] = value;
+      }
+    }
+
+    // Parse gateway port from entry's runtimeUrl, fall back to default
+    let gatewayPort = GATEWAY_INTERNAL_PORT;
+    try {
+      const parsed = new URL(entry.runtimeUrl);
+      const port = parseInt(parsed.port, 10);
+      if (!isNaN(port)) {
+        gatewayPort = port;
+      }
+    } catch {
+      // use default
     }
-  } catch {
-    // use default
-  }
 
-  // Notify connected clients that a rollback is about to begin (best-effort)
-  console.log("📢 Notifying connected clients...");
-  await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
-    type: "starting",
-    targetVersion: entry.previousServiceGroupVersion,
-    expectedDowntimeSeconds: 60,
-  });
-  // Brief pause to allow SSE delivery before containers stop.
-  await new Promise((r) => setTimeout(r, 500));
-
-  console.log("🛑 Stopping existing containers...");
-  await stopContainers(res);
-  console.log("✅ Containers stopped\n");
-
-  // Run security file migrations and signing key cleanup
-  console.log("🔄 Migrating security files to gateway volume...");
-  await migrateGatewaySecurityFiles(res, (msg) => console.log(msg));
-
-  console.log("🔄 Migrating credential files to CES security volume...");
-  await migrateCesSecurityFiles(res, (msg) => console.log(msg));
-
-  console.log("🔑 Clearing signing key bootstrap lock...");
-  await clearSigningKeyBootstrapLock(res);
-
-  console.log("🚀 Starting containers with previous version...");
-  await startContainers(
-    {
-      bootstrapSecret,
-      cesServiceToken,
-      extraAssistantEnv,
-      gatewayPort,
-      imageTags: previousImageRefs,
-      instanceName,
-      res,
-    },
-    (msg) => console.log(msg),
-  );
-  console.log("✅ Containers started\n");
-
-  console.log("Waiting for assistant to become ready...");
-  const ready = await waitForReady(entry.runtimeUrl);
-
-  if (ready) {
-    // Capture new digests from the rolled-back containers
-    const newDigests = await captureImageRefs(res);
-
-    // Swap current/previous state to enable "rollback the rollback"
-    const updatedEntry: AssistantEntry = {
-      ...entry,
-      serviceGroupVersion: entry.previousServiceGroupVersion,
-      containerInfo: {
-        assistantImage: prev.assistantImage ?? previousImageRefs.assistant,
-        gatewayImage: prev.gatewayImage ?? previousImageRefs.gateway,
-        cesImage: prev.cesImage ?? previousImageRefs["credential-executor"],
-        assistantDigest: newDigests?.assistant,
-        gatewayDigest: newDigests?.gateway,
-        cesDigest: newDigests?.["credential-executor"],
-        networkName: res.network,
+    // Notify connected clients that a rollback is about to begin (best-effort)
+    console.log("📢 Notifying connected clients...");
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildStartingEvent(entry.previousServiceGroupVersion),
+    );
+    // Brief pause to allow SSE delivery before containers stop.
+    await new Promise((r) => setTimeout(r, 500));
+
+    // Roll back migrations to pre-upgrade state (must happen before containers stop)
+    if (
+      entry.previousDbMigrationVersion !== undefined ||
+      entry.previousWorkspaceMigrationId !== undefined
+    ) {
+      console.log("🔄 Reverting database changes...");
+      await broadcastUpgradeEvent(
+        entry.runtimeUrl,
+        entry.assistantId,
+        buildProgressEvent(UPGRADE_PROGRESS.REVERTING_MIGRATIONS),
+      );
+      await rollbackMigrations(
+        entry.runtimeUrl,
+        entry.assistantId,
+        entry.previousDbMigrationVersion,
+        entry.previousWorkspaceMigrationId,
+      );
+    }
+
+    // Progress: switching version (must be sent BEFORE stopContainers)
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildProgressEvent(UPGRADE_PROGRESS.SWITCHING),
+    );
+
+    console.log("🛑 Stopping existing containers...");
+    await stopContainers(res);
+    console.log("✅ Containers stopped\n");
+
+    // Run security file migrations and signing key cleanup
+    console.log("🔄 Migrating security files to gateway volume...");
+    await migrateGatewaySecurityFiles(res, (msg) => console.log(msg));
+
+    console.log("🔄 Migrating credential files to CES security volume...");
+    await migrateCesSecurityFiles(res, (msg) => console.log(msg));
+
+    console.log("🚀 Starting containers with previous version...");
+    await startContainers(
+      {
+        signingKey,
+        bootstrapSecret,
+        cesServiceToken,
+        extraAssistantEnv,
+        gatewayPort,
+        imageTags: previousImageRefs,
+        instanceName,
+        res,
       },
-      previousServiceGroupVersion: entry.serviceGroupVersion,
-      previousContainerInfo: entry.containerInfo,
-    };
-    saveAssistantEntry(updatedEntry);
-
-    // Notify clients that the rollback succeeded
-    await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
-      type: "complete",
-      installedVersion: entry.previousServiceGroupVersion,
-      success: true,
-    });
+      (msg) => console.log(msg),
+    );
+    console.log("✅ Containers started\n");
+
+    console.log("Waiting for assistant to become ready...");
+    const ready = await waitForReady(entry.runtimeUrl);
+
+    if (ready) {
+      // Restore data from the backup created for the specific upgrade being
+      // rolled back. We use the persisted preUpgradeBackupPath rather than
+      // scanning for the latest backup on disk — if the most recent upgrade's
+      // backup failed, a global scan would find a stale backup from a prior
+      // cycle and overwrite newer user data.
+      const backupPath = entry.preUpgradeBackupPath as string | undefined;
+      if (backupPath) {
+        // Progress: restoring data (gateway is back up at this point)
+        await broadcastUpgradeEvent(
+          entry.runtimeUrl,
+          entry.assistantId,
+          buildProgressEvent(UPGRADE_PROGRESS.RESTORING),
+        );
+
+        console.log(`📦 Restoring data from pre-upgrade backup...`);
+        console.log(`   Source: ${backupPath}`);
+        const restored = await restoreBackup(
+          entry.runtimeUrl,
+          entry.assistantId,
+          backupPath,
+        );
+        if (restored) {
+          console.log("   ✅ Data restored successfully\n");
+        } else {
+          console.warn(
+            "   ⚠️  Data restore failed (rollback continues without data restoration)\n",
+          );
+        }
+      } else {
+        console.log(
+          "ℹ️  No pre-upgrade backup was created for this upgrade, skipping data restoration\n",
+        );
+      }
+
+      // Capture new digests from the rolled-back containers
+      const newDigests = await captureImageRefs(res);
+
+      // Swap current/previous state to enable "rollback the rollback"
+      const updatedEntry: AssistantEntry = {
+        ...entry,
+        serviceGroupVersion: entry.previousServiceGroupVersion,
+        containerInfo: {
+          assistantImage: prev.assistantImage ?? previousImageRefs.assistant,
+          gatewayImage: prev.gatewayImage ?? previousImageRefs.gateway,
+          cesImage: prev.cesImage ?? previousImageRefs["credential-executor"],
+          assistantDigest: newDigests?.assistant,
+          gatewayDigest: newDigests?.gateway,
+          cesDigest: newDigests?.["credential-executor"],
+          networkName: res.network,
+        },
+        previousServiceGroupVersion: entry.serviceGroupVersion,
+        previousContainerInfo: entry.containerInfo,
+        // Clear the backup path — it belonged to the upgrade we just rolled back
+        preUpgradeBackupPath: undefined,
+        previousDbMigrationVersion: undefined,
+        previousWorkspaceMigrationId: undefined,
+      };
+      saveAssistantEntry(updatedEntry);
+
+      // Notify clients that the rollback succeeded
+      await broadcastUpgradeEvent(
+        entry.runtimeUrl,
+        entry.assistantId,
+        buildCompleteEvent(entry.previousServiceGroupVersion, true),
+      );
 
-    console.log(
-      `\n✅ Docker assistant '${instanceName}' rolled back to ${entry.previousServiceGroupVersion}.`,
+      // Record successful rollback in workspace git history
+      if (workspaceDir) {
+        try {
+          await commitWorkspaceState(
+            workspaceDir,
+            buildUpgradeCommitMessage({
+              action: "rollback",
+              phase: "complete",
+              from: entry.serviceGroupVersion ?? "unknown",
+              to: entry.previousServiceGroupVersion ?? "unknown",
+              topology: "docker",
+              assistantId: entry.assistantId,
+              result: "success",
+            }),
+          );
+        } catch (err) {
+          console.warn(
+            `⚠️  Failed to create post-rollback workspace commit: ${err instanceof Error ? err.message : String(err)}`,
+          );
+        }
+      }
+
+      console.log(
+        `\n✅ Docker assistant '${instanceName}' rolled back to ${entry.previousServiceGroupVersion}.`,
+      );
+    } else {
+      console.error(
+        `\n❌ Containers failed to become ready within the timeout.`,
+      );
+      console.log(
+        `   Check logs with: docker logs -f ${res.assistantContainer}`,
+      );
+      await broadcastUpgradeEvent(
+        entry.runtimeUrl,
+        entry.assistantId,
+        buildCompleteEvent(
+          entry.previousServiceGroupVersion ?? "unknown",
+          false,
+        ),
+      );
+      emitCliError(
+        "READINESS_TIMEOUT",
+        "Rolled-back containers failed to become ready within the timeout.",
+      );
+      process.exit(1);
+    }
+  } catch (err) {
+    const detail = err instanceof Error ? err.message : String(err);
+    console.error(`\n❌ Rollback failed: ${detail}`);
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
     );
-  } else {
-    console.error(`\n❌ Containers failed to become ready within the timeout.`);
-    console.log(`   Check logs with: docker logs -f ${res.assistantContainer}`);
+    emitCliError(categorizeUpgradeError(err), "Rollback failed", detail);
     process.exit(1);
   }
 }