@vellumai/cli 0.10.6 → 0.10.7-dev.202607072325.09ad030

package/src/commands/wake.ts

@@ -12,7 +12,15 @@ import {
   resetGuardianBootstrap,
   seedGuardianTokenFromSiblingEnv,
 } from "../lib/guardian-token.js";
-import { resolveProcessState, stopProcessByPidFile, isProcessAlive } from "../lib/process";
+import {
+  probeDaemonReadinessWithRetry,
+  waitForDaemonMigrationsReady,
+} from "../lib/http-client.js";
+import {
+  isProcessAlive,
+  resolveProcessState,
+  stopProcessByPidFile,
+} from "../lib/process";
 import {
   generateLocalSigningKey,
   isAssistantWatchModeAvailable,
@@ -109,13 +117,23 @@ export async function wake(): Promise<void> {
 
   const pidFile = getDaemonPidPath(resources);
 
+  // Budget anchor for the migration-coordination wait below: the host
+  // wrapper (packages/local-mode/src/wake.ts WAKE_TIMEOUT_MS) SIGTERMs wake
+  // after 180s, and the hung-daemon health wait + post-spawn wait + gateway
+  // start can already consume most of that.
+  const wakeStartedAt = Date.now();
+
   let daemonRunning = false;
+  let daemonUnready = false;
+  let daemonMigrationsFailed = false;
   const daemonState = await resolveProcessState(
     pidFile,
     resources.daemonPort,
     "Assistant",
+    60_000,
+    "readyz",
   );
-  if (daemonState.status === "healthy") {
+  if (daemonState.status !== "needs_start") {
     if (watch && isAssistantWatchModeAvailable()) {
       console.log(
         `Assistant running (pid ${daemonState.pid}) — restarting in watch mode...`,
@@ -123,10 +141,20 @@ export async function wake(): Promise<void> {
       await stopProcessByPidFile(pidFile, "assistant");
     } else {
       daemonRunning = true;
+      daemonUnready = daemonState.status !== "healthy";
+      daemonMigrationsFailed = daemonState.status === "migration_failed";
       if (watch) {
         console.log(
           `Assistant running (pid ${daemonState.pid}) — watch mode not available (no source files). Keeping existing process.`,
         );
+      } else if (daemonMigrationsFailed) {
+        console.log(
+          `Assistant running (pid ${daemonState.pid}) but its database migrations failed.`,
+        );
+      } else if (daemonUnready) {
+        console.log(
+          `Assistant running (pid ${daemonState.pid}) — database migrations still running.`,
+        );
       } else {
         console.log(`Assistant already running (pid ${daemonState.pid}).`);
       }
@@ -189,6 +217,13 @@ export async function wake(): Promise<void> {
       startCes(watch, resources),
       startLocalDaemon(watch, resources, { foreground, signingKey }),
     ]);
+    // startLocalDaemon's post-spawn wait is bounded (60s) — a longer
+    // migration outlives it. Classify the fresh spawn the same way the
+    // attach path does, so the gateway-coordination wait below applies to
+    // both paths and wake's closing summary stays honest.
+    const readiness = await probeDaemonReadinessWithRetry(resources.daemonPort);
+    daemonUnready = readiness === "migrating";
+    daemonMigrationsFailed = readiness === "failed";
   } else {
     // Self-heal: the daemon is already healthy, but the CES sibling may have
     // died independently (crash, OOM kill). A dead ces.pid under a live daemon
@@ -205,6 +240,7 @@ export async function wake(): Promise<void> {
   }
 
   // Start gateway
+  let gatewayStarted = false;
   {
     const vellumDir = join(resources.instanceDir, ".vellum");
     const gatewayPidFile = join(vellumDir, "gateway.pid");
@@ -231,6 +267,7 @@ export async function wake(): Promise<void> {
         signingKey,
         bootstrapSecret,
       });
+      gatewayStarted = true;
     } else if (gatewayAlive) {
       if (watch && isGatewayWatchModeAvailable()) {
         console.log(
@@ -238,6 +275,7 @@ export async function wake(): Promise<void> {
         );
         await stopProcessByPidFile(gatewayPidFile, "gateway");
         await startGateway(watch, resources, { signingKey, bootstrapSecret });
+        gatewayStarted = true;
       } else {
         if (watch) {
           console.log(
@@ -249,6 +287,32 @@ export async function wake(): Promise<void> {
       }
     } else {
       await startGateway(watch, resources, { signingKey, bootstrapSecret });
+      gatewayStarted = true;
+    }
+  }
+
+  // A freshly-(re)started gateway refuses all non-probe traffic until the
+  // daemon reports migration readiness, so consumers that act right after
+  // wake (the web connect-repair retry, the guardian re-provision below)
+  // would hit its closed gate. When the daemon is migrating AND the gateway
+  // was just started (or a guardian repair was requested), wait out the
+  // migration — capped at 60s and at whatever budget remains before the
+  // 180s host-wrapper SIGTERM (30s headroom kept), since earlier bounded
+  // waits may already have consumed most of it. Fast path (gateway already
+  // serving, no repair) stays ~1s.
+  if (
+    daemonUnready &&
+    !daemonMigrationsFailed &&
+    (gatewayStarted || repairGuardian)
+  ) {
+    const waitBudgetMs = Math.min(60_000, wakeStartedAt + 150_000 - Date.now());
+    if (waitBudgetMs > 0) {
+      const readiness = await waitForDaemonMigrationsReady(
+        resources.daemonPort,
+        Date.now() + waitBudgetMs,
+      );
+      daemonUnready = readiness !== "ready";
+      daemonMigrationsFailed = readiness === "failed";
     }
   }
 
@@ -295,6 +359,11 @@ export async function wake(): Promise<void> {
           console.warn(
             `   Guardian token re-provision failed after ${maxAttempts} attempts: ${err}`,
           );
+          // The user explicitly confirmed this destructive repair — a
+          // success exit here would make callers (the web recovery flow)
+          // treat a repair that never ran as done and drop their cached
+          // gateway token. Surface the failure through the exit code.
+          process.exitCode = 1;
         }
       }
     }
@@ -311,6 +380,15 @@ export async function wake(): Promise<void> {
     writeFileSync(ngrokPidFile, String(ngrokChild.pid));
   }
 
+  if (daemonMigrationsFailed) {
+    console.log(
+      "Assistant database migrations FAILED — DB-backed routes will return 503 until the assistant is restarted. Check the daemon logs.",
+    );
+  } else if (daemonUnready) {
+    console.log(
+      "Assistant is still running database migrations; DB-backed routes return 503 until they finish.",
+    );
+  }
   console.log("Wake complete.");
 
   if (foreground) {

package/src/lib/__tests__/local-ces.test.ts

@@ -0,0 +1,104 @@
+import { mkdtempSync, rmSync, writeFileSync, existsSync, mkdirSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test";
+
+import { startCes } from "../local.js";
+
+// ---------------------------------------------------------------------------
+// Mocks
+// ---------------------------------------------------------------------------
+
+// Capture spawn calls so we can assert on cmd/env.
+let lastSpawnCall: {
+  cmd: string[];
+  options: { detached?: boolean; env?: Record<string, string | undefined>; cwd?: string };
+} | null = null;
+
+mock.module("node:child_process", () => ({
+  spawn: mock((cmd: string, args: string[], options: object) => {
+    lastSpawnCall = { cmd: [cmd, ...args], options };
+    // Return a fake subprocess with a pid and no-op methods.
+    return {
+      pid: 42,
+      unref: () => {},
+      stdout: { on: () => {} },
+      stderr: { on: () => {} },
+      on: () => {},
+    };
+  }),
+  execSync: () => "",
+  execFileSync: () => "",
+  spawnSync: () => ({ status: 0, stdout: "", stderr: "" }),
+}));
+
+// Mock xdg-log so we don't open real log files.
+mock.module("../xdg-log.js", () => ({
+  openLogFile: mock(() => 42),
+  pipeToLogFile: mock(() => {}),
+}));
+
+// Mock process helpers so stopProcessByPidFile is a no-op.
+mock.module("../process.js", () => ({
+  stopProcessByPidFile: mock(async () => {}),
+  isProcessAlive: mock(() => false),
+  stopProcessGracefully: mock(async () => {}),
+}));
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe("startCes", () => {
+  let tempDir: string;
+
+  beforeAll(() => {
+    tempDir = mkdtempSync(join(tmpdir(), "ces-test-"));
+  });
+
+  afterAll(() => {
+    rmSync(tempDir, { recursive: true, force: true });
+  });
+
+  test("spawns CES with correct env vars and writes PID file", async () => {
+    // Create the socket path before calling startCes so the wait loop exits.
+    // startCes unlinks it first, then waits for it to reappear. We create it
+    // after a short delay to simulate CES binding the socket.
+    const resources = {
+      instanceDir: tempDir,
+      name: "test-assistant",
+    } as unknown as Parameters<typeof startCes>[1];
+
+    // Create the socket file asynchronously after startCes unlinks it.
+    // We use a small setTimeout to create it during the wait loop.
+    const vellumDir = join(tempDir, ".vellum");
+    mkdirSync(vellumDir, { recursive: true });
+
+    // Pre-create the socket so it exists when startCes checks after unlinking.
+    // startCes unlinks the stale socket, then polls for it. We create it with
+    // a slight delay so the poll catches it.
+    setTimeout(() => {
+      const socketDir = join(vellumDir, "workspace");
+      mkdirSync(socketDir, { recursive: true });
+      writeFileSync(join(socketDir, "ces.sock"), "");
+    }, 50);
+
+    lastSpawnCall = null;
+    await startCes(false, resources);
+
+    // Verify spawn was called
+    expect(lastSpawnCall).not.toBeNull();
+    expect(lastSpawnCall!.options.detached).toBe(true);
+
+    // Verify env vars
+    const env = lastSpawnCall!.options.env!;
+    expect(env["CES_STANDALONE"]).toBe("1");
+    expect(env["CES_LOCAL_SOCKET"]).toBeDefined();
+    expect(env["CREDENTIAL_SECURITY_DIR"]).toBeDefined();
+    expect(env["VELLUM_WORKSPACE_DIR"]).toBeDefined();
+
+    // Verify PID file was written
+    const cesPidFile = join(vellumDir, "ces.pid");
+    expect(existsSync(cesPidFile)).toBe(true);
+  }, 15_000);
+});

package/src/lib/backup-ops.ts

@@ -9,10 +9,7 @@ import {
 import { homedir } from "os";
 import { dirname, join } from "path";
 
-import {
-  loadGuardianToken,
-  refreshGuardianToken,
-} from "./guardian-token.js";
+import { loadGuardianToken, refreshGuardianToken } from "./guardian-token.js";
 import { loopbackSafeFetch } from "./loopback-fetch.js";
 
 /** Default backup directory following XDG convention */
@@ -63,28 +60,35 @@ export async function createBackup(
   try {
     let accessToken = await getGuardianAccessToken(runtimeUrl, assistantId);
     if (!accessToken) {
-      console.warn("Warning: backup skipped — no valid guardian token available");
+      console.warn(
+        "Warning: backup skipped — no valid guardian token available",
+      );
       return null;
     }
 
-    let response = await loopbackSafeFetch(`${runtimeUrl}/v1/migrations/export`, {
-      method: "POST",
-      headers: {
-        Authorization: `Bearer ${accessToken}`,
-        "Content-Type": "application/json",
+    let response = await loopbackSafeFetch(
+      `${runtimeUrl}/v1/migrations/export`,
+      {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${accessToken}`,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          description: options?.description ?? "CLI backup",
+        }),
+        signal: AbortSignal.timeout(120_000),
       },
-      body: JSON.stringify({
-        description: options?.description ?? "CLI backup",
-      }),
-      signal: AbortSignal.timeout(120_000),
-    });
+    );
 
     // Retry once with a refreshed token on 401 — the cached token may be
     // stale after a container restart that regenerated the gateway signing key.
     if (response.status === 401) {
       const refreshed = await refreshGuardianToken(runtimeUrl, assistantId);
       if (!refreshed) {
-        console.warn(`Warning: backup export failed (401) and token refresh failed`);
+        console.warn(
+          `Warning: backup export failed (401) and token refresh failed`,
+        );
         return null;
       }
       accessToken = refreshed.accessToken;
@@ -149,19 +153,24 @@ export async function restoreBackup(
     const bundleData = readFileSync(backupPath);
     let accessToken = await getGuardianAccessToken(runtimeUrl, assistantId);
     if (!accessToken) {
-      console.warn("Warning: restore skipped — no valid guardian token available");
+      console.warn(
+        "Warning: restore skipped — no valid guardian token available",
+      );
       return false;
     }
 
-    let response = await loopbackSafeFetch(`${runtimeUrl}/v1/migrations/import`, {
-      method: "POST",
-      headers: {
-        Authorization: `Bearer ${accessToken}`,
-        "Content-Type": "application/octet-stream",
-      },
-      body: bundleData,
-      signal: AbortSignal.timeout(120_000),
-    });
+    const postImport = () =>
+      loopbackSafeFetch(`${runtimeUrl}/v1/migrations/import`, {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${accessToken}`,
+          "Content-Type": "application/octet-stream",
+        },
+        body: bundleData,
+        signal: AbortSignal.timeout(120_000),
+      });
+
+    let response = await postImport();
 
     // Retry once with a refreshed token on 401 — the cached token may be
     // stale after a container restart that regenerated the gateway signing key.
@@ -172,19 +181,33 @@ export async function restoreBackup(
         return false;
       }
       accessToken = refreshed.accessToken;
-      response = await loopbackSafeFetch(`${runtimeUrl}/v1/migrations/import`, {
-        method: "POST",
-        headers: {
-          Authorization: `Bearer ${accessToken}`,
-          "Content-Type": "application/octet-stream",
-        },
-        body: bundleData,
-        signal: AbortSignal.timeout(120_000),
-      });
+      response = await postImport();
     }
 
-    if (!response.ok) {
+    // A freshly-(re)started gateway 503s {status:"starting"} for a couple of
+    // seconds until its own assistant poll observes the migration state and
+    // opens the traffic gate — retry through that window rather than failing
+    // the recovery. The daemon's running-migrations 503 carries a `reason`
+    // field and is excluded: the import must not race an in-flight migration.
+    const startingGateDeadline = Date.now() + 15_000;
+    while (!response.ok) {
       const body = await response.text();
+      let gatewayStarting = false;
+      try {
+        const parsed = JSON.parse(body) as {
+          status?: string;
+          reason?: string;
+        };
+        gatewayStarting =
+          parsed.status === "starting" && parsed.reason === undefined;
+      } catch {
+        // Non-JSON error body — not the starting gate.
+      }
+      if (gatewayStarting && Date.now() < startingGateDeadline) {
+        await new Promise((r) => setTimeout(r, 1_000));
+        response = await postImport();
+        continue;
+      }
       console.warn(`Warning: restore failed (${response.status}): ${body}`);
       return false;
     }

package/src/lib/docker.ts

@@ -23,7 +23,10 @@ import { buildHatchConfigValues, writeInitialConfig } from "./config-utils";
 import { buildServiceRunArgs } from "./statefulset.js";
 import type { Species } from "./constants";
 import { getOrCreateHostDeviceId } from "./device-id.js";
-import { ASSISTANT_INTERNAL_PORT, getDefaultPorts } from "./environments/paths.js";
+import {
+  ASSISTANT_INTERNAL_PORT,
+  getDefaultPorts,
+} from "./environments/paths.js";
 import { getCurrentEnvironment } from "./environments/resolve.js";
 import { leaseGuardianToken } from "./guardian-token";
 import { logHatchNextSteps } from "./hatch-next-steps.js";
@@ -68,6 +71,7 @@ export {
   ASSISTANT_INTERNAL_PORT,
   GATEWAY_INTERNAL_PORT,
 } from "./environments/paths.js";
+import { classifyReadyzResponse } from "./http-client.js";
 import { loopbackSafeFetch } from "./loopback-fetch.js";
 
 /** Max time to wait for the assistant container to emit the readiness sentinel. */
@@ -953,7 +957,8 @@ function startFileWatcher(opts: {
   netnsContainer?: string;
   assistantCaCertPath?: string;
 }): () => void {
-  const { gatewayPort, assistantPort, imageTags, instanceName, repoRoot, res } = opts;
+  const { gatewayPort, assistantPort, imageTags, instanceName, repoRoot, res } =
+    opts;
 
   const { dirs: watchDirs, files: watchFiles } = collectWatchTargets(repoRoot);
 
@@ -1487,19 +1492,24 @@ export async function hatchDocker(params: HatchDockerParams): Promise<void> {
 
     emitProgress(6, 6, "Waiting for services...");
     const waitDetached = watch ? false : detached;
-    const { ready, guardianAccessToken } = await waitForGatewayAndLease({
-      bootstrapSecret: ownSecret,
-      containerName: res.assistantContainer,
-      detached: waitDetached,
-      instanceName,
-      logFd,
-      runtimeUrl,
-      containersUpAt,
-      analyze: params.analyze ?? false,
-    });
+    const { ready, guardianAccessToken, notReadyReason } =
+      await waitForGatewayAndLease({
+        bootstrapSecret: ownSecret,
+        containerName: res.assistantContainer,
+        detached: waitDetached,
+        instanceName,
+        logFd,
+        runtimeUrl,
+        containersUpAt,
+        analyze: params.analyze ?? false,
+      });
 
     if (!ready && !(watch && repoRoot)) {
-      throw new Error("Timed out waiting for assistant to become ready");
+      throw new Error(
+        notReadyReason === "migrations_failed"
+          ? "Assistant database migrations failed — the container is running but DB-backed routes are unavailable; check its logs and restore a backup or retry the migration"
+          : "Timed out waiting for assistant to become ready",
+      );
     }
 
     if (ready) {
@@ -1601,7 +1611,13 @@ async function waitForGatewayAndLease(opts: {
   runtimeUrl: string;
   containersUpAt: number;
   analyze: boolean;
-}): Promise<{ ready: boolean; guardianAccessToken?: string }> {
+}): Promise<{
+  ready: boolean;
+  guardianAccessToken?: string;
+  /** Present when readiness was not reached because migrations terminally
+   * failed (returned in seconds, unlike an actual timeout). */
+  notReadyReason?: "migrations_failed";
+}> {
   const {
     bootstrapSecret,
     containerName,
@@ -1675,24 +1691,34 @@ async function waitForGatewayAndLease(opts: {
   const readyUrl = `${runtimeUrl}/readyz`;
   const start = containersUpAt;
   let ready = false;
+  let migrationsFailed = false;
 
+  // Readiness is classified from the /readyz BODY: the assistant returns 200
+  // with `ready: false` while DB migrations run (the k8s keep-the-pod
+  // contract), so `resp.ok` alone would declare readiness mid-migration and
+  // the guardian token lease below would burn its shared budget against the
+  // gateway's still-closed traffic gate.
   while (Date.now() - start < DOCKER_READY_TIMEOUT_MS) {
     try {
       const resp = await loopbackSafeFetch(readyUrl, {
         signal: AbortSignal.timeout(5000),
       });
-      if (resp.ok) {
+      const body = (await resp.json().catch(() => null)) as {
+        status?: string;
+        upstream?: number;
+      } | null;
+      const readiness = classifyReadyzResponse(resp.ok, body);
+      if (readiness === "ready") {
         ready = true;
         break;
       }
-      const body = await resp.text();
-      let detail = "";
-      try {
-        const json = JSON.parse(body);
-        const parts = [json.status];
-        if (json.upstream != null) parts.push(`upstream=${json.upstream}`);
-        detail = ` — ${parts.join(", ")}`;
-      } catch {}
+      if (readiness === "failed") {
+        migrationsFailed = true;
+        break;
+      }
+      const parts = [body?.status].filter(Boolean);
+      if (body?.upstream != null) parts.push(`upstream=${body.upstream}`);
+      const detail = parts.length > 0 ? ` — ${parts.join(", ")}` : "";
       log(`Readiness check: ${resp.status}${detail} (retrying...)`);
     } catch {
       // Connection refused / timeout — not up yet
@@ -1702,11 +1728,20 @@ async function waitForGatewayAndLease(opts: {
 
   if (!ready) {
     log("");
-    log(`   \u26a0\ufe0f  Timed out waiting for assistant to become ready.`);
+    if (migrationsFailed) {
+      log(`   \u26a0\ufe0f  Assistant database migrations FAILED.`);
+    } else {
+      log(`   \u26a0\ufe0f  Timed out waiting for assistant to become ready.`);
+    }
     log(`   The container is still running.`);
     log(`   Check logs with: docker logs -f ${containerName}`);
     log("");
-    return { ready: false };
+    return {
+      ready: false,
+      ...(migrationsFailed
+        ? { notReadyReason: "migrations_failed" as const }
+        : {}),
+    };
   }
 
   const readyAt = Date.now();