@vellumai/cli 0.5.6 → 0.5.7

package/src/commands/upgrade.ts

@@ -1,4 +1,5 @@
 import { randomBytes } from "crypto";
+import { join } from "node:path";
 
 import cliPkg from "../../package.json";
 
@@ -11,9 +12,6 @@ import {
 import type { AssistantEntry } from "../lib/assistant-config";
 import {
   captureImageRefs,
-  clearSigningKeyBootstrapLock,
-  DOCKERHUB_IMAGES,
-  DOCKER_READY_TIMEOUT_MS,
   GATEWAY_INTERNAL_PORT,
   dockerResourceNames,
   migrateCesSecurityFiles,
@@ -21,14 +19,37 @@ import {
   startContainers,
   stopContainers,
 } from "../lib/docker";
-import type { ServiceName } from "../lib/docker";
+import { resolveImageRefs } from "../lib/platform-releases";
 import {
   fetchOrganizationId,
   getPlatformUrl,
   readPlatformToken,
 } from "../lib/platform-client";
-import { loadBootstrapSecret, loadGuardianToken } from "../lib/guardian-token";
-import { exec, execOutput } from "../lib/step-runner";
+import {
+  loadBootstrapSecret,
+  saveBootstrapSecret,
+} from "../lib/guardian-token";
+import {
+  createBackup,
+  pruneOldBackups,
+  restoreBackup,
+} from "../lib/backup-ops.js";
+import { emitCliError, categorizeUpgradeError } from "../lib/cli-error.js";
+import { exec } from "../lib/step-runner.js";
+import {
+  broadcastUpgradeEvent,
+  buildCompleteEvent,
+  buildProgressEvent,
+  buildStartingEvent,
+  buildUpgradeCommitMessage,
+  captureContainerEnv,
+  CONTAINER_ENV_EXCLUDE_KEYS,
+  rollbackMigrations,
+  UPGRADE_PROGRESS,
+  waitForReady,
+} from "../lib/upgrade-lifecycle.js";
+import { parseVersion } from "../lib/version-compat.js";
+import { commitWorkspaceState } from "../lib/workspace-git.js";
 
 interface UpgradeArgs {
   name: string | null;
@@ -72,6 +93,7 @@ function parseArgs(): UpgradeArgs {
       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;
@@ -80,6 +102,7 @@ function parseArgs(): UpgradeArgs {
       name = arg;
     } else {
       console.error(`Error: Unknown option '${arg}'.`);
+      emitCliError("UNKNOWN", `Unknown option '${arg}'`);
       process.exit(1);
     }
   }
@@ -111,6 +134,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;
@@ -126,124 +153,32 @@ 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);
 }
 
-/**
- * Capture environment variables from a running Docker container so they
- * can be replayed onto the replacement container after upgrade.
- */
-export async function captureContainerEnv(
-  containerName: string,
-): Promise<Record<string, string>> {
-  const captured: Record<string, string> = {};
-  try {
-    const raw = await execOutput("docker", [
-      "inspect",
-      "--format",
-      "{{json .Config.Env}}",
-      containerName,
-    ]);
-    const entries = JSON.parse(raw) as string[];
-    for (const entry of entries) {
-      const eqIdx = entry.indexOf("=");
-      if (eqIdx > 0) {
-        captured[entry.slice(0, eqIdx)] = entry.slice(eqIdx + 1);
-      }
-    }
-  } catch {
-    // Container may not exist or not be inspectable
-  }
-  return captured;
-}
-
-/**
- * Poll the gateway `/readyz` endpoint until it returns 200 or the timeout
- * elapses. Returns whether the assistant became ready.
- */
-export async function waitForReady(runtimeUrl: string): Promise<boolean> {
-  const readyUrl = `${runtimeUrl}/readyz`;
-  const start = Date.now();
-
-  while (Date.now() - start < DOCKER_READY_TIMEOUT_MS) {
-    try {
-      const resp = await fetch(readyUrl, {
-        signal: AbortSignal.timeout(5000),
-      });
-      if (resp.ok) {
-        const elapsedSec = ((Date.now() - start) / 1000).toFixed(1);
-        console.log(`Assistant ready after ${elapsedSec}s`);
-        return true;
-      }
-      let detail = "";
-      try {
-        const body = await resp.text();
-        const json = JSON.parse(body);
-        const parts = [json.status];
-        if (json.upstream != null) parts.push(`upstream=${json.upstream}`);
-        detail = ` — ${parts.join(", ")}`;
-      } catch {
-        // ignore parse errors
-      }
-      console.log(`Readiness check: ${resp.status}${detail} (retrying...)`);
-    } catch {
-      // Connection refused / timeout — not up yet
-    }
-    await new Promise((r) => setTimeout(r, 1000));
-  }
-
-  return false;
-}
-
-/**
- * Best-effort broadcast of an upgrade lifecycle event to connected clients
- * via the gateway's upgrade-broadcast proxy. Uses guardian token auth.
- * Failures are logged but never block the upgrade flow.
- */
-export async function broadcastUpgradeEvent(
-  gatewayUrl: string,
-  assistantId: string,
-  event: Record<string, unknown>,
-): Promise<void> {
-  try {
-    const token = loadGuardianToken(assistantId);
-    const headers: Record<string, string> = {
-      "Content-Type": "application/json",
-    };
-    if (token?.accessToken) {
-      headers["Authorization"] = `Bearer ${token.accessToken}`;
-    }
-    await fetch(`${gatewayUrl}/v1/admin/upgrade-broadcast`, {
-      method: "POST",
-      headers,
-      body: JSON.stringify(event),
-      signal: AbortSignal.timeout(3000),
-    });
-  } catch {
-    // Best-effort — gateway/daemon may already be shutting down or not yet ready
-  }
-}
-
 async function upgradeDocker(
   entry: AssistantEntry,
   version: string | null,
 ): Promise<void> {
   const instanceName = entry.assistantId;
   const res = dockerResourceNames(instanceName);
+  const workspaceDir = entry.resources
+    ? join(entry.resources.instanceDir, ".vellum", "workspace")
+    : null;
 
   const versionTag =
     version ?? (cliPkg.version ? `v${cliPkg.version}` : "latest");
-  const imageTags: Record<ServiceName, string> = {
-    assistant: `${DOCKERHUB_IMAGES.assistant}:${versionTag}`,
-    "credential-executor": `${DOCKERHUB_IMAGES["credential-executor"]}:${versionTag}`,
-    gateway: `${DOCKERHUB_IMAGES.gateway}:${versionTag}`,
-  };
+  console.log("🔍 Resolving image references...");
+  const { imageTags } = await resolveImageRefs(versionTag);
 
   console.log(
     `🔄 Upgrading Docker assistant '${instanceName}' to ${versionTag}...\n`,
@@ -265,6 +200,86 @@ async function upgradeDocker(
     );
   }
 
+  // Capture current migration state for rollback targeting.
+  // Must happen while daemon is still running (before containers are stopped).
+  let preMigrationState: {
+    dbVersion?: number;
+    lastWorkspaceMigrationId?: string;
+  } = {};
+  try {
+    const healthResp = await fetch(
+      `${entry.runtimeUrl}/healthz?include=migrations`,
+      {
+        signal: AbortSignal.timeout(5000),
+      },
+    );
+    if (healthResp.ok) {
+      const health = (await healthResp.json()) as {
+        migrations?: { dbVersion?: number; lastWorkspaceMigrationId?: string };
+      };
+      preMigrationState = health.migrations ?? {};
+    }
+  } catch {
+    // Best-effort — if we can't get migration state, rollback will skip migration reversal
+  }
+
+  // Detect if this upgrade is actually a downgrade (user picked an older
+  // version via the version picker). Used after readiness succeeds to align
+  // the DB schema with the now-running old daemon.
+  const currentVersion = entry.serviceGroupVersion;
+  const isDowngrade =
+    currentVersion &&
+    versionTag &&
+    (() => {
+      const current = parseVersion(currentVersion);
+      const target = parseVersion(versionTag);
+      if (!current || !target) return false;
+      if (target.major !== current.major) return target.major < current.major;
+      if (target.minor !== current.minor) return target.minor < current.minor;
+      return target.patch < current.patch;
+    })();
+
+  // For downgrades, fetch the target version's migration ceiling from the
+  // releases API. This tells us exactly which DB migration version and
+  // workspace migration the target version expects, enabling a precise
+  // rollback on the CURRENT (newer) daemon before swapping containers.
+  let targetMigrationCeiling: {
+    dbVersion?: number;
+    workspaceMigrationId?: string;
+  } = {};
+  if (isDowngrade) {
+    try {
+      const platformUrl = getPlatformUrl();
+      const releasesResp = await fetch(
+        `${platformUrl}/v1/releases/?stable=true`,
+        { signal: AbortSignal.timeout(10000) },
+      );
+      if (releasesResp.ok) {
+        const releases = (await releasesResp.json()) as Array<{
+          version: string;
+          db_migration_version?: number | null;
+          last_workspace_migration_id?: string;
+        }>;
+        const normalizedTag = versionTag.replace(/^v/, "");
+        const targetRelease = releases.find(
+          (r) => r.version?.replace(/^v/, "") === normalizedTag,
+        );
+        if (
+          targetRelease?.db_migration_version != null ||
+          targetRelease?.last_workspace_migration_id
+        ) {
+          targetMigrationCeiling = {
+            dbVersion: targetRelease.db_migration_version ?? undefined,
+            workspaceMigrationId:
+              targetRelease.last_workspace_migration_id || undefined,
+          };
+        }
+      }
+    } catch {
+      // Best-effort — fall back to rollbackToRegistryCeiling post-swap
+    }
+  }
+
   // Persist rollback state to lockfile BEFORE any destructive changes.
   // This enables the `vellum rollback` command to restore the previous version.
   if (entry.serviceGroupVersion && entry.containerInfo) {
@@ -272,36 +287,80 @@ async function upgradeDocker(
       ...entry,
       previousServiceGroupVersion: entry.serviceGroupVersion,
       previousContainerInfo: { ...entry.containerInfo },
+      previousDbMigrationVersion: preMigrationState.dbVersion,
+      previousWorkspaceMigrationId: preMigrationState.lastWorkspaceMigrationId,
     };
     saveAssistantEntry(rollbackEntry);
     console.log(`   Saved rollback state: ${entry.serviceGroupVersion}\n`);
   }
 
+  // Record version transition start in workspace git history
+  if (workspaceDir) {
+    try {
+      await commitWorkspaceState(
+        workspaceDir,
+        buildUpgradeCommitMessage({
+          action: "upgrade",
+          phase: "starting",
+          from: entry.serviceGroupVersion ?? "unknown",
+          to: versionTag,
+          topology: "docker",
+          assistantId: entry.assistantId,
+        }),
+      );
+    } catch (err) {
+      console.warn(
+        `⚠️  Failed to create pre-upgrade workspace commit: ${err instanceof Error ? err.message : String(err)}`,
+      );
+    }
+  }
+
   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("📦 Pulling new Docker images...");
-  await exec("docker", ["pull", imageTags.assistant]);
-  await exec("docker", ["pull", imageTags.gateway]);
-  await exec("docker", ["pull", imageTags["credential-executor"]]);
-  console.log("✅ Docker images pulled\n");
-
   // Notify connected clients that an upgrade is about to begin.
+  // This must fire BEFORE any progress broadcasts so the UI sets
+  // isUpdateInProgress = true and starts displaying status messages.
   console.log("📢 Notifying connected clients...");
-  await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
-    type: "starting",
-    targetVersion: versionTag,
-    expectedDowntimeSeconds: 60,
-  });
-  // Brief pause to allow SSE delivery before containers stop.
+  await broadcastUpgradeEvent(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildStartingEvent(versionTag),
+  );
+  // Brief pause to allow SSE delivery before progress events.
   await new Promise((r) => setTimeout(r, 500));
 
-  console.log("🛑 Stopping existing containers...");
-  await stopContainers(res);
-  console.log("✅ Containers stopped\n");
+  await broadcastUpgradeEvent(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildProgressEvent(UPGRADE_PROGRESS.DOWNLOADING),
+  );
+  console.log("📦 Pulling new Docker images...");
+  const pullImages: Array<[string, string]> = [
+    ["assistant", imageTags.assistant],
+    ["gateway", imageTags.gateway],
+    ["credential-executor", imageTags["credential-executor"]],
+  ];
+  try {
+    for (const [service, image] of pullImages) {
+      console.log(`   Pulling ${service}: ${image}`);
+      await exec("docker", ["pull", image]);
+    }
+  } catch (pullErr) {
+    const detail = pullErr instanceof Error ? pullErr.message : String(pullErr);
+    console.error(`\n❌ Failed to pull Docker images: ${detail}`);
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
+    );
+    emitCliError("IMAGE_PULL_FAILED", "Failed to pull Docker images", detail);
+    process.exit(1);
+  }
+  console.log("✅ Docker images pulled\n");
 
   // Parse gateway port from entry's runtimeUrl, fall back to default
   let gatewayPort = GATEWAY_INTERNAL_PORT;
@@ -324,18 +383,91 @@ async function upgradeDocker(
 
   // Retrieve or generate a bootstrap secret for the gateway. The secret was
   // persisted to disk during hatch; older instances won't have one yet.
-  const bootstrapSecret =
-    loadBootstrapSecret(instanceName) || randomBytes(32).toString("hex");
+  // This runs BEFORE stopping containers so a write failure (disk full,
+  // permissions) doesn't leave the assistant offline.
+  const loadedSecret = loadBootstrapSecret(instanceName);
+  const bootstrapSecret = loadedSecret || randomBytes(32).toString("hex");
+  if (!loadedSecret) {
+    saveBootstrapSecret(instanceName, bootstrapSecret);
+  }
+
+  // Extract or generate the shared JWT signing key. Pre-env-var instances
+  // won't have it in capturedEnv, so generate fresh in that case.
+  const signingKey =
+    capturedEnv["ACTOR_TOKEN_SIGNING_KEY"] || randomBytes(32).toString("hex");
+
+  // Create pre-upgrade backup (best-effort, daemon must be running)
+  await broadcastUpgradeEvent(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildProgressEvent(UPGRADE_PROGRESS.BACKING_UP),
+  );
+  console.log("📦 Creating pre-upgrade backup...");
+  const backupPath = await createBackup(entry.runtimeUrl, entry.assistantId, {
+    prefix: `${entry.assistantId}-pre-upgrade`,
+    description: `Pre-upgrade snapshot before ${entry.serviceGroupVersion ?? "unknown"} → ${versionTag}`,
+  });
+  if (backupPath) {
+    console.log(`   Backup saved: ${backupPath}\n`);
+    // Clean up old pre-upgrade backups, keep last 3
+    pruneOldBackups(entry.assistantId, 3);
+  } else {
+    console.warn("⚠️  Pre-upgrade backup failed (continuing with upgrade)\n");
+  }
+
+  // Persist the backup path so `vellum rollback` can restore the exact backup
+  // created for this upgrade attempt — never a stale backup from a prior cycle.
+  // Re-read the entry to pick up the rollback state saved earlier.
+  {
+    const current = findAssistantByName(entry.assistantId);
+    if (current) {
+      saveAssistantEntry({
+        ...current,
+        preUpgradeBackupPath: backupPath ?? undefined,
+      });
+    }
+  }
+
+  await broadcastUpgradeEvent(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildProgressEvent(UPGRADE_PROGRESS.INSTALLING),
+  );
+
+  // If we have the target version's migration ceiling, run a PRECISE
+  // rollback on the CURRENT (newer) daemon before stopping it. The current
+  // daemon has the `down()` code for all migrations it applied, so it can
+  // cleanly revert to the target version's ceiling. This is critical for
+  // multi-version downgrades where the old daemon wouldn't know about
+  // migrations introduced after its release.
+  let preSwapRollbackOk = true;
+  if (
+    isDowngrade &&
+    (targetMigrationCeiling.dbVersion !== undefined ||
+      targetMigrationCeiling.workspaceMigrationId !== undefined)
+  ) {
+    console.log("🔄 Reverting database changes for downgrade...");
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildProgressEvent(UPGRADE_PROGRESS.REVERTING_MIGRATIONS),
+    );
+    preSwapRollbackOk = await rollbackMigrations(
+      entry.runtimeUrl,
+      entry.assistantId,
+      targetMigrationCeiling.dbVersion,
+      targetMigrationCeiling.workspaceMigrationId,
+    );
+  }
+
+  console.log("🛑 Stopping existing containers...");
+  await stopContainers(res);
+  console.log("✅ Containers stopped\n");
 
   // Build the set of extra env vars to replay on the new assistant container.
   // Captured env vars serve as the base; keys already managed by
   // serviceDockerRunArgs are excluded to avoid duplicates.
-  const envKeysSetByRunArgs = new Set([
-    "CES_SERVICE_TOKEN",
-    "VELLUM_ASSISTANT_NAME",
-    "RUNTIME_HTTP_HOST",
-    "PATH",
-  ]);
+  const envKeysSetByRunArgs = new Set(CONTAINER_ENV_EXCLUDE_KEYS);
   // Only exclude keys that serviceDockerRunArgs will actually set
   for (const envVar of ["ANTHROPIC_API_KEY", "VELLUM_PLATFORM_URL"]) {
     if (process.env[envVar]) {
@@ -355,12 +487,10 @@ async function upgradeDocker(
   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 upgraded containers...");
   await startContainers(
     {
+      signingKey,
       bootstrapSecret,
       cesServiceToken,
       extraAssistantEnv,
@@ -392,15 +522,62 @@ async function upgradeDocker(
       },
       previousServiceGroupVersion: entry.serviceGroupVersion,
       previousContainerInfo: entry.containerInfo,
+      previousDbMigrationVersion: preMigrationState.dbVersion,
+      previousWorkspaceMigrationId: preMigrationState.lastWorkspaceMigrationId,
+      // Preserve the backup path so `vellum rollback` can restore it later
+      preUpgradeBackupPath: backupPath ?? undefined,
     };
     saveAssistantEntry(updatedEntry);
 
+    // After a downgrade, fall back to asking the now-running old daemon
+    // to roll back migrations above its own registry ceiling when either:
+    // (a) no release metadata was available for a precise pre-swap rollback, or
+    // (b) the precise pre-swap rollback failed (timeout, daemon crash, etc.).
+    // This is a no-op for multi-version jumps where the old daemon doesn't
+    // know about the newer migrations, but correct for single-step rollbacks.
+    if (
+      isDowngrade &&
+      (!preSwapRollbackOk ||
+        (targetMigrationCeiling.dbVersion === undefined &&
+          targetMigrationCeiling.workspaceMigrationId === undefined))
+    ) {
+      await rollbackMigrations(
+        entry.runtimeUrl,
+        entry.assistantId,
+        undefined,
+        undefined,
+        true,
+      );
+    }
+
     // Notify clients on the new service group that the upgrade succeeded.
-    await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
-      type: "complete",
-      installedVersion: versionTag,
-      success: true,
-    });
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildCompleteEvent(versionTag, true),
+    );
+
+    // Record successful upgrade in workspace git history
+    if (workspaceDir) {
+      try {
+        await commitWorkspaceState(
+          workspaceDir,
+          buildUpgradeCommitMessage({
+            action: "upgrade",
+            phase: "complete",
+            from: entry.serviceGroupVersion ?? "unknown",
+            to: versionTag,
+            topology: "docker",
+            assistantId: entry.assistantId,
+            result: "success",
+          }),
+        );
+      } catch (err) {
+        console.warn(
+          `⚠️  Failed to create post-upgrade workspace commit: ${err instanceof Error ? err.message : String(err)}`,
+        );
+      }
+    }
 
     console.log(
       `\n✅ Docker assistant '${instanceName}' upgraded to ${versionTag}.`,
@@ -409,12 +586,41 @@ async function upgradeDocker(
     console.error(`\n❌ Containers failed to become ready within the timeout.`);
 
     if (previousImageRefs) {
+      await broadcastUpgradeEvent(
+        entry.runtimeUrl,
+        entry.assistantId,
+        buildProgressEvent(UPGRADE_PROGRESS.REVERTING),
+      );
       console.log(`\n🔄 Rolling back to previous images...`);
       try {
+        // Attempt to roll back migrations before swapping containers.
+        // The new daemon may be partially up — try best-effort.
+        if (
+          preMigrationState.dbVersion !== undefined ||
+          preMigrationState.lastWorkspaceMigrationId !== undefined
+        ) {
+          console.log("🔄 Reverting database changes...");
+          await broadcastUpgradeEvent(
+            entry.runtimeUrl,
+            entry.assistantId,
+            buildProgressEvent(UPGRADE_PROGRESS.REVERTING_MIGRATIONS),
+          );
+          await rollbackMigrations(
+            entry.runtimeUrl,
+            entry.assistantId,
+            preMigrationState.dbVersion,
+            preMigrationState.lastWorkspaceMigrationId,
+          );
+        }
+
         await stopContainers(res);
 
+        await migrateGatewaySecurityFiles(res, (msg) => console.log(msg));
+        await migrateCesSecurityFiles(res, (msg) => console.log(msg));
+
         await startContainers(
           {
+            signingKey,
             bootstrapSecret,
             cesServiceToken,
             extraAssistantEnv,
@@ -428,46 +634,90 @@ async function upgradeDocker(
 
         const rollbackReady = await waitForReady(entry.runtimeUrl);
         if (rollbackReady) {
-          // Restore previous container info in lockfile after rollback.
-          // previousImageRefs contains sha256 digests from `docker inspect
-          // --format {{.Image}}`.  The *Image fields should hold
-          // human-readable image:tag names, so prefer the pre-upgrade
-          // containerInfo values and store digests in the *Digest fields.
-          if (previousImageRefs) {
-            const rolledBackEntry: AssistantEntry = {
-              ...entry,
-              containerInfo: {
-                assistantImage:
-                  entry.containerInfo?.assistantImage ??
-                  previousImageRefs.assistant,
-                gatewayImage:
-                  entry.containerInfo?.gatewayImage ??
-                  previousImageRefs.gateway,
-                cesImage:
-                  entry.containerInfo?.cesImage ??
-                  previousImageRefs["credential-executor"],
-                assistantDigest: previousImageRefs.assistant,
-                gatewayDigest: previousImageRefs.gateway,
-                cesDigest: previousImageRefs["credential-executor"],
-                networkName: res.network,
-              },
-              previousServiceGroupVersion: undefined,
-              previousContainerInfo: undefined,
-            };
-            saveAssistantEntry(rolledBackEntry);
+          // Restore data from the backup created for THIS upgrade attempt.
+          // Only use the specific backupPath — never scan for the latest
+          // backup on disk, which could be from a previous upgrade cycle
+          // and contain stale data.
+          if (backupPath) {
+            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 attempt, skipping data restoration\n",
+            );
           }
 
+          // Capture fresh digests from the now-running rolled-back containers.
+          const rollbackDigests = await captureImageRefs(res);
+
+          // Restore previous container info in lockfile after rollback.
+          // The *Image fields hold human-readable image:tag names from the
+          // pre-upgrade containerInfo; *Digest fields get fresh values from
+          // the running containers (or fall back to previousImageRefs).
+          const rolledBackEntry: AssistantEntry = {
+            ...entry,
+            containerInfo: {
+              assistantImage:
+                entry.containerInfo?.assistantImage ??
+                previousImageRefs.assistant,
+              gatewayImage:
+                entry.containerInfo?.gatewayImage ?? previousImageRefs.gateway,
+              cesImage:
+                entry.containerInfo?.cesImage ??
+                previousImageRefs["credential-executor"],
+              assistantDigest:
+                rollbackDigests?.assistant ?? previousImageRefs.assistant,
+              gatewayDigest:
+                rollbackDigests?.gateway ?? previousImageRefs.gateway,
+              cesDigest:
+                rollbackDigests?.["credential-executor"] ??
+                previousImageRefs["credential-executor"],
+              networkName: res.network,
+            },
+            previousServiceGroupVersion: undefined,
+            previousContainerInfo: undefined,
+            previousDbMigrationVersion: undefined,
+            previousWorkspaceMigrationId: undefined,
+            // Clear the backup path — the upgrade that created it just failed
+            preUpgradeBackupPath: undefined,
+          };
+          saveAssistantEntry(rolledBackEntry);
+
           // Notify clients that the upgrade failed and rolled back.
-          await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
-            type: "complete",
-            installedVersion: entry.serviceGroupVersion ?? "unknown",
-            success: false,
-            rolledBackToVersion: entry.serviceGroupVersion,
-          });
+          await broadcastUpgradeEvent(
+            entry.runtimeUrl,
+            entry.assistantId,
+            buildCompleteEvent(
+              entry.serviceGroupVersion ?? "unknown",
+              false,
+              entry.serviceGroupVersion,
+            ),
+          );
 
           console.log(
             `\n⚠️  Rolled back to previous version. Upgrade to ${versionTag} failed.`,
           );
+          emitCliError(
+            "READINESS_TIMEOUT",
+            `Upgrade to ${versionTag} failed: containers did not become ready. Rolled back to previous version.`,
+          );
         } else {
           console.error(
             `\n❌ Rollback also failed. Manual intervention required.`,
@@ -475,21 +725,51 @@ async function upgradeDocker(
           console.log(
             `   Check logs with: docker logs -f ${res.assistantContainer}`,
           );
+          await broadcastUpgradeEvent(
+            entry.runtimeUrl,
+            entry.assistantId,
+            buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
+          );
+          emitCliError(
+            "ROLLBACK_FAILED",
+            "Rollback also failed after readiness timeout. Manual intervention required.",
+          );
         }
       } catch (rollbackErr) {
-        console.error(
-          `\n❌ Rollback failed: ${rollbackErr instanceof Error ? rollbackErr.message : String(rollbackErr)}`,
-        );
+        const rollbackDetail =
+          rollbackErr instanceof Error
+            ? rollbackErr.message
+            : String(rollbackErr);
+        console.error(`\n❌ Rollback failed: ${rollbackDetail}`);
         console.error(`   Manual intervention required.`);
         console.log(
           `   Check logs with: docker logs -f ${res.assistantContainer}`,
         );
+        await broadcastUpgradeEvent(
+          entry.runtimeUrl,
+          entry.assistantId,
+          buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
+        );
+        emitCliError(
+          "ROLLBACK_FAILED",
+          "Auto-rollback failed after readiness timeout. Manual intervention required.",
+          rollbackDetail,
+        );
       }
     } else {
       console.log(`   No previous images available for rollback.`);
       console.log(
         `   Check logs with: docker logs -f ${res.assistantContainer}`,
       );
+      await broadcastUpgradeEvent(
+        entry.runtimeUrl,
+        entry.assistantId,
+        buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
+      );
+      emitCliError(
+        "ROLLBACK_NO_STATE",
+        "Containers failed to become ready and no previous images available for rollback.",
+      );
     }
 
     process.exit(1);
@@ -505,15 +785,41 @@ async function upgradePlatform(
   entry: AssistantEntry,
   version: string | null,
 ): Promise<void> {
+  const workspaceDir = entry.resources
+    ? join(entry.resources.instanceDir, ".vellum", "workspace")
+    : null;
+
+  // Record version transition start in workspace git history
+  if (workspaceDir) {
+    try {
+      await commitWorkspaceState(
+        workspaceDir,
+        buildUpgradeCommitMessage({
+          action: "upgrade",
+          phase: "starting",
+          from: entry.serviceGroupVersion ?? "unknown",
+          to: version ?? "latest",
+          topology: "managed",
+          assistantId: entry.assistantId,
+        }),
+      );
+    } catch (err) {
+      console.warn(
+        `⚠️  Failed to create pre-upgrade workspace commit: ${err instanceof Error ? err.message : String(err)}`,
+      );
+    }
+  }
+
   console.log(
     `🔄 Upgrading platform-hosted assistant '${entry.assistantId}'...\n`,
   );
 
   const token = readPlatformToken();
   if (!token) {
-    console.error(
-      "Error: Not logged in. Run `vellum login --token <token>` first.",
-    );
+    const msg =
+      "Error: Not logged in. Run `vellum login --token <token>` first.";
+    console.error(msg);
+    emitCliError("AUTH_FAILED", msg);
     process.exit(1);
   }
 
@@ -530,11 +836,11 @@ async function upgradePlatform(
   // Notify connected clients that an upgrade is about to begin.
   const targetVersion = version ?? `v${cliPkg.version}`;
   console.log("📢 Notifying connected clients...");
-  await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
-    type: "starting",
-    targetVersion,
-    expectedDowntimeSeconds: 90,
-  });
+  await broadcastUpgradeEvent(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildStartingEvent(targetVersion, 90),
+  );
 
   const response = await fetch(url, {
     method: "POST",
@@ -551,11 +857,16 @@ async function upgradePlatform(
     console.error(
       `Error: Platform upgrade failed (${response.status}): ${text}`,
     );
-    await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
-      type: "complete",
-      installedVersion: entry.serviceGroupVersion ?? "unknown",
-      success: false,
-    });
+    emitCliError(
+      "PLATFORM_API_ERROR",
+      `Platform upgrade failed (${response.status})`,
+      text,
+    );
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
+    );
     process.exit(1);
   }
 
@@ -568,6 +879,28 @@ async function upgradePlatform(
   // version-change detection (DaemonConnection.swift) once the new
   // version actually appears after the platform restarts the service group.
 
+  // Record successful upgrade in workspace git history
+  if (workspaceDir) {
+    try {
+      await commitWorkspaceState(
+        workspaceDir,
+        buildUpgradeCommitMessage({
+          action: "upgrade",
+          phase: "complete",
+          from: entry.serviceGroupVersion ?? "unknown",
+          to: version ?? "latest",
+          topology: "managed",
+          assistantId: entry.assistantId,
+          result: "success",
+        }),
+      );
+    } catch (err) {
+      console.warn(
+        `⚠️  Failed to create post-upgrade workspace commit: ${err instanceof Error ? err.message : String(err)}`,
+      );
+    }
+  }
+
   console.log(`✅ ${result.detail}`);
   if (result.version) {
     console.log(`   Version: ${result.version}`);
@@ -579,18 +912,33 @@ export async function upgrade(): Promise<void> {
   const entry = resolveTargetAssistant(name);
   const cloud = resolveCloud(entry);
 
-  if (cloud === "docker") {
-    await upgradeDocker(entry, version);
-    return;
-  }
+  try {
+    if (cloud === "docker") {
+      await upgradeDocker(entry, version);
+      return;
+    }
 
-  if (cloud === "vellum") {
-    await upgradePlatform(entry, version);
-    return;
+    if (cloud === "vellum") {
+      await upgradePlatform(entry, version);
+      return;
+    }
+  } catch (err) {
+    const detail = err instanceof Error ? err.message : String(err);
+    console.error(`\n❌ Upgrade failed: ${detail}`);
+    // Best-effort: notify connected clients that the upgrade failed.
+    // A `starting` event may have been sent inside upgradeDocker/upgradePlatform
+    // before the error was thrown, so we must close with `complete`.
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
+    );
+    emitCliError(categorizeUpgradeError(err), "Upgrade failed", detail);
+    process.exit(1);
   }
 
-  console.error(
-    `Error: Upgrade is not supported for '${cloud}' assistants. Only 'docker' and 'vellum' assistants can be upgraded via the CLI.`,
-  );
+  const msg = `Error: Upgrade is not supported for '${cloud}' assistants. Only 'docker' and 'vellum' assistants can be upgraded via the CLI.`;
+  console.error(msg);
+  emitCliError("UNSUPPORTED_TOPOLOGY", msg);
   process.exit(1);
 }