feat(cli): add --memory-limit flag for configurable container memory (#1243)

Reduce default agent container memory from 4GB to 2GB for better DoS
protection in shared environments. Add --memory-limit flag to override
(e.g., --memory-limit 8g for AI agent workloads).

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Mossaka

src/cli.test.ts

@@ -1,5 +1,5 @@
 import { Command } from 'commander';
-import { parseEnvironmentVariables, parseDomains, parseDomainsFile, escapeShellArg, joinShellArgs, parseVolumeMounts, isValidIPv4, isValidIPv6, parseDnsServers, validateAgentImage, isAgentImagePreset, AGENT_IMAGE_PRESETS, processAgentImageOption, processLocalhostKeyword, validateSkipPullWithBuildLocal, validateAllowHostPorts, validateFormat, validateApiProxyConfig, buildRateLimitConfig, validateRateLimitFlags } from './cli';
+import { parseEnvironmentVariables, parseDomains, parseDomainsFile, escapeShellArg, joinShellArgs, parseVolumeMounts, isValidIPv4, isValidIPv6, parseDnsServers, validateAgentImage, isAgentImagePreset, AGENT_IMAGE_PRESETS, processAgentImageOption, processLocalhostKeyword, validateSkipPullWithBuildLocal, validateAllowHostPorts, parseMemoryLimit, validateFormat, validateApiProxyConfig, buildRateLimitConfig, validateRateLimitFlags } from './cli';
 import { redactSecrets } from './redact-secrets';
 import * as fs from 'fs';
 import * as path from 'path';
@@ -1539,4 +1539,26 @@ describe('cli', () => {
       expect(result.error).toBeUndefined();
     });
   });
+
+  describe('parseMemoryLimit', () => {
+    it('accepts valid memory limits', () => {
+      expect(parseMemoryLimit('2g')).toEqual({ value: '2g' });
+      expect(parseMemoryLimit('4g')).toEqual({ value: '4g' });
+      expect(parseMemoryLimit('512m')).toEqual({ value: '512m' });
+      expect(parseMemoryLimit('1024k')).toEqual({ value: '1024k' });
+      expect(parseMemoryLimit('8G')).toEqual({ value: '8g' });
+    });
+
+    it('rejects invalid formats', () => {
+      expect(parseMemoryLimit('abc')).toHaveProperty('error');
+      expect(parseMemoryLimit('-1g')).toHaveProperty('error');
+      expect(parseMemoryLimit('2x')).toHaveProperty('error');
+      expect(parseMemoryLimit('')).toHaveProperty('error');
+      expect(parseMemoryLimit('g')).toHaveProperty('error');
+    });
+
+    it('rejects zero', () => {
+      expect(parseMemoryLimit('0g')).toHaveProperty('error');
+    });
+  });
 });

src/cli.ts

@@ -422,6 +422,23 @@ export function validateAllowHostPorts(
   return { valid: true };
 }
 
+/**
+ * Parses and validates a Docker memory limit string.
+ * Valid formats: positive integer followed by b, k, m, or g (e.g., "2g", "512m", "4g").
+ */
+export function parseMemoryLimit(input: string): { value: string; error?: undefined } | { value?: undefined; error: string } {
+  const pattern = /^(\d+)([bkmg])$/i;
+  const match = input.match(pattern);
+  if (!match) {
+    return { error: `Invalid --memory-limit value "${input}". Expected format: <number><unit> (e.g., 2g, 512m, 4g)` };
+  }
+  const num = parseInt(match[1], 10);
+  if (num <= 0) {
+    return { error: `Invalid --memory-limit value "${input}". Memory limit must be a positive number.` };
+  }
+  return { value: input.toLowerCase() };
+}
+
 /**
  * Parses and validates DNS servers from a comma-separated string
  * @param input - Comma-separated DNS server string (e.g., "8.8.8.8,1.1.1.1")
@@ -780,6 +797,11 @@ program
     '--container-workdir <dir>',
     'Working directory inside the container (should match GITHUB_WORKSPACE for path consistency)'
   )
+  .option(
+    '--memory-limit <limit>',
+    'Memory limit for the agent container (e.g., 1g, 2g, 4g, 512m). Default: 2g',
+    '2g'
+  )
   .option(
     '--dns-servers <servers>',
     'Comma-separated list of trusted DNS servers. DNS traffic is ONLY allowed to these servers (default: 8.8.8.8,8.8.4.4)',
@@ -1067,6 +1089,13 @@ program
       logger.warn('⚠️  SSL Bump intercepts HTTPS traffic. Only use for trusted workloads.');
     }
 
+    // Validate memory limit
+    const memoryLimit = parseMemoryLimit(options.memoryLimit);
+    if (memoryLimit.error) {
+      logger.error(memoryLimit.error);
+      process.exit(1);
+    }
+
     // Validate agent image option
     const agentImageResult = processAgentImageOption(options.agentImage, options.buildLocal);
     if (agentImageResult.error) {
@@ -1096,6 +1125,7 @@ program
       volumeMounts,
       containerWorkDir: options.containerWorkdir,
       dnsServers,
+      memoryLimit: memoryLimit.value,
       proxyLogsDir: options.proxyLogsDir,
       enableHostAccess: options.enableHostAccess,
       allowHostPorts: options.allowHostPorts,

src/docker-manager.test.ts

@@ -1022,12 +1022,21 @@ describe('docker-manager', () => {
       expect(agent.security_opt).toContain('no-new-privileges:true');
 
       // Verify resource limits
-      expect(agent.mem_limit).toBe('4g');
-      expect(agent.memswap_limit).toBe('4g');
+      expect(agent.mem_limit).toBe('2g');
+      expect(agent.memswap_limit).toBe('2g');
       expect(agent.pids_limit).toBe(1000);
       expect(agent.cpu_shares).toBe(1024);
     });
 
+    it('should use custom memory limit when specified', () => {
+      const customConfig = { ...mockConfig, memoryLimit: '8g' };
+      const result = generateDockerCompose(customConfig, mockNetworkConfig);
+      const agent = result.services.agent;
+
+      expect(agent.mem_limit).toBe('8g');
+      expect(agent.memswap_limit).toBe('8g');
+    });
+
     it('should disable TTY by default to prevent ANSI escape sequences', () => {
       const result = generateDockerCompose(mockConfig, mockNetworkConfig);
       const agent = result.services.agent;

src/docker-manager.ts

@@ -912,8 +912,8 @@ export function generateDockerCompose(
       'apparmor:unconfined',
     ],
     // Resource limits to prevent DoS attacks (conservative defaults)
-    mem_limit: '4g',           // 4GB memory limit
-    memswap_limit: '4g',       // No swap (same as mem_limit)
+    mem_limit: config.memoryLimit || '2g',
+    memswap_limit: config.memoryLimit || '2g',  // No swap (same as mem_limit)
     pids_limit: 1000,          // Max 1000 processes
     cpu_shares: 1024,          // Default CPU share
     stdin_open: true,

src/types.ts

@@ -292,6 +292,18 @@ export interface WrapperConfig {
    */
   dnsServers?: string[];
 
+  /**
+   * Memory limit for the agent execution container
+   *
+   * Accepts Docker memory format: a positive integer followed by a unit suffix
+   * (b, k, m, g). Controls the maximum amount of memory the container can use.
+   *
+   * @default '2g'
+   * @example '4g'
+   * @example '512m'
+   */
+  memoryLimit?: string;
+
   /**
    * Custom directory for Squid proxy logs (written directly during runtime)
    *