@vellumai/cli 0.10.12-dev.202607270141.42f7d3c → 0.10.12-dev.202607270523.ec1a36e

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@vellumai/cli",
-  "version": "0.10.12-dev.202607270141.42f7d3c",
+  "version": "0.10.12-dev.202607270523.ec1a36e",
   "description": "CLI tools for vellum-assistant",
   "type": "module",
   "exports": {

package/src/__tests__/drain.test.ts

@@ -0,0 +1,337 @@
+import { describe, expect, test } from "bun:test";
+
+import { drainAssistant, parseWaitDuration } from "../lib/drain.js";
+
+const BASE = "http://127.0.0.1:7823";
+
+interface RecordedCall {
+  url: string;
+  method: string;
+}
+
+/**
+ * Fetch stub: routes requests through `handler` and records every call.
+ * Throwing from the handler simulates a network failure.
+ */
+function fetchStub(
+  handler: (url: string, method: string) => Response | Promise<Response>,
+): {
+  impl: (url: string, init?: RequestInit) => Promise<Response>;
+  calls: RecordedCall[];
+} {
+  const calls: RecordedCall[] = [];
+  const impl = async (url: string, init?: RequestInit): Promise<Response> => {
+    const method = init?.method ?? "GET";
+    calls.push({ url, method });
+    return handler(url, method);
+  };
+  return { impl, calls };
+}
+
+function jsonResponse(body: unknown, status = 200): Response {
+  return new Response(JSON.stringify(body), {
+    status,
+    headers: { "Content-Type": "application/json" },
+  });
+}
+
+function busyStatus(): unknown {
+  return {
+    quiescedUntil: Date.now() + 60_000,
+    idle: false,
+    activeConversations: [],
+    memoryJobs: [
+      { id: "job-1", type: "memory_retrospective", startedAt: Date.now() },
+    ],
+    scheduleRuns: [],
+    workflowRuns: [],
+    heartbeatRuns: [],
+  };
+}
+
+function workflowBusyStatus(): unknown {
+  return {
+    quiescedUntil: Date.now() + 60_000,
+    idle: false,
+    activeConversations: [],
+    memoryJobs: [],
+    scheduleRuns: [],
+    workflowRuns: [
+      { runId: "wf-1", name: "nightly-report", startedAt: Date.now() },
+    ],
+    heartbeatRuns: [],
+  };
+}
+
+function idleStatus(): unknown {
+  return {
+    quiescedUntil: Date.now() + 60_000,
+    idle: true,
+    activeConversations: [],
+    memoryJobs: [],
+    scheduleRuns: [],
+    workflowRuns: [],
+    heartbeatRuns: [],
+  };
+}
+
+function idleExpiredLeaseStatus(): unknown {
+  return { ...(idleStatus() as Record<string, unknown>), quiescedUntil: null };
+}
+
+describe("parseWaitDuration", () => {
+  test("parses bare seconds, seconds, and minutes", () => {
+    expect(parseWaitDuration("90")).toBe(90_000);
+    expect(parseWaitDuration("90s")).toBe(90_000);
+    expect(parseWaitDuration("10m")).toBe(600_000);
+    expect(parseWaitDuration(" 5s ")).toBe(5_000);
+  });
+
+  test("rejects invalid durations", () => {
+    expect(parseWaitDuration("0")).toBeNull();
+    expect(parseWaitDuration("abc")).toBeNull();
+    expect(parseWaitDuration("5x")).toBeNull();
+    expect(parseWaitDuration("-5")).toBeNull();
+    expect(parseWaitDuration("")).toBeNull();
+  });
+});
+
+describe("drainAssistant", () => {
+  test("arms the lease, narrates, and returns drained when work finishes", async () => {
+    let statusCalls = 0;
+    const { impl, calls } = fetchStub((url, method) => {
+      if (url.endsWith("/v1/lifecycle/quiesce") && method === "POST") {
+        return jsonResponse({ quiescedUntil: Date.now() + 60_000 });
+      }
+      statusCalls += 1;
+      return jsonResponse(statusCalls < 3 ? busyStatus() : idleStatus());
+    });
+    const lines: string[] = [];
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      log: (line) => lines.push(line),
+    });
+
+    expect(outcome).toBe("drained");
+    expect(calls[0]).toEqual({
+      url: `${BASE}/v1/lifecycle/quiesce`,
+      method: "POST",
+    });
+    expect(lines.join("\n")).toContain("memory memory_retrospective");
+  });
+
+  test("a running workflow keeps the drain busy and is narrated", async () => {
+    let statusCalls = 0;
+    const { impl } = fetchStub((_url, method) => {
+      if (method === "POST") {
+        return jsonResponse({ quiescedUntil: Date.now() + 60_000 });
+      }
+      statusCalls += 1;
+      return jsonResponse(
+        statusCalls < 2 ? workflowBusyStatus() : idleStatus(),
+      );
+    });
+    const lines: string[] = [];
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      log: (line) => lines.push(line),
+    });
+
+    expect(outcome).toBe("drained");
+    expect(lines.join("\n")).toContain('workflow "nightly-report"');
+  });
+
+  test("a single idle snapshot is not enough — drained needs two in a row", async () => {
+    let statusCalls = 0;
+    const sequence = [idleStatus(), busyStatus(), idleStatus(), idleStatus()];
+    const { impl } = fetchStub((_url, method) => {
+      if (method === "POST") {
+        return jsonResponse({ quiescedUntil: Date.now() + 60_000 });
+      }
+      const status = sequence[Math.min(statusCalls, sequence.length - 1)];
+      statusCalls += 1;
+      return jsonResponse(status);
+    });
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("drained");
+    // idle, busy (streak reset), idle, idle — four reads, not one.
+    expect(statusCalls).toBe(4);
+  });
+
+  test("idle without a live lease re-arms instead of counting toward drained", async () => {
+    let statusCalls = 0;
+    let armCalls = 0;
+    const { impl } = fetchStub((_url, method) => {
+      if (method === "POST") {
+        armCalls += 1;
+        return jsonResponse({ quiescedUntil: Date.now() + 60_000 });
+      }
+      statusCalls += 1;
+      // The lease has lapsed on the first idle read; once re-armed, the
+      // subsequent idle reads carry a live lease.
+      return jsonResponse(
+        statusCalls === 1 ? idleExpiredLeaseStatus() : idleStatus(),
+      );
+    });
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("drained");
+    // Initial arm + the re-arm triggered by the lapsed lease.
+    expect(armCalls).toBeGreaterThanOrEqual(2);
+    // The expired-lease idle read did not count: two live-lease reads follow.
+    expect(statusCalls).toBeGreaterThanOrEqual(3);
+  });
+
+  test("an unmaintainable lease fails the drain instead of trusting idle", async () => {
+    let armCalls = 0;
+    const { impl } = fetchStub((_url, method) => {
+      if (method === "POST") {
+        armCalls += 1;
+        if (armCalls === 1) {
+          return jsonResponse({ quiescedUntil: Date.now() + 60_000 });
+        }
+        return jsonResponse({ error: "boom" }, 500);
+      }
+      return jsonResponse(idleExpiredLeaseStatus());
+    });
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("unreachable");
+  });
+
+  test("returns unsupported when the assistant lacks the quiesce route", async () => {
+    const { impl } = fetchStub(() => jsonResponse({ error: "nope" }, 404));
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      fetchImpl: impl,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("unsupported");
+  });
+
+  test("returns unreachable when the initial lease request fails", async () => {
+    const { impl } = fetchStub(() => {
+      throw new Error("connection refused");
+    });
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      fetchImpl: impl,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("unreachable");
+  });
+
+  test("returns timeout when the deadline passes while still busy", async () => {
+    const { impl } = fetchStub((_url, method) =>
+      method === "POST"
+        ? jsonResponse({ quiescedUntil: Date.now() + 60_000 })
+        : jsonResponse(busyStatus()),
+    );
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: Date.now() + 40,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("timeout");
+  });
+
+  test("returns unreachable after sustained status failures", async () => {
+    const { impl } = fetchStub((_url, method) => {
+      if (method === "POST") {
+        return jsonResponse({ quiescedUntil: Date.now() + 60_000 });
+      }
+      throw new Error("boom");
+    });
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("unreachable");
+  });
+
+  test("cancel releases the lease and returns cancelled", async () => {
+    const controller = new AbortController();
+    const { impl, calls } = fetchStub((_url, method) => {
+      if (method === "POST") {
+        return jsonResponse({ quiescedUntil: Date.now() + 60_000 });
+      }
+      if (method === "DELETE") {
+        return jsonResponse({ released: true });
+      }
+      // First busy status arrives, then the user hits Ctrl-C.
+      controller.abort();
+      return jsonResponse(busyStatus());
+    });
+
+    const outcome = await drainAssistant({
+      baseUrl: BASE,
+      token: "tok",
+      deadlineAt: null,
+      pollIntervalMs: 1,
+      fetchImpl: impl,
+      signal: controller.signal,
+      log: () => {},
+    });
+
+    expect(outcome).toBe("cancelled");
+    expect(
+      calls.some(
+        (c) => c.method === "DELETE" && c.url.endsWith("/v1/lifecycle/quiesce"),
+      ),
+    ).toBe(true);
+  });
+});

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/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/lib/drain.ts

@@ -0,0 +1,346 @@
+/**
+ * Drain phase for `vellum sleep --wait`.
+ *
+ * Arms the assistant's quiesce lease (pausing the START of new background
+ * work — heartbeats, schedule claims, memory-job claims), then polls the
+ * drain-status endpoint until nothing is in flight. The lease is TTL-bound
+ * and refreshed while polling, so a killed CLI can never leave background
+ * work paused. The wait narrates what it is waiting on — background jobs can
+ * legitimately run for many minutes, and a silent wait is indistinguishable
+ * from a hang.
+ */
+
+import { loopbackSafeFetch } from "./loopback-fetch.js";
+
+const QUIESCE_TTL_MS = 60_000;
+const QUIESCE_REFRESH_INTERVAL_MS = 20_000;
+const REQUEST_TIMEOUT_MS = 2_500;
+const DEFAULT_POLL_INTERVAL_MS = 1_000;
+const NARRATE_HEARTBEAT_MS = 15_000;
+/** Consecutive failed polls before giving up on the drain. Transient
+ * failures (a brief daemon event-loop stall) recover; a wall of failures
+ * means the assistant cannot report status and waiting is pointless. */
+const MAX_CONSECUTIVE_FAILURES = 10;
+/** Consecutive failed lease arms/refreshes before giving up. Without a live
+ * lease the claim gates are open and an idle snapshot proves nothing, so an
+ * unmaintainable lease makes the drain unsound. */
+const MAX_LEASE_FAILURES = 3;
+
+export type DrainOutcome =
+  | "drained"
+  | "timeout"
+  | "unsupported"
+  | "unreachable"
+  | "cancelled";
+
+interface DrainStatusResponse {
+  quiescedUntil: number | null;
+  idle: boolean;
+  activeConversations: Array<{
+    conversationId: string;
+    title: string | null;
+    originChannel: string | null;
+    originInterface: string | null;
+    processingStartedAt: number;
+  }>;
+  memoryJobs: Array<{
+    id: string;
+    type: string;
+    startedAt: number | null;
+  }>;
+  scheduleRuns: Array<{
+    runId: string;
+    scheduleName: string | null;
+    startedAt: number;
+  }>;
+  workflowRuns?: Array<{
+    runId: string;
+    name: string | null;
+    startedAt: number | null;
+  }>;
+  heartbeatRuns?: Array<{
+    runId: string;
+    startedAt: number;
+  }>;
+}
+
+export interface DrainOptions {
+  baseUrl: string;
+  token: string;
+  /** Epoch ms to stop waiting at, or null to wait as long as it takes. */
+  deadlineAt: number | null;
+  /** Aborted on Ctrl-C — releases the lease and returns "cancelled". */
+  signal?: AbortSignal;
+  pollIntervalMs?: number;
+  fetchImpl?: (url: string, init?: RequestInit) => Promise<Response>;
+  log?: (line: string) => void;
+}
+
+/**
+ * Parse a `--wait` duration: bare seconds (`90`), seconds (`90s`), or
+ * minutes (`10m`). Returns milliseconds, or null when not a valid duration.
+ */
+export function parseWaitDuration(raw: string): number | null {
+  const match = /^(\d+)(s|m)?$/.exec(raw.trim());
+  if (!match) {
+    return null;
+  }
+  const value = Number.parseInt(match[1]!, 10);
+  if (!Number.isFinite(value) || value < 1) {
+    return null;
+  }
+  return match[2] === "m" ? value * 60_000 : value * 1_000;
+}
+
+function formatElapsed(sinceMs: number | null, now: number): string {
+  if (sinceMs == null) {
+    return "just started";
+  }
+  const totalSeconds = Math.max(0, Math.floor((now - sinceMs) / 1000));
+  const minutes = Math.floor(totalSeconds / 60);
+  const seconds = totalSeconds % 60;
+  return minutes > 0
+    ? `${minutes}m${seconds.toString().padStart(2, "0")}s`
+    : `${seconds}s`;
+}
+
+function describeBusyWork(status: DrainStatusResponse, now: number): string[] {
+  const items: string[] = [];
+  for (const conv of status.activeConversations) {
+    const label =
+      conv.title != null && conv.title.length > 0
+        ? `conversation "${conv.title}"`
+        : `background turn (${conv.originInterface ?? conv.originChannel ?? "unknown origin"})`;
+    items.push(`${label} (${formatElapsed(conv.processingStartedAt, now)})`);
+  }
+  for (const job of status.memoryJobs) {
+    items.push(`memory ${job.type} (${formatElapsed(job.startedAt, now)})`);
+  }
+  for (const run of status.scheduleRuns) {
+    const label = run.scheduleName ?? run.runId.slice(0, 8);
+    items.push(`schedule "${label}" (${formatElapsed(run.startedAt, now)})`);
+  }
+  for (const run of status.workflowRuns ?? []) {
+    const label = run.name ?? run.runId.slice(0, 8);
+    items.push(`workflow "${label}" (${formatElapsed(run.startedAt, now)})`);
+  }
+  for (const run of status.heartbeatRuns ?? []) {
+    items.push(`heartbeat (${formatElapsed(run.startedAt, now)})`);
+  }
+  return items;
+}
+
+function busySignature(status: DrainStatusResponse): string {
+  return [
+    ...status.activeConversations.map((c) => c.conversationId),
+    ...status.memoryJobs.map((j) => j.id),
+    ...status.scheduleRuns.map((r) => r.runId),
+    ...(status.workflowRuns ?? []).map((r) => r.runId),
+    ...(status.heartbeatRuns ?? []).map((r) => r.runId),
+  ].join("|");
+}
+
+function sleepMs(ms: number, signal?: AbortSignal): Promise<void> {
+  return new Promise((resolve) => {
+    const timer = setTimeout(done, ms);
+    function done(): void {
+      clearTimeout(timer);
+      signal?.removeEventListener("abort", done);
+      resolve();
+    }
+    signal?.addEventListener("abort", done, { once: true });
+  });
+}
+
+/**
+ * Wait for the assistant's in-flight background work to finish.
+ *
+ * Never throws: every failure mode maps onto a {@link DrainOutcome} so the
+ * caller can decide how to proceed with the stop.
+ */
+export async function drainAssistant(
+  opts: DrainOptions,
+): Promise<DrainOutcome> {
+  const fetchImpl = opts.fetchImpl ?? loopbackSafeFetch;
+  const log = opts.log ?? console.log;
+  const pollIntervalMs = opts.pollIntervalMs ?? DEFAULT_POLL_INTERVAL_MS;
+  const base = opts.baseUrl.replace(/\/+$/, "");
+
+  const request = async (
+    path: string,
+    init: { method: string; body?: string },
+  ): Promise<Response> => {
+    // Per-request controller aborted by either the request timeout or the
+    // caller's cancel signal; listeners are removed per request so a long
+    // wait never accumulates them on the shared cancel signal.
+    const controller = new AbortController();
+    const timer = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
+    const onOuterAbort = (): void => controller.abort();
+    opts.signal?.addEventListener("abort", onOuterAbort, { once: true });
+    try {
+      return await fetchImpl(`${base}${path}`, {
+        method: init.method,
+        headers: {
+          Authorization: `Bearer ${opts.token}`,
+          "Content-Type": "application/json",
+        },
+        ...(init.body !== undefined && { body: init.body }),
+        signal: controller.signal,
+      });
+    } finally {
+      clearTimeout(timer);
+      opts.signal?.removeEventListener("abort", onOuterAbort);
+    }
+  };
+
+  const armLease = () =>
+    request("/v1/lifecycle/quiesce", {
+      method: "POST",
+      body: JSON.stringify({ ttlMs: QUIESCE_TTL_MS }),
+    });
+
+  const releaseLease = async (): Promise<void> => {
+    try {
+      await request("/v1/lifecycle/quiesce", { method: "DELETE" });
+    } catch {
+      // Best-effort — the lease self-expires.
+    }
+  };
+
+  // Arm the lease before anything else: without it a busy queue keeps
+  // claiming new work and the wait never converges.
+  let response: Response;
+  try {
+    response = await armLease();
+  } catch {
+    if (opts.signal?.aborted) {
+      // The abort may have landed after the server persisted the lease, so
+      // its outcome is ambiguous — release best-effort rather than leaving
+      // background work paused for a TTL on a cancelled sleep.
+      await releaseLease();
+      return "cancelled";
+    }
+    return "unreachable";
+  }
+  if (response.status === 404 || response.status === 405) {
+    return "unsupported";
+  }
+  if (!response.ok) {
+    return "unreachable";
+  }
+
+  let lastLeaseAt = Date.now();
+  let consecutiveFailures = 0;
+  let leaseFailureStreak = 0;
+  let lastSignature: string | null = null;
+  let lastNarratedAt = 0;
+  let idleStreak = 0;
+
+  for (;;) {
+    if (opts.signal?.aborted) {
+      await releaseLease();
+      return "cancelled";
+    }
+    if (opts.deadlineAt != null && Date.now() >= opts.deadlineAt) {
+      return "timeout";
+    }
+
+    try {
+      const statusResponse = await request("/v1/lifecycle/drain-status", {
+        method: "GET",
+      });
+      if (!statusResponse.ok) {
+        throw new Error(`drain-status returned ${statusResponse.status}`);
+      }
+      const status = (await statusResponse.json()) as DrainStatusResponse;
+      consecutiveFailures = 0;
+
+      if (status.idle) {
+        // Idle only counts under a LIVE lease: with the lease expired the
+        // claim gates are open again, so new work invisible to this snapshot
+        // could be starting. Re-arm and retry rather than trusting the read.
+        const leaseActive =
+          status.quiescedUntil != null && status.quiescedUntil > Date.now();
+        if (!leaseActive) {
+          idleStreak = 0;
+          try {
+            const rearm = await armLease();
+            if (rearm.ok) {
+              lastLeaseAt = Date.now();
+              leaseFailureStreak = 0;
+            } else {
+              leaseFailureStreak += 1;
+            }
+          } catch {
+            if (opts.signal?.aborted) {
+              await releaseLease();
+              return "cancelled";
+            }
+            leaseFailureStreak += 1;
+          }
+          if (leaseFailureStreak >= MAX_LEASE_FAILURES) {
+            return "unreachable";
+          }
+          await sleepMs(pollIntervalMs, opts.signal);
+          continue;
+        }
+        // Require two consecutive idle snapshots: work that raced past a
+        // pre-lease gate check lands its run row within a tick, so a second
+        // read one poll later rules out a stop landing on that window.
+        idleStreak += 1;
+        if (idleStreak >= 2) {
+          return "drained";
+        }
+        await sleepMs(pollIntervalMs, opts.signal);
+        continue;
+      }
+      idleStreak = 0;
+
+      const now = Date.now();
+      const signature = busySignature(status);
+      if (
+        signature !== lastSignature ||
+        now - lastNarratedAt >= NARRATE_HEARTBEAT_MS
+      ) {
+        const items = describeBusyWork(status, now);
+        const shown = items.slice(0, 4);
+        const more = items.length - shown.length;
+        log(
+          `Waiting for the assistant to finish: ${shown.join(", ")}${
+            more > 0 ? ` and ${more} more` : ""
+          }…`,
+        );
+        lastSignature = signature;
+        lastNarratedAt = now;
+      }
+    } catch {
+      if (opts.signal?.aborted) {
+        await releaseLease();
+        return "cancelled";
+      }
+      consecutiveFailures += 1;
+      if (consecutiveFailures >= MAX_CONSECUTIVE_FAILURES) {
+        return "unreachable";
+      }
+    }
+
+    if (Date.now() - lastLeaseAt >= QUIESCE_REFRESH_INTERVAL_MS) {
+      try {
+        const refresh = await armLease();
+        if (refresh.ok) {
+          lastLeaseAt = Date.now();
+          leaseFailureStreak = 0;
+        } else {
+          leaseFailureStreak += 1;
+        }
+      } catch {
+        leaseFailureStreak += 1;
+      }
+      if (leaseFailureStreak >= MAX_LEASE_FAILURES) {
+        return "unreachable";
+      }
+    }
+
+    await sleepMs(pollIntervalMs, opts.signal);
+  }
+}