@vellumai/assistant 0.7.0 → 0.7.2

package/src/__tests__/mcp-cli.test.ts

@@ -15,6 +15,41 @@ import { Command } from "commander";
 
 // ── Module mocks (must precede imports that pull them in) ──────────────
 
+type MockIpcResult = { ok: boolean; result?: unknown; error?: string };
+let mockCliIpcCallFn: ReturnType<
+  typeof mock<
+    (
+      method: string,
+      params?: Record<string, unknown>,
+      opts?: { timeoutMs?: number },
+    ) => Promise<MockIpcResult>
+  >
+> = mock(
+  (
+    _method: string,
+    _params?: Record<string, unknown>,
+    _opts?: { timeoutMs?: number },
+  ): Promise<MockIpcResult> =>
+    Promise.resolve({
+      ok: false,
+      error: "Could not connect to assistant daemon. Is it running?",
+    }),
+);
+
+mock.module("../ipc/cli-client.js", () => ({
+  cliIpcCall: (
+    method: string,
+    params?: Record<string, unknown>,
+    opts?: { timeoutMs?: number },
+  ) => mockCliIpcCallFn(method, params, opts),
+}));
+
+let mockOpenInHostBrowserFn = mock(async (_url: string) => {});
+
+mock.module("../util/browser.js", () => ({
+  openInHostBrowser: (url: string) => mockOpenInHostBrowserFn(url),
+}));
+
 let stdoutLines: string[] = [];
 let stderrLines: string[] = [];
 
@@ -26,7 +61,9 @@ mock.module("../util/logger.js", () => ({
     error: (...args: unknown[]) => {
       stderrLines.push(args.map(String).join(" "));
     },
-    warn: () => {},
+    warn: (...args: unknown[]) => {
+      stderrLines.push(args.map(String).join(" "));
+    },
     debug: () => {},
     trace: () => {},
     fatal: () => {},
@@ -39,7 +76,6 @@ mock.module("../util/logger.js", () => ({
 
 mock.module("../mcp/mcp-oauth-provider.js", () => ({
   deleteMcpOAuthCredentials: async () => {},
-  McpOAuthProvider: class {},
 }));
 
 mock.module("../mcp/client.js", () => ({
@@ -392,6 +428,23 @@ describe("assistant mcp add", () => {
     const server = servers?.["default-risk"] as Record<string, unknown>;
     expect(server.defaultRiskLevel).toBe("high");
   });
+
+  test("calls cliIpcCall with internal_mcp_reload after saving server", async () => {
+    mockCliIpcCallFn = mock(() => Promise.resolve({ ok: true }));
+
+    await runMcpAdd("reload-test-server", [
+      "-t",
+      "sse",
+      "-u",
+      "https://example.com/sse",
+    ]);
+
+    const reloadCall = mockCliIpcCallFn.mock.calls.find(
+      (c) => c[0] === "internal_mcp_reload",
+    );
+    expect(reloadCall).toBeDefined();
+    expect(reloadCall![1]).toEqual({ body: {} });
+  });
 });
 
 describe("assistant mcp remove", () => {
@@ -475,3 +528,286 @@ describe("assistant mcp remove", () => {
     expect(servers?.["keep-me"]).toBeDefined();
   });
 });
+
+describe("assistant mcp reload", () => {
+  beforeAll(() => {
+    testDataDir = join(
+      tmpdir(),
+      `vellum-mcp-reload-test-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+    );
+    mkdirSync(testDataDir, { recursive: true });
+    configPath = join(testDataDir, "config.json");
+    writeConfig({});
+  });
+
+  afterAll(() => {
+    rmSync(testDataDir, { recursive: true, force: true });
+  });
+
+  beforeEach(() => {
+    mockCliIpcCallFn = mock(() => Promise.resolve({ ok: true }));
+  });
+
+  test("calls cliIpcCall with internal_mcp_reload", async () => {
+    const { stdout } = await runMcp("reload");
+
+    const reloadCall = mockCliIpcCallFn.mock.calls.find(
+      (c) => c[0] === "internal_mcp_reload",
+    );
+    expect(reloadCall).toBeDefined();
+    expect(reloadCall![1]).toEqual({ body: {} });
+    expect(stdout).toContain("MCP reload signal sent");
+  });
+
+  test("warns but does not fail when daemon is unreachable", async () => {
+    mockCliIpcCallFn = mock(() =>
+      Promise.resolve({
+        ok: false,
+        error: "Could not connect to assistant daemon",
+      }),
+    );
+
+    const { exitCode, stderr } = await runMcp("reload");
+
+    expect(exitCode).toBe(0); // best-effort, not fatal
+    expect(stderr).toContain("Could not signal reload");
+  });
+});
+
+describe("assistant mcp auth — IPC path", () => {
+  let ipcTestDataDir: string;
+  let ipcConfigPath: string;
+
+  beforeAll(() => {
+    ipcTestDataDir = join(
+      tmpdir(),
+      `vellum-mcp-auth-ipc-test-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+    );
+    mkdirSync(ipcTestDataDir, { recursive: true });
+    ipcConfigPath = join(ipcTestDataDir, "config.json");
+    writeFileSync(
+      ipcConfigPath,
+      JSON.stringify({
+        mcp: {
+          servers: {
+            srv: {
+              transport: { type: "sse", url: "https://mcp.example.com/sse" },
+              enabled: true,
+              defaultRiskLevel: "high",
+            },
+          },
+        },
+      }),
+      "utf-8",
+    );
+  });
+
+  afterAll(() => {
+    rmSync(ipcTestDataDir, { recursive: true, force: true });
+  });
+
+  beforeEach(() => {
+    // runMcp() sets VELLUM_WORKSPACE_DIR = testDataDir; align the module-level
+    // variable with the IPC-specific directory so the config loader sees the
+    // SSE server config created in beforeAll.
+    testDataDir = ipcTestDataDir;
+    mockCliIpcCallFn = mock(() =>
+      Promise.resolve({
+        ok: false,
+        error: "Could not connect to assistant daemon. Is it running?",
+      }),
+    );
+    mockOpenInHostBrowserFn = mock(async (_url: string) => {});
+    stdoutLines = [];
+    stderrLines = [];
+    process.exitCode = 0;
+  });
+
+  test("mcp auth calls IPC internal_mcp_auth_start first", async () => {
+    let ipcCallIndex = 0;
+    mockCliIpcCallFn = mock(() => {
+      ipcCallIndex++;
+      if (ipcCallIndex === 1) {
+        // start call
+        return Promise.resolve({
+          ok: true,
+          result: { auth_url: "https://auth.example.com", state: "srv" },
+        });
+      }
+      // poll calls → complete
+      return Promise.resolve({
+        ok: true,
+        result: { status: "complete" },
+      });
+    });
+
+    await runMcp("auth", ["srv"]);
+
+    expect(mockCliIpcCallFn.mock.calls[0][0]).toBe("internal_mcp_auth_start");
+    expect(mockCliIpcCallFn.mock.calls[0][1]).toEqual({
+      body: { serverId: "srv" },
+    });
+    expect(mockOpenInHostBrowserFn).toHaveBeenCalledWith(
+      "https://auth.example.com",
+    );
+  });
+
+  test("IPC start returns ok=false (daemon unavailable) → exits 1 with helpful error", async () => {
+    mockCliIpcCallFn = mock(() =>
+      Promise.resolve({
+        ok: false,
+        error: "Could not connect to assistant daemon. Is it running?",
+      }),
+    );
+
+    const { exitCode, stderr } = await runMcp("auth", ["srv"]);
+
+    expect(exitCode).toBe(1);
+    expect(stderr).toContain("Is it running?");
+  });
+
+  test("IPC start returns ok=false (Unknown method) → exits 1 with helpful error", async () => {
+    mockCliIpcCallFn = mock(() =>
+      Promise.resolve({
+        ok: false,
+        error: "Unknown method: internal_mcp_auth_start",
+      }),
+    );
+
+    const { exitCode, stderr } = await runMcp("auth", ["srv"]);
+
+    expect(exitCode).toBe(1);
+    expect(stderr).toContain("Is it running?");
+  });
+
+  test("IPC start returns ok=false with a real daemon error → exits 1 without falling back to loopback", async () => {
+    mockCliIpcCallFn = mock(() =>
+      Promise.resolve({
+        ok: false,
+        error: "MCP server not configured for OAuth",
+      }),
+    );
+
+    await runMcp("auth", ["srv"]);
+
+    // Should have surfaced the daemon error and set exit code 1, not fallen
+    // back to loopback (which would have constructed the OAuth provider and
+    // possibly opened a browser via the loopback path).
+    expect(process.exitCode).toBe(1);
+    expect(mockOpenInHostBrowserFn).not.toHaveBeenCalled();
+    // Exactly one IPC call (the start) — no polling, no retry.
+    expect(mockCliIpcCallFn.mock.calls.length).toBe(1);
+    expect(mockCliIpcCallFn.mock.calls[0][0]).toBe("internal_mcp_auth_start");
+  });
+
+  test("polling complete → exits 0 (daemon-side reload triggers itself; CLI no longer writes file signal)", async () => {
+    let ipcCallIndex = 0;
+    mockCliIpcCallFn = mock(() => {
+      ipcCallIndex++;
+      if (ipcCallIndex === 1) {
+        return Promise.resolve({
+          ok: true,
+          result: { auth_url: "https://auth.example.com", state: "srv" },
+        });
+      }
+      return Promise.resolve({
+        ok: true,
+        result: { status: "complete" },
+      });
+    });
+
+    const { exitCode, stdout } = await runMcp("auth", ["srv"]);
+
+    expect(exitCode).toBe(0);
+    expect(stdout).toContain("Authentication successful");
+    // The IPC success path relies on the daemon-side orchestrator to call
+    // reloadMcpServers() itself on completion — no CLI-side signal needed.
+  });
+
+  test("polling error → exits 1 with error message", async () => {
+    let ipcCallIndex = 0;
+    mockCliIpcCallFn = mock(() => {
+      ipcCallIndex++;
+      if (ipcCallIndex === 1) {
+        return Promise.resolve({
+          ok: true,
+          result: { auth_url: "https://auth.example.com", state: "srv" },
+        });
+      }
+      return Promise.resolve({
+        ok: true,
+        result: { status: "error", error: "access_denied" },
+      });
+    });
+
+    const { exitCode, stderr } = await runMcp("auth", ["srv"]);
+
+    expect(exitCode).toBe(1);
+    expect(stderr).toMatch(/cancelled|access_denied|OAuth failed/);
+  });
+
+  test("polling gets NotFoundError (daemon restarted) → exits 1 with helpful message immediately", async () => {
+    let ipcCallIndex = 0;
+    mockCliIpcCallFn = mock(() => {
+      ipcCallIndex++;
+      if (ipcCallIndex === 1) {
+        return Promise.resolve({
+          ok: true,
+          result: { auth_url: "https://auth.example.com", state: "srv" },
+        });
+      }
+      // Daemon restarted — state lost
+      return Promise.resolve({
+        ok: false,
+        error: 'No active OAuth flow for server "srv"',
+      });
+    });
+
+    const { exitCode, stderr } = await runMcp("auth", ["srv"]);
+
+    expect(exitCode).toBe(1);
+    expect(stderr).toMatch(/assistant may have restarted|OAuth flow was lost/);
+  });
+
+  test("IPC start returns ok=true with already_authenticated → exits 0 without OAuth flow", async () => {
+    mockCliIpcCallFn = mock(() =>
+      Promise.resolve({
+        ok: true,
+        result: { auth_url: "", state: "srv", already_authenticated: true },
+      }),
+    );
+
+    const { exitCode, stdout } = await runMcp("auth", ["srv"]);
+
+    expect(exitCode).toBe(0);
+    expect(stdout).toContain("already authenticated");
+    // Should not have opened the browser or started polling
+    expect(mockOpenInHostBrowserFn).not.toHaveBeenCalled();
+    expect(mockCliIpcCallFn.mock.calls.length).toBe(1);
+  });
+
+  test("polling gets a non-notfound IPC error → exits 1 immediately with the error message", async () => {
+    let ipcCallIndex = 0;
+    mockCliIpcCallFn = mock(() => {
+      ipcCallIndex++;
+      if (ipcCallIndex === 1) {
+        return Promise.resolve({
+          ok: true,
+          result: { auth_url: "https://auth.example.com", state: "srv" },
+        });
+      }
+      // Unexpected IPC error (not "No active OAuth flow")
+      return Promise.resolve({
+        ok: false,
+        error: "Internal server error during polling",
+      });
+    });
+
+    const { exitCode, stderr } = await runMcp("auth", ["srv"]);
+
+    expect(exitCode).toBe(1);
+    expect(stderr).toMatch(/Internal server error during polling/);
+    // Should fail fast — only 1 start call + 1 poll call, not the full timeout
+    expect(mockCliIpcCallFn.mock.calls.length).toBe(2);
+  });
+});

package/src/__tests__/memory-jobs-worker-lanes.test.ts

@@ -0,0 +1,188 @@
+/**
+ * Head-of-line repro for the per-lane scheduler in `runMemoryJobsOnce`.
+ *
+ * Before this fix, all non-embed jobs ran through a single bounded worker
+ * pool, so a long-running `graph_consolidate` LLM call would pin every slot
+ * and starve fast-lane jobs (e.g. `embed_concept_page` consolidation pages)
+ * for the duration of that call.
+ *
+ * The new scheduler runs slow / fast / embed lanes in parallel pools, each
+ * with its own concurrency budget. This test enqueues a wave of slow LLM
+ * jobs alongside fast jobs and asserts that every fast job completes before
+ * any slow job's promise resolves — proving the lanes are truly independent.
+ */
+import { beforeAll, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import { eq } from "drizzle-orm";
+
+import { DEFAULT_CONFIG } from "../config/defaults.js";
+import type { AssistantConfig } from "../config/types.js";
+
+// ── Mocks (must precede imports of tested module) ──────────────────
+
+mock.module("../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// Per-lane caps: 1 slow slot (so only 1 of the 5 enqueued slow jobs runs in
+// this tick) and a generous fast cap so every fast job both gets claimed
+// and gets a slot in the lane pool. The OLD shared-pool scheduler — which
+// claimed jobs without lane awareness and ran them through a single
+// workerConcurrency-sized pool — would pin its slots on the first claimed
+// slow jobs and force the fast jobs to queue behind a 200ms LLM call.
+const TEST_CONFIG: AssistantConfig = {
+  ...DEFAULT_CONFIG,
+  memory: {
+    ...DEFAULT_CONFIG.memory,
+    enabled: true,
+    jobs: {
+      ...DEFAULT_CONFIG.memory.jobs,
+      slowLlmConcurrency: 1,
+      fastConcurrency: 5,
+      embedConcurrency: 1,
+      workerConcurrency: 2,
+    },
+  },
+};
+
+mock.module("../config/loader.js", () => ({
+  getConfig: () => TEST_CONFIG,
+  loadConfig: () => TEST_CONFIG,
+  invalidateConfigCache: () => {},
+}));
+
+// ── Track timestamps so we can assert ordering ─────────────────────
+
+type CompletionRecord = {
+  type: string;
+  conversationId: string;
+  completedAt: number;
+};
+const completions: CompletionRecord[] = [];
+
+// Slow-lane handler: blocks for SLOW_DELAY_MS. The test asserts every fast
+// job completes before any slow job — a single 200ms window is plenty.
+const SLOW_DELAY_MS = 200;
+
+mock.module("../memory/graph/consolidation.js", () => ({
+  runConsolidation: async (
+    scopeId: string,
+  ): Promise<{
+    totalUpdated: number;
+    totalDeleted: number;
+    totalMergeEdges: number;
+  }> => {
+    await new Promise((resolve) => setTimeout(resolve, SLOW_DELAY_MS));
+    completions.push({
+      type: "graph_consolidate",
+      conversationId: scopeId,
+      completedAt: Date.now(),
+    });
+    return { totalUpdated: 0, totalDeleted: 0, totalMergeEdges: 0 };
+  },
+}));
+
+// Fast-lane handler: resolves on the next microtask. The test fires this
+// many times in parallel; nothing should block.
+mock.module("../memory/jobs/embed-concept-page.js", () => ({
+  embedConceptPageJob: async (job: {
+    payload: { slug?: string };
+  }): Promise<void> => {
+    completions.push({
+      type: "embed_concept_page",
+      conversationId: job.payload.slug ?? "",
+      completedAt: Date.now(),
+    });
+  },
+}));
+
+// Stub remaining heavy boundaries that we never exercise but that get pulled
+// in transitively through jobs-worker's eager imports. These aren't strictly
+// required if the host machine can resolve them, but mocking them keeps the
+// test hermetic and fast under `bun test`.
+mock.module("../memory/db-maintenance.js", () => ({
+  maybeRunDbMaintenance: () => {},
+}));
+
+import { getDb } from "../memory/db-connection.js";
+import { initializeDb } from "../memory/db-init.js";
+import { enqueueMemoryJob } from "../memory/jobs-store.js";
+import { runMemoryJobsOnce } from "../memory/jobs-worker.js";
+import { _resetQdrantBreaker } from "../memory/qdrant-circuit-breaker.js";
+import { memoryJobs } from "../memory/schema.js";
+
+describe("memory jobs worker lane scheduling", () => {
+  beforeAll(() => {
+    initializeDb();
+  });
+
+  beforeEach(() => {
+    const db = getDb();
+    db.run("DELETE FROM memory_jobs");
+    completions.length = 0;
+    _resetQdrantBreaker();
+  });
+
+  test("fast lane completes before slow lane releases its slot", async () => {
+    // 5 slow `graph_consolidate` jobs across distinct scopes (so they would
+    // serialize behind a single shared pool) plus 5 fast `embed_concept_page`
+    // jobs across distinct slugs.
+    for (let i = 0; i < 5; i++) {
+      enqueueMemoryJob("graph_consolidate", { scopeId: `slow-${i}` });
+    }
+    for (let i = 0; i < 5; i++) {
+      enqueueMemoryJob("embed_concept_page", { slug: `fast-${i}` });
+    }
+
+    await runMemoryJobsOnce();
+
+    const fastDone = completions.filter((c) => c.type === "embed_concept_page");
+    const slowDone = completions.filter((c) => c.type === "graph_consolidate");
+
+    // Slow lane is capped at 1 in this test, so only 1 slow job ran in this
+    // tick. Fast lane has cap 2, but with 5 fast jobs it runs all 5 because
+    // each handler resolves on the next microtask.
+    expect(fastDone).toHaveLength(5);
+    expect(slowDone).toHaveLength(1);
+
+    // Head-of-line guarantee: every fast completion timestamp must precede
+    // the (single) slow completion timestamp. Under the old shared pool with
+    // workerConcurrency=2 and 1 slow slot occupied, this still held only if
+    // a fast slot freed up first — but with 2 slow jobs in flight (the old
+    // claim path would have claimed multiple slow jobs into the shared pool)
+    // both slots would be pinned for SLOW_DELAY_MS and fast work would queue
+    // behind. With the lane scheduler, fast work has its own pool.
+    const earliestSlow = Math.min(...slowDone.map((c) => c.completedAt));
+    for (const fast of fastDone) {
+      expect(fast.completedAt).toBeLessThan(earliestSlow);
+    }
+
+    // Sanity: exactly 1 of the 5 enqueued slow jobs reached `completed` (the
+    // slow lane's per-tick budget). The other 4 are still pending and will be
+    // picked up on subsequent ticks. (FIFO ordering inside the slow lane is
+    // covered separately in jobs-store-qdrant-breaker.test.ts.)
+    //
+    // Note: a sixth `graph_consolidate` row may also appear here — the
+    // `maybeEnqueueGraphMaintenanceJobs` tail of `runMemoryJobsOnce` enqueues
+    // its own maintenance job whose checkpoint is missing in this fresh DB.
+    // That row is irrelevant; we only care about the completed-vs-pending
+    // counts of jobs we explicitly seeded.
+    const completedSlow = countSlowByStatus("completed");
+    expect(completedSlow).toBe(1);
+  });
+});
+
+function countSlowByStatus(
+  status: "pending" | "running" | "completed",
+): number {
+  const db = getDb();
+  return db
+    .select()
+    .from(memoryJobs)
+    .where(eq(memoryJobs.type, "graph_consolidate"))
+    .all()
+    .filter((row) => row.status === status).length;
+}

package/src/__tests__/memory-v2-static-injector.test.ts

@@ -0,0 +1,95 @@
+/**
+ * Tests for the `memory-v2-static` runtime injector.
+ *
+ * Covers:
+ *   - Returns null when `memoryV2Static` is missing/empty.
+ *   - Returns null when `mode === "minimal"`.
+ *   - Wraps content in `<memory>...</memory>` and uses
+ *     `after-memory-prefix` placement.
+ *   - Escapes any `</memory>` substring inside the authored content so the
+ *     wrapper cannot be broken out of.
+ *
+ * Hermetic: drives the injector's `produce()` directly with a synthesized
+ * `TurnContext` — no daemon, no filesystem.
+ */
+
+import { describe, expect, test } from "bun:test";
+
+import { defaultInjectorsPlugin } from "../plugins/defaults/injectors.js";
+import type { Injector, TurnContext } from "../plugins/types.js";
+
+function findInjector(name: string): Injector {
+  const injector = defaultInjectorsPlugin.injectors?.find(
+    (i) => i.name === name,
+  );
+  if (!injector) {
+    throw new Error(`injector '${name}' not registered`);
+  }
+  return injector;
+}
+
+function makeContext(overrides: Partial<TurnContext> = {}): TurnContext {
+  return {
+    requestId: "req-test",
+    conversationId: "conv-test",
+    turnIndex: 0,
+    trust: { sourceChannel: "vellum", trustClass: "guardian" },
+    ...overrides,
+  };
+}
+
+const memoryV2StaticInjector = findInjector("memory-v2-static");
+
+describe("memory-v2-static injector", () => {
+  test("returns null when memoryV2Static is undefined", async () => {
+    const ctx = makeContext({ injectionInputs: {} });
+    expect(await memoryV2StaticInjector.produce(ctx)).toBeNull();
+  });
+
+  test("returns null when memoryV2Static is null", async () => {
+    const ctx = makeContext({ injectionInputs: { memoryV2Static: null } });
+    expect(await memoryV2StaticInjector.produce(ctx)).toBeNull();
+  });
+
+  test("returns null when memoryV2Static is an empty string", async () => {
+    const ctx = makeContext({ injectionInputs: { memoryV2Static: "" } });
+    expect(await memoryV2StaticInjector.produce(ctx)).toBeNull();
+  });
+
+  test("returns null in minimal mode even with content", async () => {
+    const ctx = makeContext({
+      injectionInputs: {
+        mode: "minimal",
+        memoryV2Static: "## Essentials\n\nAlice prefers VS Code.",
+      },
+    });
+    expect(await memoryV2StaticInjector.produce(ctx)).toBeNull();
+  });
+
+  test("wraps content in <memory>...</memory> with after-memory-prefix placement", async () => {
+    const content =
+      "## Essentials\n\nAlice prefers VS Code.\n\n## Threads\n\nOpen: ship PR.";
+    const ctx = makeContext({
+      injectionInputs: { memoryV2Static: content },
+    });
+
+    const block = await memoryV2StaticInjector.produce(ctx);
+    expect(block).not.toBeNull();
+    expect(block!.id).toBe("memory-v2-static");
+    expect(block!.placement).toBe("after-memory-prefix");
+    expect(block!.text).toBe(`<memory>\n${content}\n</memory>`);
+  });
+
+  test("escapes inner </memory> closing tags so the wrapper cannot be broken out of", async () => {
+    const content = "## Essentials\n\nText with </memory> embedded.";
+    const ctx = makeContext({
+      injectionInputs: { memoryV2Static: content },
+    });
+
+    const block = await memoryV2StaticInjector.produce(ctx);
+    expect(block).not.toBeNull();
+    expect(block!.text).toBe(
+      "<memory>\n## Essentials\n\nText with &lt;/memory&gt; embedded.\n</memory>",
+    );
+  });
+});