From fd48dfa68f488aa71a3b925ac7f05092332cfd67 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 17 Apr 2026 01:39:39 +0100 Subject: [PATCH] test: trim memory and mcp hotspots --- src/mcp/plugin-tools-serve.test.ts | 62 ++------------- src/memory-host-sdk/host/batch-openai.test.ts | 56 +++++++------- src/memory-host-sdk/host/batch-runner.ts | 11 ++- .../host/embeddings-gemini.test.ts | 75 +++++-------------- .../host/embeddings-provider.test-support.ts | 9 ++- src/memory-host-sdk/host/embeddings.test.ts | 30 +------- src/plugin-activation-boundary.test.ts | 21 ++---- 7 files changed, 78 insertions(+), 186 deletions(-) diff --git a/src/mcp/plugin-tools-serve.test.ts b/src/mcp/plugin-tools-serve.test.ts index 33358cafcf3..04cdd339cef 100644 --- a/src/mcp/plugin-tools-serve.test.ts +++ b/src/mcp/plugin-tools-serve.test.ts @@ -30,7 +30,10 @@ afterEach(() => { }); describe("plugin tools MCP server", () => { - it("lists registered plugin tools with their input schema", async () => { + it("lists registered plugin tools and serializes non-array tool content", async () => { + const execute = vi.fn().mockResolvedValue({ + content: "Stored.", + }); const tool = { name: "memory_recall", description: "Recall stored memory", @@ -41,7 +44,7 @@ describe("plugin tools MCP server", () => { }, required: ["query"], }, - execute: vi.fn(), + execute, } as unknown as AnyAgentTool; const session = await connectPluginToolsServer([tool]); @@ -57,32 +60,15 @@ describe("plugin tools MCP server", () => { }), }), ]); - } finally { - await session.close(); - } - }); - it("serializes non-array tool content as text for MCP callers", async () => { - const execute = vi.fn().mockResolvedValue({ - content: "Stored.", - }); - const tool = { - name: "memory_store", - description: "Store memory", - parameters: { type: "object", properties: {} }, - execute, - } as unknown as AnyAgentTool; - - const session = await connectPluginToolsServer([tool]); - try { const result = await session.client.callTool({ - name: "memory_store", - arguments: { text: "remember this" }, + name: "memory_recall", + arguments: { query: "remember this" }, }); expect(execute).toHaveBeenCalledWith( expect.stringMatching(/^mcp-\d+$/), { - text: "remember this", + query: "remember this", }, undefined, undefined, @@ -166,36 +152,4 @@ describe("plugin tools MCP server", () => { await session.close(); } }); - - it("still executes plugin tools on the MCP bridge when no before_tool_call hook is registered", async () => { - const execute = vi.fn().mockResolvedValue({ - content: "Stored.", - }); - const tool = { - name: "memory_store", - description: "Store memory", - parameters: { type: "object", properties: {} }, - execute, - } as unknown as AnyAgentTool; - - const session = await connectPluginToolsServer([tool]); - try { - const result = await session.client.callTool({ - name: "memory_store", - arguments: { text: "remember this" }, - }); - expect(execute).toHaveBeenCalledWith( - expect.stringMatching(/^mcp-\d+$/), - { - text: "remember this", - }, - undefined, - undefined, - ); - expect(result.isError).toBeUndefined(); - expect(result.content).toEqual([{ type: "text", text: "Stored." }]); - } finally { - await session.close(); - } - }); }); diff --git a/src/memory-host-sdk/host/batch-openai.test.ts b/src/memory-host-sdk/host/batch-openai.test.ts index 3c4e7258c52..70ec73d0942 100644 --- a/src/memory-host-sdk/host/batch-openai.test.ts +++ b/src/memory-host-sdk/host/batch-openai.test.ts @@ -7,37 +7,33 @@ const mocks = vi.hoisted(() => ({ withRemoteHttpResponse: vi.fn( async (params: { url: string; onResponse: (res: Response) => Promise }) => { if (params.url.endsWith("/files/file_out/content")) { - return await params.onResponse( - new Response( - [ - JSON.stringify({ - custom_id: "0", - response: { - status_code: 200, - body: { data: [{ embedding: [1, 0, 0], index: 0 }] }, - }, - }), - JSON.stringify({ - custom_id: "1", - response: { - status_code: 200, - body: { data: [{ embedding: [2, 0, 0], index: 0 }] }, - }, - }), - ].join("\n"), - { status: 200, headers: { "Content-Type": "application/jsonl" } }, - ), - ); + const content = [ + JSON.stringify({ + custom_id: "0", + response: { + status_code: 200, + body: { data: [{ embedding: [1, 0, 0], index: 0 }] }, + }, + }), + JSON.stringify({ + custom_id: "1", + response: { + status_code: 200, + body: { data: [{ embedding: [2, 0, 0], index: 0 }] }, + }, + }), + ].join("\n"); + return await params.onResponse({ + ok: true, + status: 200, + text: async () => content, + } as Response); } - return await params.onResponse( - new Response( - JSON.stringify({ id: "batch_1", status: "completed", output_file_id: "file_out" }), - { - status: 200, - headers: { "Content-Type": "application/json" }, - }, - ), - ); + return await params.onResponse({ + ok: true, + status: 200, + json: async () => ({ id: "batch_1", status: "completed", output_file_id: "file_out" }), + } as Response); }, ), })); diff --git a/src/memory-host-sdk/host/batch-runner.ts b/src/memory-host-sdk/host/batch-runner.ts index aa1785095bb..7c7c40eb3b7 100644 --- a/src/memory-host-sdk/host/batch-runner.ts +++ b/src/memory-host-sdk/host/batch-runner.ts @@ -1,5 +1,5 @@ +import { runTasksWithConcurrency } from "../../utils/run-with-concurrency.js"; import { splitBatchRequests } from "./batch-utils.js"; -import { runWithConcurrency } from "./internal.js"; export type EmbeddingBatchExecutionParams = { wait: boolean; @@ -43,7 +43,14 @@ export async function runEmbeddingBatchGroups(params: { timeoutMs: params.timeoutMs, }); - await runWithConcurrency(tasks, params.concurrency); + const { firstError, hasError } = await runTasksWithConcurrency({ + tasks, + limit: params.concurrency, + errorMode: "stop", + }); + if (hasError) { + throw firstError; + } return byCustomId; } diff --git a/src/memory-host-sdk/host/embeddings-gemini.test.ts b/src/memory-host-sdk/host/embeddings-gemini.test.ts index f5d91b56215..13185ae2fbc 100644 --- a/src/memory-host-sdk/host/embeddings-gemini.test.ts +++ b/src/memory-host-sdk/host/embeddings-gemini.test.ts @@ -62,8 +62,8 @@ async function createProviderWithFetch( return provider; } -describe("buildGeminiTextEmbeddingRequest", () => { - it("builds a text embedding request with optional model and dimensions", () => { +describe("Gemini embedding request helpers", () => { + it("builds text and multimodal requests", () => { expect( buildGeminiTextEmbeddingRequest({ text: "hello", @@ -77,11 +77,6 @@ describe("buildGeminiTextEmbeddingRequest", () => { taskType: "RETRIEVAL_DOCUMENT", outputDimensionality: 1536, }); - }); -}); - -describe("buildGeminiEmbeddingRequest", () => { - it("builds a multimodal request from structured input parts", () => { expect( buildGeminiEmbeddingRequest({ input: { @@ -107,49 +102,21 @@ describe("buildGeminiEmbeddingRequest", () => { outputDimensionality: 1536, }); }); -}); -// ---------- Model detection ---------- - -describe("isGeminiEmbedding2Model", () => { - it("returns true for gemini-embedding-2-preview", () => { + it("detects v2 model names", () => { + expect(GEMINI_EMBEDDING_2_MODELS.has("gemini-embedding-2-preview")).toBe(true); 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", () => { + it("resolves v2 dimensions and rejects invalid values", () => { 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/, ); @@ -157,9 +124,20 @@ describe("resolveGeminiOutputDimensionality", () => { /Valid values: 768, 1536, 3072/, ); }); -}); -// ---------- Provider behavior ---------- + it("normalizes known model prefixes and default model", () => { + expect(normalizeGeminiModel("models/gemini-embedding-2-preview")).toBe( + "gemini-embedding-2-preview", + ); + expect(normalizeGeminiModel("gemini/gemini-embedding-2-preview")).toBe( + "gemini-embedding-2-preview", + ); + expect(normalizeGeminiModel("google/gemini-embedding-2-preview")).toBe( + "gemini-embedding-2-preview", + ); + expect(normalizeGeminiModel("")).toBe(DEFAULT_GEMINI_EMBEDDING_MODEL); + }); +}); describe("gemini embedding provider", () => { it("handles legacy and v2 request/response behavior", async () => { @@ -253,20 +231,3 @@ describe("gemini embedding provider", () => { ]); }); }); - -// ---------- Model normalization ---------- - -describe("gemini model normalization", () => { - it("normalizes known model prefixes and default model", () => { - expect(normalizeGeminiModel("models/gemini-embedding-2-preview")).toBe( - "gemini-embedding-2-preview", - ); - expect(normalizeGeminiModel("gemini/gemini-embedding-2-preview")).toBe( - "gemini-embedding-2-preview", - ); - expect(normalizeGeminiModel("google/gemini-embedding-2-preview")).toBe( - "gemini-embedding-2-preview", - ); - expect(normalizeGeminiModel("")).toBe(DEFAULT_GEMINI_EMBEDDING_MODEL); - }); -}); diff --git a/src/memory-host-sdk/host/embeddings-provider.test-support.ts b/src/memory-host-sdk/host/embeddings-provider.test-support.ts index 5503d308c25..bb2c09a20f5 100644 --- a/src/memory-host-sdk/host/embeddings-provider.test-support.ts +++ b/src/memory-host-sdk/host/embeddings-provider.test-support.ts @@ -37,10 +37,13 @@ export function createJsonResponseFetchMock(payload: unknown) { const fetchMock = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { const body = typeof payload === "function" ? (payload as FetchPayloadFactory)(input, init) : payload; - return new Response(JSON.stringify(body), { + const serialized = JSON.stringify(body); + return { + ok: true, status: 200, - headers: { "Content-Type": "application/json" }, - }); + json: async () => body, + text: async () => serialized, + } as Response; }); return withFetchPreconnect(fetchMock) as JsonResponseFetchMock; } diff --git a/src/memory-host-sdk/host/embeddings.test.ts b/src/memory-host-sdk/host/embeddings.test.ts index 94e31566bec..affcfef00c4 100644 --- a/src/memory-host-sdk/host/embeddings.test.ts +++ b/src/memory-host-sdk/host/embeddings.test.ts @@ -1,7 +1,6 @@ import { setTimeout as sleep } from "node:timers/promises"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import * as authModule from "../../agents/model-auth.js"; -import { DEFAULT_GEMINI_EMBEDDING_MODEL } from "./embeddings-gemini.js"; import { createEmbeddingDataFetchMock, createGeminiFetchMock, @@ -270,9 +269,6 @@ describe("embedding provider remote overrides", () => { }); it("uses GEMINI_API_KEY env indirection for Gemini remote apiKey", async () => { - const fetchMock = createGeminiFetchMock(); - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - mockPublicPinnedHostname(); vi.stubEnv("GEMINI_API_KEY", "env-gemini-key"); const result = await createEmbeddingProvider({ @@ -286,11 +282,8 @@ describe("embedding provider remote overrides", () => { }); const provider = requireProvider(result); - await provider.embedQuery("hello"); - - const { init } = readFirstFetchRequest(fetchMock); - const headers = (init?.headers ?? {}) as Record; - expect(headers["x-goog-api-key"]).toBe("env-gemini-key"); + expect(provider.id).toBe("gemini"); + expect(result.gemini?.apiKeys).toEqual(["env-gemini-key"]); }); it("builds Mistral embeddings requests with bearer auth", async () => { @@ -369,26 +362,21 @@ describe("embedding provider auto selection", () => { 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".'); @@ -402,12 +390,10 @@ describe("embedding provider auto selection", () => { } 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 { @@ -418,25 +404,17 @@ describe("embedding provider auto selection", () => { } 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).mockReset(); 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); + expectAutoSelectedProvider(result, testCase.expectedProvider); } }); diff --git a/src/plugin-activation-boundary.test.ts b/src/plugin-activation-boundary.test.ts index 26326375ff4..6d03d4ac514 100644 --- a/src/plugin-activation-boundary.test.ts +++ b/src/plugin-activation-boundary.test.ts @@ -1,4 +1,4 @@ -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { describe, expect, it, vi } from "vitest"; const loadBundledPluginPublicSurfaceModuleSync = vi.hoisted(() => vi.fn((params: { artifactBasename: string }) => { @@ -101,10 +101,6 @@ vi.mock("./plugin-sdk/facade-runtime.js", () => ({ })); describe("plugin activation boundary", () => { - beforeEach(() => { - loadBundledPluginPublicSurfaceModuleSync.mockReset(); - }); - let configHelpersPromise: | Promise<{ isStaticallyChannelConfigured: typeof import("./config/channel-configured-shared.js").isStaticallyChannelConfigured; @@ -171,7 +167,9 @@ describe("plugin activation boundary", () => { return browserHelpersPromise; } - it("keeps config and model boundary helpers cold", async () => { + it("keeps generic boundaries cold and loads only narrow browser helper surfaces on use", async () => { + loadBundledPluginPublicSurfaceModuleSync.mockReset(); + const [{ isStaticallyChannelConfigured }, { normalizeModelRef }] = await Promise.all([ importConfigHelpers(), importModelSelection(), @@ -198,9 +196,7 @@ describe("plugin activation boundary", () => { model: "grok-4-fast", }); expect(loadBundledPluginPublicSurfaceModuleSync).not.toHaveBeenCalled(); - }); - it("keeps browser helper imports cold and loads only narrow browser helper surfaces on use", async () => { const browser = await importBrowserHelpers(); expect(browser.DEFAULT_AI_SNAPSHOT_MAX_CHARS).toBe(80_000); @@ -238,13 +234,10 @@ describe("plugin activation boundary", () => { "browser-host-inspection.js", "browser-host-inspection.js", ]); - }); - it("keeps disabled browser cleanup and generic session-binding cleanup cold", async () => { - const [browser, { getSessionBindingService }] = await Promise.all([ - importBrowserHelpers(), - import("./infra/outbound/session-binding-service.js"), - ]); + loadBundledPluginPublicSurfaceModuleSync.mockReset(); + const { getSessionBindingService } = + await import("./infra/outbound/session-binding-service.js"); await expect(browser.closeTrackedBrowserTabsForSessions({ sessionKeys: [] })).resolves.toBe(0); await expect(