diff --git a/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts b/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts index 358c3ccd315..44bd508d3b7 100644 --- a/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts +++ b/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts @@ -57,11 +57,8 @@ function magnitude(values: number[]) { } let buildGeminiEmbeddingRequest: typeof import("./embeddings-gemini.js").buildGeminiEmbeddingRequest; -let buildGeminiTextEmbeddingRequest: typeof import("./embeddings-gemini.js").buildGeminiTextEmbeddingRequest; let createGeminiEmbeddingProvider: typeof import("./embeddings-gemini.js").createGeminiEmbeddingProvider; let DEFAULT_GEMINI_EMBEDDING_MODEL: typeof import("./embeddings-gemini.js").DEFAULT_GEMINI_EMBEDDING_MODEL; -let GEMINI_EMBEDDING_2_MODELS: typeof import("./embeddings-gemini.js").GEMINI_EMBEDDING_2_MODELS; -let isGeminiEmbedding2Model: typeof import("./embeddings-gemini.js").isGeminiEmbedding2Model; let normalizeGeminiModel: typeof import("./embeddings-gemini.js").normalizeGeminiModel; let resolveGeminiOutputDimensionality: typeof import("./embeddings-gemini.js").resolveGeminiOutputDimensionality; @@ -69,11 +66,8 @@ beforeAll(async () => { vi.doUnmock("undici"); ({ buildGeminiEmbeddingRequest, - buildGeminiTextEmbeddingRequest, createGeminiEmbeddingProvider, DEFAULT_GEMINI_EMBEDDING_MODEL, - GEMINI_EMBEDDING_2_MODELS, - isGeminiEmbedding2Model, normalizeGeminiModel, resolveGeminiOutputDimensionality, } = await import("./embeddings-gemini.js")); @@ -124,26 +118,8 @@ function expectNormalizedThreeFourVector(embedding: number[]) { expect(magnitude(embedding)).toBeCloseTo(1, 5); } -describe("buildGeminiTextEmbeddingRequest", () => { - it("builds a text embedding request with optional model and dimensions", () => { - expect( - buildGeminiTextEmbeddingRequest({ - text: "hello", - taskType: "RETRIEVAL_DOCUMENT", - modelPath: "models/gemini-embedding-2-preview", - outputDimensionality: 1536, - }), - ).toEqual({ - model: "models/gemini-embedding-2-preview", - content: { parts: [{ text: "hello" }] }, - taskType: "RETRIEVAL_DOCUMENT", - outputDimensionality: 1536, - }); - }); -}); - -describe("buildGeminiEmbeddingRequest", () => { - it("builds a multimodal request from structured input parts", () => { +describe("package Gemini embedding provider smoke", () => { + it("builds multimodal v2 requests and resolves dimensions", () => { expect( buildGeminiEmbeddingRequest({ input: { @@ -168,62 +144,14 @@ describe("buildGeminiEmbeddingRequest", () => { taskType: "RETRIEVAL_DOCUMENT", outputDimensionality: 1536, }); - }); -}); - -// ---------- Model detection ---------- - -describe("isGeminiEmbedding2Model", () => { - it("returns true for gemini-embedding-2-preview", () => { - expect(isGeminiEmbedding2Model("gemini-embedding-2-preview")).toBe(true); - }); - - it("returns false for gemini-embedding-001", () => { - expect(isGeminiEmbedding2Model("gemini-embedding-001")).toBe(false); - }); - - it("returns false for text-embedding-004", () => { - expect(isGeminiEmbedding2Model("text-embedding-004")).toBe(false); - }); -}); - -describe("GEMINI_EMBEDDING_2_MODELS", () => { - it("contains gemini-embedding-2-preview", () => { - expect(GEMINI_EMBEDDING_2_MODELS.has("gemini-embedding-2-preview")).toBe(true); - }); -}); - -// ---------- Dimension resolution ---------- - -describe("resolveGeminiOutputDimensionality", () => { - it("returns undefined for non-v2 models", () => { expect(resolveGeminiOutputDimensionality("gemini-embedding-001")).toBeUndefined(); - expect(resolveGeminiOutputDimensionality("text-embedding-004")).toBeUndefined(); - }); - - it("returns 3072 by default for v2 models", () => { expect(resolveGeminiOutputDimensionality("gemini-embedding-2-preview")).toBe(3072); - }); - - it("accepts valid dimension values", () => { expect(resolveGeminiOutputDimensionality("gemini-embedding-2-preview", 768)).toBe(768); - expect(resolveGeminiOutputDimensionality("gemini-embedding-2-preview", 1536)).toBe(1536); - expect(resolveGeminiOutputDimensionality("gemini-embedding-2-preview", 3072)).toBe(3072); - }); - - it("throws for invalid dimension values", () => { expect(() => resolveGeminiOutputDimensionality("gemini-embedding-2-preview", 512)).toThrow( /Invalid outputDimensionality 512/, ); - expect(() => resolveGeminiOutputDimensionality("gemini-embedding-2-preview", 1024)).toThrow( - /Valid values: 768, 1536, 3072/, - ); }); -}); -// ---------- Provider behavior smoke coverage ---------- - -describe("gemini embedding provider", () => { it("handles legacy and v2 request/response behavior", async () => { const legacyFetch = createGeminiBatchFetchMock(2); const legacyProvider = await createProviderWithFetch(legacyFetch, { @@ -271,12 +199,8 @@ describe("gemini embedding provider", () => { expect.objectContaining({ outputDimensionality: 768 }), ]); }); -}); -// ---------- Model normalization ---------- - -describe("gemini model normalization", () => { - it("normalizes known model prefixes and default model", () => { + it("normalizes known model prefixes and the default model", () => { expect(normalizeGeminiModel("models/gemini-embedding-2-preview")).toBe( "gemini-embedding-2-preview", ); diff --git a/packages/memory-host-sdk/src/host/embeddings.test.ts b/packages/memory-host-sdk/src/host/embeddings.test.ts index 732723fda09..1efb4e05e21 100644 --- a/packages/memory-host-sdk/src/host/embeddings.test.ts +++ b/packages/memory-host-sdk/src/host/embeddings.test.ts @@ -1,20 +1,14 @@ -import { setTimeout as sleep } from "node:timers/promises"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import * as authModule from "../../../../src/agents/model-auth.js"; -import { DEFAULT_GEMINI_EMBEDDING_MODEL } from "./embeddings-gemini.js"; import { createEmbeddingProvider, DEFAULT_LOCAL_MODEL } from "./embeddings.js"; import * as nodeLlamaModule from "./node-llama.js"; import { mockPublicPinnedHostname } from "./test-helpers/ssrf.js"; -const { createOllamaEmbeddingProviderMock } = vi.hoisted(() => ({ - createOllamaEmbeddingProviderMock: vi.fn(async () => { - throw new Error("Unexpected ollama provider in embeddings.test.ts"); - }), -})); - -const { hasAwsCredentialsMock } = vi.hoisted(() => ({ - hasAwsCredentialsMock: vi.fn(async () => false), -})); +vi.mock("../../../../src/agents/model-auth.js", async () => { + const { createModelAuthMockModule } = + await import("../../../../src/test-utils/model-auth-mock.js"); + return createModelAuthMockModule(); +}); vi.mock("../../../../src/infra/net/fetch-guard.js", () => ({ fetchWithSsrFGuard: async (params: { @@ -35,20 +29,7 @@ vi.mock("../../../../src/infra/net/fetch-guard.js", () => ({ }, })); -vi.mock("./embeddings-ollama.js", () => ({ - createOllamaEmbeddingProvider: createOllamaEmbeddingProviderMock, -})); - -vi.mock("./embeddings-bedrock.js", async () => { - const actual = - await vi.importActual("./embeddings-bedrock.js"); - return { - ...actual, - hasAwsCredentials: hasAwsCredentialsMock, - }; -}); - -const createFetchMock = () => +const createEmbeddingDataFetchMock = () => vi.fn(async (_input?: unknown, _init?: unknown) => ({ ok: true, status: 200, @@ -62,6 +43,16 @@ const createGeminiFetchMock = () => json: async () => ({ embedding: { values: [1, 2, 3] } }), })); +beforeEach(() => { + vi.spyOn(authModule, "resolveApiKeyForProvider"); + vi.spyOn(nodeLlamaModule, "importNodeLlamaCpp"); +}); + +afterEach(() => { + vi.resetAllMocks(); + vi.unstubAllGlobals(); +}); + function installFetchMock(fetchMock: typeof globalThis.fetch) { vi.stubGlobal("fetch", fetchMock); } @@ -71,22 +62,6 @@ function readFirstFetchRequest(fetchMock: { mock: { calls: unknown[][] } }) { return { url, init: init as RequestInit | undefined }; } -type ResolvedProviderAuth = Awaited>; - -beforeEach(() => { - vi.spyOn(authModule, "resolveApiKeyForProvider"); - vi.spyOn(nodeLlamaModule, "importNodeLlamaCpp"); -}); - -beforeEach(() => { - vi.useRealTimers(); -}); - -afterEach(() => { - vi.resetAllMocks(); - vi.unstubAllGlobals(); -}); - function requireProvider(result: Awaited>) { if (!result.provider) { throw new Error("Expected embedding provider"); @@ -102,183 +77,46 @@ function mockResolvedProviderKey(apiKey = "provider-key") { }); } -function mockMissingLocalEmbeddingDependency() { - vi.mocked(nodeLlamaModule.importNodeLlamaCpp).mockRejectedValue( - Object.assign(new Error("Cannot find package 'node-llama-cpp'"), { - code: "ERR_MODULE_NOT_FOUND", - }), - ); -} - -function createLocalProvider(options?: { fallback?: "none" | "openai" }) { - return createEmbeddingProvider({ - config: {} as never, - provider: "local", - model: "text-embedding-3-small", - fallback: options?.fallback ?? "none", - }); -} - -function expectAutoSelectedProvider( - result: Awaited>, - expectedId: "openai" | "gemini" | "mistral", -) { - expect(result.requestedProvider).toBe("auto"); - const provider = requireProvider(result); - expect(provider.id).toBe(expectedId); - return provider; -} - -function createAutoProvider(model = "") { - return createEmbeddingProvider({ - config: {} as never, - provider: "auto", - model, - fallback: "none", - }); -} - -describe("embedding provider remote overrides", () => { - it("uses remote baseUrl/apiKey and merges headers", async () => { - const fetchMock = createFetchMock(); +describe("package embedding provider smoke", () => { + it("uses remote OpenAI baseUrl/apiKey and merges headers", async () => { + const fetchMock = createEmbeddingDataFetchMock(); installFetchMock(fetchMock as unknown as typeof globalThis.fetch); mockPublicPinnedHostname(); mockResolvedProviderKey("provider-key"); - const cfg = { - models: { - providers: { - openai: { - baseUrl: "https://api.openai.com/v1", - headers: { - "X-Provider": "p", - "X-Shared": "provider", + const result = await createEmbeddingProvider({ + config: { + models: { + providers: { + openai: { + baseUrl: "https://api.openai.com/v1", + headers: { "X-Provider": "p", "X-Shared": "provider" }, }, }, }, - }, - }; - - const result = await createEmbeddingProvider({ - config: cfg as never, + } as never, provider: "openai", remote: { baseUrl: "https://example.com/v1", apiKey: " remote-key ", - headers: { - "X-Shared": "remote", - "X-Remote": "r", - }, + headers: { "X-Shared": "remote", "X-Remote": "r" }, }, model: "text-embedding-3-small", fallback: "openai", }); - const provider = requireProvider(result); - await provider.embedQuery("hello"); + await requireProvider(result).embedQuery("hello"); expect(authModule.resolveApiKeyForProvider).not.toHaveBeenCalled(); - const url = fetchMock.mock.calls[0]?.[0]; - const init = fetchMock.mock.calls[0]?.[1] as RequestInit | undefined; + const { url, init } = readFirstFetchRequest(fetchMock); expect(url).toBe("https://example.com/v1/embeddings"); const headers = (init?.headers ?? {}) as Record; expect(headers.Authorization).toBe("Bearer remote-key"); - expect(headers["Content-Type"]).toBe("application/json"); expect(headers["X-Provider"]).toBe("p"); expect(headers["X-Shared"]).toBe("remote"); expect(headers["X-Remote"]).toBe("r"); }); - it("falls back to resolved api key when remote apiKey is blank", async () => { - const fetchMock = createFetchMock(); - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - mockPublicPinnedHostname(); - mockResolvedProviderKey("provider-key"); - - const cfg = { - models: { - providers: { - openai: { - baseUrl: "https://api.openai.com/v1", - }, - }, - }, - }; - - const result = await createEmbeddingProvider({ - config: cfg as never, - provider: "openai", - remote: { - baseUrl: "https://example.com/v1", - apiKey: " ", - }, - model: "text-embedding-3-small", - fallback: "openai", - }); - - const provider = requireProvider(result); - await provider.embedQuery("hello"); - - expect(authModule.resolveApiKeyForProvider).toHaveBeenCalledTimes(1); - const init = fetchMock.mock.calls[0]?.[1] as RequestInit | undefined; - const headers = (init?.headers as Record) ?? {}; - expect(headers.Authorization).toBe("Bearer provider-key"); - }); - - it("builds Gemini embeddings requests with api key header", async () => { - const fetchMock = createGeminiFetchMock(); - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - mockPublicPinnedHostname(); - mockResolvedProviderKey("provider-key"); - - const cfg = { - models: { - providers: { - google: { - baseUrl: "https://generativelanguage.googleapis.com/v1beta", - }, - }, - }, - }; - - const result = await createEmbeddingProvider({ - config: cfg as never, - provider: "gemini", - remote: { - apiKey: "gemini-key", - }, - model: "text-embedding-004", - fallback: "openai", - }); - - const provider = requireProvider(result); - await provider.embedQuery("hello"); - - const { url, init } = readFirstFetchRequest(fetchMock); - expect(url).toBe( - "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:embedContent", - ); - const headers = (init?.headers ?? {}) as Record; - expect(headers["x-goog-api-key"]).toBe("gemini-key"); - expect(headers["Content-Type"]).toBe("application/json"); - }); - - it("fails fast when Gemini remote apiKey is an unresolved SecretRef", async () => { - vi.stubEnv("GEMINI_API_KEY", ""); - - await expect( - createEmbeddingProvider({ - config: {} as never, - provider: "gemini", - remote: { - apiKey: { source: "env", provider: "default", id: "GEMINI_API_KEY" }, - }, - model: "text-embedding-004", - fallback: "openai", - }), - ).rejects.toThrow(/agents\.\*\.memorySearch\.remote\.apiKey:/i); - }); - it("uses GEMINI_API_KEY env indirection for Gemini remote apiKey", async () => { const fetchMock = createGeminiFetchMock(); installFetchMock(fetchMock as unknown as typeof globalThis.fetch); @@ -295,341 +133,26 @@ describe("embedding provider remote overrides", () => { fallback: "openai", }); - const provider = requireProvider(result); - await provider.embedQuery("hello"); + await requireProvider(result).embedQuery("hello"); const { init } = readFirstFetchRequest(fetchMock); const headers = (init?.headers ?? {}) as Record; expect(headers["x-goog-api-key"]).toBe("env-gemini-key"); }); - it("builds Mistral embeddings requests with bearer auth", async () => { - const fetchMock = createFetchMock(); - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - mockPublicPinnedHostname(); - mockResolvedProviderKey("provider-key"); - - const cfg = { - models: { - providers: { - mistral: { - baseUrl: "https://api.mistral.ai/v1", - }, - }, - }, - }; - - const result = await createEmbeddingProvider({ - config: cfg as never, - provider: "mistral", - remote: { - apiKey: "mistral-key", // pragma: allowlist secret - }, - model: "mistral/mistral-embed", - fallback: "none", - }); - - const provider = requireProvider(result); - await provider.embedQuery("hello"); - - const { url, init } = readFirstFetchRequest(fetchMock); - expect(url).toBe("https://api.mistral.ai/v1/embeddings"); - const headers = (init?.headers ?? {}) as Record; - expect(headers.Authorization).toBe("Bearer mistral-key"); - const payload = JSON.parse((init?.body as string | undefined) ?? "{}") as { model?: string }; - expect(payload.model).toBe("mistral-embed"); - }); -}); - -describe("embedding provider auto selection", () => { - it("keeps explicit model when openai is selected", async () => { - const fetchMock = vi.fn(async (_input?: unknown, _init?: unknown) => ({ - ok: true, - status: 200, - json: async () => ({ data: [{ embedding: [1, 2, 3] }] }), - })); - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - mockPublicPinnedHostname(); - vi.mocked(authModule.resolveApiKeyForProvider).mockImplementation(async ({ provider }) => { - if (provider === "openai") { - return { apiKey: "openai-key", source: "env: OPENAI_API_KEY", mode: "api-key" }; - } - throw new Error(`Unexpected provider ${provider}`); - }); - - const result = await createEmbeddingProvider({ - config: {} as never, - provider: "auto", - model: "text-embedding-3-small", - fallback: "none", - }); - - expect(result.requestedProvider).toBe("auto"); - const provider = requireProvider(result); - expect(provider.id).toBe("openai"); - await provider.embedQuery("hello"); - const url = fetchMock.mock.calls[0]?.[0]; - const init = fetchMock.mock.calls[0]?.[1] as RequestInit | undefined; - expect(url).toBe("https://api.openai.com/v1/embeddings"); - const payload = JSON.parse(init?.body as string) as { model?: string }; - expect(payload.model).toBe("text-embedding-3-small"); - }); - - it("selects the first available remote provider in auto mode", async () => { - const cases: Array<{ - name: string; - expectedProvider: "openai" | "gemini" | "mistral"; - fetchMockFactory: typeof createFetchMock | typeof createGeminiFetchMock; - resolveApiKey: (provider: string) => ResolvedProviderAuth; - expectedUrl: string; - }> = [ - { - name: "openai first", - expectedProvider: "openai" as const, - fetchMockFactory: createFetchMock, - resolveApiKey(provider: string): ResolvedProviderAuth { - if (provider === "openai") { - return { apiKey: "openai-key", source: "env: OPENAI_API_KEY", mode: "api-key" }; - } - throw new Error(`No API key found for provider "${provider}".`); - }, - expectedUrl: "https://api.openai.com/v1/embeddings", - }, - { - name: "gemini fallback", - expectedProvider: "gemini" as const, - fetchMockFactory: createGeminiFetchMock, - resolveApiKey(provider: string): ResolvedProviderAuth { - if (provider === "openai") { - throw new Error('No API key found for provider "openai".'); - } - if (provider === "google") { - return { - apiKey: "gemini-key", - source: "env: GEMINI_API_KEY", - mode: "api-key" as const, - }; - } - throw new Error(`Unexpected provider ${provider}`); - }, - expectedUrl: `https://generativelanguage.googleapis.com/v1beta/models/${DEFAULT_GEMINI_EMBEDDING_MODEL}:embedContent`, - }, - { - name: "mistral after earlier misses", - expectedProvider: "mistral" as const, - fetchMockFactory: createFetchMock, - resolveApiKey(provider: string): ResolvedProviderAuth { - if (provider === "mistral") { - return { - apiKey: "mistral-key", - source: "env: MISTRAL_API_KEY", - mode: "api-key" as const, - }; - } - throw new Error(`No API key found for provider "${provider}".`); - }, - expectedUrl: "https://api.mistral.ai/v1/embeddings", - }, - ]; - - for (const testCase of cases) { - vi.resetAllMocks(); - vi.unstubAllGlobals(); - const fetchMock = testCase.fetchMockFactory(); - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - mockPublicPinnedHostname(); - vi.mocked(authModule.resolveApiKeyForProvider).mockImplementation(async ({ provider }) => - testCase.resolveApiKey(provider), - ); - - const result = await createAutoProvider(); - const provider = expectAutoSelectedProvider(result, testCase.expectedProvider); - await provider.embedQuery("hello"); - const [url] = fetchMock.mock.calls[0] ?? []; - expect(url, testCase.name).toBe(testCase.expectedUrl); - } - }); -}); - -describe("embedding provider local fallback", () => { - it("falls back to openai when node-llama-cpp is missing", async () => { - mockMissingLocalEmbeddingDependency(); - - const fetchMock = createFetchMock(); - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - - mockResolvedProviderKey("provider-key"); - - const result = await createLocalProvider({ fallback: "openai" }); - - const provider = requireProvider(result); - expect(provider.id).toBe("openai"); - expect(result.fallbackFrom).toBe("local"); - expect(result.fallbackReason).toContain("node-llama-cpp"); - }); - - it("throws a helpful error when local is requested and fallback is none", async () => { - mockMissingLocalEmbeddingDependency(); - await expect(createLocalProvider()).rejects.toThrow(/optional dependency node-llama-cpp/i); - }); - - it("mentions every remote provider in local setup guidance", async () => { - mockMissingLocalEmbeddingDependency(); - await expect(createLocalProvider()).rejects.toThrow(/provider = "gemini"/i); - await expect(createLocalProvider()).rejects.toThrow(/provider = "mistral"/i); - }); -}); - -describe("local embedding normalization", () => { - async function createLocalProviderForTest() { - return createEmbeddingProvider({ - config: {} as never, - provider: "local", - model: "", - fallback: "none", - }); - } - - function mockSingleLocalEmbeddingVector( - vector: number[], - resolveModelFile: (modelPath: string, modelDirectory?: string) => Promise = async () => - "/fake/model.gguf", - ): void { + it("normalizes local embeddings and resolves the default local model", async () => { + const resolveModelFileMock = vi.fn(async () => "/fake/model.gguf"); vi.mocked(nodeLlamaModule.importNodeLlamaCpp).mockResolvedValue({ getLlama: async () => ({ loadModel: vi.fn().mockResolvedValue({ createEmbeddingContext: vi.fn().mockResolvedValue({ getEmbeddingFor: vi.fn().mockResolvedValue({ - vector: new Float32Array(vector), + vector: new Float32Array([2.35, 3.45, 0.63, 4.3]), }), }), }), }), - resolveModelFile, - LlamaLogLevel: { error: 0 }, - } as never); - } - - it("normalizes local embeddings to magnitude ~1.0", async () => { - const unnormalizedVector = [2.35, 3.45, 0.63, 4.3, 1.2, 5.1, 2.8, 3.9]; - const resolveModelFileMock = vi.fn(async () => "/fake/model.gguf"); - - mockSingleLocalEmbeddingVector(unnormalizedVector, resolveModelFileMock); - - const result = await createLocalProviderForTest(); - - const provider = requireProvider(result); - const embedding = await provider.embedQuery("test query"); - - const magnitude = Math.sqrt(embedding.reduce((sum, x) => sum + x * x, 0)); - - expect(magnitude).toBeCloseTo(1.0, 5); - expect(resolveModelFileMock).toHaveBeenCalledWith(DEFAULT_LOCAL_MODEL, undefined); - }); - - it("handles zero vector without division by zero", async () => { - const zeroVector = [0, 0, 0, 0]; - - mockSingleLocalEmbeddingVector(zeroVector); - - const result = await createLocalProviderForTest(); - - const provider = requireProvider(result); - const embedding = await provider.embedQuery("test"); - - expect(embedding).toEqual([0, 0, 0, 0]); - expect(embedding.every((value) => Number.isFinite(value))).toBe(true); - }); - - it("sanitizes non-finite values before normalization", async () => { - const nonFiniteVector = [1, Number.NaN, Number.POSITIVE_INFINITY, Number.NEGATIVE_INFINITY]; - - mockSingleLocalEmbeddingVector(nonFiniteVector); - - const result = await createLocalProviderForTest(); - - const provider = requireProvider(result); - const embedding = await provider.embedQuery("test"); - - expect(embedding).toEqual([1, 0, 0, 0]); - expect(embedding.every((value) => Number.isFinite(value))).toBe(true); - }); - - it("normalizes batch embeddings to magnitude ~1.0", async () => { - const unnormalizedVectors = [ - [2.35, 3.45, 0.63, 4.3], - [10.0, 0.0, 0.0, 0.0], - [1.0, 1.0, 1.0, 1.0], - ]; - - vi.mocked(nodeLlamaModule.importNodeLlamaCpp).mockResolvedValue({ - getLlama: async () => ({ - loadModel: vi.fn().mockResolvedValue({ - createEmbeddingContext: vi.fn().mockResolvedValue({ - getEmbeddingFor: vi - .fn() - .mockResolvedValueOnce({ vector: new Float32Array(unnormalizedVectors[0]) }) - .mockResolvedValueOnce({ vector: new Float32Array(unnormalizedVectors[1]) }) - .mockResolvedValueOnce({ vector: new Float32Array(unnormalizedVectors[2]) }), - }), - }), - }), - resolveModelFile: async () => "/fake/model.gguf", - LlamaLogLevel: { error: 0 }, - } as never); - - const result = await createLocalProviderForTest(); - - const provider = requireProvider(result); - const embeddings = await provider.embedBatch(["text1", "text2", "text3"]); - - for (const embedding of embeddings) { - const magnitude = Math.sqrt(embedding.reduce((sum, x) => sum + x * x, 0)); - expect(magnitude).toBeCloseTo(1.0, 5); - } - }); -}); - -describe("local embedding ensureContext concurrency", () => { - async function setupLocalProviderWithMockedInit(params?: { - initializationDelayMs?: number; - failFirstGetLlama?: boolean; - }) { - const getLlamaSpy = vi.fn(); - const loadModelSpy = vi.fn(); - const createContextSpy = vi.fn(); - let shouldFail = params?.failFirstGetLlama ?? false; - - vi.spyOn(nodeLlamaModule, "importNodeLlamaCpp").mockResolvedValue({ - getLlama: async (...args: unknown[]) => { - getLlamaSpy(...args); - if (shouldFail) { - shouldFail = false; - throw new Error("transient init failure"); - } - if (params?.initializationDelayMs) { - await sleep(params.initializationDelayMs); - } - return { - loadModel: async (...modelArgs: unknown[]) => { - loadModelSpy(...modelArgs); - if (params?.initializationDelayMs) { - await sleep(params.initializationDelayMs); - } - return { - createEmbeddingContext: async () => { - createContextSpy(); - return { - getEmbeddingFor: vi.fn().mockResolvedValue({ - vector: new Float32Array([1, 0, 0, 0]), - }), - }; - }, - }; - }, - }; - }, - resolveModelFile: async () => "/fake/model.gguf", + resolveModelFile: resolveModelFileMock, LlamaLogLevel: { error: 0 }, } as never); @@ -640,128 +163,27 @@ describe("local embedding ensureContext concurrency", () => { fallback: "none", }); - return { - provider: requireProvider(result), - getLlamaSpy, - loadModelSpy, - createContextSpy, - }; - } - - it("loads the model only once when embedBatch is called concurrently", async () => { - const { provider, getLlamaSpy, loadModelSpy, createContextSpy } = - await setupLocalProviderWithMockedInit({ - initializationDelayMs: 5, - }); - - const results = await Promise.all([ - provider.embedBatch(["text1"]), - provider.embedBatch(["text2"]), - provider.embedBatch(["text3"]), - provider.embedBatch(["text4"]), - ]); - - expect(results).toHaveLength(4); - for (const embeddings of results) { - expect(embeddings).toHaveLength(1); - expect(embeddings[0]).toHaveLength(4); - } - - expect(getLlamaSpy).toHaveBeenCalledTimes(1); - expect(loadModelSpy).toHaveBeenCalledTimes(1); - expect(createContextSpy).toHaveBeenCalledTimes(1); + const embedding = await requireProvider(result).embedQuery("test query"); + const magnitude = Math.sqrt(embedding.reduce((sum, value) => sum + value * value, 0)); + expect(magnitude).toBeCloseTo(1, 5); + expect(resolveModelFileMock).toHaveBeenCalledWith(DEFAULT_LOCAL_MODEL, undefined); }); - it("retries initialization after a transient ensureContext failure", async () => { - const { provider, getLlamaSpy, loadModelSpy, createContextSpy } = - await setupLocalProviderWithMockedInit({ - failFirstGetLlama: true, - }); - - await expect(provider.embedBatch(["first"])).rejects.toThrow("transient init failure"); - - const recovered = await provider.embedBatch(["second"]); - expect(recovered).toHaveLength(1); - expect(recovered[0]).toHaveLength(4); - - expect(getLlamaSpy).toHaveBeenCalledTimes(2); - expect(loadModelSpy).toHaveBeenCalledTimes(1); - expect(createContextSpy).toHaveBeenCalledTimes(1); - }); - - it("shares initialization when embedQuery and embedBatch start concurrently", async () => { - const { provider, getLlamaSpy, loadModelSpy, createContextSpy } = - await setupLocalProviderWithMockedInit({ - initializationDelayMs: 5, - }); - - const [queryA, batch, queryB] = await Promise.all([ - provider.embedQuery("query-a"), - provider.embedBatch(["batch-a", "batch-b"]), - provider.embedQuery("query-b"), - ]); - - expect(queryA).toHaveLength(4); - expect(batch).toHaveLength(2); - expect(queryB).toHaveLength(4); - expect(batch[0]).toHaveLength(4); - expect(batch[1]).toHaveLength(4); - - expect(getLlamaSpy).toHaveBeenCalledTimes(1); - expect(loadModelSpy).toHaveBeenCalledTimes(1); - expect(createContextSpy).toHaveBeenCalledTimes(1); - }); -}); - -describe("FTS-only fallback when no provider available", () => { - it("returns null provider when all requested auth paths fail", async () => { + it("returns null provider when explicit primary and fallback auth paths fail", async () => { vi.mocked(authModule.resolveApiKeyForProvider).mockRejectedValue( new Error("No API key found for provider"), ); - for (const testCase of [ - { - name: "auto mode", - options: { - config: {} as never, - provider: "auto" as const, - model: "", - fallback: "none" as const, - }, - requestedProvider: "auto", - fallbackFrom: undefined, - reasonIncludes: "No API key", - }, - { - name: "explicit provider only", - options: { - config: {} as never, - provider: "openai" as const, - model: "text-embedding-3-small", - fallback: "none" as const, - }, - requestedProvider: "openai", - fallbackFrom: undefined, - reasonIncludes: "No API key", - }, - { - name: "primary and fallback", - options: { - config: {} as never, - provider: "openai" as const, - model: "text-embedding-3-small", - fallback: "gemini" as const, - }, - requestedProvider: "openai", - fallbackFrom: "openai", - reasonIncludes: "Fallback to gemini failed", - }, - ]) { - const result = await createEmbeddingProvider(testCase.options); - expect(result.provider, testCase.name).toBeNull(); - expect(result.requestedProvider, testCase.name).toBe(testCase.requestedProvider); - expect(result.fallbackFrom, testCase.name).toBe(testCase.fallbackFrom); - expect(result.providerUnavailableReason, testCase.name).toContain(testCase.reasonIncludes); - } + const result = await createEmbeddingProvider({ + config: {} as never, + provider: "openai", + model: "text-embedding-3-small", + fallback: "gemini", + }); + + expect(result.provider).toBeNull(); + expect(result.requestedProvider).toBe("openai"); + expect(result.fallbackFrom).toBe("openai"); + expect(result.providerUnavailableReason).toContain("Fallback to gemini failed"); }); }); diff --git a/src/context-engine/context-engine.test.ts b/src/context-engine/context-engine.test.ts index d99ade872f1..df028042650 100644 --- a/src/context-engine/context-engine.test.ts +++ b/src/context-engine/context-engine.test.ts @@ -532,19 +532,6 @@ describe("Registry tests", () => { existingOwner: "core", }); }); - - it("shares registered engines across duplicate module copies", async () => { - const registryUrl = new URL("./registry.ts", import.meta.url).href; - const suffix = Date.now().toString(36); - const first = await import(/* @vite-ignore */ `${registryUrl}?copy=${suffix}-a`); - const second = await import(/* @vite-ignore */ `${registryUrl}?copy=${suffix}-b`); - - const engineId = `dup-copy-${suffix}`; - const factory = () => new MockContextEngine(); - first.registerContextEngine(engineId, factory); - - expect(second.getContextEngineFactory(engineId)).toBe(factory); - }); }); // ═══════════════════════════════════════════════════════════════════════════ @@ -1019,34 +1006,15 @@ describe("Initialization guard", () => { // ═══════════════════════════════════════════════════════════════════════════ describe("Bundle chunk isolation (#40096)", () => { - it("Symbol.for key is stable across independently loaded modules", async () => { - // Simulate two distinct bundle chunks by loading the registry module - // twice with different query strings (forces separate module instances - // in Vite/esbuild but shares globalThis). + it("shares registrations and resolves engines across independently loaded chunks", async () => { const ts = Date.now().toString(36); const registryUrl = new URL("./registry.ts", import.meta.url).href; const chunkA = await import(/* @vite-ignore */ `${registryUrl}?chunk=a-${ts}`); const chunkB = await import(/* @vite-ignore */ `${registryUrl}?chunk=b-${ts}`); - // Chunk A registers an engine const engineId = `cross-chunk-${ts}`; - chunkA.registerContextEngine(engineId, () => new MockContextEngine()); - - // Chunk B must see it - expect(chunkB.getContextEngineFactory(engineId)).toBeDefined(); - expect(chunkB.listContextEngineIds()).toContain(engineId); - }); - - it("resolveContextEngine from chunk B finds engine registered in chunk A", async () => { - const ts = Date.now().toString(36); - const registryUrl = new URL("./registry.ts", import.meta.url).href; - - const chunkA = await import(/* @vite-ignore */ `${registryUrl}?chunk=resolve-a-${ts}`); - const chunkB = await import(/* @vite-ignore */ `${registryUrl}?chunk=resolve-b-${ts}`); - - const engineId = `resolve-cross-${ts}`; - chunkA.registerContextEngine(engineId, () => ({ + const factory = () => ({ info: { id: engineId, name: "Cross-chunk Engine", version: "0.0.1" }, async ingest() { return { ingested: true }; @@ -1057,9 +1025,11 @@ describe("Bundle chunk isolation (#40096)", () => { async compact() { return { ok: true, compacted: false }; }, - })); + }); + chunkA.registerContextEngine(engineId, factory); - // Resolve from chunk B using a config that points to this engine + expect(chunkB.getContextEngineFactory(engineId)).toBe(factory); + expect(chunkB.listContextEngineIds()).toContain(engineId); const engine = await chunkB.resolveContextEngine(configWithSlot(engineId)); expect(engine.info.id).toBe(engineId); }); diff --git a/src/memory-host-sdk/host/embeddings.test.ts b/src/memory-host-sdk/host/embeddings.test.ts index b762ea4a6ca..e45c1088892 100644 --- a/src/memory-host-sdk/host/embeddings.test.ts +++ b/src/memory-host-sdk/host/embeddings.test.ts @@ -27,6 +27,11 @@ const { resolveCredentialsMock: vi.fn(), })); +vi.mock("../../agents/model-auth.js", async () => { + const { createModelAuthMockModule } = await import("../../test-utils/model-auth-mock.js"); + return createModelAuthMockModule(); +}); + vi.mock("./embeddings-ollama.js", () => ({ createOllamaEmbeddingProvider: createOllamaEmbeddingProviderMock, })); diff --git a/src/plugin-activation-boundary.test.ts b/src/plugin-activation-boundary.test.ts index b2fb255fef6..26326375ff4 100644 --- a/src/plugin-activation-boundary.test.ts +++ b/src/plugin-activation-boundary.test.ts @@ -108,7 +108,6 @@ describe("plugin activation boundary", () => { let configHelpersPromise: | Promise<{ isStaticallyChannelConfigured: typeof import("./config/channel-configured-shared.js").isStaticallyChannelConfigured; - resolveEnvApiKey: typeof import("./agents/model-auth-env.js").resolveEnvApiKey; }> | undefined; let modelSelectionPromise: @@ -134,13 +133,11 @@ describe("plugin activation boundary", () => { }> | undefined; function importConfigHelpers() { - configHelpersPromise ??= Promise.all([ - import("./config/channel-configured-shared.js"), - import("./agents/model-auth-env.js"), - ]).then(([channelConfigured, modelAuthEnv]) => ({ - isStaticallyChannelConfigured: channelConfigured.isStaticallyChannelConfigured, - resolveEnvApiKey: modelAuthEnv.resolveEnvApiKey, - })); + configHelpersPromise ??= import("./config/channel-configured-shared.js").then( + (channelConfigured) => ({ + isStaticallyChannelConfigured: channelConfigured.isStaticallyChannelConfigured, + }), + ); return configHelpersPromise; } @@ -175,8 +172,10 @@ describe("plugin activation boundary", () => { } it("keeps config and model boundary helpers cold", async () => { - const [{ isStaticallyChannelConfigured, resolveEnvApiKey }, { normalizeModelRef }] = - await Promise.all([importConfigHelpers(), importModelSelection()]); + const [{ isStaticallyChannelConfigured }, { normalizeModelRef }] = await Promise.all([ + importConfigHelpers(), + importModelSelection(), + ]); expect(isStaticallyChannelConfigured({}, "telegram", { TELEGRAM_BOT_TOKEN: "token" })).toBe( true, @@ -190,14 +189,6 @@ describe("plugin activation boundary", () => { }), ).toBe(true); expect(isStaticallyChannelConfigured({}, "whatsapp", {})).toBe(false); - expect( - resolveEnvApiKey("anthropic-vertex", { - ANTHROPIC_VERTEX_USE_GCP_METADATA: "true", - }), - ).toEqual({ - apiKey: "gcp-vertex-credentials", - source: "gcloud adc", - }); expect(normalizeModelRef("google", "gemini-3.1-pro")).toEqual({ provider: "google", model: "gemini-3.1-pro-preview",