Jacob314/max old space (#1314)

3 files changed, 78 insertions, 6 deletions

diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts
index 882df403..36b16446 100644
--- a/packages/cli/src/config/settings.ts
+++ b/packages/cli/src/config/settings.ts
@@ -53,6 +53,7 @@ export interface Settings {
   preferredEditor?: string;
   bugCommand?: BugCommandSettings;
   checkpointing?: CheckpointingSettings;
+  autoConfigureMaxOldSpaceSize?: boolean;
 
   // Git-aware file filtering settings
   fileFiltering?: {
diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx
index 293d50a5..cd8e23fc 100644
--- a/packages/cli/src/gemini.tsx
+++ b/packages/cli/src/gemini.tsx
@@ -10,6 +10,9 @@ import { AppWrapper } from './ui/App.js';
 import { loadCliConfig } from './config/config.js';
 import { readStdin } from './utils/readStdin.js';
 import { basename } from 'node:path';
+import v8 from 'node:v8';
+import os from 'node:os';
+import { spawn } from 'node:child_process';
 import { start_sandbox } from './utils/sandbox.js';
 import {
   LoadedSettings,
@@ -34,6 +37,50 @@ import {
 import { validateAuthMethod } from './config/auth.js';
 import { setMaxSizedBoxDebugging } from './ui/components/shared/MaxSizedBox.js';
 
+function getNodeMemoryArgs(config: Config): string[] {
+  const totalMemoryMB = os.totalmem() / (1024 * 1024);
+  const heapStats = v8.getHeapStatistics();
+  const currentMaxOldSpaceSizeMb = Math.floor(
+    heapStats.heap_size_limit / 1024 / 1024,
+  );
+
+  // Set target to 50% of total memory
+  const targetMaxOldSpaceSizeInMB = Math.floor(totalMemoryMB * 0.5);
+  if (config.getDebugMode()) {
+    console.debug(
+      `Current heap size ${currentMaxOldSpaceSizeMb.toFixed(2)} MB`,
+    );
+  }
+
+  if (process.env.GEMINI_CLI_NO_RELAUNCH) {
+    return [];
+  }
+
+  if (targetMaxOldSpaceSizeInMB > currentMaxOldSpaceSizeMb) {
+    if (config.getDebugMode()) {
+      console.debug(
+        `Need to relaunch with more memory: ${targetMaxOldSpaceSizeInMB.toFixed(2)} MB`,
+      );
+    }
+    return [`--max-old-space-size=${targetMaxOldSpaceSizeInMB}`];
+  }
+
+  return [];
+}
+
+async function relaunchWithAdditionalArgs(additionalArgs: string[]) {
+  const nodeArgs = [...additionalArgs, ...process.argv.slice(1)];
+  const newEnv = { ...process.env, GEMINI_CLI_NO_RELAUNCH: 'true' };
+
+  const child = spawn(process.execPath, nodeArgs, {
+    stdio: 'inherit',
+    env: newEnv,
+  });
+
+  await new Promise((resolve) => child.on('close', resolve));
+  process.exit(0);
+}
+
 export async function main() {
   const workspaceRoot = process.cwd();
   const settings = loadSettings(workspaceRoot);
@@ -84,6 +131,10 @@ export async function main() {
     }
   }
 
+  const memoryArgs = settings.merged.autoConfigureMaxOldSpaceSize
+    ? getNodeMemoryArgs(config)
+    : [];
+
   // hop into sandbox if we are outside and sandboxing is enabled
   if (!process.env.SANDBOX) {
     const sandboxConfig = config.getSandbox();
@@ -97,11 +148,17 @@ export async function main() {
         }
         await config.refreshAuth(settings.merged.selectedAuthType);
       }
-      await start_sandbox(sandboxConfig);
+      await start_sandbox(sandboxConfig, memoryArgs);
       process.exit(0);
+    } else {
+      // Not in a sandbox and not entering one, so relaunch with additional
+      // arguments to control memory usage if needed.
+      if (memoryArgs.length > 0) {
+        await relaunchWithAdditionalArgs(memoryArgs);
+        process.exit(0);
+      }
     }
   }
-
   let input = config.getQuestion();
   const startupWarnings = await getStartupWarnings();
 
diff --git a/packages/cli/src/utils/sandbox.ts b/packages/cli/src/utils/sandbox.ts
index 6a08edef..ce374ceb 100644
--- a/packages/cli/src/utils/sandbox.ts
+++ b/packages/cli/src/utils/sandbox.ts
@@ -180,7 +180,10 @@ function entrypoint(workdir: string): string[] {
   return ['bash', '-c', args.join(' ')];
 }
 
-export async function start_sandbox(config: SandboxConfig) {
+export async function start_sandbox(
+  config: SandboxConfig,
+  nodeArgs: string[] = [],
+) {
   if (config.command === 'sandbox-exec') {
     // disallow BUILD_SANDBOX
     if (process.env.BUILD_SANDBOX) {
@@ -206,6 +209,11 @@ export async function start_sandbox(config: SandboxConfig) {
     // Log on STDERR so it doesn't clutter the output on STDOUT
     console.error(`using macos seatbelt (profile: ${profile}) ...`);
     // if DEBUG is set, convert to --inspect-brk in NODE_OPTIONS
+    const nodeOptions = [
+      ...(process.env.DEBUG ? ['--inspect-brk'] : []),
+      ...nodeArgs,
+    ].join(' ');
+
     const args = [
       '-D',
       `TARGET_DIR=${fs.realpathSync(process.cwd())}`,
@@ -221,7 +229,7 @@ export async function start_sandbox(config: SandboxConfig) {
       '-c',
       [
         `SANDBOX=sandbox-exec`,
-        `NODE_OPTIONS="${process.env.DEBUG ? `--inspect-brk` : ''}"`,
+        `NODE_OPTIONS="${nodeOptions}"`,
         ...process.argv.map((arg) => quote([arg])),
       ].join(' '),
     ];
@@ -588,8 +596,14 @@ export async function start_sandbox(config: SandboxConfig) {
   }
 
   // copy NODE_OPTIONS
-  if (process.env.NODE_OPTIONS) {
-    args.push('--env', `NODE_OPTIONS="${process.env.NODE_OPTIONS}"`);
+  const existingNodeOptions = process.env.NODE_OPTIONS || '';
+  const allNodeOptions = [
+    ...(existingNodeOptions ? [existingNodeOptions] : []),
+    ...nodeArgs,
+  ].join(' ');
+
+  if (allNodeOptions.length > 0) {
+    args.push('--env', `NODE_OPTIONS="${allNodeOptions}"`);
   }
 
   // set SANDBOX as container name