@vellumai/cli 0.5.5 → 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,8 +12,6 @@ import {
 import type { AssistantEntry } from "../lib/assistant-config";
 import {
   captureImageRefs,
-  DOCKERHUB_IMAGES,
-  DOCKER_READY_TIMEOUT_MS,
   GATEWAY_INTERNAL_PORT,
   dockerResourceNames,
   migrateCesSecurityFiles,
@@ -20,13 +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 { 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;
@@ -70,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;
@@ -78,6 +102,7 @@ function parseArgs(): UpgradeArgs {
       name = arg;
     } else {
       console.error(`Error: Unknown option '${arg}'.`);
+      emitCliError("UNKNOWN", `Unknown option '${arg}'`);
       process.exit(1);
     }
   }
@@ -109,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;
@@ -124,95 +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.
- */
-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.
- */
-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;
-}
-
 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`,
@@ -234,22 +200,168 @@ 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) {
+    const rollbackEntry: AssistantEntry = {
+      ...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`,
   );
 
+  // 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,
+    buildStartingEvent(versionTag),
+  );
+  // Brief pause to allow SSE delivery before progress events.
+  await new Promise((r) => setTimeout(r, 500));
+
+  await broadcastUpgradeEvent(
+    entry.runtimeUrl,
+    entry.assistantId,
+    buildProgressEvent(UPGRADE_PROGRESS.DOWNLOADING),
+  );
   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"]]);
+  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");
 
-  console.log("🛑 Stopping existing containers...");
-  await stopContainers(res);
-  console.log("✅ Containers stopped\n");
-
   // Parse gateway port from entry's runtimeUrl, fall back to default
   let gatewayPort = GATEWAY_INTERNAL_PORT;
   try {
@@ -269,15 +381,93 @@ async function upgradeDocker(
   const cesServiceToken =
     capturedEnv["CES_SERVICE_TOKEN"] || randomBytes(32).toString("hex");
 
+  // Retrieve or generate a bootstrap secret for the gateway. The secret was
+  // persisted to disk during hatch; older instances won't have one yet.
+  // 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]) {
@@ -300,6 +490,8 @@ async function upgradeDocker(
   console.log("🚀 Starting upgraded containers...");
   await startContainers(
     {
+      signingKey,
+      bootstrapSecret,
       cesServiceToken,
       extraAssistantEnv,
       gatewayPort,
@@ -328,9 +520,65 @@ async function upgradeDocker(
         cesDigest: newDigests?.["credential-executor"],
         networkName: res.network,
       },
+      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,
+      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}.`,
     );
@@ -338,12 +586,42 @@ 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,
             gatewayPort,
@@ -356,22 +634,90 @@ async function upgradeDocker(
 
         const rollbackReady = await waitForReady(entry.runtimeUrl);
         if (rollbackReady) {
-          // Restore previous container info in lockfile after rollback
-          if (previousImageRefs) {
-            const rolledBackEntry: AssistantEntry = {
-              ...entry,
-              containerInfo: {
-                assistantImage: previousImageRefs.assistant,
-                gatewayImage: previousImageRefs.gateway,
-                cesImage: previousImageRefs["credential-executor"],
-                networkName: res.network,
-              },
-            };
-            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,
+            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.`,
@@ -379,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);
@@ -409,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);
   }
 
@@ -431,6 +833,15 @@ async function upgradePlatform(
     body.version = version;
   }
 
+  // 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,
+    buildStartingEvent(targetVersion, 90),
+  );
+
   const response = await fetch(url, {
     method: "POST",
     headers: {
@@ -446,10 +857,50 @@ async function upgradePlatform(
     console.error(
       `Error: Platform upgrade failed (${response.status}): ${text}`,
     );
+    emitCliError(
+      "PLATFORM_API_ERROR",
+      `Platform upgrade failed (${response.status})`,
+      text,
+    );
+    await broadcastUpgradeEvent(
+      entry.runtimeUrl,
+      entry.assistantId,
+      buildCompleteEvent(entry.serviceGroupVersion ?? "unknown", false),
+    );
     process.exit(1);
   }
 
   const result = (await response.json()) as UpgradeApiResponse;
+
+  // NOTE: We intentionally do NOT broadcast a "complete" event here.
+  // The platform API returning 200 only means "upgrade request accepted" —
+  // the service group has not yet restarted with the new version.  The
+  // completion signal will come from the client's health-check
+  // 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}`);
@@ -461,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);
 }