perf(memory): implement simulation script and identify growth sources

2026-05-16 23:02:51 -07:00 · 2026-04-03 02:29:16 +00:00
parent e293424bb4
commit 21ecd3ca9f
3 changed files with 291 additions and 0 deletions
@@ -0,0 +1,117 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as crypto from 'node:crypto';
+import { GeminiChat } from '../packages/core/src/core/geminiChat.js';
+import { Config } from '../packages/core/src/config/config.js';
+import { ToolRegistry } from '../packages/core/src/tools/tool-registry.js';
+import { MessageBus } from '../packages/core/src/confirmation-bus/message-bus.js';
+import { PromptRegistry } from '../packages/core/src/prompts/prompt-registry.js';
+import { ResourceRegistry } from '../packages/core/src/resources/resource-registry.js';
+import { NoopSandboxManager } from '../packages/core/src/services/sandboxManager.js';
+import type { AgentLoopContext } from '../packages/core/src/config/agent-loop-context.js';
+
+// Helper to force GC if run with --expose-gc
+const runGC = () => {
+  if (global.gc) {
+    global.gc();
+  }
+};
+
+const printMemory = (turn: number) => {
+  runGC();
+  const usage = process.memoryUsage();
+  console.log(
+    `Turn ${turn} - RSS: ${(usage.rss / 1024 / 1024).toFixed(2)} MB, ` +
+      `HeapUsed: ${(usage.heapUsed / 1024 / 1024).toFixed(2)} MB, ` +
+      `External: ${(usage.external / 1024 / 1024).toFixed(2)} MB`,
+  );
+};
+
+async function runReproduction() {
+  console.log(
+    'Starting memory growth reproduction (ChatRecordingService focus)...',
+  );
+
+  const config = new Config({
+    sessionId: 'reproduction-session',
+    targetDir: process.cwd(),
+    cwd: process.cwd(),
+    debugMode: false,
+    model: 'gemini-2.0-flash',
+  });
+  await config.initialize();
+
+  const context: AgentLoopContext = {
+    config,
+    promptId: 'reproduction-session',
+    toolRegistry: new ToolRegistry(config),
+    promptRegistry: new PromptRegistry(),
+    resourceRegistry: new ResourceRegistry(),
+    messageBus: new MessageBus(),
+    sandboxManager: new NoopSandboxManager(),
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    geminiClient: null as any,
+  };
+
+  const chat = new GeminiChat(context);
+
+  for (let i = 1; i <= 200; i++) {
+    const LARGE_STRING = crypto.randomBytes(512 * 1024).toString('hex'); // 1MB string
+    // 1. User message
+    chat.addHistory({
+      role: 'user',
+      parts: [{ text: `Turn ${i}: Get the large data.` }],
+    });
+
+    // 2. Model message with tool call
+    chat.addHistory({
+      role: 'model',
+      parts: [
+        {
+          functionCall: {
+            name: 'get_large_data',
+            args: {},
+          },
+        },
+      ],
+    });
+
+    // 3. User message with tool response (LARGE)
+    chat.addHistory({
+      role: 'user',
+      parts: [
+        {
+          functionResponse: {
+            name: 'get_large_data',
+            response: { output: LARGE_STRING },
+          },
+        },
+      ],
+    });
+
+    // 4. Model message with final text
+    chat.addHistory({
+      role: 'model',
+      parts: [{ text: 'I have processed the large data.' }],
+    });
+
+    // Trigger ChatRecordingService update (as GeminiClient does)
+    await chat
+      .getChatRecordingService()
+      ?.updateMessagesFromHistory(chat.getHistory());
+
+    if (i % 20 === 0) {
+      const history = chat.getHistory();
+      console.log(`Turn ${i}: History size: ${history.length}`);
+      printMemory(i);
+    }
+  }
+
+  console.log('Reproduction complete.');
+}
+
+runReproduction().catch(console.error);