@vellumai/cli 0.5.4 → 0.5.6

package/src/commands/upgrade.ts

@@ -1,17 +1,23 @@
+import { randomBytes } from "crypto";
+
 import cliPkg from "../../package.json";
 
 import {
   findAssistantByName,
   getActiveAssistant,
   loadAllAssistants,
+  saveAssistantEntry,
 } from "../lib/assistant-config";
 import type { AssistantEntry } from "../lib/assistant-config";
 import {
   captureImageRefs,
+  clearSigningKeyBootstrapLock,
   DOCKERHUB_IMAGES,
   DOCKER_READY_TIMEOUT_MS,
   GATEWAY_INTERNAL_PORT,
   dockerResourceNames,
+  migrateCesSecurityFiles,
+  migrateGatewaySecurityFiles,
   startContainers,
   stopContainers,
 } from "../lib/docker";
@@ -21,6 +27,7 @@ import {
   getPlatformUrl,
   readPlatformToken,
 } from "../lib/platform-client";
+import { loadBootstrapSecret, loadGuardianToken } from "../lib/guardian-token";
 import { exec, execOutput } from "../lib/step-runner";
 
 interface UpgradeArgs {
@@ -132,7 +139,7 @@ function resolveTargetAssistant(nameArg: string | null): AssistantEntry {
  * Capture environment variables from a running Docker container so they
  * can be replayed onto the replacement container after upgrade.
  */
-async function captureContainerEnv(
+export async function captureContainerEnv(
   containerName: string,
 ): Promise<Record<string, string>> {
   const captured: Record<string, string> = {};
@@ -160,7 +167,7 @@ async function captureContainerEnv(
  * 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> {
+export async function waitForReady(runtimeUrl: string): Promise<boolean> {
   const readyUrl = `${runtimeUrl}/readyz`;
   const start = Date.now();
 
@@ -194,6 +201,35 @@ async function waitForReady(runtimeUrl: string): Promise<boolean> {
   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,
@@ -229,6 +265,18 @@ async function upgradeDocker(
     );
   }
 
+  // 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 },
+    };
+    saveAssistantEntry(rollbackEntry);
+    console.log(`   Saved rollback state: ${entry.serviceGroupVersion}\n`);
+  }
+
   console.log("💾 Capturing existing container environment...");
   const capturedEnv = await captureContainerEnv(res.assistantContainer);
   console.log(
@@ -241,6 +289,16 @@ async function upgradeDocker(
   await exec("docker", ["pull", imageTags["credential-executor"]]);
   console.log("✅ Docker images pulled\n");
 
+  // Notify connected clients that an upgrade is about to begin.
+  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 new Promise((r) => setTimeout(r, 500));
+
   console.log("🛑 Stopping existing containers...");
   await stopContainers(res);
   console.log("✅ Containers stopped\n");
@@ -257,10 +315,23 @@ async function upgradeDocker(
     // use default
   }
 
+  // Extract CES_SERVICE_TOKEN from the captured env so it can be passed via
+  // the dedicated cesServiceToken parameter (which propagates it to all three
+  // containers). If the old instance predates CES_SERVICE_TOKEN, generate a
+  // fresh one so gateway and CES can authenticate.
+  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.
+  const bootstrapSecret =
+    loadBootstrapSecret(instanceName) || randomBytes(32).toString("hex");
+
   // 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",
@@ -278,9 +349,20 @@ async function upgradeDocker(
     }
   }
 
+  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 upgraded containers...");
   await startContainers(
     {
+      bootstrapSecret,
+      cesServiceToken,
       extraAssistantEnv,
       gatewayPort,
       imageTags,
@@ -294,6 +376,32 @@ async function upgradeDocker(
   console.log("Waiting for assistant to become ready...");
   const ready = await waitForReady(entry.runtimeUrl);
   if (ready) {
+    // Update lockfile with new service group topology
+    const newDigests = await captureImageRefs(res);
+    const updatedEntry: AssistantEntry = {
+      ...entry,
+      serviceGroupVersion: versionTag,
+      containerInfo: {
+        assistantImage: imageTags.assistant,
+        gatewayImage: imageTags.gateway,
+        cesImage: imageTags["credential-executor"],
+        assistantDigest: newDigests?.assistant,
+        gatewayDigest: newDigests?.gateway,
+        cesDigest: newDigests?.["credential-executor"],
+        networkName: res.network,
+      },
+      previousServiceGroupVersion: entry.serviceGroupVersion,
+      previousContainerInfo: entry.containerInfo,
+    };
+    saveAssistantEntry(updatedEntry);
+
+    // Notify clients on the new service group that the upgrade succeeded.
+    await broadcastUpgradeEvent(entry.runtimeUrl, entry.assistantId, {
+      type: "complete",
+      installedVersion: versionTag,
+      success: true,
+    });
+
     console.log(
       `\n✅ Docker assistant '${instanceName}' upgraded to ${versionTag}.`,
     );
@@ -307,6 +415,8 @@ async function upgradeDocker(
 
         await startContainers(
           {
+            bootstrapSecret,
+            cesServiceToken,
             extraAssistantEnv,
             gatewayPort,
             imageTags: previousImageRefs,
@@ -318,6 +428,43 @@ 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);
+          }
+
+          // 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,
+          });
+
           console.log(
             `\n⚠️  Rolled back to previous version. Upgrade to ${versionTag} failed.`,
           );
@@ -380,6 +527,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, {
+    type: "starting",
+    targetVersion,
+    expectedDowntimeSeconds: 90,
+  });
+
   const response = await fetch(url, {
     method: "POST",
     headers: {
@@ -395,10 +551,23 @@ 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,
+    });
     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.
+
   console.log(`✅ ${result.detail}`);
   if (result.version) {
     console.log(`   Version: ${result.version}`);

package/src/commands/wake.ts

@@ -40,6 +40,14 @@ export async function wake(): Promise<void> {
   const entry = resolveTargetAssistant(nameArg);
 
   if (entry.cloud === "docker") {
+    if (watch || foreground) {
+      const ignored = [watch && "--watch", foreground && "--foreground"]
+        .filter(Boolean)
+        .join(" and ");
+      console.warn(
+        `Warning: ${ignored} ignored for Docker instances (not supported).`,
+      );
+    }
     const res = dockerResourceNames(entry.assistantId);
     await wakeContainers(res);
     console.log("Docker containers started.");

package/src/index.ts

@@ -1,6 +1,7 @@
 #!/usr/bin/env bun
 
 import cliPkg from "../package.json";
+import { backup } from "./commands/backup";
 import { clean } from "./commands/clean";
 import { client } from "./commands/client";
 import { hatch } from "./commands/hatch";
@@ -8,7 +9,9 @@ import { login, logout, whoami } from "./commands/login";
 import { pair } from "./commands/pair";
 import { ps } from "./commands/ps";
 import { recover } from "./commands/recover";
+import { restore } from "./commands/restore";
 import { retire } from "./commands/retire";
+import { rollback } from "./commands/rollback";
 import { setup } from "./commands/setup";
 import { sleep } from "./commands/sleep";
 import { ssh } from "./commands/ssh";
@@ -26,6 +29,7 @@ import { loadGuardianToken } from "./lib/guardian-token";
 import { checkHealth } from "./lib/health-check";
 
 const commands = {
+  backup,
   clean,
   client,
   hatch,
@@ -34,7 +38,9 @@ const commands = {
   pair,
   ps,
   recover,
+  restore,
   retire,
+  rollback,
   setup,
   sleep,
   ssh,
@@ -51,6 +57,7 @@ function printHelp(): void {
   console.log("Usage: vellum <command> [options]");
   console.log("");
   console.log("Commands:");
+  console.log("  backup   Export a backup of a running assistant");
   console.log("  clean    Kill orphaned vellum processes");
   console.log("  client   Connect to a hatched assistant");
   console.log("  hatch    Create a new assistant instance");
@@ -61,7 +68,11 @@ function printHelp(): void {
     "  ps       List assistants (or processes for a specific assistant)",
   );
   console.log("  recover  Restore a previously retired local assistant");
+  console.log("  restore  Restore a .vbundle backup into a running assistant");
   console.log("  retire   Delete an assistant instance");
+  console.log(
+    "  rollback  Roll back a Docker assistant to the previous version",
+  );
   console.log("  setup    Configure API keys interactively");
   console.log("  sleep    Stop the assistant process");
   console.log("  ssh      SSH into a remote assistant instance");

package/src/lib/assistant-config.ts

@@ -4,6 +4,7 @@ import { join } from "path";
 
 import {
   DAEMON_INTERNAL_ASSISTANT_ID,
+  DEFAULT_CES_PORT,
   DEFAULT_DAEMON_PORT,
   DEFAULT_GATEWAY_PORT,
   DEFAULT_QDRANT_PORT,
@@ -29,11 +30,28 @@ export interface LocalInstanceResources {
   gatewayPort: number;
   /** HTTP port for the Qdrant vector store */
   qdrantPort: number;
+  /** HTTP port for the CES (Claude Extension Server) */
+  cesPort: number;
   /** Absolute path to the daemon PID file */
   pidFile: string;
   [key: string]: unknown;
 }
 
+/** Docker image metadata for the service group. Enables rollback to known-good digests. */
+export interface ContainerInfo {
+  assistantImage: string;
+  gatewayImage: string;
+  cesImage: string;
+  /** sha256 digest of the assistant image at time of hatch/upgrade */
+  assistantDigest?: string;
+  /** sha256 digest of the gateway image at time of hatch/upgrade */
+  gatewayDigest?: string;
+  /** sha256 digest of the CES image at time of hatch/upgrade */
+  cesDigest?: string;
+  /** Docker network name for the service group */
+  networkName?: string;
+}
+
 export interface AssistantEntry {
   assistantId: string;
   runtimeUrl: string;
@@ -56,6 +74,14 @@ export interface AssistantEntry {
   resources?: LocalInstanceResources;
   /** PID of the file watcher process for docker instances hatched with --watch. */
   watcherPid?: number;
+  /** Last-known version of the service group, populated at hatch and updated by health checks. */
+  serviceGroupVersion?: string;
+  /** Docker image metadata for rollback. Only present for docker topology entries. */
+  containerInfo?: ContainerInfo;
+  /** The service group version that was running before the last upgrade. */
+  previousServiceGroupVersion?: string;
+  /** Docker image metadata from before the last upgrade. Enables rollback to the prior version. */
+  previousContainerInfo?: ContainerInfo;
   [key: string]: unknown;
 }
 
@@ -166,6 +192,7 @@ export function migrateLegacyEntry(raw: Record<string, unknown>): boolean {
       daemonPort: DEFAULT_DAEMON_PORT,
       gatewayPort,
       qdrantPort: DEFAULT_QDRANT_PORT,
+      cesPort: DEFAULT_CES_PORT,
       pidFile: join(instanceDir, ".vellum", "vellum.pid"),
     };
     mutated = true;
@@ -198,6 +225,10 @@ export function migrateLegacyEntry(raw: Record<string, unknown>): boolean {
       res.qdrantPort = DEFAULT_QDRANT_PORT;
       mutated = true;
     }
+    if (typeof res.cesPort !== "number") {
+      res.cesPort = DEFAULT_CES_PORT;
+      mutated = true;
+    }
     if (typeof res.pidFile !== "string") {
       res.pidFile = join(res.instanceDir as string, ".vellum", "vellum.pid");
       mutated = true;
@@ -333,6 +364,23 @@ export function saveAssistantEntry(entry: AssistantEntry): void {
   writeAssistants(entries);
 }
 
+/**
+ * Update just the serviceGroupVersion field on a lockfile entry.
+ * Reads the current entry, updates the version if changed, and writes back.
+ * No-op if the entry doesn't exist or the version hasn't changed.
+ */
+export function updateServiceGroupVersion(
+  assistantId: string,
+  version: string,
+): void {
+  const entries = readAssistants();
+  const entry = entries.find((e) => e.assistantId === assistantId);
+  if (!entry) return;
+  if (entry.serviceGroupVersion === version) return;
+  entry.serviceGroupVersion = version;
+  writeAssistants(entries);
+}
+
 /**
  * Scan upward from `basePort` to find an available port. A port is considered
  * available when `probePort()` returns false (nothing listening). Scans up to
@@ -373,6 +421,7 @@ export async function allocateLocalResources(
       daemonPort: DEFAULT_DAEMON_PORT,
       gatewayPort: DEFAULT_GATEWAY_PORT,
       qdrantPort: DEFAULT_QDRANT_PORT,
+      cesPort: DEFAULT_CES_PORT,
       pidFile: join(vellumDir, "vellum.pid"),
     };
   }
@@ -398,6 +447,7 @@ export async function allocateLocalResources(
         entry.resources.daemonPort,
         entry.resources.gatewayPort,
         entry.resources.qdrantPort,
+        entry.resources.cesPort,
       );
     }
   }
@@ -417,12 +467,19 @@ export async function allocateLocalResources(
     daemonPort,
     gatewayPort,
   ]);
+  const cesPort = await findAvailablePort(DEFAULT_CES_PORT, [
+    ...reservedPorts,
+    daemonPort,
+    gatewayPort,
+    qdrantPort,
+  ]);
 
   return {
     instanceDir,
     daemonPort,
     gatewayPort,
     qdrantPort,
+    cesPort,
     pidFile: join(instanceDir, ".vellum", "vellum.pid"),
   };
 }

package/src/lib/aws.ts

@@ -6,7 +6,7 @@ import { buildStartupScript, watchHatching } from "../commands/hatch";
 import type { PollResult } from "../commands/hatch";
 import { saveAssistantEntry, setActiveAssistant } from "./assistant-config";
 import type { AssistantEntry } from "./assistant-config";
-import { GATEWAY_PORT } from "./constants";
+import { GATEWAY_PORT, PROVIDER_ENV_VAR_NAMES } from "./constants";
 import type { Species } from "./constants";
 import { leaseGuardianToken } from "./guardian-token";
 import { generateInstanceName } from "./random-name";
@@ -410,10 +410,18 @@ export async function hatchAws(
 
     const sshUser = userInfo().username;
     const hatchedBy = process.env.VELLUM_HATCHED_BY;
-    const anthropicApiKey = process.env.ANTHROPIC_API_KEY;
-    if (!anthropicApiKey) {
+    const providerApiKeys: Record<string, string> = {};
+    for (const [, envVar] of Object.entries(PROVIDER_ENV_VAR_NAMES)) {
+      const value = process.env[envVar];
+      if (value) {
+        providerApiKeys[envVar] = value;
+      }
+    }
+    if (Object.keys(providerApiKeys).length === 0) {
       console.error(
-        "Error: ANTHROPIC_API_KEY environment variable is not set.",
+        "Error: No provider API key environment variable is set. " +
+          "Set at least one of: " +
+          Object.values(PROVIDER_ENV_VAR_NAMES).join(", "),
       );
       process.exit(1);
     }
@@ -437,7 +445,7 @@ export async function hatchAws(
     const startupScript = await buildStartupScript(
       species,
       sshUser,
-      anthropicApiKey,
+      providerApiKeys,
       instanceName,
       "aws",
     );

package/src/lib/constants.ts

@@ -1,3 +1,5 @@
+import providerEnvVarsRegistry from "../../../meta/provider-env-vars.json";
+
 /**
  * Canonical internal assistant ID used as the default/fallback across the CLI
  * and daemon. Mirrors `DAEMON_INTERNAL_ASSISTANT_ID` from
@@ -14,6 +16,16 @@ export const GATEWAY_PORT = process.env.GATEWAY_PORT
 export const DEFAULT_DAEMON_PORT = 7821;
 export const DEFAULT_GATEWAY_PORT = 7830;
 export const DEFAULT_QDRANT_PORT = 6333;
+export const DEFAULT_CES_PORT = 8090;
+
+/**
+ * Environment variable names for provider API keys, keyed by provider ID.
+ * Loaded from the shared registry at `meta/provider-env-vars.json` — the
+ * single source of truth also consumed by the assistant runtime and the
+ * macOS client.
+ */
+export const PROVIDER_ENV_VAR_NAMES: Record<string, string> =
+  providerEnvVarsRegistry.providers;
 
 export const VALID_REMOTE_HOSTS = [
   "local",