From 184160ebb845d529a96c57c01250821fd8c83245 Mon Sep 17 00:00:00 2001 From: James Feng <47167674+GhostDragon124@users.noreply.github.com> Date: Mon, 1 Jun 2026 20:37:23 +0800 Subject: [PATCH] Apply predictive compact memory optimizations --- src/query.ts | 60 ++++++++++++-- src/screens/REPL.tsx | 10 ++- src/services/compact/autoCompact.ts | 31 ++++++- src/services/compact/compact.ts | 13 ++- src/services/compact/reactiveCompact.ts | 105 ++++++++++++++++++++---- src/utils/readFileInRange.ts | 14 ++++ 6 files changed, 201 insertions(+), 32 deletions(-) diff --git a/src/query.ts b/src/query.ts index 7d712e9dd..58aaf78df 100644 --- a/src/query.ts +++ b/src/query.ts @@ -7,6 +7,9 @@ import type { CanUseToolFn } from './hooks/useCanUseTool.js' import { FallbackTriggeredError } from './services/api/withRetry.js' import { calculateTokenWarningState, + estimateMaxTurnGrowth, + getAutoCompactThreshold, + getEffectiveContextWindowSize, isAutoCompactEnabled, type AutoCompactTrackingState, } from './services/compact/autoCompact.js' @@ -434,7 +437,7 @@ async function* queryLoop( queryTracking, } - let messagesForQuery = [...getMessagesAfterCompactBoundary(messages)] + let messagesForQuery = getMessagesAfterCompactBoundary(messages) let tracking = autoCompactTracking @@ -719,6 +722,48 @@ async function* queryLoop( } } + // Predictive autocompact: estimate if this turn's growth will push + // us past the context window. Uses effectiveContextWindow directly + // (without the autocompact buffer) to avoid double-reserving with + // getAutoCompactThreshold which already subtracts buffer. + if (!compactionResult && isAutoCompactEnabled()) { + const model = toolUseContext.options.mainLoopModel + const currentTokens = + tokenCountWithEstimation(messagesForQuery) - snipTokensFreed + const estimatedGrowth = estimateMaxTurnGrowth(model) + const predictiveThreshold = + getEffectiveContextWindowSize(model) - estimatedGrowth + if (currentTokens > predictiveThreshold) { + const predictiveResult = await deps.autocompact( + messagesForQuery, + toolUseContext, + { + systemPrompt, + userContext, + systemContext, + toolUseContext, + forkContextMessages: messagesForQuery, + }, + querySource, + tracking, + snipTokensFreed, + ) + if (predictiveResult.compactionResult) { + messagesForQuery = buildPostCompactMessages( + predictiveResult.compactionResult, + ) + snipTokensFreed = 0 + tracking = tracking + ? { + ...tracking, + compacted: true, + consecutiveFailures: predictiveResult.consecutiveFailures ?? 0, + } + : tracking + } + } + } + let attemptWithFallback = true queryCheckpoint('query_api_loop_start') @@ -1075,7 +1120,7 @@ async function* queryLoop( // Execute post-sampling hooks after model response is complete if (assistantMessages.length > 0) { void executePostSamplingHooks( - [...messagesForQuery, ...assistantMessages], + messagesForQuery.concat(assistantMessages), systemPrompt, userContext, systemContext, @@ -1768,11 +1813,10 @@ async function* queryLoop( userContext, systemContext, toolUseContext, - forkContextMessages: [ - ...messagesForQuery, - ...assistantMessages, - ...toolResults, - ], + forkContextMessages: messagesForQuery.concat( + assistantMessages, + toolResults, + ), }) } } @@ -1789,7 +1833,7 @@ async function* queryLoop( queryCheckpoint('query_recursive_call') const next: State = { - messages: [...messagesForQuery, ...assistantMessages, ...toolResults], + messages: messagesForQuery.concat(assistantMessages, toolResults), toolUseContext: toolUseContextWithQueryTracking, autoCompactTracking: tracking, turnCount: nextTurnCount, diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index 31c210231..9e63d436e 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -1553,7 +1553,15 @@ export function REPL({ // Deferred messages for the Messages component — renders at transition // priority so the reconciler yields every 5ms, keeping input responsive // while the expensive message processing pipeline runs. - const deferredMessages = useDeferredValue(messages); + // Cap at 500 messages to limit memory double-buffering. The bypass + // at display-time uses sync messages during streaming and non-loading, + // so this cap only affects reduced-motion scenarios. + const DEFERRED_CAP = 500; + const cappedMessages = React.useMemo( + () => (messages.length > DEFERRED_CAP ? messages.slice(-DEFERRED_CAP) : messages), + [messages], + ); + const deferredMessages = useDeferredValue(cappedMessages); const deferredBehind = messages.length - deferredMessages.length; if (deferredBehind > 0) { logForDebugging( diff --git a/src/services/compact/autoCompact.ts b/src/services/compact/autoCompact.ts index 4025897ce..b77fb8b61 100644 --- a/src/services/compact/autoCompact.ts +++ b/src/services/compact/autoCompact.ts @@ -64,6 +64,35 @@ export const WARNING_THRESHOLD_BUFFER_TOKENS = 20_000 export const ERROR_THRESHOLD_BUFFER_TOKENS = 20_000 export const MANUAL_COMPACT_BUFFER_TOKENS = 3_000 +// Conservative estimate for tool result growth per turn. +// Typical tool results (file reads, grep, bash) average ~5-10K tokens; +// occasional large reads can spike to 20K+. +const TOOL_RESULT_GROWTH_ESTIMATE = 15_000 + +/** + * Context-aware autocompact buffer. Larger context windows need more + * headroom because a single turn can produce proportionally more tokens + * (longer model outputs + larger tool results). + */ +export function getAutocompactBufferTokens(model: string): number { + const effectiveWindow = getEffectiveContextWindowSize(model) + if (effectiveWindow >= 800_000) return 50_000 + if (effectiveWindow >= 400_000) return 30_000 + return AUTOCOMPACT_BUFFER_TOKENS +} + +/** + * Estimate the maximum token growth a single turn can produce. + * Used for predictive autocompact checks before the API call. + */ +export function estimateMaxTurnGrowth(model: string): number { + const maxOutput = Math.min( + getMaxOutputTokensForModel(model), + MAX_OUTPUT_TOKENS_FOR_SUMMARY, + ) + return maxOutput + TOOL_RESULT_GROWTH_ESTIMATE +} + // Stop trying autocompact after this many consecutive failures. // BQ 2026-03-10: 1,279 sessions had 50+ consecutive failures (up to 3,272) // in a single session, wasting ~250K API calls/day globally. @@ -73,7 +102,7 @@ export function getAutoCompactThreshold(model: string): number { const effectiveContextWindow = getEffectiveContextWindowSize(model) const autocompactThreshold = - effectiveContextWindow - AUTOCOMPACT_BUFFER_TOKENS + effectiveContextWindow - getAutocompactBufferTokens(model) // Override for easier testing of autocompact const envPercent = process.env.CLAUDE_AUTOCOMPACT_PCT_OVERRIDE diff --git a/src/services/compact/compact.ts b/src/services/compact/compact.ts index 710701dcb..c0f195e07 100644 --- a/src/services/compact/compact.ts +++ b/src/services/compact/compact.ts @@ -332,13 +332,12 @@ export type RecompactionInfo = { * Order: boundaryMarker, summaryMessages, messagesToKeep, attachments, hookResults */ export function buildPostCompactMessages(result: CompactionResult): Message[] { - return [ - result.boundaryMarker, - ...result.summaryMessages, - ...(result.messagesToKeep ?? []), - ...result.attachments, - ...result.hookResults, - ] + return ([result.boundaryMarker] as Message[]).concat( + result.summaryMessages, + result.messagesToKeep ?? [], + result.attachments, + result.hookResults, + ) } /** diff --git a/src/services/compact/reactiveCompact.ts b/src/services/compact/reactiveCompact.ts index 67c872cda..8045f3f3b 100644 --- a/src/services/compact/reactiveCompact.ts +++ b/src/services/compact/reactiveCompact.ts @@ -1,22 +1,97 @@ -// Auto-generated stub — replace with real implementation -export {}; +import { isEnvTruthy } from '../../utils/envUtils.js' +import { + isMediaSizeErrorMessage, + isPromptTooLongMessage, +} from '../api/errors.js' +import type { AssistantMessage, Message } from '../../types/message.js' +import { type CompactionResult, compactConversation } from './compact.js' +import { logError } from '../../utils/log.js' +import { logForDebugging } from '../../utils/debug.js' +import type { CacheSafeParams } from '../../utils/forkedAgent.js' -import type { Message } from 'src/types/message'; -import type { CompactionResult } from './compact.js'; +export const isReactiveOnlyMode: () => boolean = () => false -export const isReactiveOnlyMode: () => boolean = () => false; export const reactiveCompactOnPromptTooLong: ( messages: Message[], cacheSafeParams: Record, options: { customInstructions?: string; trigger?: string }, -) => Promise<{ ok: boolean; reason?: string; result?: CompactionResult }> = async () => ({ ok: false }); -export const isReactiveCompactEnabled: () => boolean = () => false; -export const isWithheldPromptTooLong: (message: Message) => boolean = () => false; -export const isWithheldMediaSizeError: (message: Message) => boolean = () => false; +) => Promise<{ ok: boolean; reason?: string; result?: CompactionResult }> = + async (messages, cacheSafeParams, options) => { + const params = cacheSafeParams as unknown as CacheSafeParams + try { + const result = await compactConversation( + messages, + params.toolUseContext, + params, + true, + options.customInstructions, + true, + { + isRecompactionInChain: false, + turnsSincePreviousCompact: 0, + autoCompactThreshold: 0, + querySource: 'compact', + }, + ) + return { ok: true, result } + } catch (error) { + logError(error) + return { ok: false, reason: String(error) } + } + } + +export const isReactiveCompactEnabled: () => boolean = () => { + if (isEnvTruthy(process.env.DISABLE_COMPACT)) return false + return true +} + +export const isWithheldPromptTooLong: (message: Message) => boolean = + message => { + if (message.type !== 'assistant' || !message.isApiErrorMessage) return false + return isPromptTooLongMessage(message as AssistantMessage) + } + +export const isWithheldMediaSizeError: (message: Message) => boolean = + message => { + if (message.type !== 'assistant' || !message.isApiErrorMessage) return false + return isMediaSizeErrorMessage(message as AssistantMessage) + } + export const tryReactiveCompact: (params: { - hasAttempted: boolean; - querySource: string; - aborted: boolean; - messages: Message[]; - cacheSafeParams: Record; -}) => Promise = async () => null; + hasAttempted: boolean + querySource: string + aborted: boolean + messages: Message[] + cacheSafeParams: Record +}) => Promise = async ({ + hasAttempted, + aborted, + messages, + cacheSafeParams, +}) => { + if (hasAttempted || aborted) return null + const params = cacheSafeParams as unknown as CacheSafeParams + try { + const result = await compactConversation( + messages, + params.toolUseContext, + params, + true, + undefined, + true, + { + isRecompactionInChain: false, + turnsSincePreviousCompact: 0, + autoCompactThreshold: 0, + }, + ) + return result + } catch (error) { + logForDebugging( + `reactiveCompact: emergency compaction failed — ${String(error)}`, + { level: 'warn' }, + ) + logError(error) + return null + } +} diff --git a/src/utils/readFileInRange.ts b/src/utils/readFileInRange.ts index de12d2b49..18086135c 100644 --- a/src/utils/readFileInRange.ts +++ b/src/utils/readFileInRange.ts @@ -101,6 +101,20 @@ export async function readFileInRange( throw new FileTooLargeError(stats.size, maxBytes) } + // For targeted reads of moderately large files, prefer streaming to + // avoid loading the full file into memory when only a slice is needed. + const isTargetedRead = offset > 0 || maxLines !== undefined + if (isTargetedRead && stats.size > FAST_PATH_MAX_SIZE / 4) { + return readFileInRangeStreaming( + filePath, + offset, + maxLines, + maxBytes, + truncateOnByteLimit, + signal, + ) + } + const text = await readFile(filePath, { encoding: 'utf8', signal }) return readFileInRangeFast( text,