From 28bbec49481ba625fda57d99418dbc50f67e30a2 Mon Sep 17 00:00:00 2001 From: Bogdan Date: Thu, 25 Jun 2026 00:39:05 +0200 Subject: [PATCH] fix(memory): bound memory-directory scanning work (#1757) * fix(memory): bound memory-directory scanning work * fix(memory): harden bounded memory scanning follow-up --- src/memdir/memoryScan.test.ts | 510 +++++++++++++++++++++++++++++++--- src/memdir/memoryScan.ts | 235 +++++++++++++--- 2 files changed, 671 insertions(+), 74 deletions(-) diff --git a/src/memdir/memoryScan.test.ts b/src/memdir/memoryScan.test.ts index e5b66eca0..bdcd06f07 100644 --- a/src/memdir/memoryScan.test.ts +++ b/src/memdir/memoryScan.test.ts @@ -1,59 +1,507 @@ import { afterEach, expect, test } from 'bun:test' -import { mkdtemp, mkdir, writeFile, rm } from 'fs/promises' -import { join } from 'path' +import { mkdir, mkdtemp, rm, writeFile } from 'fs/promises' import { tmpdir } from 'os' -import { scanMemoryFiles } from './memoryScan.ts' +import { join } from 'path' +import { __test, scanMemoryFiles } from './memoryScan.ts' -// Finding #42-3: readdir({ recursive: true }) has no depth limit. -// A deeply nested directory in the memory dir causes a full unbounded walk. +let tempDir: string | undefined -let tempDir: string +type TestDirent = { + name: string + isFile(): boolean + isDirectory(): boolean + isSymbolicLink(): boolean +} + +type FakeFile = { + content?: string + mtimeMs?: number + error?: unknown +} + +type FakeReadCall = { + filePath: string + offset: number + maxLines?: number + maxBytes?: number + truncateOnByteLimit?: boolean +} + +function file(name: string): TestDirent { + return { + name, + isFile: () => true, + isDirectory: () => false, + isSymbolicLink: () => false, + } +} + +function dir(name: string): TestDirent { + return { + name, + isFile: () => false, + isDirectory: () => true, + isSymbolicLink: () => false, + } +} + +function symlink(name: string): TestDirent { + return { + name, + isFile: () => false, + isDirectory: () => false, + isSymbolicLink: () => true, + } +} + +function deferred(): { + promise: Promise + resolve: () => void + reject: (error: unknown) => void +} { + let resolve!: () => void + let reject!: (error: unknown) => void + const promise = new Promise((res, rej) => { + resolve = res + reject = rej + }) + return { promise, resolve, reject } +} + +async function waitFor( + condition: () => boolean, + message: string, +): Promise { + const deadline = Date.now() + 250 + while (Date.now() < deadline) { + if (condition()) return + await new Promise(resolve => setTimeout(resolve, 1)) + } + throw new Error(message) +} + +function createFakeDeps({ + tree, + files = {}, + onRead, + throwAfterRead = true, +}: { + tree: Record + files?: Record + onRead?: (filePath: string, signal: AbortSignal) => Promise | void + throwAfterRead?: boolean +}) { + const openedDirs: string[] = [] + const readPaths: string[] = [] + const readCalls: FakeReadCall[] = [] + + return { + openedDirs, + readCalls, + readPaths, + deps: { + readdir: async (dirPath: string) => { + openedDirs.push(dirPath) + const entries = tree[dirPath] + if (!entries) { + throw Object.assign(new Error(`ENOENT: ${dirPath}`), { + code: 'ENOENT', + }) + } + return entries + }, + readFileInRange: async ( + filePath: string, + _offset = 0, + _maxLines?: number, + _maxBytes?: number, + signal?: AbortSignal, + options?: { truncateOnByteLimit?: boolean }, + ) => { + readCalls.push({ + filePath, + offset: _offset, + maxLines: _maxLines, + maxBytes: _maxBytes, + truncateOnByteLimit: options?.truncateOnByteLimit, + }) + readPaths.push(filePath) + signal?.throwIfAborted() + await onRead?.(filePath, signal ?? new AbortController().signal) + if (throwAfterRead) { + signal?.throwIfAborted() + } + + const fakeFile = files[filePath] + if (fakeFile?.error) throw fakeFile.error + + return { + content: + fakeFile?.content ?? + '---\ndescription: fake memory\ntype: user\n---\nBody', + lineCount: 4, + totalLines: 4, + totalBytes: 0, + readBytes: 0, + mtimeMs: fakeFile?.mtimeMs ?? 0, + } + }, + }, + } +} + +async function writeMemoryFile(path: string): Promise { + await writeFile(path, '---\ndescription: test\ntype: user\n---\nContent') +} afterEach(async () => { if (tempDir) { await rm(tempDir, { recursive: true, force: true }) + tempDir = undefined } }) -test('scanMemoryFiles finds .md files at shallow depth', async () => { +test('scanMemoryFiles returns markdown files within the current allowed depth', async () => { tempDir = await mkdtemp(join(tmpdir(), 'memoryScan-')) - await writeFile(join(tempDir, 'note.md'), '---\nname: test\ntype: user\n---\nContent') + await mkdir(join(tempDir, 'one', 'two', 'three'), { recursive: true }) + await writeMemoryFile(join(tempDir, 'root.md')) + await writeMemoryFile(join(tempDir, 'one', 'one.md')) + await writeMemoryFile(join(tempDir, 'one', 'two', 'two.md')) + await writeMemoryFile(join(tempDir, 'one', 'two', 'three', 'three.md')) - const controller = new AbortController() - const result = await scanMemoryFiles(tempDir, controller.signal) + const result = await scanMemoryFiles(tempDir, new AbortController().signal) - expect(result.length).toBe(1) - expect(result[0].filename).toBe('note.md') + const filenames = result.map(r => r.filename).sort() + expect(filenames).toEqual([ + join('one', 'one.md'), + join('one', 'two', 'two.md'), + 'root.md', + ]) }) -test('scanMemoryFiles ignores MEMORY.md', async () => { +test('scanMemoryFiles does not open directories beyond the current allowed depth', async () => { + const root = '/memory' + const one = join(root, 'one') + const two = join(one, 'two') + const tooDeep = join(two, 'three') + const { deps, openedDirs } = createFakeDeps({ + tree: { + [root]: [dir('one')], + [one]: [dir('two')], + [two]: [dir('three'), file('two.md')], + [tooDeep]: [file('three.md')], + }, + }) + + const result = await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + expect(result.map(r => r.filename)).toEqual([join('one', 'two', 'two.md')]) + expect(openedDirs).toEqual([root, one, two]) +}) + +test('scanMemoryFiles does not follow symlinked directories', async () => { + const root = '/memory' + const linked = join(root, 'linked') + const { deps, openedDirs } = createFakeDeps({ + tree: { + [root]: [symlink('linked'), file('root.md')], + [linked]: [file('linked.md')], + }, + }) + + const result = await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + expect(result.map(r => r.filename)).toEqual(['root.md']) + expect(openedDirs).toEqual([root]) +}) + +test('scanMemoryFiles keeps symlinked markdown file candidates', async () => { + const root = '/memory' + const { deps, openedDirs } = createFakeDeps({ + tree: { + [root]: [symlink('linked.md')], + }, + files: { + [join(root, 'linked.md')]: { mtimeMs: 1 }, + }, + }) + + const result = await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + expect(result.map(r => r.filename)).toEqual(['linked.md']) + expect(openedDirs).toEqual([root]) +}) + +test('scanMemoryFiles returns the newest 200 markdown files', async () => { + const root = '/memory' + const tree = { + [root]: Array.from({ length: 205 }, (_, i) => file(`file-${i}.md`)), + } + const files = Object.fromEntries( + Array.from({ length: 205 }, (_, i) => [ + join(root, `file-${i}.md`), + { mtimeMs: i }, + ]), + ) + + const { deps } = createFakeDeps({ tree, files }) + const result = await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + expect(result).toHaveLength(200) + expect(result[0].filename).toBe('file-204.md') + expect(result.at(-1)?.filename).toBe('file-5.md') + expect(result.some(r => r.filename === 'file-4.md')).toBe(false) +}) + +test('scanMemoryFiles never exceeds header read concurrency', async () => { + const root = '/memory' + const gate = deferred() + let activeReads = 0 + let maxActiveReads = 0 + const { deps } = createFakeDeps({ + tree: { + [root]: Array.from( + { length: __test.HEADER_READ_CONCURRENCY * 3 }, + (_, i) => file(`file-${i}.md`), + ), + }, + onRead: async () => { + activeReads++ + maxActiveReads = Math.max(maxActiveReads, activeReads) + await gate.promise + activeReads-- + }, + }) + + const promise = __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + await waitFor( + () => activeReads === __test.HEADER_READ_CONCURRENCY, + 'expected the first read batch to start', + ) + expect(maxActiveReads).toBe(__test.HEADER_READ_CONCURRENCY) + + gate.resolve() + await promise + expect(maxActiveReads).toBeLessThanOrEqual(__test.HEADER_READ_CONCURRENCY) +}) + +test('scanMemoryFiles does not schedule every file in a broad directory at once', async () => { + const root = '/memory' + const gate = deferred() + let readsStarted = 0 + const { deps } = createFakeDeps({ + tree: { + [root]: Array.from({ length: 500 }, (_, i) => file(`file-${i}.md`)), + }, + onRead: async () => { + readsStarted++ + await gate.promise + }, + }) + + const promise = __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + await waitFor( + () => readsStarted >= __test.HEADER_READ_CONCURRENCY, + 'expected bounded reads to start', + ) + expect(readsStarted).toBe(__test.HEADER_READ_CONCURRENCY) + + gate.resolve() + await promise +}) + +test('scanMemoryFiles bounds header reads by lines and bytes', async () => { + const root = '/memory' + const { deps, readCalls } = createFakeDeps({ + tree: { + [root]: [file('note.md')], + }, + }) + + await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + expect(readCalls).toEqual([ + { + filePath: join(root, 'note.md'), + offset: 0, + maxLines: __test.FRONTMATTER_MAX_LINES, + maxBytes: __test.FRONTMATTER_MAX_BYTES, + truncateOnByteLimit: true, + }, + ]) +}) + +test('scanMemoryFiles excludes MEMORY.md with the current case-sensitive basename rule', async () => { tempDir = await mkdtemp(join(tmpdir(), 'memoryScan-')) await writeFile(join(tempDir, 'MEMORY.md'), '# index') - await writeFile(join(tempDir, 'user_role.md'), '---\nname: role\ntype: user\n---\nContent') + await writeMemoryFile(join(tempDir, 'user_role.md')) - const controller = new AbortController() - const result = await scanMemoryFiles(tempDir, controller.signal) + const result = await scanMemoryFiles(tempDir, new AbortController().signal) - expect(result.length).toBe(1) - expect(result[0].filename).toBe('user_role.md') + expect(result.map(r => r.filename)).toEqual(['user_role.md']) }) -test('scanMemoryFiles does not return .md files nested beyond max depth', async () => { - tempDir = await mkdtemp(join(tmpdir(), 'memoryScan-')) +test('scanMemoryFiles preserves case-sensitive MEMORY.md exclusion semantics', async () => { + const root = '/memory' + const { deps, readPaths } = createFakeDeps({ + tree: { + [root]: [file('MEMORY.md'), file('memory.md'), file('user_role.md')], + }, + }) - // Shallow file - should be found - await writeFile(join(tempDir, 'shallow.md'), '---\nname: shallow\ntype: user\n---\nContent') + const result = await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) - // Deeply nested file (depth 5) - should be excluded - const deepDir = join(tempDir, 'd1', 'd2', 'd3', 'd4', 'd5') - await mkdir(deepDir, { recursive: true }) - await writeFile(join(deepDir, 'deep.md'), '---\nname: deep\ntype: user\n---\nContent') + expect(result.map(r => r.filename).sort()).toEqual([ + 'memory.md', + 'user_role.md', + ]) + expect(readPaths).not.toContain(join(root, 'MEMORY.md')) +}) +test('scanMemoryFiles treats non-string descriptions as absent', async () => { + const root = '/memory' + const { deps } = createFakeDeps({ + tree: { + [root]: [file('number-description.md')], + }, + files: { + [join(root, 'number-description.md')]: { + content: '---\ndescription: 123\ntype: user\n---\nBody', + }, + }, + }) + + const result = await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + expect(result).toHaveLength(1) + expect(result[0]?.description).toBeNull() +}) + +test('scanMemoryFiles skips unreadable files without discarding valid siblings', async () => { + const root = '/memory' + const { deps } = createFakeDeps({ + tree: { + [root]: [file('good-a.md'), file('bad.md'), file('good-b.md')], + }, + files: { + [join(root, 'good-a.md')]: { mtimeMs: 3 }, + [join(root, 'bad.md')]: { error: new Error('unreadable') }, + [join(root, 'good-b.md')]: { mtimeMs: 2 }, + }, + }) + + const result = await __test.scanMemoryFilesWithDependencies( + root, + new AbortController().signal, + deps, + ) + + expect(result.map(r => r.filename)).toEqual(['good-a.md', 'good-b.md']) +}) + +test('scanMemoryFiles returns promptly when the signal is already aborted', async () => { + const root = '/memory' + const { deps, openedDirs, readPaths } = createFakeDeps({ + tree: { + [root]: [file('note.md')], + }, + }) const controller = new AbortController() - const result = await scanMemoryFiles(tempDir, controller.signal) + controller.abort() + + const result = await __test.scanMemoryFilesWithDependencies( + root, + controller.signal, + deps, + ) + + expect(result).toEqual([]) + expect(openedDirs).toEqual([]) + expect(readPaths).toEqual([]) +}) + +test('scanMemoryFiles drops headers when the signal aborts after a read', async () => { + const root = '/memory' + const controller = new AbortController() + const { deps, readPaths } = createFakeDeps({ + tree: { + [root]: [file('late-abort.md')], + }, + onRead: () => { + controller.abort() + }, + throwAfterRead: false, + }) + + const result = await __test.scanMemoryFilesWithDependencies( + root, + controller.signal, + deps, + ) + + expect(readPaths).toEqual([join(root, 'late-abort.md')]) + expect(result).toEqual([]) +}) + +test('scanMemoryFiles stops scheduling additional reads after abort', async () => { + const root = '/memory' + const controller = new AbortController() + let readsStarted = 0 + const { deps } = createFakeDeps({ + tree: { + [root]: Array.from({ length: 50 }, (_, i) => file(`file-${i}.md`)), + }, + onRead: () => { + readsStarted++ + controller.abort() + }, + }) + + const result = await __test.scanMemoryFilesWithDependencies( + root, + controller.signal, + deps, + ) - const filenames = result.map(r => r.filename) - expect(filenames).toContain('shallow.md') - // The deeply nested file must not appear - expect(filenames.some(f => f.includes('deep.md'))).toBe(false) + expect(result).toEqual([]) + expect(readsStarted).toBeGreaterThan(0) + expect(readsStarted).toBeLessThanOrEqual(__test.HEADER_READ_CONCURRENCY) }) diff --git a/src/memdir/memoryScan.ts b/src/memdir/memoryScan.ts index 5440f30e7..a15e2aa32 100644 --- a/src/memdir/memoryScan.ts +++ b/src/memdir/memoryScan.ts @@ -4,9 +4,11 @@ * the API-client chain (which closed a cycle through memdir.ts — #25372). */ +import type { Dirent } from 'fs' import { readdir } from 'fs/promises' -import { basename, join } from 'path' +import { join } from 'path' import { parseFrontmatter } from '../utils/frontmatterParser.js' +import type { ReadFileRangeResult } from '../utils/readFileInRange.js' import { readFileInRange } from '../utils/readFileInRange.js' import { type MemoryType, parseMemoryType } from './memoryTypes.js' @@ -20,6 +22,36 @@ export type MemoryHeader = { const MAX_MEMORY_FILES = 200 const FRONTMATTER_MAX_LINES = 30 +const FRONTMATTER_MAX_BYTES = 64 * 1024 +const MAX_DEPTH = 3 +const HEADER_READ_CONCURRENCY = 8 + +type MemoryScanDirent = Pick< + Dirent, + 'name' | 'isFile' | 'isDirectory' | 'isSymbolicLink' +> + +type MemoryScanDependencies = { + readdir: (dir: string) => Promise + readFileInRange: ( + filePath: string, + offset: number, + maxLines: number, + maxBytes: number, + signal: AbortSignal, + options: { truncateOnByteLimit: true }, + ) => Promise> +} + +type RankedMemoryHeader = { + header: MemoryHeader + order: number +} + +const defaultDependencies: MemoryScanDependencies = { + readdir: dir => readdir(dir, { withFileTypes: true }), + readFileInRange, +} /** * Scan a memory directory for .md files, read their frontmatter, and return @@ -27,62 +59,179 @@ const FRONTMATTER_MAX_LINES = 30 * findRelevantMemories (query-time recall) and extractMemories (pre-injects * the listing so the extraction agent doesn't spend a turn on `ls`). * - * Single-pass: readFileInRange stats internally and returns mtimeMs, so we - * read-then-sort rather than stat-sort-read. For the common case (N ≤ 200) - * this halves syscalls vs a separate stat round; for large N we read a few - * extra small files but still avoid the double-stat on the surviving 200. + * Traversal is depth-bounded before opening child directories. Header reads + * run through a small worker pool, and only the newest MAX_MEMORY_FILES + * parsed headers are retained while scanning. */ export async function scanMemoryFiles( memoryDir: string, signal: AbortSignal, +): Promise { + return scanMemoryFilesWithDependencies(memoryDir, signal, defaultDependencies) +} + +async function scanMemoryFilesWithDependencies( + memoryDir: string, + signal: AbortSignal, + deps: MemoryScanDependencies, ): Promise { try { - const entries = await readdir(memoryDir, { recursive: true }) - // Limit depth to 3 levels to prevent DoS from deep/symlinked directory trees. - // Relative paths from readdir use the OS separator, so count separators. - const sep = require('path').sep as string - const MAX_DEPTH = 3 - const mdFiles = entries.filter( - f => - f.endsWith('.md') && - basename(f) !== 'MEMORY.md' && - (f.split(sep).length - 1) < MAX_DEPTH, - ) + signal.throwIfAborted() + const topHeaders: RankedMemoryHeader[] = [] + const fileIterator = walkMarkdownFiles( + memoryDir, + signal, + deps, + )[Symbol.asyncIterator]() + let nextOrder = 0 - const headerResults = await Promise.allSettled( - mdFiles.map(async (relativePath): Promise => { - const filePath = join(memoryDir, relativePath) - const { content, mtimeMs } = await readFileInRange( - filePath, - 0, - FRONTMATTER_MAX_LINES, - undefined, - signal, - ) - const { frontmatter } = parseFrontmatter(content, filePath) - return { - filename: relativePath, - filePath, - mtimeMs, - description: frontmatter.description || null, - type: parseMemoryType(frontmatter.type), + const workers = Array.from({ length: HEADER_READ_CONCURRENCY }, async () => { + while (!signal.aborted) { + let next: IteratorResult + try { + next = await fileIterator.next() + } catch (error) { + if (signal.aborted) return + throw error } - }), - ) + if (next.done) return + const order = nextOrder++ - return headerResults - .filter( - (r): r is PromiseFulfilledResult => - r.status === 'fulfilled', - ) - .map(r => r.value) - .sort((a, b) => b.mtimeMs - a.mtimeMs) - .slice(0, MAX_MEMORY_FILES) + try { + const header = await readMemoryHeader( + memoryDir, + next.value, + signal, + deps, + ) + if (signal.aborted) return + insertNewestHeader(topHeaders, { header, order }) + } catch { + if (signal.aborted) return + } + } + }) + + await Promise.all(workers) + return signal.aborted ? [] : topHeaders.map(entry => entry.header) } catch { return [] } } +async function* walkMarkdownFiles( + memoryDir: string, + signal: AbortSignal, + deps: MemoryScanDependencies, +): AsyncGenerator { + const pendingDirs: Array<{ + absolutePath: string + relativePath: string + depth: number + }> = [{ absolutePath: memoryDir, relativePath: '', depth: 0 }] + + while (pendingDirs.length > 0) { + signal.throwIfAborted() + const current = pendingDirs.pop()! + let entries: MemoryScanDirent[] + try { + entries = await deps.readdir(current.absolutePath) + } catch { + continue + } + + for (const entry of entries) { + signal.throwIfAborted() + const relativePath = current.relativePath + ? join(current.relativePath, entry.name) + : entry.name + const absolutePath = join(memoryDir, relativePath) + const isMarkdownMemoryFile = + entry.name.endsWith('.md') && entry.name !== 'MEMORY.md' + + if (entry.isSymbolicLink()) { + if (isMarkdownMemoryFile) { + yield relativePath + } + continue + } + + if (entry.isDirectory()) { + const nextDepth = current.depth + 1 + if (nextDepth < MAX_DEPTH) { + pendingDirs.push({ absolutePath, relativePath, depth: nextDepth }) + } + continue + } + + if (entry.isFile() && isMarkdownMemoryFile) { + yield relativePath + } + } + } +} + +async function readMemoryHeader( + memoryDir: string, + relativePath: string, + signal: AbortSignal, + deps: MemoryScanDependencies, +): Promise { + signal.throwIfAborted() + const filePath = join(memoryDir, relativePath) + const { content, mtimeMs } = await deps.readFileInRange( + filePath, + 0, + FRONTMATTER_MAX_LINES, + FRONTMATTER_MAX_BYTES, + signal, + { truncateOnByteLimit: true }, + ) + const { frontmatter } = parseFrontmatter(content, filePath) + const description = + typeof frontmatter.description === 'string' && frontmatter.description + ? frontmatter.description + : null + + return { + filename: relativePath, + filePath, + mtimeMs, + description, + type: parseMemoryType(frontmatter.type), + } +} + +function insertNewestHeader( + headers: RankedMemoryHeader[], + entry: RankedMemoryHeader, +): void { + const index = headers.findIndex( + existing => + entry.header.mtimeMs > existing.header.mtimeMs || + (entry.header.mtimeMs === existing.header.mtimeMs && + entry.order < existing.order), + ) + if (index === -1) { + if (headers.length < MAX_MEMORY_FILES) { + headers.push(entry) + } + return + } + + headers.splice(index, 0, entry) + if (headers.length > MAX_MEMORY_FILES) { + headers.length = MAX_MEMORY_FILES + } +} + +export const __test = { + FRONTMATTER_MAX_BYTES, + FRONTMATTER_MAX_LINES, + HEADER_READ_CONCURRENCY, + scanMemoryFilesWithDependencies, +} + /** * Format memory headers as a text manifest: one line per file with * [type] filename (timestamp): description. Used by both the recall