test: trim memory and mcp hotspots

This commit is contained in:
Peter Steinberger
2026-04-17 01:39:39 +01:00
parent 2e08c77582
commit fd48dfa68f
7 changed files with 78 additions and 186 deletions

View File

@@ -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();
}
});
});

View File

@@ -7,37 +7,33 @@ const mocks = vi.hoisted(() => ({
withRemoteHttpResponse: vi.fn(
async (params: { url: string; onResponse: (res: Response) => Promise<unknown> }) => {
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);
},
),
}));

View File

@@ -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<TRequest>(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;
}

View File

@@ -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);
});
});

View File

@@ -37,10 +37,13 @@ export function createJsonResponseFetchMock(payload: unknown) {
const fetchMock = vi.fn<FetchMock>(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;
}

View File

@@ -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<string, string>;
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);
}
});

View File

@@ -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(