@vellumai/assistant 0.6.1 → 0.6.3

package/src/memory/embedding-backend.test.ts

@@ -0,0 +1,75 @@
+import { afterEach, describe, expect, mock, test } from "bun:test";
+
+import type { AssistantConfig } from "../config/types.js";
+import {
+  clearEmbeddingBackendCache,
+  resetLocalEmbeddingFailureState,
+  selectEmbeddingBackend,
+} from "./embedding-backend.js";
+
+const LOCAL_CONFIG = {
+  memory: {
+    embeddings: {
+      provider: "local",
+      localModel: "BAAI/bge-small-en-v1.5",
+    },
+  },
+} as unknown as AssistantConfig;
+
+describe("embedding backend cache invalidation", () => {
+  afterEach(() => {
+    clearEmbeddingBackendCache();
+  });
+
+  test("clearEmbeddingBackendCache disposes cached backends before clearing", async () => {
+    const firstSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(firstSelection.backend).not.toBeNull();
+
+    const dispose = mock();
+    (firstSelection.backend as { dispose?: () => void }).dispose = dispose;
+
+    clearEmbeddingBackendCache();
+
+    expect(dispose).toHaveBeenCalledTimes(1);
+
+    const secondSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(secondSelection.backend).not.toBe(firstSelection.backend);
+  });
+
+  test("resetLocalEmbeddingFailureState preserves live cached backends", async () => {
+    const firstSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(firstSelection.backend).not.toBeNull();
+
+    const dispose = mock();
+    (firstSelection.backend as { dispose?: () => void }).dispose = dispose;
+
+    resetLocalEmbeddingFailureState();
+
+    expect(dispose).not.toHaveBeenCalled();
+
+    const secondSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(secondSelection.backend).toBe(firstSelection.backend);
+  });
+
+  test("resetLocalEmbeddingFailureState clears poisoned local backend retry state", async () => {
+    const firstSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(firstSelection.backend).not.toBeNull();
+
+    const poisonedInitPromise = Promise.reject(new Error("poisoned"));
+    poisonedInitPromise.catch(() => {});
+
+    const backend = firstSelection.backend as unknown as {
+      delegate: unknown;
+      initPromise: Promise<unknown> | null;
+    };
+    backend.delegate = null;
+    backend.initPromise = poisonedInitPromise;
+
+    resetLocalEmbeddingFailureState();
+
+    expect(backend.initPromise).toBeNull();
+
+    const secondSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(secondSelection.backend).toBe(firstSelection.backend);
+  });
+});

package/src/memory/embedding-backend.ts

@@ -1,7 +1,10 @@
 import { createHash } from "node:crypto";
 
+import { isAssistantFeatureFlagEnabled } from "../config/assistant-feature-flags.js";
 import { getOllamaBaseUrlEnv } from "../config/env.js";
 import type { AssistantConfig } from "../config/types.js";
+import { MANAGED_PROVIDER_META } from "../providers/managed-proxy/constants.js";
+import { resolveManagedProxyContext } from "../providers/managed-proxy/context.js";
 import { getProviderKeyAsync } from "../security/secure-keys.js";
 import { getLogger } from "../util/logger.js";
 import { GeminiEmbeddingBackend } from "./embedding-gemini.js";
@@ -67,6 +70,16 @@ class LazyLocalEmbeddingBackend implements EmbeddingBackend {
     }
   }
 
+  dispose(): void {
+    this.delegate?.dispose?.();
+  }
+
+  resetForRetry(): void {
+    if (!this.delegate) {
+      this.initPromise = null;
+    }
+  }
+
   private async getDelegate(): Promise<EmbeddingBackend> {
     if (this.delegate) return this.delegate;
     if (!this.initPromise) {
@@ -176,12 +189,32 @@ function putInVectorCache(
 
 /** Clear cached embedding backends and the in-memory vector cache. */
 export function clearEmbeddingBackendCache(): void {
+  for (const backend of new Set(backendCache.values())) {
+    try {
+      backend.dispose?.();
+    } catch (err) {
+      log.warn(
+        { err, provider: backend.provider, model: backend.model },
+        "Failed to dispose embedding backend during cache clear",
+      );
+    }
+  }
   backendCache.clear();
   vectorCache.clear();
   vectorCacheBytes = 0;
   localBackendBroken = false;
 }
 
+/** Reset the sticky local-backend failure flag without evicting live backends. */
+export function resetLocalEmbeddingFailureState(): void {
+  localBackendBroken = false;
+  for (const backend of new Set(backendCache.values())) {
+    if (backend instanceof LazyLocalEmbeddingBackend) {
+      backend.resetForRetry();
+    }
+  }
+}
+
 function cacheKey(provider: string, model: string, extras?: string[]): string {
   if (extras && extras.length > 0) {
     return `${provider}:${model}:${extras.join(":")}`;
@@ -243,6 +276,7 @@ export interface EmbeddingBackend {
     inputs: EmbeddingInput[],
     options?: EmbeddingRequestOptions,
   ): Promise<number[][]>;
+  dispose?(): void;
 }
 
 export interface EmbeddingBackendSelection {
@@ -280,6 +314,44 @@ export async function selectEmbeddingBackend(
     };
   }
 
+  // When the managed-gemini-embeddings-enabled flag is on AND managed proxy
+  // prerequisites are satisfied, insert managed-proxy Gemini at the front of
+  // the auto chain so platform assistants use Vellum-managed Gemini embeddings.
+  if (
+    (requested === "auto" || requested === "gemini") &&
+    isAssistantFeatureFlagEnabled("managed-gemini-embeddings-enabled", config)
+  ) {
+    const proxyCtx = await resolveManagedProxyContext();
+    if (proxyCtx.enabled) {
+      const meta = MANAGED_PROVIDER_META["gemini"];
+      if (meta?.managed && meta.proxyPath) {
+        const managedBaseUrl = `${proxyCtx.platformBaseUrl}${meta.proxyPath}`;
+        const managedModel = config.memory.embeddings.geminiModel;
+        const managedDimensions =
+          config.memory.embeddings.geminiDimensions ?? 3072;
+        const extras = geminiCacheExtras(config);
+        return {
+          backend: getCachedOrCreate(
+            "gemini",
+            managedModel,
+            () =>
+              new GeminiEmbeddingBackend(
+                proxyCtx.assistantApiKey,
+                managedModel,
+                {
+                  taskType: config.memory.embeddings.geminiTaskType,
+                  dimensions: managedDimensions,
+                  managedBaseUrl,
+                },
+              ),
+            [...extras, "managed"],
+          ),
+          reason: null,
+        };
+      }
+    }
+  }
+
   // Auto order: local → openai → gemini → ollama
   const order: EmbeddingProviderName[] =
     requested === "auto"
@@ -329,11 +401,18 @@ export async function selectEmbeddingBackend(
       case "gemini": {
         const geminiKey = await getProviderKeyAsync("gemini");
         if (!geminiKey) {
-          const cached = getCached(
-            "gemini",
-            config.memory.embeddings.geminiModel,
-            geminiCacheExtras(config),
-          );
+          // Check managed cache variant first so a warm managed backend
+          // survives transient proxy-context blips, then non-managed.
+          const cached =
+            getCached("gemini", config.memory.embeddings.geminiModel, [
+              ...geminiCacheExtras(config),
+              "managed",
+            ]) ??
+            getCached(
+              "gemini",
+              config.memory.embeddings.geminiModel,
+              geminiCacheExtras(config),
+            );
           if (cached) return { backend: cached, reason: null };
           continue;
         }
@@ -614,6 +693,53 @@ async function selectFallbackBackends(
               geminiCacheExtras(config),
             ),
           );
+        } else if (
+          isAssistantFeatureFlagEnabled(
+            "managed-gemini-embeddings-enabled",
+            config,
+          )
+        ) {
+          // Try managed proxy Gemini as fallback when no direct key exists.
+          const proxyCtx = await resolveManagedProxyContext();
+          const meta = MANAGED_PROVIDER_META["gemini"];
+          if (proxyCtx.enabled && meta?.managed && meta.proxyPath) {
+            const managedBaseUrl = `${proxyCtx.platformBaseUrl}${meta.proxyPath}`;
+            const managedModel = config.memory.embeddings.geminiModel;
+            const managedDimensions =
+              config.memory.embeddings.geminiDimensions ?? 3072;
+            const extras = geminiCacheExtras(config);
+            backends.push(
+              getCachedOrCreate(
+                "gemini",
+                managedModel,
+                () =>
+                  new GeminiEmbeddingBackend(
+                    proxyCtx.assistantApiKey,
+                    managedModel,
+                    {
+                      taskType: config.memory.embeddings.geminiTaskType,
+                      dimensions: managedDimensions,
+                      managedBaseUrl,
+                    },
+                  ),
+                [...extras, "managed"],
+              ),
+            );
+          } else {
+            // Check managed cache variant first, then non-managed, so a warm
+            // managed backend survives transient proxy-context blips.
+            const cached =
+              getCached("gemini", config.memory.embeddings.geminiModel, [
+                ...geminiCacheExtras(config),
+                "managed",
+              ]) ??
+              getCached(
+                "gemini",
+                config.memory.embeddings.geminiModel,
+                geminiCacheExtras(config),
+              );
+            if (cached) backends.push(cached);
+          }
         } else {
           // Preserve cached backend on transient credential-store failures.
           const cached = getCached(

package/src/memory/embedding-gemini.test.ts

@@ -253,4 +253,58 @@ describe("GeminiEmbeddingBackend", () => {
       expect(result[1]).toEqual([0.2, 0.4]);
     });
   });
+
+  describe("managed proxy transport", () => {
+    test("routes through managed proxy base URL when managedBaseUrl is set", async () => {
+      const backend = new GeminiEmbeddingBackend(
+        "ast-managed-key",
+        "gemini-embedding-2-preview",
+        {
+          managedBaseUrl:
+            "https://platform.example.com/v1/runtime-proxy/gemini",
+        },
+      );
+      await backend.embed(["hello"]);
+
+      expect(mockFetch).toHaveBeenCalledTimes(1);
+      const [url, init] = mockFetch.mock.calls[0] as [string, RequestInit];
+      expect(url).toBe(
+        "https://platform.example.com/v1/runtime-proxy/gemini/v1beta/models/gemini-embedding-2-preview:embedContent",
+      );
+      // Should NOT have key= query param
+      expect(url).not.toContain("key=");
+      // Should have Bearer auth header
+      const headers = init.headers as Record<string, string>;
+      expect(headers["Authorization"]).toBe("Bearer ast-managed-key");
+    });
+
+    test("uses direct Google API URL when managedBaseUrl is not set", async () => {
+      const backend = new GeminiEmbeddingBackend("direct-key", "test-model");
+      await backend.embed(["hello"]);
+
+      const [url, init] = mockFetch.mock.calls[0] as [string, RequestInit];
+      expect(url).toContain("generativelanguage.googleapis.com");
+      expect(url).toContain("key=direct-key");
+      // Should NOT have Authorization header
+      const headers = init.headers as Record<string, string>;
+      expect(headers["Authorization"]).toBeUndefined();
+    });
+
+    test("includes outputDimensionality with managed proxy", async () => {
+      const backend = new GeminiEmbeddingBackend(
+        "ast-managed-key",
+        "gemini-embedding-2-preview",
+        {
+          managedBaseUrl:
+            "https://platform.example.com/v1/runtime-proxy/gemini",
+          dimensions: 3072,
+        },
+      );
+      await backend.embed(["hello"]);
+
+      const [, init] = mockFetch.mock.calls[0] as [string, RequestInit];
+      const body = JSON.parse(init.body as string);
+      expect(body.outputDimensionality).toBe(3072);
+    });
+  });
 });

package/src/memory/embedding-gemini.ts

@@ -15,22 +15,27 @@ interface GeminiEmbedResponse {
   };
 }
 
+export interface GeminiEmbeddingOptions {
+  taskType?: EmbeddingTaskType;
+  dimensions?: number;
+  /** When set, routes requests through the managed proxy at this base URL. */
+  managedBaseUrl?: string;
+}
+
 export class GeminiEmbeddingBackend implements EmbeddingBackend {
   readonly provider = "gemini" as const;
   readonly model: string;
   private readonly apiKey: string;
   private readonly taskType?: EmbeddingTaskType;
   private readonly dimensions?: number;
+  private readonly managedBaseUrl?: string;
 
-  constructor(
-    apiKey: string,
-    model: string,
-    options?: { taskType?: EmbeddingTaskType; dimensions?: number },
-  ) {
+  constructor(apiKey: string, model: string, options?: GeminiEmbeddingOptions) {
     this.apiKey = apiKey;
     this.model = model;
     this.taskType = options?.taskType;
     this.dimensions = options?.dimensions;
+    this.managedBaseUrl = options?.managedBaseUrl;
   }
 
   async embed(
@@ -59,12 +64,18 @@ export class GeminiEmbeddingBackend implements EmbeddingBackend {
     if (this.taskType) body.taskType = this.taskType;
     if (this.dimensions) body.outputDimensionality = this.dimensions;
 
-    const url = `https://generativelanguage.googleapis.com/v1beta/models/${encodeURIComponent(
-      this.model,
-    )}:embedContent?key=${encodeURIComponent(this.apiKey)}`;
+    const url = this.managedBaseUrl
+      ? `${this.managedBaseUrl}/v1beta/models/${encodeURIComponent(this.model)}:embedContent`
+      : `https://generativelanguage.googleapis.com/v1beta/models/${encodeURIComponent(this.model)}:embedContent?key=${encodeURIComponent(this.apiKey)}`;
+    const headers: Record<string, string> = {
+      "Content-Type": "application/json",
+    };
+    if (this.managedBaseUrl) {
+      headers["Authorization"] = `Bearer ${this.apiKey}`;
+    }
     const response = await fetch(url, {
       method: "POST",
-      headers: { "Content-Type": "application/json" },
+      headers,
       body: JSON.stringify(body),
       signal: options?.signal,
     });

package/src/memory/embedding-local.ts

@@ -1,4 +1,10 @@
-import { existsSync, rmSync, unlinkSync, writeFileSync } from "node:fs";
+import {
+  existsSync,
+  readFileSync,
+  rmSync,
+  unlinkSync,
+  writeFileSync,
+} from "node:fs";
 import { join } from "node:path";
 
 import { getIsContainerized } from "../config/env-registry.js";
@@ -83,6 +89,8 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
     }
   >();
   private stdoutReaderActive = false;
+  private activeEmbeds = 0;
+  private disposeRequested = false;
 
   private readonly initGuard = new PromiseGuard<void>();
 
@@ -94,6 +102,9 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
     inputs: EmbeddingInput[],
     options?: EmbeddingRequestOptions,
   ): Promise<number[][]> {
+    if (this.disposeRequested) {
+      throw new Error("Local embedding backend is shutting down");
+    }
     if (inputs.length === 0) return [];
 
     const texts = inputs.map((i) => {
@@ -106,24 +117,30 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
     if (options?.signal?.aborted)
       throw new DOMException("Aborted", "AbortError");
 
-    await this.ensureInitialized();
-
-    const results: number[][] = [];
-    const batchSize = 32;
-    for (let i = 0; i < texts.length; i += batchSize) {
-      if (options?.signal?.aborted)
-        throw new DOMException("Aborted", "AbortError");
-      const batch = texts.slice(i, i + batchSize);
-      const response = await this.sendRequest(batch);
-      if (response.error) {
-        throw new Error(`Embedding worker error: ${response.error}`);
-      }
-      if (!response.vectors) {
-        throw new Error("Embedding worker returned no vectors");
+    this.activeEmbeds++;
+    try {
+      await this.ensureInitialized();
+
+      const results: number[][] = [];
+      const batchSize = 32;
+      for (let i = 0; i < texts.length; i += batchSize) {
+        if (options?.signal?.aborted)
+          throw new DOMException("Aborted", "AbortError");
+        const batch = texts.slice(i, i + batchSize);
+        const response = await this.sendRequest(batch);
+        if (response.error) {
+          throw new Error(`Embedding worker error: ${response.error}`);
+        }
+        if (!response.vectors) {
+          throw new Error("Embedding worker returned no vectors");
+        }
+        results.push(...response.vectors);
       }
-      results.push(...response.vectors);
+      return results;
+    } finally {
+      this.activeEmbeds--;
+      this.disposeIfIdle();
     }
-    return results;
   }
 
   private sendRequest(texts: string[]): Promise<WorkerResponse> {
@@ -148,6 +165,11 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
     await this.initGuard.run(() => this.initialize());
   }
 
+  dispose(): void {
+    this.disposeRequested = true;
+    this.disposeIfIdle();
+  }
+
   private async initialize(): Promise<void> {
     log.info({ model: this.model }, "Initializing local embedding backend");
 
@@ -199,13 +221,19 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
     const embeddingModelsDir = getEmbeddingModelsDir();
     const modelCacheDir = `${embeddingModelsDir}/model-cache`;
 
+    // Singleton guard: an orphaned embed worker from a previous daemon
+    // (e.g. one that crashed without cleanup) may still be running and
+    // holding the workspace's PID file. Detect and reclaim it before
+    // spawning so we never leave duplicate workers eating CPU/memory.
+    this.reclaimStaleWorker(workerPath);
+
     log.info(
       { bunPath, workerPath, model: this.model },
       "Spawning embedding worker process",
     );
 
     const proc = Bun.spawn({
-      cmd: [bunPath, workerPath, this.model, modelCacheDir],
+      cmd: [bunPath, "--smol", workerPath, this.model, modelCacheDir],
       stdin: "pipe",
       stdout: "pipe",
       stderr: "pipe",
@@ -255,6 +283,8 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
       { pid: proc.pid, model: this.model },
       "Embedding worker process started",
     );
+
+    this.disposeIfIdle();
   }
 
   private drainStderr(stderr: ReadableStream<Uint8Array>): void {
@@ -355,6 +385,7 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
         if (pending) {
           this.pendingRequests.delete(msg.id);
           pending.resolve(msg);
+          this.disposeIfIdle();
         }
       }
     }
@@ -425,4 +456,132 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
       // Best-effort
     }
   }
+
+  /** Read the PID from the on-disk PID file, or null if missing/invalid. */
+  private readPidFile(): number | null {
+    const path = this.getPidFilePath();
+    if (!existsSync(path)) return null;
+    try {
+      const pid = parseInt(readFileSync(path, "utf-8").trim(), 10);
+      return Number.isFinite(pid) && pid > 0 ? pid : null;
+    } catch {
+      return null;
+    }
+  }
+
+  /**
+   * Verify a PID belongs to this workspace's embed worker before sending
+   * signals — defends against PID reuse killing an unrelated process if the
+   * original worker exited and the OS recycled the PID.
+   *
+   * Matching `embed-worker` alone would also match a sibling assistant
+   * instance's worker (different VELLUM_WORKSPACE_DIR), so we match against
+   * the absolute worker script path, which lives under THIS workspace's
+   * embedding-models directory and is therefore unique per instance.
+   */
+  private isOurEmbedWorker(pid: number, workerPath: string): boolean {
+    try {
+      // `-ww` disables column-width truncation. Without it, macOS `ps` clips
+      // the command field to the terminal width, which can cut off the
+      // workerPath argument and cause this check to spuriously return false
+      // for genuine orphans. Same flag is used by daemon-control.ts:123 for
+      // exactly this reason.
+      const result = Bun.spawnSync({
+        cmd: ["ps", "-ww", "-p", String(pid), "-o", "command="],
+        stdout: "pipe",
+        stderr: "ignore",
+      });
+      if (result.exitCode !== 0) return false;
+      const cmd = new TextDecoder().decode(result.stdout).trim();
+      if (!cmd) return false;
+      return cmd.includes(workerPath);
+    } catch {
+      return false;
+    }
+  }
+
+  /**
+   * If a previous embed worker is still running for this workspace (orphaned
+   * by a crashed daemon, for example), terminate it before spawning a new one
+   * so we never end up with duplicate workers competing for the same workspace.
+   *
+   * Stale PID files (process no longer exists) are silently cleaned up.
+   * PIDs that have been recycled to unrelated processes — including embed
+   * workers belonging to *other* assistant instances — are left untouched.
+   */
+  private reclaimStaleWorker(workerPath: string): void {
+    const pid = this.readPidFile();
+    if (pid == null) return;
+
+    // Never signal ourselves — should not happen since the worker is a child
+    // process, but guard against logic bugs that would deadlock the daemon.
+    if (pid === process.pid) {
+      this.removePidFile();
+      return;
+    }
+
+    let isAlive = false;
+    try {
+      // Signal 0 just probes for liveness without delivering a signal.
+      process.kill(pid, 0);
+      isAlive = true;
+    } catch {
+      // ESRCH — no such process. PID file is stale.
+    }
+
+    if (!isAlive) {
+      log.info(
+        { pid, model: this.model },
+        "Removing stale embed worker PID file (process no longer exists)",
+      );
+      this.removePidFile();
+      return;
+    }
+
+    if (!this.isOurEmbedWorker(pid, workerPath)) {
+      // PID points to something that isn't this workspace's embed worker —
+      // either an unrelated process (PID reuse after the original worker
+      // exited) or another assistant instance's worker. Either way, don't
+      // signal it; just drop the stale file so the new worker can claim it.
+      log.warn(
+        { pid, model: this.model },
+        "PID file points to a process that is not this workspace's embed worker; clearing without killing",
+      );
+      this.removePidFile();
+      return;
+    }
+
+    log.warn(
+      { pid, model: this.model },
+      "Found orphaned embed worker from a previous daemon, terminating it",
+    );
+    try {
+      process.kill(pid, "SIGTERM");
+    } catch {
+      // Race: it exited between the liveness check and the kill — fine.
+    }
+    this.removePidFile();
+  }
+
+  private disposeIfIdle(): void {
+    if (!this.disposeRequested) return;
+    if (this.activeEmbeds > 0) return;
+    if (this.pendingRequests.size > 0) return;
+    if (this.readyResolve || this.readyReject) return;
+
+    const proc = this.workerProc;
+    this.workerProc = null;
+    this.stdoutReaderActive = false;
+    this.stdoutBuffer = "";
+    this.initGuard.reset();
+    this.removePidFile();
+
+    if (!proc) return;
+
+    try {
+      proc.kill();
+    } catch {
+      // Worker may already be exiting
+    }
+  }
 }

package/src/memory/graph/capability-seed.ts

@@ -127,9 +127,7 @@ export function seedSkillGraphNodes(): void {
       // skip catalog-based pruning to avoid incorrectly marking valid
       // uninstalled catalog nodes as gone. But still prune locally disabled
       // skills so stale capability nodes don't linger after cold start.
-      log.info(
-        "Catalog cache is cold — pruning only locally disabled skills",
-      );
+      log.info("Catalog cache is cold — pruning only locally disabled skills");
       const disabled = resolved.filter((r) => r.state !== "enabled");
       for (const { summary } of disabled) {
         deleteSkillCapabilityNode(summary.id);
@@ -158,9 +156,9 @@ export function seedSkillGraphNodes(): void {
  * Seed graph nodes for all CLI commands.
  * Prunes stale nodes whose commands are no longer registered.
  */
-export function seedCliGraphNodes(): void {
+export async function seedCliGraphNodes(): Promise<void> {
   try {
-    const program = buildCliProgram();
+    const program = await buildCliProgram();
 
     const seenKeys = new Set<string>();
     for (const cmd of program.commands) {