chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,528 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, beforeAll, afterAll, afterEach } from 'vitest';
|
||||
import { TestRig, MemoryTestHarness } from '@google/gemini-cli-test-utils';
|
||||
import { join, dirname } from 'node:path';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
import {
|
||||
createWriteStream,
|
||||
copyFileSync,
|
||||
readFileSync,
|
||||
existsSync,
|
||||
mkdirSync,
|
||||
rmSync,
|
||||
} from 'node:fs';
|
||||
import { randomUUID, createHash } from 'node:crypto';
|
||||
|
||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
const BASELINES_PATH = join(__dirname, 'baselines.json');
|
||||
const UPDATE_BASELINES = process.env['UPDATE_MEMORY_BASELINES'] === 'true';
|
||||
function getProjectHash(projectRoot: string): string {
|
||||
return createHash('sha256').update(projectRoot).digest('hex');
|
||||
}
|
||||
const TOLERANCE_PERCENT = 10;
|
||||
|
||||
// Fake API key for tests using fake responses
|
||||
const TEST_ENV = {
|
||||
GEMINI_API_KEY: 'fake-memory-test-key',
|
||||
GEMINI_MEMORY_MONITOR_INTERVAL: '100',
|
||||
};
|
||||
|
||||
describe('Memory Usage Tests', () => {
|
||||
let harness: MemoryTestHarness;
|
||||
let rig: TestRig;
|
||||
|
||||
beforeAll(() => {
|
||||
harness = new MemoryTestHarness({
|
||||
baselinesPath: BASELINES_PATH,
|
||||
defaultTolerancePercent: TOLERANCE_PERCENT,
|
||||
gcCycles: 3,
|
||||
gcDelayMs: 100,
|
||||
sampleCount: 3,
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await rig.cleanup();
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
// Generate the summary report after all tests
|
||||
await harness.generateReport();
|
||||
});
|
||||
|
||||
it('idle-session-startup: memory usage within baseline', async () => {
|
||||
rig = new TestRig();
|
||||
rig.setup('memory-idle-startup', {
|
||||
fakeResponsesPath: join(__dirname, 'memory.idle-startup.responses'),
|
||||
});
|
||||
|
||||
const result = await harness.runScenario(
|
||||
rig,
|
||||
'idle-session-startup',
|
||||
async (recordSnapshot) => {
|
||||
await rig.run({
|
||||
args: ['hello'],
|
||||
timeout: 120000,
|
||||
env: TEST_ENV,
|
||||
});
|
||||
|
||||
await recordSnapshot('after-startup');
|
||||
},
|
||||
);
|
||||
|
||||
if (UPDATE_BASELINES) {
|
||||
harness.updateScenarioBaseline(result);
|
||||
console.log(
|
||||
`Updated baseline for idle-session-startup: ${(result.finalHeapUsed / (1024 * 1024)).toFixed(1)} MB`,
|
||||
);
|
||||
} else {
|
||||
harness.assertWithinBaseline(result);
|
||||
}
|
||||
});
|
||||
|
||||
it('simple-prompt-response: memory usage within baseline', async () => {
|
||||
rig = new TestRig();
|
||||
rig.setup('memory-simple-prompt', {
|
||||
fakeResponsesPath: join(__dirname, 'memory.simple-prompt.responses'),
|
||||
});
|
||||
|
||||
const result = await harness.runScenario(
|
||||
rig,
|
||||
'simple-prompt-response',
|
||||
async (recordSnapshot) => {
|
||||
await rig.run({
|
||||
args: ['What is the capital of France?'],
|
||||
timeout: 120000,
|
||||
env: TEST_ENV,
|
||||
});
|
||||
|
||||
await recordSnapshot('after-response');
|
||||
},
|
||||
);
|
||||
|
||||
if (UPDATE_BASELINES) {
|
||||
harness.updateScenarioBaseline(result);
|
||||
console.log(
|
||||
`Updated baseline for simple-prompt-response: ${(result.finalHeapUsed / (1024 * 1024)).toFixed(1)} MB`,
|
||||
);
|
||||
} else {
|
||||
harness.assertWithinBaseline(result);
|
||||
}
|
||||
});
|
||||
|
||||
it('multi-turn-conversation: memory remains stable over turns', async () => {
|
||||
rig = new TestRig();
|
||||
rig.setup('memory-multi-turn', {
|
||||
fakeResponsesPath: join(__dirname, 'memory.multi-turn.responses'),
|
||||
});
|
||||
|
||||
const prompts = [
|
||||
'Hello, what can you help me with?',
|
||||
'Tell me about JavaScript',
|
||||
'How is TypeScript different?',
|
||||
'Can you write a simple TypeScript function?',
|
||||
'What are some TypeScript best practices?',
|
||||
];
|
||||
|
||||
const result = await harness.runScenario(
|
||||
rig,
|
||||
'multi-turn-conversation',
|
||||
async (recordSnapshot) => {
|
||||
// Run through all turns as a piped sequence
|
||||
const stdinContent = prompts.join('\n');
|
||||
await rig.run({
|
||||
stdin: stdinContent,
|
||||
timeout: 120000,
|
||||
env: TEST_ENV,
|
||||
});
|
||||
|
||||
// Take snapshots after the conversation completes
|
||||
await recordSnapshot('after-all-turns');
|
||||
},
|
||||
);
|
||||
|
||||
if (UPDATE_BASELINES) {
|
||||
harness.updateScenarioBaseline(result);
|
||||
console.log(
|
||||
`Updated baseline for multi-turn-conversation: ${(result.finalHeapUsed / (1024 * 1024)).toFixed(1)} MB`,
|
||||
);
|
||||
} else {
|
||||
harness.assertWithinBaseline(result);
|
||||
harness.assertMemoryReturnsToBaseline(result.snapshots, 20);
|
||||
const { leaked, message } = harness.analyzeSnapshots(result.snapshots);
|
||||
if (leaked) console.warn(`⚠ ${message}`);
|
||||
}
|
||||
});
|
||||
|
||||
it('multi-function-call-repo-search: memory after tool use', async () => {
|
||||
rig = new TestRig();
|
||||
rig.setup('memory-multi-func-call', {
|
||||
fakeResponsesPath: join(
|
||||
__dirname,
|
||||
'memory.multi-function-call.responses',
|
||||
),
|
||||
});
|
||||
|
||||
// Create directories first, then files in the workspace so the tools have targets
|
||||
rig.mkdir('packages/core/src/telemetry');
|
||||
rig.createFile(
|
||||
'packages/core/src/telemetry/memory-monitor.ts',
|
||||
'export class MemoryMonitor { constructor() {} }',
|
||||
);
|
||||
rig.createFile(
|
||||
'packages/core/src/telemetry/metrics.ts',
|
||||
'export function recordMemoryUsage() {}',
|
||||
);
|
||||
|
||||
const result = await harness.runScenario(
|
||||
rig,
|
||||
'multi-function-call-repo-search',
|
||||
async (recordSnapshot) => {
|
||||
await rig.run({
|
||||
args: [
|
||||
'Search this repository for MemoryMonitor and tell me what it does',
|
||||
],
|
||||
timeout: 120000,
|
||||
env: TEST_ENV,
|
||||
});
|
||||
|
||||
await recordSnapshot('after-tool-calls');
|
||||
},
|
||||
);
|
||||
|
||||
if (UPDATE_BASELINES) {
|
||||
harness.updateScenarioBaseline(result);
|
||||
console.log(
|
||||
`Updated baseline for multi-function-call-repo-search: ${(result.finalHeapUsed / (1024 * 1024)).toFixed(1)} MB`,
|
||||
);
|
||||
} else {
|
||||
harness.assertWithinBaseline(result);
|
||||
harness.assertMemoryReturnsToBaseline(result.snapshots, 20);
|
||||
}
|
||||
});
|
||||
|
||||
describe('Large Chat Scenarios', () => {
|
||||
let sharedResumeResponsesPath: string;
|
||||
let sharedActiveResponsesPath: string;
|
||||
let sharedHistoryPath: string;
|
||||
let sharedPrompts: string;
|
||||
let tempDir: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
tempDir = join(__dirname, `large-chat-tmp-${randomUUID()}`);
|
||||
mkdirSync(tempDir, { recursive: true });
|
||||
|
||||
const { resumeResponsesPath, activeResponsesPath, historyPath, prompts } =
|
||||
await generateSharedLargeChatData(tempDir);
|
||||
sharedActiveResponsesPath = activeResponsesPath;
|
||||
sharedResumeResponsesPath = resumeResponsesPath;
|
||||
sharedHistoryPath = historyPath;
|
||||
sharedPrompts = prompts;
|
||||
}, 60000);
|
||||
|
||||
afterAll(() => {
|
||||
if (existsSync(tempDir)) {
|
||||
rmSync(tempDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await rig.cleanup();
|
||||
});
|
||||
|
||||
it('large-chat: memory usage within baseline', async () => {
|
||||
rig = new TestRig();
|
||||
rig.setup('memory-large-chat', {
|
||||
fakeResponsesPath: sharedActiveResponsesPath,
|
||||
});
|
||||
|
||||
const result = await harness.runScenario(
|
||||
rig,
|
||||
'large-chat',
|
||||
async (recordSnapshot) => {
|
||||
await rig.run({
|
||||
stdin: sharedPrompts,
|
||||
timeout: 600000,
|
||||
env: TEST_ENV,
|
||||
});
|
||||
|
||||
await recordSnapshot('after-large-chat');
|
||||
},
|
||||
);
|
||||
|
||||
if (UPDATE_BASELINES) {
|
||||
harness.updateScenarioBaseline(result);
|
||||
console.log(
|
||||
`Updated baseline for large-chat: ${(result.finalHeapUsed / (1024 * 1024)).toFixed(1)} MB`,
|
||||
);
|
||||
} else {
|
||||
harness.assertWithinBaseline(result);
|
||||
}
|
||||
});
|
||||
|
||||
it('resume-large-chat: memory usage within baseline', async () => {
|
||||
rig = new TestRig();
|
||||
rig.setup('memory-resume-large-chat', {
|
||||
fakeResponsesPath: sharedResumeResponsesPath,
|
||||
});
|
||||
|
||||
const result = await harness.runScenario(
|
||||
rig,
|
||||
'resume-large-chat',
|
||||
async (recordSnapshot) => {
|
||||
// Ensure the history file is linked
|
||||
const targetChatsDir = join(
|
||||
rig.homeDir!,
|
||||
'.gemini',
|
||||
'tmp',
|
||||
getProjectHash(rig.testDir!),
|
||||
'chats',
|
||||
);
|
||||
mkdirSync(targetChatsDir, { recursive: true });
|
||||
const targetHistoryPath = join(
|
||||
targetChatsDir,
|
||||
'session-large-chat.json',
|
||||
);
|
||||
if (existsSync(targetHistoryPath)) rmSync(targetHistoryPath);
|
||||
copyFileSync(sharedHistoryPath, targetHistoryPath);
|
||||
|
||||
await rig.run({
|
||||
// add a prompt to make sure it does not hang there and exits immediately
|
||||
args: ['--resume', 'latest', '--prompt', 'hello'],
|
||||
timeout: 600000,
|
||||
env: TEST_ENV,
|
||||
});
|
||||
|
||||
await recordSnapshot('after-resume-large-chat');
|
||||
},
|
||||
);
|
||||
|
||||
if (UPDATE_BASELINES) {
|
||||
harness.updateScenarioBaseline(result);
|
||||
console.log(
|
||||
`Updated baseline for resume-large-chat: ${(result.finalHeapUsed / (1024 * 1024)).toFixed(1)} MB`,
|
||||
);
|
||||
} else {
|
||||
harness.assertWithinBaseline(result);
|
||||
}
|
||||
});
|
||||
|
||||
it('resume-large-chat-with-messages: memory usage within baseline', async () => {
|
||||
rig = new TestRig();
|
||||
rig.setup('memory-resume-large-chat-msgs', {
|
||||
fakeResponsesPath: sharedResumeResponsesPath,
|
||||
});
|
||||
|
||||
const result = await harness.runScenario(
|
||||
rig,
|
||||
'resume-large-chat-with-messages',
|
||||
async (recordSnapshot) => {
|
||||
// Ensure the history file is linked
|
||||
const targetChatsDir = join(
|
||||
rig.homeDir!,
|
||||
'.gemini',
|
||||
'tmp',
|
||||
getProjectHash(rig.testDir!),
|
||||
'chats',
|
||||
);
|
||||
mkdirSync(targetChatsDir, { recursive: true });
|
||||
const targetHistoryPath = join(
|
||||
targetChatsDir,
|
||||
'session-large-chat.json',
|
||||
);
|
||||
if (existsSync(targetHistoryPath)) rmSync(targetHistoryPath);
|
||||
copyFileSync(sharedHistoryPath, targetHistoryPath);
|
||||
|
||||
const stdinContent = 'new prompt 1\nnew prompt 2\n';
|
||||
|
||||
await rig.run({
|
||||
args: ['--resume', 'latest'],
|
||||
stdin: stdinContent,
|
||||
timeout: 600000,
|
||||
env: TEST_ENV,
|
||||
});
|
||||
|
||||
await recordSnapshot('after-resume-and-append');
|
||||
},
|
||||
);
|
||||
|
||||
if (UPDATE_BASELINES) {
|
||||
harness.updateScenarioBaseline(result);
|
||||
console.log(
|
||||
`Updated baseline for resume-large-chat-with-messages: ${(result.finalHeapUsed / (1024 * 1024)).toFixed(1)} MB`,
|
||||
);
|
||||
} else {
|
||||
harness.assertWithinBaseline(result);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
async function generateSharedLargeChatData(tempDir: string) {
|
||||
const resumeResponsesPath = join(tempDir, 'large-chat-resume-chat.responses');
|
||||
const activeResponsesPath = join(tempDir, 'large-chat-active-chat.responses');
|
||||
const historyPath = join(tempDir, 'large-chat-history.json');
|
||||
const sourceSessionPath = join(__dirname, 'large-chat-session.json');
|
||||
|
||||
const session = JSON.parse(readFileSync(sourceSessionPath, 'utf8'));
|
||||
const messages = session.messages;
|
||||
|
||||
copyFileSync(sourceSessionPath, historyPath);
|
||||
|
||||
// Generate fake responses for active chat
|
||||
const promptsList: string[] = [];
|
||||
const activeResponsesStream = createWriteStream(activeResponsesPath);
|
||||
const complexityResponse = {
|
||||
method: 'generateContent',
|
||||
response: {
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [
|
||||
{
|
||||
text: '{"complexity_reasoning":"simple","complexity_score":1}',
|
||||
},
|
||||
],
|
||||
role: 'model',
|
||||
},
|
||||
finishReason: 'STOP',
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
const summaryResponse = {
|
||||
method: 'generateContent',
|
||||
response: {
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [
|
||||
{ text: '{"originalSummary":"large chat summary","events":[]}' },
|
||||
],
|
||||
role: 'model',
|
||||
},
|
||||
finishReason: 'STOP',
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
|
||||
for (let i = 0; i < messages.length; i++) {
|
||||
const msg = messages[i];
|
||||
if (msg.type === 'user') {
|
||||
promptsList.push(msg.content[0].text);
|
||||
|
||||
// Start of a new turn
|
||||
activeResponsesStream.write(JSON.stringify(complexityResponse) + '\n');
|
||||
|
||||
// Find all subsequent gemini messages until the next user message
|
||||
let j = i + 1;
|
||||
while (j < messages.length && messages[j].type === 'gemini') {
|
||||
const geminiMsg = messages[j];
|
||||
const parts = [];
|
||||
if (geminiMsg.content) {
|
||||
parts.push({ text: geminiMsg.content });
|
||||
}
|
||||
if (geminiMsg.toolCalls) {
|
||||
for (const tc of geminiMsg.toolCalls) {
|
||||
parts.push({
|
||||
functionCall: {
|
||||
name: tc.name,
|
||||
args: tc.args,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
activeResponsesStream.write(
|
||||
JSON.stringify({
|
||||
method: 'generateContentStream',
|
||||
response: [
|
||||
{
|
||||
candidates: [
|
||||
{
|
||||
content: { parts, role: 'model' },
|
||||
finishReason: 'STOP',
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 100,
|
||||
candidatesTokenCount: 100,
|
||||
totalTokenCount: 200,
|
||||
promptTokensDetails: [{ modality: 'TEXT', tokenCount: 100 }],
|
||||
},
|
||||
},
|
||||
],
|
||||
}) + '\n',
|
||||
);
|
||||
j++;
|
||||
}
|
||||
// End of turn
|
||||
activeResponsesStream.write(JSON.stringify(summaryResponse) + '\n');
|
||||
// Skip the gemini messages we just processed
|
||||
i = j - 1;
|
||||
}
|
||||
}
|
||||
activeResponsesStream.end();
|
||||
|
||||
// Generate responses for resumed chat
|
||||
const resumeResponsesStream = createWriteStream(resumeResponsesPath);
|
||||
for (let i = 0; i < 5; i++) {
|
||||
// Doubling up on non-streaming responses to satisfy classifier and complexity checks
|
||||
resumeResponsesStream.write(JSON.stringify(complexityResponse) + '\n');
|
||||
resumeResponsesStream.write(JSON.stringify(summaryResponse) + '\n');
|
||||
resumeResponsesStream.write(JSON.stringify(complexityResponse) + '\n');
|
||||
resumeResponsesStream.write(
|
||||
JSON.stringify({
|
||||
method: 'generateContentStream',
|
||||
response: [
|
||||
{
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
parts: [{ text: `Resume response ${i}` }],
|
||||
role: 'model',
|
||||
},
|
||||
finishReason: 'STOP',
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 10,
|
||||
candidatesTokenCount: 10,
|
||||
totalTokenCount: 20,
|
||||
promptTokensDetails: [{ modality: 'TEXT', tokenCount: 10 }],
|
||||
},
|
||||
},
|
||||
],
|
||||
}) + '\n',
|
||||
);
|
||||
resumeResponsesStream.write(JSON.stringify(summaryResponse) + '\n');
|
||||
}
|
||||
resumeResponsesStream.end();
|
||||
|
||||
// Wait for streams to finish
|
||||
await Promise.all([
|
||||
new Promise((res) =>
|
||||
activeResponsesStream.on('finish', () => res(undefined)),
|
||||
),
|
||||
new Promise((res) =>
|
||||
resumeResponsesStream.on('finish', () => res(undefined)),
|
||||
),
|
||||
]);
|
||||
|
||||
return {
|
||||
resumeResponsesPath,
|
||||
activeResponsesPath,
|
||||
historyPath,
|
||||
prompts: promptsList.join('\n'),
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user