Skip to content

Commit 236d99e

Browse files
claude-code-bestclaude-code-best
authored
Merge pull request #1288 from Ericcc-Ma/fix/session-messages-cache-bounded
 fix(session-storage): bound getSessionMessages cache to prevent unbounded growth
2 parents 837ba00 + 562da4c commit 236d99e

2 files changed

Lines changed: 193 additions & 13 deletions

File tree

‎src/utils/__tests__/sessionStorage.test.ts‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
2+
import { existsSync, mkdirSync, rmSync, writeFileSync } from 'node:fs'
3+
import { join } from 'node:path'
4+
import { tmpdir } from 'node:os'
5+
6+
const MAX_CACHED_ENTRIES = 200 // mirrors MAX_CACHED_SESSION_FILES in sessionStorage.ts
7+
8+
const {
9+
getSessionMessages,
10+
getSessionMessagesCache,
11+
clearSessionMessagesCache,
12+
} = await import('../sessionStorage.js')
13+
14+
function asUuid(s: string): any {
15+
return s as unknown as any
16+
}
17+
18+
let tempDir: string
19+
let originalConfigDir: string | undefined
20+
21+
beforeEach(() => {
22+
tempDir = join(
23+
tmpdir(),
24+
`claude-session-test-${Date.now()}-${Math.random().toString(36).slice(2)}`,
25+
)
26+
mkdirSync(tempDir, { recursive: true })
27+
// `getProjectsDir()` returns `${CLAUDE_CONFIG_DIR}/projects`, and
28+
// loadSessionFile reads from `${getProjectsDir()}/${sessionId}.jsonl`.
29+
// Pre-create the projects subdir so writeFileSync doesn't fail.
30+
mkdirSync(join(tempDir, 'projects'), { recursive: true })
31+
// Pin session-file lookups to a temp dir via CLAUDE_CONFIG_DIR.
32+
// Restoring in afterEach keeps tests hermetic.
33+
originalConfigDir = process.env.CLAUDE_CONFIG_DIR
34+
process.env.CLAUDE_CONFIG_DIR = tempDir
35+
})
36+
37+
afterEach(() => {
38+
clearSessionMessagesCache()
39+
if (originalConfigDir === undefined) {
40+
delete process.env.CLAUDE_CONFIG_DIR
41+
} else {
42+
process.env.CLAUDE_CONFIG_DIR = originalConfigDir
43+
}
44+
if (tempDir && existsSync(tempDir)) {
45+
rmSync(tempDir, { recursive: true, force: true })
46+
}
47+
})
48+
49+
function sessionFilePath(sessionId: string): string {
50+
// Mirror sessionStorage.ts's path computation:
51+
// getSessionProjectDir() ?? getProjectDir(getOriginalCwd())
52+
// With CLAUDE_CONFIG_DIR=tempDir and getSessionProjectDir() returning
53+
// null in tests, files live at `${tempDir}/projects/${sessionId}.jsonl`.
54+
return join(tempDir, 'projects', `${sessionId}.jsonl`)
55+
}
56+
57+
describe('getSessionMessagesCache', () => {
58+
test('returns the same Map instance across calls', () => {
59+
// Cache identity must be stable — `getLastSessionLog` uses
60+
// `getSessionMessagesCache()` directly to prime entries, so a
61+
// different instance each call would break that priming.
62+
expect(getSessionMessagesCache()).toBe(getSessionMessagesCache())
63+
})
64+
65+
test('clearSessionMessagesCache empties a populated cache', async () => {
66+
const cache = getSessionMessagesCache()
67+
writeFileSync(sessionFilePath('id-1'), '')
68+
writeFileSync(sessionFilePath('id-2'), '')
69+
await getSessionMessages(asUuid('id-1'))
70+
await getSessionMessages(asUuid('id-2'))
71+
expect(cache.size).toBeGreaterThan(0)
72+
73+
clearSessionMessagesCache()
74+
expect(cache.size).toBe(0)
75+
})
76+
77+
test('clearSessionMessagesCache is a no-op on empty cache', () => {
78+
const cache = getSessionMessagesCache()
79+
expect(cache.size).toBe(0)
80+
clearSessionMessagesCache()
81+
expect(cache.size).toBe(0)
82+
})
83+
84+
test('getSessionMessages dedups concurrent calls for the same sessionId', async () => {
85+
const cache = getSessionMessagesCache()
86+
const id = asUuid('same-id')
87+
writeFileSync(sessionFilePath('same-id'), '')
88+
const [a, b, c] = await Promise.all([
89+
getSessionMessages(id),
90+
getSessionMessages(id),
91+
getSessionMessages(id),
92+
])
93+
expect(a).toBe(b)
94+
expect(b).toBe(c)
95+
expect(cache.size).toBe(1)
96+
})
97+
})
98+
99+
describe('getSessionMessages bounded cache (memory leak fix)', () => {
100+
test('cache size stays at MAX_CACHED_ENTRIES after many distinct sessionIds', async () => {
101+
// Bounded cache — calling getSessionMessages with N distinct
102+
// sessionIds must NOT grow the cache beyond MAX_CACHED_ENTRIES.
103+
// Pre-fix: lodash memoize grew unbounded. Post-fix: Map-based
104+
// cache evicts oldest entry when at capacity.
105+
const cache = getSessionMessagesCache()
106+
const total = MAX_CACHED_ENTRIES * 3 // 600 distinct sessionIds
107+
for (let i = 0; i < total; i++) {
108+
writeFileSync(sessionFilePath(`id-${i}`), '')
109+
await getSessionMessages(asUuid(`id-${i}`))
110+
}
111+
expect(cache.size).toBe(MAX_CACHED_ENTRIES)
112+
})
113+
114+
test('FIFO eviction: oldest entry is removed first', async () => {
115+
// Fill cache to MAX with sequential ids. The first inserted
116+
// (`oldest`) should be evicted on the (MAX+1)th insertion.
117+
const cache = getSessionMessagesCache()
118+
const oldestId = asUuid('id-0')
119+
writeFileSync(sessionFilePath('id-0'), '')
120+
await getSessionMessages(oldestId)
121+
for (let i = 1; i < MAX_CACHED_ENTRIES; i++) {
122+
writeFileSync(sessionFilePath(`id-${i}`), '')
123+
await getSessionMessages(asUuid(`id-${i}`))
124+
}
125+
expect(cache.size).toBe(MAX_CACHED_ENTRIES)
126+
expect(cache.has(oldestId)).toBe(true)
127+
128+
writeFileSync(sessionFilePath('id-overflow'), '')
129+
await getSessionMessages(asUuid('id-overflow'))
130+
expect(cache.size).toBe(MAX_CACHED_ENTRIES)
131+
expect(cache.has(oldestId)).toBe(false)
132+
})
133+
134+
test('cleared cache can be refilled without leaking entries', async () => {
135+
const cache = getSessionMessagesCache()
136+
for (let i = 0; i < MAX_CACHED_ENTRIES; i++) {
137+
writeFileSync(sessionFilePath(`id-${i}`), '')
138+
await getSessionMessages(asUuid(`id-${i}`))
139+
}
140+
expect(cache.size).toBe(MAX_CACHED_ENTRIES)
141+
142+
clearSessionMessagesCache()
143+
expect(cache.size).toBe(0)
144+
145+
for (let i = 0; i < MAX_CACHED_ENTRIES + 5; i++) {
146+
writeFileSync(sessionFilePath(`refill-${i}`), '')
147+
await getSessionMessages(asUuid(`refill-${i}`))
148+
}
149+
expect(cache.size).toBe(MAX_CACHED_ENTRIES)
150+
})
151+
})

‎src/utils/sessionStorage.ts‎

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3941,24 +3941,55 @@ async function loadSessionFile(sessionId: UUID): Promise<{
39413941
return loadTranscriptFile(sessionFile)
39423942
}
39433943

3944+
/**
3945+
* Bounded cache for {@link getSessionMessages}. Bounded at
3946+
* {@link MAX_CACHED_SESSION_FILES} to prevent unbounded Map growth in
3947+
* long-running daemon / swarm sessions that spawn many agents — mirrors
3948+
* the existingSessionFiles pattern above. Entries are evicted in FIFO
3949+
* order via `Map` insertion order.
3950+
*/
3951+
const sessionMessagesCache = new Map<UUID, Promise<Set<UUID>>>()
3952+
39443953
/**
39453954
* Gets message UUIDs for a specific session without loading all sessions.
3946-
* Memoized to avoid re-reading the same session file multiple times.
3955+
* Cached (bounded) to avoid re-reading the same session file multiple
3956+
* times. Concurrent calls for the same `sessionId` share one in-flight
3957+
* load promise.
3958+
*
3959+
* Exported so tests can verify cache eviction behavior directly; not
3960+
* intended as a public API — prefer {@link doesMessageExistInSession}.
39473961
*/
3948-
const getSessionMessages = memoize(
3949-
async (sessionId: UUID): Promise<Set<UUID>> => {
3962+
export async function getSessionMessages(sessionId: UUID): Promise<Set<UUID>> {
3963+
const existing = sessionMessagesCache.get(sessionId)
3964+
if (existing !== undefined) {
3965+
return existing
3966+
}
3967+
// Evict oldest entry when at capacity so the Map stays bounded.
3968+
if (sessionMessagesCache.size >= MAX_CACHED_SESSION_FILES) {
3969+
const oldestKey = sessionMessagesCache.keys().next().value
3970+
if (oldestKey !== undefined) {
3971+
sessionMessagesCache.delete(oldestKey)
3972+
}
3973+
}
3974+
const promise = (async () => {
39503975
const { messages } = await loadSessionFile(sessionId)
39513976
return new Set(messages.keys())
3952-
},
3953-
(sessionId: UUID) => sessionId,
3954-
)
3977+
})()
3978+
sessionMessagesCache.set(sessionId, promise)
3979+
return promise
3980+
}
3981+
3982+
/** Underlying cache for direct manipulation (priming, clearing, tests). */
3983+
export function getSessionMessagesCache(): Map<UUID, Promise<Set<UUID>>> {
3984+
return sessionMessagesCache
3985+
}
39553986

39563987
/**
3957-
* Clear the memoized session messages cache.
3988+
* Clear the cached session messages.
39583989
* Call after compaction when old message UUIDs are no longer valid.
39593990
*/
39603991
export function clearSessionMessagesCache(): void {
3961-
getSessionMessages.cache.clear?.()
3992+
sessionMessagesCache.clear()
39623993
}
39633994

39643995
/**
@@ -3996,11 +4027,9 @@ export async function getLastSessionLog(
39964027
// Guard: only prime if cache is empty. Mid-session callers (e.g. IssueFeedback)
39974028
// may call getLastSessionLog on the current session — overwriting a live cache
39984029
// with a stale disk snapshot would lose unflushed UUIDs and break dedup.
3999-
if (!getSessionMessages.cache.has(sessionId)) {
4000-
getSessionMessages.cache.set(
4001-
sessionId,
4002-
Promise.resolve(new Set(messages.keys())),
4003-
)
4030+
const messagesCache = getSessionMessagesCache()
4031+
if (!messagesCache.has(sessionId)) {
4032+
messagesCache.set(sessionId, Promise.resolve(new Set(messages.keys())))
40044033
}
40054034

40064035
// Find the most recent non-sidechain message

0 commit comments

Comments
 (0)