@vellumai/cli 0.5.6 → 0.5.8

package/src/commands/rollback.ts

@@ -9,7 +9,6 @@ import {
 import type { AssistantEntry } from "../lib/assistant-config";
 import {
   captureImageRefs,
-  clearSigningKeyBootstrapLock,
   GATEWAY_INTERNAL_PORT,
   dockerResourceNames,
   migrateCesSecurityFiles,
@@ -18,46 +17,79 @@ import {
   stopContainers,
 } from "../lib/docker";
 import type { ServiceName } from "../lib/docker";
-import { loadBootstrapSecret } from "../lib/guardian-token";
+import { emitCliError, categorizeUpgradeError } from "../lib/cli-error.js";
+import {
+  fetchOrganizationId,
+  readPlatformToken,
+  rollbackPlatformAssistant,
+} from "../lib/platform-client.js";
 import {
   broadcastUpgradeEvent,
+  buildCompleteEvent,
+  buildProgressEvent,
+  buildStartingEvent,
+  buildUpgradeCommitMessage,
   captureContainerEnv,
+  commitWorkspaceViaGateway,
+  CONTAINER_ENV_EXCLUDE_KEYS,
+  performDockerRollback,
+  rollbackMigrations,
+  UPGRADE_PROGRESS,
   waitForReady,
-} from "./upgrade";
+} from "../lib/upgrade-lifecycle.js";
+import { parseVersion } from "../lib/version-compat.js";
 
-function parseArgs(): { name: string | null } {
+function parseArgs(): { name: string | null; version: string | null } {
   const args = process.argv.slice(3);
   let name: string | null = null;
+  let version: string | null = null;
 
   for (let i = 0; i < args.length; i++) {
     const arg = args[i];
     if (arg === "--help" || arg === "-h") {
-      console.log("Usage: vellum rollback [<name>]");
+      console.log("Usage: vellum rollback [<name>] [--version <version>]");
       console.log("");
-      console.log("Roll back a Docker assistant to the previous version.");
+      console.log(
+        "Roll back a Docker or managed assistant to a previous version.",
+      );
       console.log("");
       console.log("Arguments:");
       console.log(
-        "  <name>  Name of the assistant to roll back (default: active or only assistant)",
+        "  <name>               Name of the assistant (default: active or only assistant)",
+      );
+      console.log("");
+      console.log("Options:");
+      console.log(
+        "  --version <version>  Target version (optional for managed — omit to roll back to previous)",
       );
       console.log("");
       console.log("Examples:");
       console.log(
-        "  vellum rollback                # Roll back the active assistant",
+        "  vellum rollback my-assistant                  # Roll back to previous version (Docker or managed)",
       );
       console.log(
-        "  vellum rollback my-assistant   # Roll back a specific assistant by name",
+        "  vellum rollback my-assistant --version v1.2.3 # Roll back to a specific version",
       );
       process.exit(0);
+    } else if (arg === "--version") {
+      const next = args[i + 1];
+      if (!next || next.startsWith("-")) {
+        console.error("Error: --version requires a value");
+        emitCliError("UNKNOWN", "--version requires a value");
+        process.exit(1);
+      }
+      version = next;
+      i++;
     } else if (!arg.startsWith("-")) {
       name = arg;
     } else {
       console.error(`Error: Unknown option '${arg}'.`);
+      emitCliError("UNKNOWN", `Unknown option '${arg}'`);
       process.exit(1);
     }
   }
 
-  return { name };
+  return { name, version };
 }
 
 function resolveCloud(entry: AssistantEntry): string {
@@ -84,6 +116,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,42 +135,218 @@ 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);
 }
 
+async function rollbackPlatformViaEndpoint(
+  entry: AssistantEntry,
+  version?: string,
+): Promise<void> {
+  const currentVersion = entry.serviceGroupVersion;
+
+  // Step 1 — Version validation (only if version provided)
+  if (version && currentVersion) {
+    const current = parseVersion(currentVersion);
+    const target = parseVersion(version);
+    if (current && target) {
+      const isOlder =
+        target.major < current.major ||
+        (target.major === current.major && target.minor < current.minor) ||
+        (target.major === current.major &&
+          target.minor === current.minor &&
+          target.patch < current.patch);
+      if (!isOlder) {
+        const msg = `Target version ${version} is not older than the current version ${currentVersion}. Use \`vellum upgrade --version ${version}\` to upgrade.`;
+        console.error(msg);
+        emitCliError("VERSION_DIRECTION", msg);
+        process.exit(1);
+      }
+    }
+  }
+
+  // Step 2 — Workspace commit (starting)
+  await commitWorkspaceViaGateway(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildUpgradeCommitMessage({
+      action: "rollback",
+      phase: "starting",
+      from: currentVersion ?? "unknown",
+      to: version ?? "previous",
+      topology: "managed",
+      assistantId: entry.assistantId,
+    }),
+  );
+
+  // Step 3 — Authenticate
+  const token = readPlatformToken();
+  if (!token) {
+    const msg =
+      "Error: Not logged in. Run `vellum login --token <token>` first.";
+    console.error(msg);
+    emitCliError("AUTH_FAILED", msg);
+    process.exit(1);
+  }
+
+  let orgId: string;
+  try {
+    orgId = await fetchOrganizationId(token);
+  } catch (err) {
+    const msg = err instanceof Error ? err.message : String(err);
+    if (msg.includes("401") || msg.includes("403")) {
+      console.error("Authentication failed. Run 'vellum login' to refresh.");
+    } else {
+      console.error(`Error: ${msg}`);
+    }
+    emitCliError("AUTH_FAILED", "Failed to authenticate with platform", msg);
+    process.exit(1);
+  }
+
+  // Step 4 — Broadcast starting event
+  console.log("📢 Notifying connected clients...");
+  await broadcastUpgradeEvent(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildStartingEvent(version ?? "previous", 90),
+  );
+
+  // Step 5 — Call rollback endpoint
+  if (version) {
+    console.log(`Rolling back to ${version}...`);
+  } else {
+    console.log("Rolling back to previous version...");
+  }
+
+  let result: { detail: string; version: string | null };
+  try {
+    result = await rollbackPlatformAssistant(token, orgId, version);
+  } catch (err) {
+    const detail = err instanceof Error ? err.message : String(err);
+
+    // Map specific server error messages to actionable CLI output
+    if (detail.includes("No previous version")) {
+      console.error(
+        "No previous version available. A successful upgrade must have been performed first.",
+      );
+    } else if (detail.includes("not older")) {
+      console.error(
+        `Target version is not older than the current version. Use 'vellum upgrade --version' instead.`,
+      );
+    } else if (detail.includes("not found")) {
+      console.error(
+        version
+          ? `Version ${version} not found.`
+          : `Rollback target not found.`,
+      );
+    } else if (
+      err instanceof TypeError ||
+      detail.includes("fetch failed") ||
+      detail.includes("ECONNREFUSED")
+    ) {
+      console.error(
+        `Connection error: ${detail}\nIs the platform reachable? Try 'vellum wake' if the assistant is asleep.`,
+      );
+    } else {
+      console.error(`Error: ${detail}`);
+    }
+
+    emitCliError("PLATFORM_API_ERROR", "Platform rollback failed", detail);
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildCompleteEvent(currentVersion ?? "unknown", false),
+    );
+    process.exit(1);
+  }
+
+  const rolledBackVersion = result.version ?? version ?? "unknown";
+
+  // Step 6 — Workspace commit (complete)
+  await commitWorkspaceViaGateway(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildUpgradeCommitMessage({
+      action: "rollback",
+      phase: "complete",
+      from: currentVersion ?? "unknown",
+      to: rolledBackVersion,
+      topology: "managed",
+      assistantId: entry.assistantId,
+      result: "success",
+    }),
+  );
+
+  // Step 7 — Print success
+  console.log(`Rolled back to version ${rolledBackVersion}.`);
+  if (!version) {
+    console.log("Tip: Run 'vellum rollback' again to undo.");
+  }
+}
+
 export async function rollback(): Promise<void> {
-  const { name } = parseArgs();
+  const { name, version } = parseArgs();
   const entry = resolveTargetAssistant(name);
   const cloud = resolveCloud(entry);
 
-  // Only Docker assistants support rollback
+  // ---------- Managed (Vellum platform) rollback ----------
+  if (cloud === "vellum") {
+    await rollbackPlatformViaEndpoint(entry, version ?? undefined);
+    return;
+  }
+
+  // ---------- Unsupported topologies ----------
   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 and managed assistants.";
+    console.error(msg);
+    emitCliError("UNSUPPORTED_TOPOLOGY", msg);
     process.exit(1);
   }
 
+  // ---------- Docker: Targeted version rollback (--version specified) ----------
+  if (version) {
+    try {
+      await performDockerRollback(entry, { targetVersion: version });
+    } 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),
+      );
+      emitCliError(categorizeUpgradeError(err), "Rollback failed", detail);
+      process.exit(1);
+    }
+    return;
+  }
+
+  // ---------- Docker: Saved-state rollback (no --version) ----------
+
   // 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 +360,215 @@ 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`,
-  );
+  try {
+    // Record rollback start in workspace git history
+    await commitWorkspaceViaGateway(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildUpgradeCommitMessage({
+        action: "rollback",
+        phase: "starting",
+        from: entry.serviceGroupVersion ?? "unknown",
+        to: entry.previousServiceGroupVersion ?? "unknown",
+        topology: "docker",
+        assistantId: entry.assistantId,
+      }),
+    );
 
-  // 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`,
-  );
+    console.log(
+      `🔄 Rolling back Docker assistant '${instanceName}' to ${entry.previousServiceGroupVersion}...\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);
+    // 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");
+
+    // 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;
+    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;
+    // 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,
+        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) {
+      // 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
+      await commitWorkspaceViaGateway(
+        entry.runtimeUrl,
+        entry.assistantId,
+        buildUpgradeCommitMessage({
+          action: "rollback",
+          phase: "complete",
+          from: entry.serviceGroupVersion ?? "unknown",
+          to: entry.previousServiceGroupVersion ?? "unknown",
+          topology: "docker",
+          assistantId: entry.assistantId,
+          result: "success",
+        }),
+      );
+
+      console.log(
+        `\n✅ Docker assistant '${instanceName}' rolled back to ${entry.previousServiceGroupVersion}.`,
+      );
+      console.log(
+        "\nTip: To also restore data from before the upgrade, use `vellum restore --from <backup-path>`.",
+      );
+    } 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);
   }
 }