fix(memory): bound memory-directory scanning work (#1757)

* fix(memory): bound memory-directory scanning work

* fix(memory): harden bounded memory scanning follow-up
This commit is contained in:
Bogdan
2026-06-25 06:39:05 +08:00
committed by GitHub
parent 701b68c215
commit 28bbec4948
2 changed files with 671 additions and 74 deletions
+479 -31
View File
@@ -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<void>
resolve: () => void
reject: (error: unknown) => void
} {
let resolve!: () => void
let reject!: (error: unknown) => void
const promise = new Promise<void>((res, rej) => {
resolve = res
reject = rej
})
return { promise, resolve, reject }
}
async function waitFor(
condition: () => boolean,
message: string,
): Promise<void> {
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<string, TestDirent[]>
files?: Record<string, FakeFile>
onRead?: (filePath: string, signal: AbortSignal) => Promise<void> | 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<void> {
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)
})
+192 -43
View File
@@ -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<MemoryScanDirent[]>
readFileInRange: (
filePath: string,
offset: number,
maxLines: number,
maxBytes: number,
signal: AbortSignal,
options: { truncateOnByteLimit: true },
) => Promise<Pick<ReadFileRangeResult, 'content' | 'mtimeMs'>>
}
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<MemoryHeader[]> {
return scanMemoryFilesWithDependencies(memoryDir, signal, defaultDependencies)
}
async function scanMemoryFilesWithDependencies(
memoryDir: string,
signal: AbortSignal,
deps: MemoryScanDependencies,
): Promise<MemoryHeader[]> {
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<MemoryHeader> => {
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<string>
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<MemoryHeader> =>
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<string> {
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<MemoryHeader> {
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