diff --git a/docs/memory-leak-audit.md b/docs/memory-leak-audit.md new file mode 100644 index 000000000..e2b423f61 --- /dev/null +++ b/docs/memory-leak-audit.md @@ -0,0 +1,579 @@ +# 内存泄漏排查报告 + +> 基于官方 CHANGELOG 记录的 11 个已修复内存泄漏 + 1 个代码注释中的已知问题,对反编译代码库进行逐文件验证。 +> 审计日期:2026-04-28 + +## TODO + +- [ ] #1 图片处理无限内存增长 — 确认已实现 ✅ +- [ ] #2 /usage 命令泄漏约 2GB — 确认已实现 ✅ +- [ ] #3 长时间运行工具进度事件泄漏 — 确认已实现 ✅ +- [ ] #4 空闲重新渲染循环 — 部分实现,keepAlive 集成完整性待确认 +- [ ] #5 虚拟滚动器保留历史消息拷贝 — 确认已实现 ✅ +- [ ] #6 管道模式超宽行过度分配 — 确认已实现 ✅ +- [ ] #7 语言语法按需加载 — 已回退为静态导入,修正内存估计(~5-15MB 非 ~50MB) +- [ ] #8 NO_FLICKER 模式流状态泄漏 — resetLoadingState 已实现,StreamingToolExecutor.discard 完整性待确认 +- [ ] #9 Remote Control 权限条目保留 — 核心清理已实现,hook cleanup 完整性待确认 +- [ ] #10 MCP HTTP/SSE 缓冲区累积 — 确认已实现 ✅ +- [ ] #11 LRU 缓存键保留大 JSON — **归类需修正**:实际已完整实现(sizeCalculation + maxSize 上限),非"未实现" +- [x] #12 QueryEngine.mutableMessages 不收缩 — **已修复**:实现 snipCompactIfNeeded(按 removedUuids 过滤)+ snipProjection(边界检测 + 视图投影),28 tests + +## 总览 +--- + +## 1. 图片处理无限内存增长 (v2.1.121) + +**CHANGELOG 描述**:Fixed unbounded memory growth (multi-GB RSS) when processing many images in a session + +### 实现位置 + +- `src/utils/imageStore.ts` — 核心修复 +- `src/commands/clear/caches.ts` — 缓存清理 +- `src/screens/REPL.tsx` — UI 层释放 + +### 修复方式 + +三层防护机制: + +1. **LRU 内存缓存**:`storedImagePaths` Map 上限 200 条目(`MAX_STORED_IMAGE_PATHS`),超出自动驱逐最早条目 +2. **磁盘持久化**:图片 base64 数据写入 `~/.claude/image-cache//`,内存中仅保留路径字符串 +3. **立即释放**:`setPastedContents({})` 在消息提交/命令执行后清空 React state 中的 base64 数据 + +### 关键代码 + +```typescript +// imageStore.ts:10 +const MAX_STORED_IMAGE_PATHS = 200 + +// imageStore.ts:115-124 +function evictOldestIfAtCap(): void { + while (storedImagePaths.size >= MAX_STORED_IMAGE_PATHS) { + const oldest = storedImagePaths.keys().next().value + if (oldest !== undefined) { + storedImagePaths.delete(oldest) + } else { + break + } + } +} + +// imageStore.ts:129-167 — 清理旧会话目录 +export async function cleanupOldImageCaches(): Promise { ... } +``` + +--- + +## 2. /usage 命令泄漏约 2GB (v2.1.121) + + +**CHANGELOG 描述**:Fixed /usage leaking up to ~2GB of memory on machines with large transcript histories + +### 实现位置 + +- `src/utils/sessionStoragePortable.ts:716-792` — 核心流式读取 +- `src/utils/attribution.ts` — 调用方 + +### 修复方式 + +1. **分块流式读取**:使用 `TRANSCRIPT_READ_CHUNK_SIZE = 1MB` 固定块大小,通过 `fd.read()` 逐块处理,避免一次性加载整个 transcript +2. **字节级过滤**:在 fd 层面直接跳过 `attribution-snapshot` 类型的行(占长会话 84% 的字节空间) +3. **边界截断**:搜索 `compact_boundary` 标记,只保留边界之后的数据 +4. **缓冲区控制**:初始缓冲区限制 `Math.min(fileSize, 8MB)` + +### 关键代码 + +```typescript +// sessionStoragePortable.ts:716-792 +export async function readTranscriptForLoad( + filePath: string, + fileSize: number, +): Promise<{ + boundaryStartOffset: number + postBoundaryBuf: Buffer + hasPreservedSegment: boolean +}> { + const s: LoadState = { + out: { + buf: Buffer.allocUnsafe(Math.min(fileSize, 8 * 1024 * 1024)), + len: 0, + cap: fileSize + 1, + }, + // ... + } + const chunk = Buffer.allocUnsafe(CHUNK_SIZE) + const fd = await fsOpen(filePath, 'r') + try { + let filePos = 0 + while (filePos < fileSize) { + const { bytesRead } = await fd.read(chunk, 0, Math.min(CHUNK_SIZE, fileSize - filePos), filePos) + if (bytesRead === 0) break + filePos += bytesRead + // ... 分块处理逻辑 + } + finalizeOutput(s) + } finally { + await fd.close() + } +} +``` + +--- + +## 3. 长时间运行工具进度事件泄漏 (v2.1.121) + + +**CHANGELOG 描述**:Fixed memory leak when long-running tools fail to emit a clear progress event + +### 实现位置 + +- `src/screens/REPL.tsx:3054-3114` — progress 消息替换逻辑 +- `src/utils/sessionStorage.ts:186-196` — 临时消息类型定义 + +### 修复方式 + +1. **向后扫描替换**:从只检查最后一条消息改为向后遍历所有 progress 消息,找到匹配的 `parentToolUseID` + `type` 后替换(修复交错消息导致 13k+ 条目堆积) +2. **全屏模式硬上限**:`MAX_FULLSCREEN_SCROLLBACK = 500`,超出截断 +3. **临时消息识别**:`isEphemeralToolProgress()` 区分 `bash_progress`、`sleep_progress` 等一次性消息与需要保留的 `agent_progress` 等 + +### 关键代码 + +```typescript +// REPL.tsx:3094-3114 +setMessages(oldMessages => { + const newData = newMessage.data as Record; + // Scan backwards to find the last ephemeral progress with matching + // parentToolUseID and type. + for (let i = oldMessages.length - 1; i >= 0; i--) { + const m = oldMessages[i]! + if (m.type !== 'progress') break + const mData = m.data as Record | undefined + if ( + m.parentToolUseID === newMessage.parentToolUseID && + mData?.type === newData.type + ) { + const copy = oldMessages.slice(); + copy[i] = newMessage; + return copy; + } + } + return [...oldMessages, newMessage]; +}); + +// REPL.tsx:3058-3064 — 全屏模式硬上限 +const MAX_FULLSCREEN_SCROLLBACK = 500 +const kept = postBoundary.length > MAX_FULLSCREEN_SCROLLBACK + ? postBoundary.slice(-MAX_FULLSCREEN_SCROLLBACK) + : postBoundary +return [...kept, newMessage] +``` + +--- + +## 4. 空闲重新渲染循环 (v2.1.117) + +**状态:部分实现** + +**CHANGELOG 描述**:Fixed idle re-render loop when background tasks are present, reducing memory growth on Linux + +### 实现位置 + +- `packages/@ant/ink/src/components/ClockContext.tsx` — 核心时钟管理 + +### 已实现部分 + +`ClockContext` 的 `keepAlive` 订阅者分类机制完整存在: + +```typescript +// ClockContext.tsx:11-43 +function createClock(tickIntervalMs: number): Clock { + const subscribers = new Map<() => void, boolean>() + let interval: ReturnType | null = null + + function updateInterval(): void { + const anyKeepAlive = [...subscribers.values()].some(Boolean) + if (anyKeepAlive) { + // 有 keepAlive 订阅者时启动 interval + interval = setInterval(tick, currentTickIntervalMs) + } else if (interval) { + // 无 keepAlive 订阅者时停止 interval + clearInterval(interval) + interval = null + } + } + + return { + subscribe(onChange, keepAlive) { + subscribers.set(onChange, keepAlive) + updateInterval() + return () => { + subscribers.delete(onChange) + updateInterval() + } + }, + // ... + } +} +``` + +### 不确定部分 + +无法确认 `useAnimationFrame` hook 是否在所有使用时钟的组件中正确传递了 `keepAlive` 参数。反编译代码中调用链可能不完整。 + +--- + +## 5. 虚拟滚动器保留历史消息拷贝 (v2.1.101) + + +**CHANGELOG 描述**:Fixed a memory leak where long sessions retained dozens of historical copies of the message list in the virtual scroller + +### 实现位置 + +- `src/components/VirtualMessageList.tsx:276-296` + +### 修复方式 + +增量式键值数组:使用 `useRef` 保存 keys 数组引用,流式追加而非每次 O(n) 全量重建。 + +```typescript +// VirtualMessageList.tsx:276-296 +const keysRef = useRef([]) +const prevMessagesRef = useRef(messages) +const prevItemKeyRef = useRef(itemKey) +if ( + prevItemKeyRef.current !== itemKey || + messages.length < keysRef.current.length || + messages[0] !== prevMessagesRef.current[0] +) { + // 全量重建(仅在 itemKey 变化、数组缩短等场景) + keysRef.current = messages.map(m => itemKey(m)) +} else { + // 增量追加(正常流式场景) + for (let i = keysRef.current.length; i < messages.length; i++) { + keysRef.current.push(itemKey(messages[i]!)) + } +} +prevMessagesRef.current = messages +prevItemKeyRef.current = itemKey +const keys = keysRef.current +``` + +修复前 27k 消息时每次新消息添加产生 ~1MB 内存分配,修复后降为 O(1) 追加。 + +--- + +## 6. 管道模式超宽行过度分配 (v2.1.110) + + +**CHANGELOG 描述**:Fixed potential excessive memory allocation when piped (non-TTY) Ink output contains a single very wide line + +### 实现位置 + +- `packages/@ant/ink/src/core/output.ts:200-207` + +### 修复方式 + +在 `Output.reset()` 中当字符缓存超过 16384 条目时清空: + +```typescript +// output.ts:200-207 +reset(width: number, height: number, screen: Screen): void { + this.width = width + this.height = height + this.screen = screen + this.operations.length = 0 + resetScreen(screen, width, height) + if (this.charCache.size > 16384) this.charCache.clear() // 关键修复 +} +``` + +--- + +## 7. 语言语法按需加载 (v2.1.108) + +**状态:已回退为静态导入** + +**CHANGELOG 描述**:Reduced memory footprint for file reads, edits, and syntax highlighting by loading language grammars on demand + +### 实现位置 + +- `packages/color-diff-napi/src/index.ts:21-37` + +### 当前状态 + +延迟加载逻辑**已被移除**,改为顶层静态导入。代码注释说明原因: + +```typescript +// color-diff-napi/src/index.ts:21-37 +// Static import — createRequire(import.meta.url) fails in Bun --compile mode +// because the resolved path points to the internal bunfs binary path where +// node_modules cannot be found. A top-level import ensures the module is +// bundled and accessible at runtime. +import hljs from 'highlight.js' // 顶层静态导入 + +type HLJSApi = typeof hljs +let cachedHljs: HLJSApi | null = null +function hljsApi(): HLJSApi { + if (cachedHljs) return cachedHljs + const mod = hljs as HLJSApi & { default?: HLJSApi } + cachedHljs = 'default' in mod && mod.default ? mod.default : mod + return cachedHljs! +} +``` + +**影响**:highlight.js 包含 190+ 语言语法(约 50MB),现在在模块加载时即全部载入内存,无法按需释放。这是为了兼容 Bun `--compile` 模式做的妥协。 + +--- + +## 8. NO_FLICKER 模式流状态泄漏 (v2.1.105) + +**状态:可疑 — 框架存在但完整性不确定** + +**CHANGELOG 描述**:Fixed a NO_FLICKER mode memory leak where API retries left stale streaming state + +### 实现位置 + +- `src/screens/REPL.tsx:1841-1861` — `resetLoadingState()` +- `src/screens/REPL.tsx:3568-3578` — finally 块调用 + +### 已实现部分 + +`resetLoadingState()` 在 `onQuery` 的 finally 块中无条件调用,清理 `streamingText`、`streamingToolUses` 等: + +```typescript +// REPL.tsx:1841-1861 +const resetLoadingState = useCallback(() => { + setStreamingText(null); + setStreamingToolUses([]); + setSpinnerMessage(null); + // ... +}, [pickNewSpinnerTip]); + +// REPL.tsx:3568-3578 — finally 块 +} finally { + if (queryGuard.end(thisGeneration)) { + resetLoadingState(); // 无条件清理 + } +} +``` + +### 不确定部分 + +无法确认 `query.ts` 中 `StreamingToolExecutor.discard()` 的逻辑是否完整实现了旧工具结果的释放。 + +--- + +## 9. Remote Control 权限条目保留 (v2.1.98) + +**状态:部分实现** + +**CHANGELOG 描述**:Fixed a memory leak where Remote Control permission handler entries were retained for the lifetime of the session + +### 实现位置 + +- `src/hooks/useReplBridge.tsx:466-491` — 处理 + 删除 +- `src/hooks/useReplBridge.tsx:712-717` — 注册 + 清理函数 + +### 已实现部分 + +```typescript +// useReplBridge.tsx:466-491 +const pendingPermissionHandlers = new Map void>() + +function handlePermissionResponse(msg: SDKControlResponse): void { + const requestId = msg.response?.request_id + if (!requestId) return + const handler = pendingPermissionHandlers.get(requestId) + if (!handler) return + const parsed = parseBridgePermissionResponse(msg) + if (!parsed) return + pendingPermissionHandlers.delete(requestId) // 处理后删除 + handler(parsed) +} + +// useReplBridge.tsx:712-717 +onResponse(requestId, handler) { + pendingPermissionHandlers.set(requestId, handler) + return () => { + pendingPermissionHandlers.delete(requestId) // 取消时删除 + } +} +``` + +### 不确定部分 + +hook 的 cleanup 函数(组件卸载时的 `replBridgePermissionCallbacks = undefined`)是否完整调用。 + +--- + +## 10. MCP HTTP/SSE 缓冲区累积 (v2.1.97) + + +**CHANGELOG 描述**:Fixed MCP HTTP/SSE connections accumulating ~50 MB/hr of unreleased buffers when servers reconnect + +### 实现位置 + +- `src/services/api/claude.ts:1557-1564` — `releaseStreamResources()` +- `src/cli/transports/SSETransport.ts:419` — `reader.releaseLock()` +- `@modelcontextprotocol/sdk` (sse.js, streamableHttp.js) — `response.body?.cancel()` + +### 修复方式 + +1. **主动释放响应体**:`releaseStreamResources()` 清理 stream 和 response + +```typescript +// claude.ts:1553-1564 +// Release all stream resources to prevent native memory leaks. +// The Response object holds native TLS/socket buffers that live outside the +// V8 heap (observed on the Node.js/npm path; see GH #32920), so we must +// explicitly cancel and release it regardless of how the generator exits. +function releaseStreamResources(): void { + cleanupStream(stream) + stream = undefined + if (streamResponse) { + streamResponse.body?.cancel().catch(() => {}) + streamResponse = undefined + } +} +``` + +2. **SSE 读取器释放**: + +```typescript +// SSETransport.ts:418-419 +} finally { + reader.releaseLock() +} +``` + +3. **MCP SDK 层面**:在所有 HTTP 路径(成功/失败/重连)调用 `response.body?.cancel()` + +--- + +## 11. LRU 缓存键保留大 JSON (v2.1.89) + + +**CHANGELOG 描述**:Fixed memory leak where large JSON inputs were retained as LRU cache keys in long-running sessions + +### 实现位置 + +- `src/utils/fileStateCache.ts:37-48` — 大小计算修复 +- `src/utils/queryHelpers.ts:48-54` — 类型强制转换 + +### 修复方式 + +1. **正确计算缓存大小**:处理 `content` 为嵌套对象的情况 + +```typescript +// fileStateCache.ts:37-48 +sizeCalculation: value => { + const c = value.content + const s = + typeof c === 'string' + ? c + : c === null || c === undefined + ? '' + : typeof c === 'object' + ? JSON.stringify(c) + : String(c) + return Math.max(1, Buffer.byteLength(s, 'utf8')) +} +``` + +2. **强制类型转换**:确保 Write 工具 content 始终为字符串 + +```typescript +// queryHelpers.ts:48-54 +function coerceToolContentToString(value: unknown): string { + if (typeof value === 'string') return value + if (value === null || value === undefined) return '' + if (typeof value === 'object') return JSON.stringify(value) + return String(value) +} +``` + +--- + +## 12. QueryEngine.mutableMessages 不收缩 + +**状态:已修复** + +**代码注释描述**:`markers persist and re-trigger on every turn, and mutableMessages never shrinks (memory leak in long SDK sessions)`(`src/QueryEngine.ts:929-930`) + +### 实现位置 + +- `src/services/compact/snipCompact.ts` — **存根文件** +- `src/QueryEngine.ts:925-962` — 消息处理逻辑 + +### 问题详情 + +`mutableMessages` 数组只增不减,每轮对话 push 多条消息(assistant、progress、user、attachment 等)。清理依赖两条路径: + +**路径 1:API 返回 compact_boundary**(已实现) + +```typescript +// QueryEngine.ts:946-962 +if (msg.subtype === 'compact_boundary' && msg.compactMetadata) { + const mutableBoundaryIdx = this.mutableMessages.length - 1 + if (mutableBoundaryIdx > 0) { + this.mutableMessages.splice(0, mutableBoundaryIdx) // 清理旧消息 + } +} +``` + +**路径 2:本地 snip 压缩**(存根 — 永不执行) + +```typescript +// snipCompact.ts — 完整文件 +// Auto-generated stub — replace with real implementation +export {}; +import type { Message } from 'src/types/message'; + +export const isSnipMarkerMessage: (message: Message) => boolean = () => false; +export const snipCompactIfNeeded: ( + messages: Message[], + options?: { force?: boolean }, +) => { messages: Message[]; executed: boolean; tokensFreed: number; boundaryMessage?: Message } = (messages) => ({ + messages, + executed: false, // 永远 false — 清理从不执行 + tokensFreed: 0, +}); +export const isSnipRuntimeEnabled: () => boolean = () => false; +export const shouldNudgeForSnips: (messages: Message[]) => boolean = () => false; +export const SNIP_NUDGE_TEXT: string = ''; +``` + +`snipReplay` 回调依赖 `HISTORY_SNIP` feature flag,且调用的 `snipCompactIfNeeded` 永远返回 `executed: false`。 + +```typescript +// QueryEngine.ts:933-942 +const snipResult = this.config.snipReplay?.(msg, this.mutableMessages) +if (snipResult !== undefined) { + if (snipResult.executed) { // 永远是 false + this.mutableMessages.length = 0 + this.mutableMessages.push(...snipResult.messages) + } + break +} +``` + +### 风险评估 + +- 在长时间 SDK 会话中,如果 API 不频繁返回 `compact_boundary`,`mutableMessages` 会持续增长 +- 每条消息可能包含大量内容(工具输出、文件内容等),长时间运行可能导致 GB 级内存占用 +- 这是当前代码库中**最明确的未实现内存泄漏点** + +--- + +## 总结 + +``` +确认已实现 (7): #1 图片 #2 /usage #3 进度消息 #5 虚拟滚动器 #6 管道输出 #10 MCP缓冲区 #12 snipCompact +部分实现 (3): #4 空闲渲染 #9 RC权限 #8 NO_FLICKER流状态 +未实现/存根 (2): #7 语法加载(已回退) #11 LRU缓存键 +``` + +### 需要关注的优先级 + +1. ~~**P0 — `snipCompact.ts` 存根**:唯一完全不工作的清理路径,长时间 SDK 会话必然触发~~ **已修复** +2. **P1 — 语法按需加载回退**:highlight.js 190+ 语法常驻内存(~5-15MB),无释放时机 +3. **P2 — NO_FLICKER / 空闲渲染**:框架存在但反编译代码中集成完整性不确定 diff --git a/src/services/compact/__tests__/snipCompact.test.ts b/src/services/compact/__tests__/snipCompact.test.ts new file mode 100644 index 000000000..893f2ddf6 --- /dev/null +++ b/src/services/compact/__tests__/snipCompact.test.ts @@ -0,0 +1,222 @@ +import { describe, expect, test } from 'bun:test' +import { + isSnipMarkerMessage, + isSnipRuntimeEnabled, + shouldNudgeForSnips, + snipCompactIfNeeded, + SNIP_NUDGE_TEXT, +} from '../snipCompact.js' +import type { Message } from 'src/types/message.js' + +// --- Helpers --- + +function makeMessage(uuid: string, type: Message['type'] = 'user'): Message { + return { + type, + uuid, + message: { + role: type === 'user' ? 'user' : 'assistant', + content: `Message ${uuid}`, + }, + } as Message +} + +function makeSystemMessage( + uuid: string, + subtype?: string, + extra?: Record, +): Message { + const msg: Message = { + type: 'system', + uuid, + message: { role: 'system', content: '' }, + ...extra, + } as Message + if (subtype) { + ;(msg as Record).subtype = subtype + } + return msg +} + +function makeSnipBoundary( + uuid: string, + removedUuids: string[], +): Message { + return makeSystemMessage(uuid, 'snip_boundary', { + snipMetadata: { removedUuids }, + content: '[snip] Conversation history before this point has been snipped.', + }) +} + +// --- isSnipMarkerMessage --- + +describe('isSnipMarkerMessage', () => { + test('returns true for system message with snip_marker subtype', () => { + const msg = makeSystemMessage('m1', 'snip_marker') + expect(isSnipMarkerMessage(msg)).toBe(true) + }) + + test('returns false for system message with other subtype', () => { + const msg = makeSystemMessage('m1', 'snip_boundary') + expect(isSnipMarkerMessage(msg)).toBe(false) + }) + + test('returns false for non-system message', () => { + const msg = makeMessage('m1', 'user') + expect(isSnipMarkerMessage(msg)).toBe(false) + }) +}) + +// --- isSnipRuntimeEnabled --- + +describe('isSnipRuntimeEnabled', () => { + test('returns true (module is only loaded when HISTORY_SNIP is on)', () => { + expect(isSnipRuntimeEnabled()).toBe(true) + }) +}) + +// --- shouldNudgeForSnips --- + +describe('shouldNudgeForSnips', () => { + test('returns false for short conversation', () => { + const msgs = Array.from({ length: 10 }, (_, i) => makeMessage(`u${i}`)) + expect(shouldNudgeForSnips(msgs)).toBe(false) + }) + + test('returns true for long conversation', () => { + const msgs = Array.from({ length: 35 }, (_, i) => makeMessage(`u${i}`)) + expect(shouldNudgeForSnips(msgs)).toBe(true) + }) + + test('returns true at exact threshold', () => { + const msgs = Array.from({ length: 30 }, (_, i) => makeMessage(`u${i}`)) + expect(shouldNudgeForSnips(msgs)).toBe(true) + }) +}) + +// --- SNIP_NUDGE_TEXT --- + +describe('SNIP_NUDGE_TEXT', () => { + test('is a non-empty string', () => { + expect(typeof SNIP_NUDGE_TEXT).toBe('string') + expect(SNIP_NUDGE_TEXT.length).toBeGreaterThan(0) + }) +}) + +// --- snipCompactIfNeeded --- + +describe('snipCompactIfNeeded', () => { + test('returns messages unchanged when no snip boundary exists', () => { + const msgs = [makeMessage('a'), makeMessage('b'), makeMessage('c')] + const result = snipCompactIfNeeded(msgs) + expect(result.executed).toBe(false) + expect(result.messages).toBe(msgs) // same reference + expect(result.tokensFreed).toBe(0) + expect(result.boundaryMessage).toBeUndefined() + }) + + test('removes messages listed in removedUuids', () => { + const a = makeMessage('a') + const b = makeMessage('b') + const c = makeMessage('c') + const boundary = makeSnipBoundary('bnd', ['a', 'b']) + + const msgs = [a, b, c, boundary] + const result = snipCompactIfNeeded(msgs) + + expect(result.executed).toBe(true) + expect(result.messages).toHaveLength(2) + expect(result.messages.map((m) => m.uuid) as string[]).toEqual(['c', 'bnd']) + expect(result.tokensFreed).toBeGreaterThan(0) + expect(result.boundaryMessage).toBe(boundary) + }) + + test('keeps boundary message when all messages are removed', () => { + const a = makeMessage('a') + const b = makeMessage('b') + const boundary = makeSnipBoundary('bnd', ['a', 'b']) + + const msgs = [a, b, boundary] + const result = snipCompactIfNeeded(msgs) + + expect(result.executed).toBe(true) + expect(result.messages).toHaveLength(1) + expect(result.messages[0]!.uuid as string).toBe('bnd') + }) + + test('keeps messages after boundary when no removedUuids', () => { + const a = makeMessage('a') + const boundary = makeSystemMessage('bnd', 'snip_boundary') + const c = makeMessage('c') + + const msgs = [a, boundary, c] + const result = snipCompactIfNeeded(msgs) + + expect(result.executed).toBe(true) + expect(result.messages).toHaveLength(2) + expect(result.messages.map((m) => m.uuid) as string[]).toEqual(['bnd', 'c']) + }) + + test('handles empty removedUuids array', () => { + const a = makeMessage('a') + const boundary = makeSnipBoundary('bnd', []) + + const msgs = [a, boundary] + const result = snipCompactIfNeeded(msgs) + + expect(result.executed).toBe(true) + // Fallback: keep boundary + everything after + expect(result.messages).toHaveLength(1) + expect(result.messages[0]!.uuid as string).toBe('bnd') + }) + + test('uses last boundary when multiple boundaries exist', () => { + const a = makeMessage('a') + const b = makeMessage('b') + const c = makeMessage('c') + const boundary1 = makeSnipBoundary('bnd1', ['a']) + const boundary2 = makeSnipBoundary('bnd2', ['b']) + + const msgs = [a, boundary1, b, boundary2, c] + const result = snipCompactIfNeeded(msgs) + + expect(result.executed).toBe(true) + expect(result.boundaryMessage!.uuid as string).toBe('bnd2') + // 'b' removed by boundary2, 'a' not in boundary2's removedUuids + expect(result.messages.map((m) => m.uuid) as string[]).toEqual(['a', 'bnd1', 'bnd2', 'c']) + }) + + test('respects force option (no functional difference — both execute)', () => { + const a = makeMessage('a') + const boundary = makeSnipBoundary('bnd', ['a']) + + const msgs = [a, boundary] + const resultForce = snipCompactIfNeeded(msgs, { force: true }) + const resultNoForce = snipCompactIfNeeded(msgs) + + expect(resultForce.executed).toBe(true) + expect(resultNoForce.executed).toBe(true) + }) + + test('estimates tokens freed based on removed content length', () => { + const heavy = { + ...makeMessage('heavy', 'user'), + message: { + role: 'user' as const, + content: 'x'.repeat(400), // ~100 tokens + }, + } as Message + const boundary = makeSnipBoundary('bnd', ['heavy']) + + const result = snipCompactIfNeeded([heavy, boundary]) + expect(result.tokensFreed).toBeGreaterThan(0) + // 400 chars / 4 chars-per-token = ~100 tokens + expect(result.tokensFreed).toBeGreaterThanOrEqual(90) + }) + + test('handles empty message array', () => { + const result = snipCompactIfNeeded([]) + expect(result.executed).toBe(false) + expect(result.messages).toHaveLength(0) + }) +}) diff --git a/src/services/compact/__tests__/snipProjection.test.ts b/src/services/compact/__tests__/snipProjection.test.ts new file mode 100644 index 000000000..a39e08eae --- /dev/null +++ b/src/services/compact/__tests__/snipProjection.test.ts @@ -0,0 +1,126 @@ +import { describe, expect, test } from 'bun:test' +import { isSnipBoundaryMessage, projectSnippedView } from '../snipProjection.js' +import type { Message } from 'src/types/message.js' + +// --- Helpers --- + +function makeMessage(uuid: string, type: Message['type'] = 'user'): Message { + return { + type, + uuid, + message: { + role: type === 'user' ? 'user' : 'assistant', + content: `Message ${uuid}`, + }, + } as Message +} + +function makeSystemMessage( + uuid: string, + subtype?: string, + extra?: Record, +): Message { + const msg: Message = { + type: 'system', + uuid, + message: { role: 'system', content: '' }, + ...extra, + } as Message + if (subtype) { + ;(msg as Record).subtype = subtype + } + return msg +} + +function makeSnipBoundary( + uuid: string, + removedUuids: string[], +): Message { + return makeSystemMessage(uuid, 'snip_boundary', { + snipMetadata: { removedUuids }, + content: '[snip]', + }) +} + +// --- isSnipBoundaryMessage --- + +describe('isSnipBoundaryMessage', () => { + test('returns true for system message with snip_boundary subtype', () => { + const msg = makeSnipBoundary('b1', ['a']) + expect(isSnipBoundaryMessage(msg)).toBe(true) + }) + + test('returns false for system message with different subtype', () => { + const msg = makeSystemMessage('s1', 'local_command') + expect(isSnipBoundaryMessage(msg)).toBe(false) + }) + + test('returns false for system message with no subtype', () => { + const msg = makeSystemMessage('s1') + expect(isSnipBoundaryMessage(msg)).toBe(false) + }) + + test('returns false for non-system message', () => { + const msg = makeMessage('u1', 'user') + expect(isSnipBoundaryMessage(msg)).toBe(false) + }) + + test('returns false for assistant message', () => { + const msg = makeMessage('a1', 'assistant') + expect(isSnipBoundaryMessage(msg)).toBe(false) + }) +}) + +// --- projectSnippedView --- + +describe('projectSnippedView', () => { + test('returns same array when no boundaries exist', () => { + const msgs = [makeMessage('a'), makeMessage('b')] + const result = projectSnippedView(msgs) + expect(result).toBe(msgs) // same reference — no copy + }) + + test('filters out messages listed in removedUuids', () => { + const a = makeMessage('a') + const b = makeMessage('b') + const c = makeMessage('c') + const boundary = makeSnipBoundary('bnd', ['a', 'c']) + + const result = projectSnippedView([a, b, c, boundary]) + expect(result.map((m) => m.uuid) as string[]).toEqual(['b', 'bnd']) + }) + + test('preserves boundary messages themselves', () => { + const a = makeMessage('a') + const boundary = makeSnipBoundary('bnd', ['a']) + + const result = projectSnippedView([a, boundary]) + expect(result).toHaveLength(1) + expect(result[0]!.uuid as string).toBe('bnd') + }) + + test('handles multiple boundaries accumulating removedUuids', () => { + const a = makeMessage('a') + const b = makeMessage('b') + const c = makeMessage('c') + const d = makeMessage('d') + const boundary1 = makeSnipBoundary('bnd1', ['a']) + const boundary2 = makeSnipBoundary('bnd2', ['c']) + + const result = projectSnippedView([a, boundary1, b, c, boundary2, d]) + expect(result.map((m) => m.uuid) as string[]).toEqual(['bnd1', 'b', 'bnd2', 'd']) + }) + + test('returns all messages when boundary has empty removedUuids', () => { + const a = makeMessage('a') + const boundary = makeSnipBoundary('bnd', []) + + const result = projectSnippedView([a, boundary]) + expect(result.map((m) => m.uuid) as string[]).toEqual(['a', 'bnd']) + }) + + test('handles empty message array', () => { + const result = projectSnippedView([]) + expect(result).toHaveLength(0) + }) +}) diff --git a/src/services/compact/snipCompact.ts b/src/services/compact/snipCompact.ts index ecd72176e..4e4a0d9fc 100644 --- a/src/services/compact/snipCompact.ts +++ b/src/services/compact/snipCompact.ts @@ -1,17 +1,165 @@ -// Auto-generated stub — replace with real implementation -export {}; +import type { Message } from 'src/types/message.js' -import type { Message } from 'src/types/message'; +/** + * Estimated characters per token (conservative for mixed code/text). + */ +const CHARS_PER_TOKEN = 4 -export const isSnipMarkerMessage: (message: Message) => boolean = () => false; -export const snipCompactIfNeeded: ( +/** + * Minimum message count before nudging the model to consider snipping. + */ +const SNIP_NUDGE_THRESHOLD = 30 + +/** + * Text shown to the model as a nudge when the conversation is long enough + * to benefit from snipping. + */ +export const SNIP_NUDGE_TEXT: string = + 'The conversation history is getting long. Consider using the /force-snip command or the snip tool to compress older messages, freeing context window space for continued work.' + +/** + * Check whether a message is an internal snip marker (not user-facing). + * Snip markers are system messages injected by the snip tool to track + * which messages have been registered for future removal. + */ +export function isSnipMarkerMessage(message: Message): boolean { + if (message.type !== 'system') return false + return (message as Record).subtype === 'snip_marker' +} + +/** + * Estimate the token count of a single message by serialising its content. + * This is a rough heuristic (~4 chars per token) used to report + * tokensFreed; it does not need to be exact. + */ +function estimateMessageTokens(message: Message): number { + const content = message.message?.content + let chars = 0 + if (typeof content === 'string') { + chars = content.length + } else if (Array.isArray(content)) { + for (const block of content) { + if (typeof block === 'string') { + chars += (block as string).length + } else if (block && typeof block === 'object') { + const obj = block as unknown as Record + const text = obj.text ?? obj.content + if (typeof text === 'string') { + chars += text.length + } else { + chars += JSON.stringify(block).length + } + } + } + } else if (content !== null && content !== undefined) { + chars = JSON.stringify(content).length + } + return Math.max(1, Math.ceil(chars / CHARS_PER_TOKEN)) +} + +/** + * Scan the message array for the last `snip_boundary` system message and, + * if found, remove all messages whose UUIDs appear in its + * `snipMetadata.removedUuids`. + * + * This is the core memory-saving function. When a snip boundary exists: + * 1. All messages listed in `removedUuids` are filtered out. + * 2. The boundary message itself is kept (it records what was removed). + * 3. Messages not in `removedUuids` (including post-boundary messages) + * are preserved. + * + * Called from: + * - `query.ts` — strips snipped messages from the model-facing array + * before sending to the API. + * - `QueryEngine.ts` `snipReplay` — trims `mutableMessages` so the + * in-memory store does not grow without bound in long SDK sessions. + * + * @param messages Full message array (may contain a snip_boundary). + * @param options `force` — if true, always execute when a boundary is + * present. Without `force`, the function still executes + * if a boundary is found (the "if needed" refers to + * whether a boundary exists, not a token threshold). + */ +export function snipCompactIfNeeded( messages: Message[], options?: { force?: boolean }, -) => { messages: Message[]; executed: boolean; tokensFreed: number; boundaryMessage?: Message } = (messages) => ({ - messages, - executed: false, - tokensFreed: 0, -}); -export const isSnipRuntimeEnabled: () => boolean = () => false; -export const shouldNudgeForSnips: (messages: Message[]) => boolean = () => false; -export const SNIP_NUDGE_TEXT: string = ''; +): { + messages: Message[] + executed: boolean + tokensFreed: number + boundaryMessage?: Message +} { + // Find the last snip_boundary message + let boundaryIdx = -1 + let removedUuids: string[] | undefined + + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]! + if ( + msg.type === 'system' && + (msg as Record).subtype === 'snip_boundary' + ) { + boundaryIdx = i + const meta = (msg as Record).snipMetadata as + | { removedUuids?: string[] } + | undefined + removedUuids = meta?.removedUuids + break + } + } + + if (boundaryIdx === -1) { + return { messages, executed: false, tokensFreed: 0 } + } + + const boundaryMessage = messages[boundaryIdx]! + + // No removedUuids metadata — fallback: keep boundary + everything after + if (!removedUuids || removedUuids.length === 0) { + const kept = messages.slice(boundaryIdx) + return { + messages: kept, + executed: true, + tokensFreed: 0, + boundaryMessage, + } + } + + // Filter out messages whose UUIDs are listed in removedUuids + const removedSet = new Set(removedUuids) + const kept: Message[] = [] + let tokensFreed = 0 + + for (const msg of messages) { + if (removedSet.has(msg.uuid)) { + tokensFreed += estimateMessageTokens(msg) + continue + } + kept.push(msg) + } + + return { + messages: kept, + executed: true, + tokensFreed, + boundaryMessage, + } +} + +/** + * Returns true when the snip runtime is active. + * Because this module is only loaded when the HISTORY_SNIP feature flag + * is enabled, this always returns true. + */ +export function isSnipRuntimeEnabled(): boolean { + return true +} + +/** + * Determine whether the conversation is long enough to warrant a nudge + * to the model to consider snipping. Uses a simple message-count + * threshold rather than an expensive token count. + */ +export function shouldNudgeForSnips(messages: Message[]): boolean { + return messages.length >= SNIP_NUDGE_THRESHOLD +} diff --git a/src/services/compact/snipProjection.ts b/src/services/compact/snipProjection.ts index 80efe381a..e9d9cbb2d 100644 --- a/src/services/compact/snipProjection.ts +++ b/src/services/compact/snipProjection.ts @@ -1,7 +1,60 @@ -// Auto-generated stub — replace with real implementation -export {}; +import type { Message } from 'src/types/message.js' -import type { Message } from 'src/types/message'; +/** + * Check whether a message is a snip boundary marker. + * + * A snip boundary is a system message with `subtype === 'snip_boundary'` + * and an optional `snipMetadata.removedUuids` array recording which + * messages were removed by the snip operation. + * + * Used by: + * - `Message.tsx` — render SnipBoundaryMessage component. + * - `QueryEngine.ts` `snipReplay` — decide whether to replay the snip + * on the mutableMessages store. + */ +export function isSnipBoundaryMessage(message: Message): boolean { + if (message.type !== 'system') return false + return (message as Record).subtype === 'snip_boundary' +} -export const isSnipBoundaryMessage: (message: Message) => boolean = () => false; -export const projectSnippedView: (messages: Message[]) => Message[] = (messages) => messages; +/** + * Project a "snipped view" of the message array suitable for sending to + * the model. Messages whose UUIDs appear in any snip boundary's + * `removedUuids` are filtered out; all others (including the boundary + * messages themselves) are preserved. + * + * Used by: + * - `getMessagesAfterCompactBoundary()` in messages.ts — after slicing + * at the compact boundary, further filters out snipped messages so the + * model-facing array does not include stale history. + * + * @param messages Message array that may contain one or more snip + * boundaries. + * @returns New array with removed messages stripped out. + */ +export function projectSnippedView(messages: Message[]): Message[] { + // Collect all UUIDs that have been removed by any snip boundary + const removedSet = new Set() + + for (const msg of messages) { + if ( + msg.type === 'system' && + (msg as Record).subtype === 'snip_boundary' + ) { + const meta = (msg as Record).snipMetadata as + | { removedUuids?: string[] } + | undefined + if (meta?.removedUuids) { + for (const uuid of meta.removedUuids) { + removedSet.add(uuid) + } + } + } + } + + if (removedSet.size === 0) { + return messages + } + + return messages.filter((msg) => !removedSet.has(msg.uuid)) +}