@vellumai/cli 0.10.12 → 0.11.0-dev.202607290043.3b26b6b

package/src/__tests__/sleep.test.ts

@@ -30,10 +30,53 @@ mock.module("../lib/process.js", () => ({
   stopProcessByPidFile: stopProcessByPidFileMock,
 }));
 
-// Restore the real module once this file finishes so the mock does not leak
+// Mock the drain phase and the guardian-token loader so `--wait` tests can
+// steer outcomes without a live daemon. Both spread the real module so
+// untouched exports keep working (parseWaitDuration stays real).
+import type { DrainOptions, DrainOutcome } from "../lib/drain.js";
+
+const drainAssistantMock = mock(
+  async (_opts: DrainOptions): Promise<DrainOutcome> => "drained",
+);
+const realDrain = { ...(await import("../lib/drain.js")) };
+mock.module("../lib/drain.js", () => ({
+  ...realDrain,
+  drainAssistant: drainAssistantMock,
+}));
+
+const realGuardianToken = { ...(await import("../lib/guardian-token.js")) };
+const loadGuardianTokenMock = mock(
+  (_assistantId: string) =>
+    ({ accessToken: "test-guardian-token" }) as unknown as ReturnType<
+      typeof realGuardianToken.loadGuardianToken
+    >,
+);
+mock.module("../lib/guardian-token.js", () => ({
+  ...realGuardianToken,
+  loadGuardianToken: loadGuardianTokenMock,
+}));
+
+// Stub the token-refresh helper without importing the real client module
+// (it drags in the interactive chat client's dependency graph). Pass-through
+// by default: the stored token is returned as-is.
+const resolveFreshBearerTokenMock = mock(
+  async (
+    _runtimeUrl: string,
+    _assistantId: string,
+    bearerToken: string | undefined,
+    _cloud: string | undefined,
+  ) => bearerToken,
+);
+mock.module("../commands/client.js", () => ({
+  resolveFreshBearerToken: resolveFreshBearerTokenMock,
+}));
+
+// Restore the real modules once this file finishes so the mocks do not leak
 // into other test files in the same `bun test` run.
 afterAll(() => {
   mock.module("../lib/process.js", () => realProcess);
+  mock.module("../lib/drain.js", () => realDrain);
+  mock.module("../lib/guardian-token.js", () => realGuardianToken);
 });
 
 import { sleep } from "../commands/sleep.js";
@@ -45,7 +88,7 @@ import {
 
 // Write a lockfile entry so the real resolveTargetAssistant() finds our test
 // assistant without needing to mock the entire assistant-config module.
-function writeLockfile(): void {
+function writeLockfile(extraEntryFields: Record<string, unknown> = {}): void {
   writeFileSync(
     join(testDir, ".vellum.lock.json"),
     JSON.stringify(
@@ -61,6 +104,7 @@ function writeLockfile(): void {
               gatewayPort: DEFAULT_GATEWAY_PORT,
               qdrantPort: DEFAULT_QDRANT_PORT,
             },
+            ...extraEntryFields,
           },
         ],
         activeAssistant: "sleep-test",
@@ -99,6 +143,16 @@ describe("sleep command", () => {
     isProcessAliveMock.mockReturnValue({ alive: false, pid: null });
     stopProcessByPidFileMock.mockReset();
     stopProcessByPidFileMock.mockResolvedValue(true);
+    drainAssistantMock.mockReset();
+    drainAssistantMock.mockResolvedValue("drained");
+    loadGuardianTokenMock.mockReset();
+    loadGuardianTokenMock.mockReturnValue({
+      accessToken: "test-guardian-token",
+    } as unknown as ReturnType<typeof realGuardianToken.loadGuardianToken>);
+    resolveFreshBearerTokenMock.mockReset();
+    resolveFreshBearerTokenMock.mockImplementation(
+      async (_url, _id, bearerToken, _cloud) => bearerToken,
+    );
     rmSync(assistantRootDir, { recursive: true, force: true });
     writeLockfile();
   });
@@ -187,4 +241,119 @@ describe("sleep command", () => {
       "credential-executor",
     );
   });
+
+  describe("--wait", () => {
+    test("runs the drain phase against localUrl before stopping", async () => {
+      writeLockfile({ localUrl: "http://127.0.0.1:9999" });
+      isProcessAliveMock.mockReturnValue({ alive: true, pid: 4321 });
+      process.argv = ["bun", "vellum", "sleep", "sleep-test", "--wait"];
+
+      const consoleLog = spyOn(console, "log").mockImplementation(() => {});
+      try {
+        await sleep();
+      } finally {
+        consoleLog.mockRestore();
+      }
+
+      expect(drainAssistantMock).toHaveBeenCalledTimes(1);
+      const opts = drainAssistantMock.mock.calls[0]![0];
+      expect(opts.baseUrl).toBe("http://127.0.0.1:9999");
+      expect(opts.token).toBe("test-guardian-token");
+      // Bare --wait waits as long as it takes.
+      expect(opts.deadlineAt).toBeNull();
+      expect(stopProcessByPidFileMock).toHaveBeenCalledTimes(3);
+    });
+
+    test("a duration bounds the wait with a deadline", async () => {
+      isProcessAliveMock.mockReturnValue({ alive: true, pid: 4321 });
+      process.argv = ["bun", "vellum", "sleep", "sleep-test", "--wait", "90s"];
+
+      const consoleLog = spyOn(console, "log").mockImplementation(() => {});
+      const before = Date.now();
+      try {
+        await sleep();
+      } finally {
+        consoleLog.mockRestore();
+      }
+
+      const opts = drainAssistantMock.mock.calls[0]![0];
+      expect(opts.deadlineAt).toBeGreaterThanOrEqual(before + 89_000);
+      expect(opts.deadlineAt).toBeLessThanOrEqual(Date.now() + 91_000);
+      expect(stopProcessByPidFileMock).toHaveBeenCalledTimes(3);
+    });
+
+    test("cancelled drain exits 130 without stopping anything", async () => {
+      isProcessAliveMock.mockReturnValue({ alive: true, pid: 4321 });
+      drainAssistantMock.mockResolvedValue("cancelled");
+      process.argv = ["bun", "vellum", "sleep", "sleep-test", "--wait"];
+
+      const consoleLog = spyOn(console, "log").mockImplementation(() => {});
+      const exitMock = mock((code?: number) => {
+        throw new Error(`process.exit:${code}`);
+      });
+      const originalExit = process.exit;
+      process.exit = exitMock as unknown as typeof process.exit;
+      try {
+        await expect(sleep()).rejects.toThrow("process.exit:130");
+      } finally {
+        process.exit = originalExit;
+        consoleLog.mockRestore();
+      }
+
+      expect(stopProcessByPidFileMock).not.toHaveBeenCalled();
+    });
+
+    test("skips the drain when the assistant is not running", async () => {
+      isProcessAliveMock.mockReturnValue({ alive: false, pid: null });
+      process.argv = ["bun", "vellum", "sleep", "sleep-test", "--wait"];
+
+      const consoleLog = spyOn(console, "log").mockImplementation(() => {});
+      try {
+        await sleep();
+      } finally {
+        consoleLog.mockRestore();
+      }
+
+      expect(drainAssistantMock).not.toHaveBeenCalled();
+      expect(stopProcessByPidFileMock).toHaveBeenCalledTimes(3);
+    });
+
+    test("stops without waiting when no guardian token is available", async () => {
+      isProcessAliveMock.mockReturnValue({ alive: true, pid: 4321 });
+      loadGuardianTokenMock.mockReturnValue(null);
+      process.argv = ["bun", "vellum", "sleep", "sleep-test", "--wait"];
+
+      const consoleLog = spyOn(console, "log").mockImplementation(() => {});
+      try {
+        await sleep();
+      } finally {
+        consoleLog.mockRestore();
+      }
+
+      expect(drainAssistantMock).not.toHaveBeenCalled();
+      expect(stopProcessByPidFileMock).toHaveBeenCalledTimes(3);
+    });
+
+    test("rejects an invalid --wait duration", async () => {
+      process.argv = ["bun", "vellum", "sleep", "sleep-test", "--wait=abc"];
+
+      const consoleError = spyOn(console, "error").mockImplementation(() => {});
+      const exitMock = mock((code?: number) => {
+        throw new Error(`process.exit:${code}`);
+      });
+      const originalExit = process.exit;
+      process.exit = exitMock as unknown as typeof process.exit;
+      try {
+        await expect(sleep()).rejects.toThrow("process.exit:1");
+        expect(consoleError).toHaveBeenCalledWith(
+          expect.stringContaining("invalid --wait duration"),
+        );
+      } finally {
+        process.exit = originalExit;
+        consoleError.mockRestore();
+      }
+
+      expect(stopProcessByPidFileMock).not.toHaveBeenCalled();
+    });
+  });
 });

package/src/__tests__/wake.test.ts

@@ -179,7 +179,7 @@ mock.module("../lib/nginx-ingress.js", () => ({
   startRemoteWebIngress: startRemoteWebIngressMock,
 }));
 
-const { wake } = await import("../commands/wake.js");
+const { wake, WEB_INGRESS_FLAG_RETRY } = await import("../commands/wake.js");
 
 let tempDir: string;
 let originalArgv: string[];
@@ -408,8 +408,18 @@ describe("vellum wake — web ingress restore", () => {
     ingress: { enabled: true, publicBaseUrl: "https://assistant.example.com" },
   };
 
+  const originalRetry = { ...WEB_INGRESS_FLAG_RETRY };
+
   beforeEach(() => {
     process.argv = ["bun", "vellum", "wake", "local-assistant"];
+    // Shrink the startup-window retry so failure-path tests stay fast.
+    WEB_INGRESS_FLAG_RETRY.attempts = 3;
+    WEB_INGRESS_FLAG_RETRY.intervalMs = 5;
+  });
+
+  afterEach(() => {
+    WEB_INGRESS_FLAG_RETRY.attempts = originalRetry.attempts;
+    WEB_INGRESS_FLAG_RETRY.intervalMs = originalRetry.intervalMs;
   });
 
   test("restores the edge when ingress is enabled, the flag is on, and it is not already running", async () => {
@@ -515,8 +525,35 @@ describe("vellum wake — web ingress restore", () => {
 
     await wake();
 
+    // The probe retries through the startup window before giving up.
+    expect(isAssistantFeatureFlagEnabledMock.mock.calls.length).toBe(3);
     expect(startRemoteWebIngressMock).not.toHaveBeenCalled();
     expect(warnSpy).toHaveBeenCalled();
     expect(logSpy).toHaveBeenCalledWith("Wake complete.");
   });
+
+  test("a flag probe that fails while the gateway is starting retries and then restores", async () => {
+    loadRawConfigMock.mockReturnValue(enabledConfig);
+    isAssistantFeatureFlagEnabledMock
+      .mockRejectedValueOnce(
+        new Error(
+          'Failed to fetch feature flags: HTTP 503 {"status":"starting"}',
+        ),
+      )
+      .mockRejectedValueOnce(
+        new Error(
+          'Failed to fetch feature flags: HTTP 503 {"status":"starting"}',
+        ),
+      )
+      .mockResolvedValueOnce(true);
+
+    await wake();
+
+    expect(isAssistantFeatureFlagEnabledMock.mock.calls.length).toBe(3);
+    expect(logSpy).toHaveBeenCalledWith(
+      "   Waiting for the gateway before restoring the web ingress edge...",
+    );
+    expect(startRemoteWebIngressMock).toHaveBeenCalled();
+    expect(logSpy).toHaveBeenCalledWith("Wake complete.");
+  });
 });

package/src/commands/flags.ts

@@ -93,13 +93,13 @@ function printHelp(): void {
     "  $ vellum flags                                              # list flags for active assistant",
   );
   console.log(
-    "  $ vellum flags get voice-mode                                 # inspect one flag",
+    "  $ vellum flags get browser                                     # inspect one flag",
   );
   console.log(
-    "  $ vellum flags set voice-mode true                           # enable a flag",
+    "  $ vellum flags set browser true                               # enable a flag",
   );
   console.log(
-    "  $ vellum flags set voice-mode true --assistant eval-1       # target by name/id",
+    "  $ vellum flags set browser true     --assistant eval-1       # target by name/id",
   );
 }
 

package/src/commands/pair.ts

@@ -390,7 +390,7 @@ export async function pair(): Promise<void> {
       `Re-run with a reachable URL, e.g.:\n  vellum pair --url ${suggestion}`,
     );
     console.error(
-      "Pairing a phone by QR instead? Use: vellum pair --qr " +
+      "Pairing another device by QR instead? Use: vellum pair --qr " +
         "(needs an https URL — run `vellum tunnel --provider tailscale` first).",
     );
     process.exit(1);
@@ -478,7 +478,7 @@ export async function pair(): Promise<void> {
         }'s website, not your assistant's address`,
       };
       console.error(
-        "Error: --qr needs a public https URL the phone can open — " +
+        "Error: --qr needs a public https URL the scanning device can open — " +
           `${detailByReason[qrResult.reason]}.`,
       );
       console.error(

package/src/commands/sleep.ts

@@ -7,8 +7,15 @@ import {
 } from "../lib/assistant-config.js";
 import type { AssistantEntry } from "../lib/assistant-config.js";
 import { dockerResourceNames, sleepContainers } from "../lib/docker.js";
+import {
+  drainAssistant,
+  type DrainOutcome,
+  parseWaitDuration,
+} from "../lib/drain.js";
+import { loadGuardianToken } from "../lib/guardian-token.js";
 import { stopIngressNginx } from "../lib/nginx-ingress.js";
 import { isProcessAlive, stopProcessByPidFile } from "../lib/process";
+import { resolveFreshBearerToken } from "./client.js";
 
 const ACTIVE_CALL_LEASES_FILE = "active-call-leases.json";
 
@@ -49,7 +56,7 @@ function readActiveCallLeases(vellumDir: string): ActiveCallLease[] {
 export async function sleep(): Promise<void> {
   const args = process.argv.slice(3);
   if (args.includes("--help") || args.includes("-h")) {
-    console.log("Usage: vellum sleep [<name>] [--force]");
+    console.log("Usage: vellum sleep [<name>] [--force] [--wait [<duration>]]");
     console.log("");
     console.log("Stop the assistant and gateway processes.");
     console.log("");
@@ -62,14 +69,66 @@ export async function sleep(): Promise<void> {
     console.log(
       "  --force   Stop the assistant even if a phone call keepalive lease is active",
     );
+    console.log(
+      "  --wait    Wait for in-flight background work (heartbeats, memory jobs,",
+    );
+    console.log(
+      "            schedules) to finish before stopping. Waits as long as it",
+    );
+    console.log(
+      "            takes by default; pass a duration (e.g. --wait 90s, --wait 10m)",
+    );
+    console.log(
+      "            to bound the wait. Ctrl-C cancels and leaves the assistant running.",
+    );
     process.exit(0);
   }
 
   const force = args.includes("--force");
-  const nameArg = args.find((a) => !a.startsWith("-"));
+  let wait = false;
+  let waitMs: number | null = null;
+  let waitDurationToken: string | null = null;
+  const waitEqArg = args.find((a) => a.startsWith("--wait="));
+  if (waitEqArg) {
+    wait = true;
+    const raw = waitEqArg.slice("--wait=".length);
+    waitMs = parseWaitDuration(raw);
+    if (waitMs == null) {
+      console.error(
+        `Error: invalid --wait duration '${raw}'. Use seconds or a value like 90s or 10m.`,
+      );
+      process.exit(1);
+    }
+  } else {
+    const waitIndex = args.indexOf("--wait");
+    if (waitIndex !== -1) {
+      wait = true;
+      // A bare `--wait` waits indefinitely; a following duration token
+      // (e.g. `90s`, `10m`, `120`) bounds it. Anything else is the name arg.
+      const next = args[waitIndex + 1];
+      if (next !== undefined && /^\d+(s|m)?$/.test(next)) {
+        waitDurationToken = next;
+        waitMs = parseWaitDuration(next);
+        if (waitMs == null) {
+          console.error(
+            `Error: invalid --wait duration '${next}'. Use seconds or a value like 90s or 10m.`,
+          );
+          process.exit(1);
+        }
+      }
+    }
+  }
+  const nameArg = args.find(
+    (a) => !a.startsWith("-") && a !== waitDurationToken,
+  );
   const entry = resolveTargetAssistant(nameArg);
 
   if (entry.cloud === "docker") {
+    if (wait) {
+      console.log(
+        "--wait is not supported for Docker assistants yet — stopping containers now.",
+      );
+    }
     const res = dockerResourceNames(entry.assistantId);
     await sleepContainers(res);
     console.log("Docker containers stopped.");
@@ -135,6 +194,14 @@ export async function sleep(): Promise<void> {
     }
   }
 
+  if (wait) {
+    const outcome = await runDrainPhase(entry, assistantPidFile, waitMs);
+    if (outcome === "cancelled") {
+      console.log("Sleep cancelled — the assistant is still running.");
+      process.exit(130);
+    }
+  }
+
   // Stop assistant — use a generous timeout. On SIGTERM the daemon runs a
   // WAL checkpoint before exiting, which can take several seconds on a
   // multi-GB database. The default 2s grace in stopProcess() would SIGKILL a
@@ -186,3 +253,85 @@ export async function sleep(): Promise<void> {
     console.log("nginx ingress stopped.");
   }
 }
+
+/**
+ * Wait for the assistant's in-flight background work (heartbeats, memory
+ * jobs, schedule runs, active conversation turns) to finish before the stop.
+ * Every failure mode falls back to a normal stop with a note — a drain must
+ * never make `vellum sleep` less reliable than it is without `--wait`.
+ */
+async function runDrainPhase(
+  entry: AssistantEntry,
+  assistantPidFile: string,
+  waitMs: number | null,
+): Promise<DrainOutcome | "skipped"> {
+  if (!isProcessAlive(assistantPidFile).alive) {
+    return "skipped";
+  }
+  const baseUrl = entry.localUrl ?? entry.runtimeUrl;
+  const stored = loadGuardianToken(entry.assistantId)?.accessToken;
+  if (!baseUrl || !stored) {
+    console.log(
+      "Cannot reach the assistant API to wait for background work (missing URL or guardian token) — stopping without waiting.",
+    );
+    return "skipped";
+  }
+  // A token past its renewal point would 401 and read as "unreachable",
+  // silently skipping the wait — refresh-and-fall-back like other CLI
+  // request paths.
+  const token = await resolveFreshBearerToken(
+    baseUrl,
+    entry.assistantId,
+    stored,
+    entry.cloud,
+  );
+  if (!token) {
+    console.log(
+      "Cannot reach the assistant API to wait for background work (missing URL or guardian token) — stopping without waiting.",
+    );
+    return "skipped";
+  }
+
+  console.log(
+    waitMs == null
+      ? "Waiting for the assistant to finish background work (Ctrl-C cancels)…"
+      : `Waiting up to ${Math.round(waitMs / 1000)}s for the assistant to finish background work (Ctrl-C cancels)…`,
+  );
+
+  const controller = new AbortController();
+  const onSigint = (): void => controller.abort();
+  process.once("SIGINT", onSigint);
+  try {
+    const outcome = await drainAssistant({
+      baseUrl,
+      token,
+      deadlineAt: waitMs == null ? null : Date.now() + waitMs,
+      signal: controller.signal,
+    });
+    switch (outcome) {
+      case "drained":
+        console.log("Background work finished — stopping now.");
+        break;
+      case "timeout":
+        console.log(
+          "Wait limit reached — stopping now. Interrupted background jobs resume after the next start.",
+        );
+        break;
+      case "unsupported":
+        console.log(
+          "This assistant version does not support --wait — stopping without waiting.",
+        );
+        break;
+      case "unreachable":
+        console.log(
+          "Could not reach the assistant to check background work — stopping without waiting.",
+        );
+        break;
+      case "cancelled":
+        break;
+    }
+    return outcome;
+  } finally {
+    process.removeListener("SIGINT", onSigint);
+  }
+}

package/src/commands/wake.ts

@@ -425,6 +425,52 @@ export async function wake(): Promise<void> {
   }
 }
 
+/**
+ * Retry policy for the flag probe that gates the web-ingress restore. The
+ * gateway has typically been up for milliseconds at this point and answers
+ * `503 {"status":"starting"}` (or refuses connections) until its startup
+ * completes, so a single probe races it. Mutable so tests can shrink the
+ * window.
+ */
+export const WEB_INGRESS_FLAG_RETRY = {
+  attempts: 15,
+  intervalMs: 2_000,
+};
+
+/**
+ * Probe the `web-remote-ingress` flag, riding out the gateway's startup
+ * window: transient failures retry on an interval until the attempt budget is
+ * spent, then the last error propagates to the caller's warn path.
+ */
+async function verifyWebIngressFlagWithRetry(
+  assistantId: string,
+  gatewayPort: number,
+): Promise<boolean> {
+  let lastError: unknown;
+  for (let attempt = 1; attempt <= WEB_INGRESS_FLAG_RETRY.attempts; attempt++) {
+    try {
+      return await isAssistantFeatureFlagEnabled(
+        assistantId,
+        WEB_REMOTE_INGRESS_FLAG,
+        { runtimeUrl: `http://127.0.0.1:${gatewayPort}` },
+      );
+    } catch (err) {
+      lastError = err;
+      if (attempt === 1) {
+        console.log(
+          "   Waiting for the gateway before restoring the web ingress edge...",
+        );
+      }
+      if (attempt < WEB_INGRESS_FLAG_RETRY.attempts) {
+        await new Promise((resolve) =>
+          setTimeout(resolve, WEB_INGRESS_FLAG_RETRY.intervalMs),
+        );
+      }
+    }
+  }
+  throw lastError;
+}
+
 /**
  * Bring the nginx web ingress edge back up after a wake when the workspace
  * config still wants it. Only restores when ingress is explicitly enabled with
@@ -460,11 +506,7 @@ async function restoreWebIngressIfEnabled(
 
   let flagEnabled: boolean;
   try {
-    flagEnabled = await isAssistantFeatureFlagEnabled(
-      assistantId,
-      WEB_REMOTE_INGRESS_FLAG,
-      { runtimeUrl: `http://127.0.0.1:${gatewayPort}` },
-    );
+    flagEnabled = await verifyWebIngressFlagWithRetry(assistantId, gatewayPort);
   } catch (err) {
     console.warn(
       `   Could not verify the \`${WEB_REMOTE_INGRESS_FLAG}\` flag to restore the web ingress edge; leaving it down. Bring it up manually with \`vellum nginx-ingress up\`. ${