@vellumai/assistant 0.6.2 → 0.6.4

package/src/__tests__/embedding-managed-proxy-selection.test.ts

@@ -0,0 +1,256 @@
+/**
+ * Tests for managed proxy Gemini embedding backend selection.
+ *
+ * Verifies that selectEmbeddingBackend correctly routes through the
+ * managed proxy when the feature flag is enabled and managed proxy
+ * prerequisites are satisfied.
+ */
+
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import { credentialKey } from "../security/credential-key.js";
+
+// ---------------------------------------------------------------------------
+// Mocks — must be before importing the module under test
+// ---------------------------------------------------------------------------
+
+// Suppress logger output
+mock.module("../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// Mutable state for managed proxy context
+let mockPlatformBaseUrl = "";
+let mockAssistantApiKey: string | null = null;
+let mockProviderKeys: Record<string, string | null> = {};
+
+mock.module("../config/env.js", () => ({
+  getPlatformBaseUrl: () => mockPlatformBaseUrl,
+  getOllamaBaseUrlEnv: () => "",
+}));
+
+mock.module("../security/secure-keys.js", () => ({
+  getSecureKeyAsync: async (key: string) => {
+    if (key === credentialKey("vellum", "assistant_api_key")) {
+      return mockAssistantApiKey;
+    }
+    // Provider keys are looked up by plain name
+    return mockProviderKeys[key] ?? null;
+  },
+}));
+
+// Feature flag mock
+const mockFeatureFlags: Record<string, boolean> = {};
+
+mock.module("../config/assistant-feature-flags.js", () => ({
+  isAssistantFeatureFlagEnabled: (key: string, _config: unknown) => {
+    return mockFeatureFlags[key] ?? false;
+  },
+}));
+
+import type { AssistantConfig } from "../config/types.js";
+import {
+  clearEmbeddingBackendCache,
+  selectEmbeddingBackend,
+} from "../memory/embedding-backend.js";
+import { GeminiEmbeddingBackend } from "../memory/embedding-gemini.js";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+const PLATFORM_BASE = "https://platform.example.com";
+const MANAGED_API_KEY = "ast-managed-key-123";
+
+function enableManagedProxy() {
+  mockPlatformBaseUrl = PLATFORM_BASE;
+  mockAssistantApiKey = MANAGED_API_KEY;
+}
+
+function disableManagedProxy() {
+  mockPlatformBaseUrl = "";
+  mockAssistantApiKey = null;
+}
+
+function enableFlag() {
+  mockFeatureFlags["managed-gemini-embeddings-enabled"] = true;
+}
+
+function disableFlag() {
+  mockFeatureFlags["managed-gemini-embeddings-enabled"] = false;
+}
+
+function makeConfig(
+  overrides: {
+    provider?: string;
+    geminiModel?: string;
+    geminiDimensions?: number;
+  } = {},
+): AssistantConfig {
+  return {
+    memory: {
+      embeddings: {
+        provider: overrides.provider ?? "auto",
+        localModel: "Xenova/bge-small-en-v1.5",
+        openaiModel: "text-embedding-3-small",
+        geminiModel: overrides.geminiModel ?? "gemini-embedding-2-preview",
+        geminiDimensions: overrides.geminiDimensions,
+        ollamaModel: "nomic-embed-text",
+      },
+      qdrant: {
+        vectorSize: 384,
+      },
+    },
+    services: {
+      inference: { provider: "anthropic" },
+    },
+  } as unknown as AssistantConfig;
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+beforeEach(() => {
+  disableManagedProxy();
+  disableFlag();
+  mockProviderKeys = {};
+  clearEmbeddingBackendCache();
+});
+
+afterEach(() => {
+  clearEmbeddingBackendCache();
+});
+
+describe("managed proxy Gemini embedding selection", () => {
+  test("selects managed proxy Gemini when flag enabled and proxy context available", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend, reason } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    expect(backend!.provider).toBe("gemini");
+    expect(reason).toBeNull();
+  });
+
+  test("managed proxy backend uses default 3072 dimensions when geminiDimensions not set", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    // Access the private dimensions field to verify default
+    const dimensions = (backend as unknown as { dimensions: number })
+      .dimensions;
+    expect(dimensions).toBe(3072);
+  });
+
+  test("managed proxy backend uses explicit geminiDimensions when set", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig({ geminiDimensions: 768 });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    const dimensions = (backend as unknown as { dimensions: number })
+      .dimensions;
+    expect(dimensions).toBe(768);
+  });
+
+  test("managed proxy backend uses managedBaseUrl (not direct Google API)", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    const managedBaseUrl = (backend as unknown as { managedBaseUrl: string })
+      .managedBaseUrl;
+    expect(managedBaseUrl).toBe(`${PLATFORM_BASE}/v1/runtime-proxy/gemini`);
+  });
+
+  test("falls back to local when flag is disabled (no managed proxy)", async () => {
+    enableManagedProxy();
+    disableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    // With auto and no provider keys, falls through to local
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("local");
+  });
+
+  test("falls back to local when managed proxy context unavailable", async () => {
+    disableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("local");
+  });
+
+  test("selects managed proxy when provider is explicitly gemini", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig({ provider: "gemini" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    const managedBaseUrl = (backend as unknown as { managedBaseUrl: string })
+      .managedBaseUrl;
+    expect(managedBaseUrl).toContain("/v1/runtime-proxy/gemini");
+  });
+
+  test("does not use managed proxy when provider is explicitly local", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig({ provider: "local" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("local");
+  });
+
+  test("does not use managed proxy when provider is explicitly openai", async () => {
+    enableManagedProxy();
+    enableFlag();
+    mockProviderKeys["openai"] = "user-openai-key";
+    const config = makeConfig({ provider: "openai" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("openai");
+  });
+
+  test("direct Gemini key still works when flag is off", async () => {
+    disableManagedProxy();
+    disableFlag();
+    mockProviderKeys["gemini"] = "user-gemini-key";
+    const config = makeConfig({ provider: "gemini" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    expect(backend!.provider).toBe("gemini");
+    // Should NOT use managed proxy
+    const managedBaseUrl = (backend as unknown as { managedBaseUrl?: string })
+      .managedBaseUrl;
+    expect(managedBaseUrl).toBeUndefined();
+  });
+});

package/src/__tests__/emit-event-signal.test.ts

@@ -0,0 +1,71 @@
+/**
+ * Behavioral test for the generic CLI→daemon event bridge.
+ *
+ * Callers write a JSON-encoded {@link ServerMessage} to
+ * `<signalsDir>/emit-event`. {@link handleEmitEventSignal} reads that
+ * payload and republishes it through the {@link assistantEventHub} so
+ * SSE subscribers receive it.
+ */
+import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+import { afterEach, describe, expect, test } from "bun:test";
+
+import type { ServerMessage } from "../daemon/message-protocol.js";
+import type { AssistantEvent } from "../runtime/assistant-event.js";
+import { assistantEventHub } from "../runtime/assistant-event-hub.js";
+import { DAEMON_INTERNAL_ASSISTANT_ID } from "../runtime/assistant-scope.js";
+import { handleEmitEventSignal } from "../signals/emit-event.js";
+import { getSignalsDir } from "../util/platform.js";
+
+function signalPath(): string {
+  return join(getSignalsDir(), "emit-event");
+}
+
+const subscriptions: Array<{ dispose(): void }> = [];
+
+afterEach(() => {
+  for (const sub of subscriptions.splice(0)) {
+    sub.dispose();
+  }
+  const path = signalPath();
+  if (existsSync(path)) {
+    rmSync(path, { force: true });
+  }
+});
+
+describe("handleEmitEventSignal", () => {
+  test("reads a ServerMessage from the signal file and publishes it to the event hub", async () => {
+    mkdirSync(getSignalsDir(), { recursive: true });
+
+    const payload: ServerMessage = { type: "tasks_changed" };
+
+    writeFileSync(signalPath(), JSON.stringify(payload), "utf-8");
+
+    const received: AssistantEvent[] = [];
+    let resolveDelivered: (() => void) | null = null;
+    const delivered = new Promise<void>((resolve) => {
+      resolveDelivered = resolve;
+    });
+
+    subscriptions.push(
+      assistantEventHub.subscribe(
+        { assistantId: DAEMON_INTERNAL_ASSISTANT_ID },
+        (event) => {
+          received.push(event);
+          resolveDelivered?.();
+        },
+      ),
+    );
+
+    handleEmitEventSignal();
+
+    await delivered;
+
+    expect(received).toHaveLength(1);
+    const event = received[0];
+    expect(event.assistantId).toBe(DAEMON_INTERNAL_ASSISTANT_ID);
+    expect(event.message).toEqual(payload);
+    expect(typeof event.id).toBe("string");
+    expect(typeof event.emittedAt).toBe("string");
+  });
+});

package/src/__tests__/extension-id-sync-guard.test.ts

@@ -0,0 +1,222 @@
+/**
+ * Guard tests for Chrome extension allowlist configuration.
+ *
+ * Single source of truth:
+ *   meta/browser-extension/chrome-extension-allowlist.json
+ *
+ * This guard ensures:
+ *   1) Canonical config has a valid shape and valid extension IDs.
+ *   2) Assistant runtime allowlist mirrors canonical config.
+ *   3) The concrete extension ID literal appears only in canonical config
+ *      (not duplicated across runtime/tests/docs).
+ */
+
+import { existsSync, readdirSync, readFileSync, statSync } from "node:fs";
+import { homedir } from "node:os";
+import { join, resolve } from "node:path";
+import { describe, expect, test } from "bun:test";
+
+import { ALLOWED_EXTENSION_ORIGINS } from "../runtime/routes/browser-extension-pair-routes.js";
+
+const repoRoot = resolve(__dirname, "..", "..", "..");
+const CANONICAL_CONFIG_REL_PATH =
+  "meta/browser-extension/chrome-extension-allowlist.json";
+const CANONICAL_CONFIG_ABS_PATH = join(repoRoot, CANONICAL_CONFIG_REL_PATH);
+const LOCAL_OVERRIDE_PATH = join(
+  homedir(),
+  ".vellum",
+  "chrome-extension-allowlist.local.json",
+);
+
+const EXTENSION_ID_REGEX = /^[a-p]{32}$/;
+const PLACEHOLDER_ID_REGEX = /^TODO_[A-Z0-9_]+$/;
+
+type AllowlistConfig = {
+  version: number;
+  allowedExtensionIds: string[];
+};
+
+function parseCanonicalConfig(): AllowlistConfig {
+  const raw = readFileSync(CANONICAL_CONFIG_ABS_PATH, "utf8");
+  const parsed = JSON.parse(raw) as Partial<AllowlistConfig>;
+
+  if (!Number.isInteger(parsed.version) || (parsed.version ?? 0) <= 0) {
+    throw new Error("Invalid canonical config: version must be a positive integer");
+  }
+
+  if (!Array.isArray(parsed.allowedExtensionIds)) {
+    throw new Error("Invalid canonical config: allowedExtensionIds must be an array");
+  }
+
+  if (parsed.allowedExtensionIds.length === 0) {
+    throw new Error(
+      "Invalid canonical config: allowedExtensionIds must contain at least one id",
+    );
+  }
+
+  const seen = new Set<string>();
+  const validIds: string[] = [];
+  for (const id of parsed.allowedExtensionIds) {
+    if (typeof id !== "string") {
+      throw new Error(`Invalid canonical extension id: ${String(id)}`);
+    }
+    if (seen.has(id)) {
+      throw new Error(`Duplicate canonical extension id: ${id}`);
+    }
+    seen.add(id);
+
+    if (EXTENSION_ID_REGEX.test(id)) {
+      validIds.push(id);
+    } else if (!PLACEHOLDER_ID_REGEX.test(id)) {
+      throw new Error(`Invalid canonical extension id: ${id}`);
+    }
+  }
+
+  if (validIds.length === 0) {
+    throw new Error(
+      "Invalid canonical config: allowedExtensionIds must contain at least one real extension id",
+    );
+  }
+
+  return {
+    version: parsed.version as number,
+    allowedExtensionIds: validIds,
+  };
+}
+
+function listTextFilesRecursively(root: string): string[] {
+  const ignoredDirs = new Set([
+    ".git",
+    "node_modules",
+    "dist",
+    "build",
+    "coverage",
+    ".next",
+    ".turbo",
+    ".idea",
+    ".vscode",
+  ]);
+
+  const allowedExtensions = new Set([
+    ".ts",
+    ".tsx",
+    ".js",
+    ".jsx",
+    ".mjs",
+    ".cjs",
+    ".json",
+    ".md",
+    ".swift",
+    ".sh",
+    ".toml",
+    ".yml",
+    ".yaml",
+    ".txt",
+  ]);
+
+  const out: string[] = [];
+
+  function walk(dir: string): void {
+    let entries;
+    try {
+      entries = readdirSync(dir, { withFileTypes: true });
+    } catch (err) {
+      // Directory may have been removed by a concurrent test; skip it.
+      if ((err as NodeJS.ErrnoException).code === "ENOENT") return;
+      throw err;
+    }
+    for (const entry of entries) {
+      if (entry.name.startsWith(".DS_Store")) continue;
+      // Skip temp fixtures created by parallel tests (e.g. .test-starter-bundle-<pid>).
+      if (entry.name.startsWith(".test-")) continue;
+      const absPath = join(dir, entry.name);
+      if (entry.isDirectory()) {
+        if (ignoredDirs.has(entry.name)) continue;
+        walk(absPath);
+        continue;
+      }
+      if (!entry.isFile()) continue;
+
+      const ext = entry.name.slice(entry.name.lastIndexOf("."));
+      if (!allowedExtensions.has(ext)) continue;
+
+      // Skip large files to keep this guard lightweight.
+      let size: number;
+      try {
+        size = statSync(absPath).size;
+      } catch (err) {
+        if ((err as NodeJS.ErrnoException).code === "ENOENT") continue;
+        throw err;
+      }
+      if (size > 1_000_000) continue;
+      out.push(absPath);
+    }
+  }
+
+  walk(root);
+  return out;
+}
+
+describe("Chrome extension allowlist guard", () => {
+  test("canonical allowlist config is valid", () => {
+    const config = parseCanonicalConfig();
+    expect(config.version).toBeGreaterThan(0);
+    expect(config.allowedExtensionIds.length).toBeGreaterThan(0);
+  });
+
+  test("assistant runtime allowlist contains every canonical origin", () => {
+    // The runtime set is the union of canonical + local override
+    // (~/.vellum/chrome-extension-allowlist.local.json) + env var. A dev
+    // machine may have extras; we only assert the canonical IDs are present.
+    const config = parseCanonicalConfig();
+    for (const id of config.allowedExtensionIds) {
+      const origin = `chrome-extension://${id}/`;
+      expect(ALLOWED_EXTENSION_ORIGINS.has(origin)).toBe(true);
+    }
+  });
+
+  test("assistant runtime allowlist exactly mirrors canonical when no override sources are active", () => {
+    // Exact-equality invariant: when neither the local override file nor
+    // the env-var override is active, the runtime set must equal the
+    // canonical set — nothing extra, nothing missing. A dev machine with
+    // an unpacked-extension override will skip this deterministic check
+    // and rely on the subset assertion above.
+    if (existsSync(LOCAL_OVERRIDE_PATH)) return;
+    if (
+      process.env.VELLUM_CHROME_EXTENSION_IDS ||
+      process.env.VELLUM_CHROME_EXTENSION_ID
+    ) {
+      return;
+    }
+
+    const config = parseCanonicalConfig();
+    const expectedOrigins = new Set(
+      config.allowedExtensionIds.map((id) => `chrome-extension://${id}/`),
+    );
+    expect(new Set(ALLOWED_EXTENSION_ORIGINS)).toEqual(expectedOrigins);
+  });
+
+  test("concrete extension IDs appear only in canonical config", () => {
+    const config = parseCanonicalConfig();
+    const allFiles = listTextFilesRecursively(repoRoot);
+
+    for (const extensionId of config.allowedExtensionIds) {
+      const matches: string[] = [];
+      for (const absPath of allFiles) {
+        const relPath = absPath.replace(`${repoRoot}/`, "");
+        let content: string;
+        try {
+          content = readFileSync(absPath, "utf8");
+        } catch (err) {
+          // File may have been removed by a concurrent test between listing and reading.
+          if ((err as NodeJS.ErrnoException).code === "ENOENT") continue;
+          throw err;
+        }
+        if (content.includes(extensionId)) {
+          matches.push(relPath);
+        }
+      }
+      expect(matches).toEqual([CANONICAL_CONFIG_REL_PATH]);
+    }
+  });
+});