@vellumai/assistant 0.7.1 → 0.7.2

package/src/mcp/mcp-auth-state.ts

@@ -0,0 +1,133 @@
+/**
+ * In-memory MCP OAuth flow status map.
+ *
+ * Tracks the current state of daemon-owned MCP OAuth flows so the CLI can
+ * poll for completion via the IPC route.
+ */
+
+/**
+ * Sibling: `assistant/src/security/oauth-callback-registry.ts`. The two look
+ * similar (in-memory map, supersede semantics, ~5 min TTL) but live at
+ * different layers. The callback registry stores the deferred resolve/reject
+ * pair for a single OAuth code arrival, keyed by OAuth `state`. This file
+ * stores observable status (pending / complete / error) keyed by serverId so
+ * the polling CLI can render progress without holding a long-lived IPC
+ * connection.
+ */
+type McpAuthState =
+  | { status: "pending"; authUrl: string; attemptId: string; expiresAt: number }
+  | {
+      status: "complete";
+      serverId: string;
+      attemptId: string;
+      completedAt: number;
+    }
+  | { status: "error"; error: string; attemptId: string; failedAt: number };
+
+const activeMcpAuthFlows = new Map<string, McpAuthState>();
+
+const PENDING_TTL_MS = 5 * 60 * 1000; // 5 min — matches oauth-callback-registry.ts
+const COMPLETION_GRACE_MS = 60 * 1000; // 60s so the polling CLI gets one final read
+
+/**
+ * Record that an OAuth flow is pending authorization.
+ * Overwrites any prior state for the same serverId (supersede semantics
+ * matching registerPendingCallback). The caller must pass an `attemptId`
+ * (a unique token per attempt) so that fire-and-forget completion writes
+ * can verify they still own the slot before mutating shared state — see
+ * `setMcpAuthComplete` / `setMcpAuthError`.
+ */
+export function setMcpAuthPending(
+  serverId: string,
+  authUrl: string,
+  attemptId: string,
+): void {
+  activeMcpAuthFlows.set(serverId, {
+    status: "pending",
+    authUrl,
+    attemptId,
+    expiresAt: Date.now() + PENDING_TTL_MS,
+  });
+}
+
+/**
+ * Record that an OAuth flow completed successfully. Returns true if the
+ * write was applied; false if the attempt has been superseded by a newer
+ * one (in which case the caller's tail should silently exit without
+ * touching state).
+ */
+export function setMcpAuthComplete(
+  serverId: string,
+  attemptId: string,
+): boolean {
+  const current = activeMcpAuthFlows.get(serverId);
+  if (current && current.attemptId !== attemptId) {
+    return false; // superseded
+  }
+  activeMcpAuthFlows.set(serverId, {
+    status: "complete",
+    serverId,
+    attemptId,
+    completedAt: Date.now(),
+  });
+  return true;
+}
+
+/**
+ * Record that an OAuth flow failed. Returns true if the write was applied;
+ * false if the attempt has been superseded.
+ */
+export function setMcpAuthError(
+  serverId: string,
+  error: string,
+  attemptId: string,
+): boolean {
+  const current = activeMcpAuthFlows.get(serverId);
+  if (current && current.attemptId !== attemptId) {
+    return false; // superseded
+  }
+  activeMcpAuthFlows.set(serverId, {
+    status: "error",
+    error,
+    attemptId,
+    failedAt: Date.now(),
+  });
+  return true;
+}
+
+/**
+ * Get the current state of an OAuth flow, or null if none exists.
+ */
+export function getMcpAuthState(serverId: string): McpAuthState | null {
+  clearExpiredMcpAuthStates();
+  return activeMcpAuthFlows.get(serverId) ?? null;
+}
+
+/**
+ * Remove expired entries from the auth flow map.
+ */
+export function clearExpiredMcpAuthStates(): void {
+  const now = Date.now();
+  for (const [serverId, state] of activeMcpAuthFlows) {
+    if (state.status === "pending" && now > state.expiresAt) {
+      activeMcpAuthFlows.delete(serverId);
+    } else if (
+      state.status === "complete" &&
+      now > state.completedAt + COMPLETION_GRACE_MS
+    ) {
+      activeMcpAuthFlows.delete(serverId);
+    } else if (
+      state.status === "error" &&
+      now > state.failedAt + COMPLETION_GRACE_MS
+    ) {
+      activeMcpAuthFlows.delete(serverId);
+    }
+  }
+}
+
+/**
+ * Test-only helper — clears all auth flow state for test isolation.
+ */
+export function _clearAllMcpAuthStates(): void {
+  activeMcpAuthFlows.clear();
+}

package/src/mcp/mcp-oauth-provider.ts

@@ -61,6 +61,16 @@ export interface McpOAuthCallbackResult {
 /** Which callback transport to use for receiving the OAuth redirect. */
 export type McpOAuthCallbackTransport = "loopback" | "gateway";
 
+export interface McpOAuthProviderOptions {
+  /**
+   * If provided, called with the authorization URL during
+   * redirectToAuthorization() instead of opening a browser.
+   * Used by the daemon-side orchestrator so it can return the
+   * URL to the IPC caller (CLI / web client).
+   */
+  onAuthorizationUrl?: (url: string) => void;
+}
+
 export class McpOAuthProvider implements OAuthClientProvider {
   private readonly serverId: string;
   private readonly serverUrl: string;
@@ -75,6 +85,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
   /** Deferred resolver/rejector for the gateway code promise. */
   private _gatewayCodeResolve: ((code: string) => void) | undefined;
   private _gatewayCodeReject: ((err: Error) => void) | undefined;
+  private readonly _onAuthorizationUrl: ((url: string) => void) | undefined;
 
   /**
    * @param interactive When true (e.g. `mcp auth` CLI), opens browser for OAuth.
@@ -82,17 +93,20 @@ export class McpOAuthProvider implements OAuthClientProvider {
    * @param callbackTransport Which transport to use for the OAuth redirect.
    *   - `"loopback"` (default): localhost HTTP server — for desktop clients.
    *   - `"gateway"`: platform ingress + callback registry — for Docker/platform.
+   * @param options Additional options for the provider.
    */
   constructor(
     serverId: string,
     serverUrl: string,
     interactive = false,
     callbackTransport: McpOAuthCallbackTransport = "loopback",
+    options: McpOAuthProviderOptions = {},
   ) {
     this.serverId = serverId;
     this.serverUrl = serverUrl;
     this.interactive = interactive;
     this.callbackTransport = callbackTransport;
+    this._onAuthorizationUrl = options.onAuthorizationUrl;
   }
 
   // --- redirectUrl ---
@@ -269,6 +283,11 @@ export class McpOAuthProvider implements OAuthClientProvider {
       }
     }
 
+    if (this._onAuthorizationUrl) {
+      this._onAuthorizationUrl(url);
+      return;
+    }
+
     if (!this.interactive) {
       // Daemon mode — don't open browser, just log guidance
       log.info(

package/src/memory/__tests__/jobs-store-job-classes.test.ts

@@ -0,0 +1,24 @@
+import { describe, expect, test } from "bun:test";
+
+import { EMBED_JOB_TYPES, SLOW_LLM_JOB_TYPES } from "../jobs-store.js";
+
+describe("memory job classes", () => {
+  test("EMBED_JOB_TYPES and SLOW_LLM_JOB_TYPES are disjoint", () => {
+    const embedSet = new Set<string>(EMBED_JOB_TYPES);
+    const overlap = SLOW_LLM_JOB_TYPES.filter((t) => embedSet.has(t));
+    expect(overlap).toEqual([]);
+  });
+
+  test("SLOW_LLM_JOB_TYPES entries are non-empty strings", () => {
+    expect(SLOW_LLM_JOB_TYPES.length).toBeGreaterThan(0);
+    for (const t of SLOW_LLM_JOB_TYPES) {
+      expect(typeof t).toBe("string");
+      expect(t.length).toBeGreaterThan(0);
+    }
+  });
+
+  test("SLOW_LLM_JOB_TYPES has no duplicate entries", () => {
+    const set = new Set(SLOW_LLM_JOB_TYPES);
+    expect(set.size).toBe(SLOW_LLM_JOB_TYPES.length);
+  });
+});

package/src/memory/__tests__/qdrant-client-sentinel.test.ts

@@ -0,0 +1,49 @@
+import { describe, expect, test } from "bun:test";
+
+import { stripLegacySparseSuffix } from "../qdrant-client.js";
+
+describe("stripLegacySparseSuffix", () => {
+  test("strips a trailing :sparse-v<digit> suffix", () => {
+    expect(
+      stripLegacySparseSuffix("gemini:gemini-embedding-2-preview:sparse-v3"),
+    ).toBe("gemini:gemini-embedding-2-preview");
+  });
+
+  test("strips multi-digit version suffixes", () => {
+    expect(stripLegacySparseSuffix("openai:text-embed-3:sparse-v42")).toBe(
+      "openai:text-embed-3",
+    );
+  });
+
+  test("returns the input unchanged when no suffix is present", () => {
+    expect(stripLegacySparseSuffix("gemini:gemini-embedding-2-preview")).toBe(
+      "gemini:gemini-embedding-2-preview",
+    );
+  });
+
+  test("does not strip a non-trailing :sparse-v segment", () => {
+    expect(stripLegacySparseSuffix("foo:sparse-v3:bar")).toBe(
+      "foo:sparse-v3:bar",
+    );
+  });
+
+  test("does not strip when the version part is missing", () => {
+    expect(stripLegacySparseSuffix("provider:model:sparse-v")).toBe(
+      "provider:model:sparse-v",
+    );
+  });
+
+  test("normalizes legacy and current sentinels to the same value", () => {
+    const legacy = "gemini:gemini-embedding-2-preview:sparse-v2";
+    const current = "gemini:gemini-embedding-2-preview";
+    expect(stripLegacySparseSuffix(legacy)).toBe(
+      stripLegacySparseSuffix(current),
+    );
+  });
+
+  test("differentiates sentinels that diverge on the dense identity", () => {
+    const a = "gemini:gemini-embedding-2-preview:sparse-v2";
+    const b = "openai:text-embedding-3-small:sparse-v2";
+    expect(stripLegacySparseSuffix(a)).not.toBe(stripLegacySparseSuffix(b));
+  });
+});

package/src/memory/__tests__/sparse-tokenize.test.ts

@@ -0,0 +1,66 @@
+import { describe, expect, test } from "bun:test";
+
+import { tokenize, tokenizeStemmed } from "../sparse-tokenize.js";
+
+describe("tokenize", () => {
+  test("lowercases and splits on non-alphanumeric runs", () => {
+    expect(tokenize("Hello, World!")).toEqual(["hello", "world"]);
+  });
+
+  test("preserves alphanumeric runs as-is (no stemming)", () => {
+    expect(tokenize("running supplements taking")).toEqual([
+      "running",
+      "supplements",
+      "taking",
+    ]);
+  });
+
+  test("handles Unicode letters and digits", () => {
+    expect(tokenize("café-99 naïve")).toEqual(["café", "99", "naïve"]);
+  });
+
+  test("returns empty array for empty / whitespace input", () => {
+    expect(tokenize("")).toEqual([]);
+    expect(tokenize("   \n\t  ")).toEqual([]);
+  });
+});
+
+describe("tokenizeStemmed", () => {
+  test("collapses plurals to their singular stem", () => {
+    expect(tokenizeStemmed("supplements")).toEqual(["supplement"]);
+    expect(tokenizeStemmed("tests")).toEqual(["test"]);
+  });
+
+  test("collapses verb tenses and gerunds to a shared stem", () => {
+    expect(tokenizeStemmed("taking")).toEqual(["take"]);
+    expect(tokenizeStemmed("running")).toEqual(["run"]);
+    expect(tokenizeStemmed("testing")).toEqual(["test"]);
+  });
+
+  test("singular forms map to themselves (idempotent on stems)", () => {
+    expect(tokenizeStemmed("supplement")).toEqual(["supplement"]);
+    expect(tokenizeStemmed("run")).toEqual(["run"]);
+  });
+
+  test("query and document forms collapse to identical buckets", () => {
+    // Plural query terms must land on the same stem as singular doc terms
+    // (and vice versa). This is the property the BM25 sparse channel
+    // relies on so doc-side weights and query-side occurrences match.
+    const queryTokens = tokenizeStemmed("filing reports");
+    const docTokens = tokenizeStemmed("filed a report");
+    expect(queryTokens).toContain("file");
+    expect(queryTokens).toContain("report");
+    expect(docTokens).toContain("file");
+    expect(docTokens).toContain("report");
+  });
+
+  test("returns empty array for empty input", () => {
+    expect(tokenizeStemmed("")).toEqual([]);
+  });
+
+  test("preserves token order and length (one stem per token)", () => {
+    const original = tokenize("the quick brown foxes were jumping");
+    const stemmed = tokenizeStemmed("the quick brown foxes were jumping");
+    expect(stemmed.length).toBe(original.length);
+  });
+});

package/src/memory/anisotropy.test.ts

@@ -0,0 +1,247 @@
+import { describe, expect, test } from "bun:test";
+
+import {
+  type AnisotropyCalibration,
+  applyAnisotropyCorrection,
+  explainedVarianceRatio,
+  fitAnisotropyCalibration,
+} from "./anisotropy.js";
+
+const META = { provider: "gemini" as const, model: "test-model" };
+
+function dot(a: readonly number[], b: readonly number[]): number {
+  let s = 0;
+  for (let i = 0; i < a.length; i++) s += a[i] * b[i];
+  return s;
+}
+
+function l2Norm(v: readonly number[]): number {
+  return Math.sqrt(dot(v, v));
+}
+
+function l2Normalize(v: readonly number[]): number[] {
+  const n = l2Norm(v);
+  if (n === 0) return [...v];
+  return v.map((x) => x / n);
+}
+
+/**
+ * Build a synthetic anisotropic corpus: every sample lives close to a fixed
+ * direction `axis`, with small Gaussian-ish noise injected in the orthogonal
+ * subspace. This gives us a known top-1 PC the fit must recover.
+ */
+function buildAnisotropicCorpus(
+  n: number,
+  dim: number,
+  axis: readonly number[],
+  noiseScale: number,
+): number[][] {
+  const ax = l2Normalize(axis);
+  const out: number[][] = [];
+  let seed = 42;
+  const rand = () => {
+    // Tiny deterministic LCG so tests don't depend on Math.random.
+    seed = (seed * 1103515245 + 12345) & 0x7fffffff;
+    return seed / 0x7fffffff;
+  };
+  for (let i = 0; i < n; i++) {
+    const v = new Array<number>(dim);
+    // Strong common direction component, slightly varied per sample.
+    const along = 1 + (rand() - 0.5) * 0.1;
+    for (let j = 0; j < dim; j++) v[j] = along * ax[j];
+    // Orthogonal noise.
+    for (let j = 0; j < dim; j++) v[j] += (rand() - 0.5) * noiseScale;
+    out.push(v);
+  }
+  return out;
+}
+
+describe("fitAnisotropyCalibration", () => {
+  test("recovers the dominant direction of an anisotropic corpus", () => {
+    const dim = 16;
+    const axis = new Array<number>(dim)
+      .fill(0)
+      .map((_, i) => (i === 0 ? 1 : 0));
+    const corpus = buildAnisotropicCorpus(200, dim, axis, 0.05);
+
+    const calib = fitAnisotropyCalibration(corpus, 1, META);
+
+    expect(calib.components.length).toBe(1);
+    expect(calib.components[0].length).toBe(dim);
+    // PC1 should align with the planted axis (sign-agnostic).
+    const alignment = Math.abs(dot(calib.components[0], axis));
+    expect(alignment).toBeGreaterThan(0.95);
+  });
+
+  test("captures most variance in PC1 for a strongly anisotropic corpus", () => {
+    const dim = 16;
+    const axis = new Array<number>(dim)
+      .fill(0)
+      .map((_, i) => (i === 0 ? 1 : 0));
+    const corpus = buildAnisotropicCorpus(200, dim, axis, 0.02);
+
+    const calib = fitAnisotropyCalibration(corpus, 3, META);
+    const ratios = explainedVarianceRatio(calib);
+
+    // Spectrum is monotonically non-increasing.
+    expect(ratios[0]).toBeGreaterThanOrEqual(ratios[1]);
+    expect(ratios[1]).toBeGreaterThanOrEqual(ratios[2]);
+    // PC1 should dwarf PC2 by an order of magnitude when noise is small.
+    expect(ratios[0]).toBeGreaterThan(ratios[1] * 10);
+  });
+
+  test("returns mean and dim metadata", () => {
+    const dim = 8;
+    const corpus = [
+      [1, 1, 0, 0, 0, 0, 0, 0],
+      [3, 3, 0, 0, 0, 0, 0, 0],
+    ];
+    const calib = fitAnisotropyCalibration(corpus, 1, META);
+    expect(calib.dim).toBe(dim);
+    expect(calib.mean[0]).toBeCloseTo(2);
+    expect(calib.mean[1]).toBeCloseTo(2);
+    expect(calib.mean[2]).toBeCloseTo(0);
+    expect(calib.sampleCount).toBe(2);
+    expect(calib.provider).toBe("gemini");
+    expect(calib.model).toBe("test-model");
+  });
+
+  test("rejects empty input, bad k, and ragged rows", () => {
+    expect(() => fitAnisotropyCalibration([], 1, META)).toThrow(/no vectors/);
+    expect(() => fitAnisotropyCalibration([[1, 2, 3]], 0, META)).toThrow(
+      /positive integer/,
+    );
+    expect(() =>
+      fitAnisotropyCalibration(
+        [
+          [1, 2],
+          [3, 4, 5],
+        ],
+        1,
+        META,
+      ),
+    ).toThrow(/dim/);
+    expect(() => fitAnisotropyCalibration([[1, 2]], 5, META)).toThrow(
+      /exceeds/,
+    );
+  });
+});
+
+describe("applyAnisotropyCorrection", () => {
+  test("output has unit L2 norm", () => {
+    const corpus = buildAnisotropicCorpus(
+      100,
+      8,
+      [1, 0, 0, 0, 0, 0, 0, 0],
+      0.05,
+    );
+    const calib = fitAnisotropyCalibration(corpus, 1, META);
+
+    const corrected = applyAnisotropyCorrection(corpus[0], calib);
+    expect(l2Norm(corrected)).toBeCloseTo(1, 6);
+  });
+
+  test("removes projection onto the dominant direction", () => {
+    const dim = 8;
+    const axis = [1, 0, 0, 0, 0, 0, 0, 0];
+    const corpus = buildAnisotropicCorpus(200, dim, axis, 0.02);
+    const calib = fitAnisotropyCalibration(corpus, 1, META);
+
+    // Pick a vector that points strongly along the planted axis.
+    const sample = [...axis];
+    const corrected = applyAnisotropyCorrection(sample, calib);
+
+    // After removing PC1 (which is ~axis), the projection back onto axis
+    // should be nearly zero — the sample is in the deflated subspace.
+    expect(Math.abs(dot(corrected, axis))).toBeLessThan(0.05);
+  });
+
+  test("spreads cosine similarities for vectors in the cone", () => {
+    const dim = 16;
+    const axis = new Array<number>(dim)
+      .fill(0)
+      .map((_, i) => (i === 0 ? 1 : 0));
+    const corpus = buildAnisotropicCorpus(300, dim, axis, 0.1);
+    const calib = fitAnisotropyCalibration(corpus, 1, META);
+
+    // Compute pairwise cosines on raw vs corrected vectors. The corrected
+    // distribution should have noticeably larger spread (max - min) — the
+    // whole point of the correction.
+    function spread(vectors: number[][]): number {
+      const sims: number[] = [];
+      for (let i = 0; i < vectors.length; i++) {
+        const a = l2Normalize(vectors[i]);
+        for (let j = i + 1; j < vectors.length; j++) {
+          const b = l2Normalize(vectors[j]);
+          sims.push(dot(a, b));
+        }
+      }
+      let min = Infinity;
+      let max = -Infinity;
+      for (const s of sims) {
+        if (s < min) min = s;
+        if (s > max) max = s;
+      }
+      return max - min;
+    }
+
+    const rawSpread = spread(corpus.slice(0, 30));
+    const correctedSpread = spread(
+      corpus.slice(0, 30).map((v) => applyAnisotropyCorrection(v, calib)),
+    );
+
+    expect(correctedSpread).toBeGreaterThan(rawSpread * 2);
+  });
+
+  test("rejects mismatched dim", () => {
+    const calib: AnisotropyCalibration = {
+      provider: "gemini",
+      model: "x",
+      dim: 4,
+      mean: [0, 0, 0, 0],
+      components: [[1, 0, 0, 0]],
+      componentVariance: [1],
+      totalVariance: 1,
+      sampleCount: 1,
+      fitAt: 0,
+    };
+    expect(() => applyAnisotropyCorrection([1, 2, 3], calib)).toThrow(/dim/);
+  });
+});
+
+describe("explainedVarianceRatio", () => {
+  test("returns zeros when totalVariance is zero", () => {
+    const calib: AnisotropyCalibration = {
+      provider: "gemini",
+      model: "x",
+      dim: 2,
+      mean: [0, 0],
+      components: [[1, 0]],
+      componentVariance: [0],
+      totalVariance: 0,
+      sampleCount: 1,
+      fitAt: 0,
+    };
+    expect(explainedVarianceRatio(calib)).toEqual([0]);
+  });
+
+  test("each ratio is variance/total", () => {
+    const calib: AnisotropyCalibration = {
+      provider: "gemini",
+      model: "x",
+      dim: 2,
+      mean: [0, 0],
+      components: [
+        [1, 0],
+        [0, 1],
+      ],
+      componentVariance: [3, 1],
+      totalVariance: 4,
+      sampleCount: 100,
+      fitAt: 0,
+    };
+    const ratios = explainedVarianceRatio(calib);
+    expect(ratios[0]).toBeCloseTo(0.75);
+    expect(ratios[1]).toBeCloseTo(0.25);
+  });
+});