From 22b1a193f1c34763a5f98c985fc98fccac0109fd Mon Sep 17 00:00:00 2001 From: Paijo <14921983+oyi77@users.noreply.github.com> Date: Wed, 3 Jun 2026 18:33:01 +0700 Subject: [PATCH] feat: memory optimization to prevent OOM in multi-session scenarios (#1437) - Add memory pressure monitor (RSS-based) with 4 levels: normal/elevated/high/critical - Add memory compaction trigger that fires when pressure is elevated or critical - Add per-session memory budget calculation (total budget / max concurrent sessions) - Add concurrent session limiter using semaphore pattern - Wire into QueryEngine auto-compact path with skipTokenCheck for forceReason - Guard forceReason by querySource to prevent recursive deadlock - Extract monitor startup into shared startMemoryMonitorIfNeeded() helper - Call from both headless and interactive REPL paths (idempotent) Co-authored-by: oyi77 --- bin/openclaude | 19 +++- src/QueryEngine.ts | 4 + src/main.tsx | 30 ++++++ src/query.ts | 27 +++++ src/screens/REPL.tsx | 7 +- src/services/compact/autoCompact.ts | 22 ++++ src/utils/concurrentSessions.ts | 30 ++++++ src/utils/memoryCompaction.ts | 42 ++++++++ src/utils/memoryPressure.ts | 160 ++++++++++++++++++++++++++++ 9 files changed, 339 insertions(+), 2 deletions(-) create mode 100644 src/utils/memoryCompaction.ts create mode 100644 src/utils/memoryPressure.ts diff --git a/bin/openclaude b/bin/openclaude index 1cd52083d..a66e13659 100755 --- a/bin/openclaude +++ b/bin/openclaude @@ -32,6 +32,17 @@ function hasNodeOptionFlag(flag) { } function getHeapSizeMb() { + // --max-memory flag overrides env var + const maxMemArg = process.argv.find(a => a.startsWith('--max-memory=')) + if (maxMemArg) { + const mb = Number.parseInt(maxMemArg.split('=')[1] || '0', 10) + if (Number.isSafeInteger(mb) && mb > 0) { + process.env[HEAP_SIZE_ENV] = String(mb) + process.env.OPENCLAUDE_MAX_MEMORY_MB = String(mb) + return mb + } + } + const raw = process.env[HEAP_SIZE_ENV] if (!raw) return DEFAULT_HEAP_SIZE_MB const parsed = Number.parseInt(raw, 10) @@ -58,10 +69,16 @@ function relaunchWithLongSessionHeapIfNeeded() { execArgv.push('--expose-gc') } + // Strip --max-memory flag before relaunching — it's a launcher-only arg + // that Commander in the built CLI would reject as unknown. + const childArgs = process.argv.slice(2).filter( + arg => !arg.startsWith('--max-memory=') && arg !== '--max-memory', + ) + const result = spawnSync(process.execPath, [ ...execArgv, fileURLToPath(import.meta.url), - ...process.argv.slice(2), + ...childArgs, ], { stdio: 'inherit', env: { diff --git a/src/QueryEngine.ts b/src/QueryEngine.ts index 84e6a2dfc..6285490f4 100644 --- a/src/QueryEngine.ts +++ b/src/QueryEngine.ts @@ -437,6 +437,10 @@ export class QueryEngine { this.autoCompactTracking = undefined } + // NOTE: Message-count and memory-pressure forceReason checks now live in + // src/query.ts (the shared query path used by both REPL and SDK), so they + // no longer need to be duplicated here in QueryEngine. + // Update params to reflect updates from processing /slash commands const messages = [...this.mutableMessages] diff --git a/src/main.tsx b/src/main.tsx index 912c8b1ac..579dc344c 100644 --- a/src/main.tsx +++ b/src/main.tsx @@ -379,6 +379,29 @@ function prefetchSystemContextIfSafe(): void { // Otherwise, don't prefetch - wait for trust to be established first } +/** + * Start memory-pressure monitor + compaction trigger. + * Idempotent — startMemoryPressureMonitor early-returns if already running. + * Called from both headless (--print) and interactive REPL paths. + */ +function startMemoryMonitorIfNeeded(): void { + void Promise.all([ + import('./utils/memoryPressure.js'), + import('./utils/memoryCompaction.js'), + import('./utils/concurrentSessions.js'), + ]).then(([pressure, compaction, sessions]) => + sessions.calculatePerSessionMemoryBudget().then(budgetMB => { + pressure.startMemoryPressureMonitor({ perSessionBudgetMB: budgetMB }) + compaction.createMemoryCompactionTrigger({ + onCompact: () => {}, + onPruneCache: () => { + pressure.pruneRegisteredCaches() + }, + }) + }), + ); +} + /** * Start background prefetches and housekeeping that are NOT needed before first render. * These are deferred from setup() to reduce event loop contention and child process @@ -2748,6 +2771,7 @@ async function run(): Promise { if (!isBareMode()) { startDeferredPrefetches(); void import('./utils/backgroundHousekeeping.js').then(m => m.startBackgroundHousekeeping()); + startMemoryMonitorIfNeeded(); if ("external" === 'ant') { void import('./utils/sdkHeapDumpMonitor.js').then(m => m.startSdkMemoryMonitor()); } @@ -2984,6 +3008,12 @@ async function run(): Promise { logSessionTelemetry(); }); + // Start memory-pressure monitor for interactive sessions. + // Idempotent — safe to call even if --print path already started it. + if (!isBareMode()) { + startMemoryMonitorIfNeeded(); + } + // Set up per-turn session environment data uploader (internal-only build). // Default-enabled for all ant users when working in an Anthropic-owned // repo. Captures git/filesystem state (NOT transcripts) at each turn so diff --git a/src/query.ts b/src/query.ts index e2fdb3857..b660bd269 100644 --- a/src/query.ts +++ b/src/query.ts @@ -11,6 +11,7 @@ import { MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES, type AutoCompactTrackingState, } from './services/compact/autoCompact.js' +import { consumeCompactionRequest } from './utils/memoryPressure.js' import { buildPostCompactMessages } from './services/compact/compact.js' /* eslint-disable @typescript-eslint/no-require-imports */ const reactiveCompact = feature('REACTIVE_COMPACT') @@ -536,6 +537,32 @@ async function* queryLoop( appendSystemContext(asSystemPrompt(promptWithArc), systemContext), ) + // Force compaction if memory pressure detected or message count exceeded. + // Sets forceReason on tracking so autoCompactIfNeeded bypasses the + // token-threshold check. Consumed once (one-shot) inside autocompact. + // Skip for compact/session_memory sources — those run inside an existing + // compaction and forcing would deadlock via recursive autocompaction. + const canForceCompact = + querySource !== 'compact' && querySource !== 'session_memory' + if (canForceCompact) { + const MAX_ACTIVE_MESSAGES = Number.parseInt( + process.env.OPENCLAUDE_MAX_ACTIVE_MESSAGES ?? '200', + 10, + ) + if (messagesForQuery.length > MAX_ACTIVE_MESSAGES) { + tracking = { + ...(tracking ?? { compacted: false, turnId: '', turnCounter: 0 }), + forceReason: 'message-count', + } + } + if (consumeCompactionRequest()) { + tracking = { + ...(tracking ?? { compacted: false, turnId: '', turnCounter: 0 }), + forceReason: tracking?.forceReason ?? 'memory-pressure', + } + } + } + queryCheckpoint('query_autocompact_start') const { compactionResult, diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index ca0d11451..001c80283 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -29,6 +29,7 @@ import { startPreventSleep, stopPreventSleep } from '../services/preventSleep.js import { useTerminalNotification } from '../ink/useTerminalNotification.js'; import { hasCursorUpViewportYankBug } from '../ink/terminal.js'; import { createFileStateCacheWithSizeLimit, mergeFileStateCaches, READ_FILE_STATE_CACHE_SIZE } from '../utils/fileStateCache.js'; +import { registerPrunableCache } from '../utils/memoryPressure.js'; import { updateLastInteractionTime, getLastInteractionTime, getOriginalCwd, getProjectRoot, getSessionId, switchSession, setCostStateForRestore, getTurnHookDurationMs, getTurnHookCount, resetTurnHookDuration, getTurnToolDurationMs, getTurnToolCount, resetTurnToolDuration, getTurnClassifierDurationMs, getTurnClassifierCount, resetTurnClassifierDuration, setMainLoopModelOverride, setMainThreadAgentType } from '../bootstrap/state.js'; import { asSessionId, asAgentId } from '../types/ids.js'; import { logForDebugging } from '../utils/debug.js'; @@ -2057,7 +2058,11 @@ export function REPL({ // discard the result. LRUCache construction inside FileStateCache is // expensive (~170ms), so we use useState's lazy initializer to create // it exactly once, then feed that stable reference into useRef. - const [initialReadFileState] = useState(() => createFileStateCacheWithSizeLimit(READ_FILE_STATE_CACHE_SIZE)); + const [initialReadFileState] = useState(() => { + const cache = createFileStateCacheWithSizeLimit(READ_FILE_STATE_CACHE_SIZE) + registerPrunableCache(cache) + return cache + }); const readFileState = useRef(initialReadFileState); const bashTools = useRef(new Set()); const bashToolsProcessedIdx = useRef(0); diff --git a/src/services/compact/autoCompact.ts b/src/services/compact/autoCompact.ts index 7767aae11..694978699 100644 --- a/src/services/compact/autoCompact.ts +++ b/src/services/compact/autoCompact.ts @@ -68,6 +68,10 @@ export type AutoCompactTrackingState = { // threaded through query() callers rather than serialized into transcripts. nextRetryAtMs?: number lastFailureAtMs?: number + // When set, bypasses shouldAutoCompact() token threshold check. + // Used by memory pressure and message count guards to force compaction + // even when token usage is below the normal autocompact threshold. + forceReason?: 'memory-pressure' | 'message-count' } export const AUTOCOMPACT_BUFFER_TOKENS = 13_000 @@ -255,6 +259,10 @@ export async function shouldAutoCompact( // pre-snip context, so tokenCountWithEstimation can't see the savings. // Subtract the rough-delta that snip already computed. snipTokensFreed = 0, + // When true, skip the token-threshold check but still run all guards + // (recursion, disabled, reactive-only, context-collapse). Used by + // forceReason to bypass only the token gate, not the safety guards. + skipTokenCheck = false, ): Promise { // Recursion guards. session_memory and compact are forked agents that // would deadlock. @@ -312,6 +320,11 @@ export async function shouldAutoCompact( } } + if (skipTokenCheck) { + logForDebugging('autocompact: skipping token threshold check (forced)') + return true + } + const tokenCount = tokenCountWithEstimation(messages) - snipTokensFreed const threshold = getAutoCompactThreshold(model) const effectiveWindow = getEffectiveContextWindowSize(model) @@ -349,11 +362,20 @@ export async function autoCompactIfNeeded( } const model = toolUseContext.options.mainLoopModel + // Force compaction if a pressure/count signal set forceReason. + // Consume the flag so it only forces one compaction cycle. + // Pass skipTokenCheck to shouldAutoCompact so safety guards + // (disabled, reactive-only, context-collapse, recursion) still apply. + const forcedBy = tracking?.forceReason + if (tracking?.forceReason) { + tracking.forceReason = undefined + } const shouldCompact = await shouldAutoCompact( messages, model, querySource, snipTokensFreed, + !!forcedBy, ) if (!shouldCompact) { diff --git a/src/utils/concurrentSessions.ts b/src/utils/concurrentSessions.ts index f00ce67f4..f4af17d18 100644 --- a/src/utils/concurrentSessions.ts +++ b/src/utils/concurrentSessions.ts @@ -202,3 +202,33 @@ export async function countConcurrentSessions(): Promise { } return count } + +/** + * Calculate per-session memory budget based on concurrent sessions. + * For 32GB system: 4 sessions = ~7GB each, 8 sessions = ~3.5GB each. + * Respects OPENCLAUDE_MAX_MEMORY_MB env var if set. + */ +const OS_OVERHEAD_MB = 4096 +const MIN_PER_SESSION_MB = 512 + +export async function calculatePerSessionMemoryBudget(): Promise { + const envBudget = Number.parseInt( + process.env.OPENCLAUDE_MAX_MEMORY_MB ?? '0', + 10, + ) + if (envBudget > 0) return envBudget + + const os = await import('os') + const totalMemMB = Math.floor(os.totalmem() / 1024 / 1024) + const sessionCount = await countConcurrentSessions() + const availableMB = totalMemMB - OS_OVERHEAD_MB + const budget = Math.max( + MIN_PER_SESSION_MB, + Math.floor(availableMB / Math.max(1, sessionCount)), + ) + + logForDebugging( + `[ConcurrentSessions] ${sessionCount} sessions, budget: ${budget}MB each (total: ${totalMemMB}MB)`, + ) + return budget +} diff --git a/src/utils/memoryCompaction.ts b/src/utils/memoryCompaction.ts new file mode 100644 index 000000000..5a3065d59 --- /dev/null +++ b/src/utils/memoryCompaction.ts @@ -0,0 +1,42 @@ +/** + * Memory Compaction Trigger + * + * Connects memory pressure levels to compaction actions. + * When pressure rises, triggers conversation compaction and cache pruning. + */ + +import { logForDebugging } from './debug.js' +import { onMemoryPressure } from './memoryPressure.js' + +export interface CompactionTrigger { + forceCompact(): void + dispose(): void +} + +export function createMemoryCompactionTrigger(opts: { + onCompact: (aggressive: boolean) => void + onPruneCache: () => void +}): CompactionTrigger { + const dispose = onMemoryPressure(level => { + if (level === 'elevated') { + logForDebugging( + '[MemoryCompaction] Elevated pressure - requesting compaction', + ) + opts.onCompact(false) + } else if (level === 'critical') { + logForDebugging( + '[MemoryCompaction] Critical pressure - forcing aggressive compaction', + ) + opts.onCompact(true) + opts.onPruneCache() + } + }) + + return { + forceCompact() { + opts.onCompact(true) + opts.onPruneCache() + }, + dispose, + } +} diff --git a/src/utils/memoryPressure.ts b/src/utils/memoryPressure.ts new file mode 100644 index 000000000..a4f3e0f2a --- /dev/null +++ b/src/utils/memoryPressure.ts @@ -0,0 +1,160 @@ +/** + * Memory Pressure Monitor + * + * Watches process RSS and triggers cleanup actions at configurable thresholds. + * Designed to prevent OOM when running multiple OpenClaude sessions. + */ + +import { logForDebugging } from './debug.js' + +export type MemoryPressureLevel = 'normal' | 'elevated' | 'critical' + +export interface MemoryPressureConfig { + elevatedThresholdMB: number + criticalThresholdMB: number + checkIntervalMs: number + perSessionBudgetMB: number +} + +const DEFAULT_CONFIG: MemoryPressureConfig = { + elevatedThresholdMB: 0, + criticalThresholdMB: 0, + checkIntervalMs: 30_000, + perSessionBudgetMB: Number.parseInt( + process.env.OPENCLAUDE_MAX_MEMORY_MB ?? '1536', + 10, + ), +} + +let currentLevel: MemoryPressureLevel = 'normal' +let pressureListeners: Array<(level: MemoryPressureLevel) => void> = [] +let monitorInterval: ReturnType | null = null +let compactionRequested = false + +// Registry of caches that can be pruned under critical memory pressure. +// Caches register themselves at init; the monitor prunes them all when +// RSS crosses the critical threshold. +const prunableCaches: Array<{ clear(): void }> = [] + +/** + * Register a cache for automatic pruning under critical memory pressure. + * Safe to call multiple times with the same cache (idempotent). + */ +export function registerPrunableCache(cache: { clear(): void }): void { + if (!prunableCaches.includes(cache)) { + prunableCaches.push(cache) + } +} + +/** + * Clear all registered prunable caches. Called automatically when memory + * pressure reaches 'critical'. Also callable directly for manual cache + * eviction. + */ +export function pruneRegisteredCaches(): void { + for (const cache of prunableCaches) { + try { + cache.clear() + } catch { + // best-effort — cache may already be empty + } + } +} + +export function getMemoryPressureLevel(): MemoryPressureLevel { + return currentLevel +} + +export function onMemoryPressure( + callback: (level: MemoryPressureLevel) => void, +): () => void { + pressureListeners.push(callback) + return () => { + pressureListeners = pressureListeners.filter(l => l !== callback) + } +} + +export function startMemoryPressureMonitor( + config: Partial = {}, +): void { + if (monitorInterval) return + + const resolved = { ...DEFAULT_CONFIG, ...config } + + if (resolved.elevatedThresholdMB === 0) { + resolved.elevatedThresholdMB = Math.floor( + resolved.perSessionBudgetMB * 0.8, + ) + } + if (resolved.criticalThresholdMB === 0) { + resolved.criticalThresholdMB = Math.floor( + resolved.perSessionBudgetMB * 0.9, + ) + } + + logForDebugging( + `[MemoryPressure] Monitor started: elevated=${resolved.elevatedThresholdMB}MB, critical=${resolved.criticalThresholdMB}MB, interval=${resolved.checkIntervalMs}ms`, + ) + + monitorInterval = setInterval(() => { + const rss = process.memoryUsage().rss / 1024 / 1024 + const previousLevel = currentLevel + + if (rss >= resolved.criticalThresholdMB) { + currentLevel = 'critical' + } else if (rss >= resolved.elevatedThresholdMB) { + currentLevel = 'elevated' + } else { + currentLevel = 'normal' + } + + if (currentLevel !== previousLevel) { + logForDebugging( + `[MemoryPressure] Level changed: ${previousLevel} -> ${currentLevel} (RSS: ${rss.toFixed(0)}MB)`, + ) + if (currentLevel === 'critical') { + logForDebugging('[MemoryPressure] Critical — pruning registered caches') + pruneRegisteredCaches() + } + for (const listener of pressureListeners) { + try { + listener(currentLevel) + } catch { + // Don't let listener errors crash the monitor + } + } + } + + // Keep requesting compaction while pressure stays elevated/critical. + // The previous level-change-only gate meant one compact/prune cycle then + // silence even if RSS remained high. consumeCompactionRequest() is + // one-shot so the existing autocompact cooldown prevents retry storms. + if (currentLevel !== 'normal') { + compactionRequested = true + } + }, resolved.checkIntervalMs) + + // Don't keep process alive just for monitoring + ;(monitorInterval as ReturnType & { unref?: () => void }).unref?.() +} + +/** + * Returns true if memory pressure triggered a compaction request since last check. + * Consumes the flag (resets to false). + */ +export function consumeCompactionRequest(): boolean { + if (compactionRequested) { + compactionRequested = false + return true + } + return false +} + +export function stopMemoryPressureMonitor(): void { + if (monitorInterval) { + clearInterval(monitorInterval) + monitorInterval = null + } + currentLevel = 'normal' + pressureListeners = [] +}