Cap Kimi session index cache (#7673)
This commit is contained in:
Neil
GitHub
parent
9202f92220
commit
e538a93129
|
|
@ -0,0 +1,164 @@
|
|||
import { mkdtemp, rm, stat, utimes, writeFile } from 'node:fs/promises'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import {
|
||||
KimiSessionIndexCache,
|
||||
KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS,
|
||||
KIMI_WORK_DIR_CACHE_TTL_MS,
|
||||
type KimiSessionIndexIdentity
|
||||
} from './session-scanner-kimi-index-cache'
|
||||
import {
|
||||
clearKimiSessionIndexCache,
|
||||
hasKimiSessionIndexCacheEntryForTests,
|
||||
readKimiWorkDirBySessionId
|
||||
} from './session-scanner-kimi-paths'
|
||||
|
||||
const IDENTITY: KimiSessionIndexIdentity = {
|
||||
changeTimeMs: 1,
|
||||
mtimeMs: 1,
|
||||
sizeBytes: 1
|
||||
}
|
||||
let tempDirs: string[] = []
|
||||
|
||||
afterEach(async () => {
|
||||
vi.useRealTimers()
|
||||
clearKimiSessionIndexCache()
|
||||
await Promise.all(tempDirs.map((dir) => rm(dir, { recursive: true, force: true })))
|
||||
tempDirs = []
|
||||
})
|
||||
|
||||
describe('KimiSessionIndexCache', () => {
|
||||
it('bounds prolonged path churn and retains a reused index', async () => {
|
||||
const cache = new KimiSessionIndexCache()
|
||||
|
||||
for (let index = 0; index < KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS; index += 1) {
|
||||
const path = `home-${index}/session_index.jsonl`
|
||||
await cache.get(path, IDENTITY, cache.beginRead(), async () => new Map([[path, path]]))
|
||||
}
|
||||
const reusedPath = 'home-0/session_index.jsonl'
|
||||
await cache.get(reusedPath, IDENTITY, cache.beginRead(), async () => new Map())
|
||||
|
||||
const firstOverflowPath = `home-${KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS}/session_index.jsonl`
|
||||
await cache.get(
|
||||
firstOverflowPath,
|
||||
IDENTITY,
|
||||
cache.beginRead(),
|
||||
async () => new Map([[firstOverflowPath, firstOverflowPath]])
|
||||
)
|
||||
expect(cache.has(reusedPath)).toBe(true)
|
||||
expect(cache.has('home-1/session_index.jsonl')).toBe(false)
|
||||
|
||||
for (let index = KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS + 1; index < 640; index += 1) {
|
||||
const path = `home-${index}/session_index.jsonl`
|
||||
await cache.get(path, IDENTITY, cache.beginRead(), async () => new Map([[path, path]]))
|
||||
}
|
||||
|
||||
expect(cache.has(reusedPath)).toBe(false)
|
||||
expect(cache.size).toBe(KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS)
|
||||
expect(cache.has('home-576/session_index.jsonl')).toBe(true)
|
||||
expect(cache.has('home-575/session_index.jsonl')).toBe(false)
|
||||
cache.clear()
|
||||
})
|
||||
|
||||
it('refreshes active entries and expires them after an idle TTL', async () => {
|
||||
vi.useFakeTimers()
|
||||
const cache = new KimiSessionIndexCache()
|
||||
const path = 'active/session_index.jsonl'
|
||||
const value = new Map([['session', '/repo']])
|
||||
|
||||
await cache.get(path, IDENTITY, cache.beginRead(), async () => value)
|
||||
await vi.advanceTimersByTimeAsync(KIMI_WORK_DIR_CACHE_TTL_MS - 1)
|
||||
expect(await cache.get(path, IDENTITY, cache.beginRead(), async () => new Map())).toBe(value)
|
||||
|
||||
await vi.advanceTimersByTimeAsync(KIMI_WORK_DIR_CACHE_TTL_MS - 1)
|
||||
expect(cache.has(path)).toBe(true)
|
||||
await vi.advanceTimersByTimeAsync(1)
|
||||
expect(cache.has(path)).toBe(false)
|
||||
expect(vi.getTimerCount()).toBe(0)
|
||||
})
|
||||
|
||||
it('deduplicates concurrent reads of the same file identity', async () => {
|
||||
const cache = new KimiSessionIndexCache()
|
||||
const load = vi.fn(async () => new Map([['session', '/repo']]))
|
||||
const first = cache.get('index', IDENTITY, cache.beginRead(), load)
|
||||
const second = cache.get('index', IDENTITY, cache.beginRead(), load)
|
||||
|
||||
expect(second).toBe(first)
|
||||
await expect(second).resolves.toEqual(new Map([['session', '/repo']]))
|
||||
expect(load).toHaveBeenCalledOnce()
|
||||
cache.clear()
|
||||
})
|
||||
|
||||
it('does not let an older mutation race replace a newer identity', async () => {
|
||||
const cache = new KimiSessionIndexCache()
|
||||
const oldGeneration = cache.beginRead()
|
||||
const newGeneration = cache.beginRead()
|
||||
const newerIdentity = { ...IDENTITY, changeTimeMs: 2, mtimeMs: 2, sizeBytes: 2 }
|
||||
const newer = new Map([['session', '/new']])
|
||||
|
||||
await cache.get('index', newerIdentity, newGeneration, async () => newer)
|
||||
await cache.get('index', IDENTITY, oldGeneration, async () => new Map([['session', '/old']]))
|
||||
const reload = vi.fn(async () => new Map())
|
||||
|
||||
await expect(cache.get('index', newerIdentity, cache.beginRead(), reload)).resolves.toBe(newer)
|
||||
expect(reload).not.toHaveBeenCalled()
|
||||
cache.clear()
|
||||
})
|
||||
|
||||
it('does not repopulate after an owner clears an in-flight read', async () => {
|
||||
const cache = new KimiSessionIndexCache()
|
||||
const staleGeneration = cache.beginRead()
|
||||
cache.clear()
|
||||
|
||||
await cache.get('index', IDENTITY, staleGeneration, async () => new Map([['session', '/old']]))
|
||||
|
||||
expect(cache.has('index')).toBe(false)
|
||||
const current = new Map([['session', '/current']])
|
||||
await expect(
|
||||
cache.get('index', IDENTITY, cache.beginRead(), async () => current)
|
||||
).resolves.toBe(current)
|
||||
expect(cache.has('index')).toBe(true)
|
||||
cache.clear()
|
||||
})
|
||||
})
|
||||
|
||||
describe('Kimi session index reader cache', () => {
|
||||
it('releases a retained map when its index file disappears', async () => {
|
||||
const root = await mkdtemp(join(tmpdir(), 'orca-kimi-index-delete-'))
|
||||
tempDirs.push(root)
|
||||
const indexPath = join(root, 'session_index.jsonl')
|
||||
await writeFile(indexPath, `${JSON.stringify({ sessionId: 'session', workDir: '/repo' })}\n`)
|
||||
|
||||
await readKimiWorkDirBySessionId(indexPath)
|
||||
expect(hasKimiSessionIndexCacheEntryForTests(indexPath)).toBe(true)
|
||||
await rm(indexPath)
|
||||
|
||||
await expect(readKimiWorkDirBySessionId(indexPath)).resolves.toEqual(new Map())
|
||||
expect(hasKimiSessionIndexCacheEntryForTests(indexPath)).toBe(false)
|
||||
})
|
||||
|
||||
it('invalidates when size changes even if the mtime is restored', async () => {
|
||||
const root = await mkdtemp(join(tmpdir(), 'orca-kimi-index-mutation-'))
|
||||
tempDirs.push(root)
|
||||
const indexPath = join(root, 'session_index.jsonl')
|
||||
await writeFile(indexPath, `${JSON.stringify({ sessionId: 'session-old', workDir: '/old' })}\n`)
|
||||
const originalStat = await stat(indexPath)
|
||||
|
||||
await expect(readKimiWorkDirBySessionId(indexPath)).resolves.toEqual(
|
||||
new Map([['session-old', '/old']])
|
||||
)
|
||||
await writeFile(
|
||||
indexPath,
|
||||
`${JSON.stringify({ sessionId: 'session-old', workDir: '/old' })}\n${JSON.stringify({ sessionId: 'session-new', workDir: '/new' })}\n`
|
||||
)
|
||||
await utimes(indexPath, originalStat.atime, originalStat.mtime)
|
||||
|
||||
await expect(readKimiWorkDirBySessionId(indexPath)).resolves.toEqual(
|
||||
new Map([
|
||||
['session-old', '/old'],
|
||||
['session-new', '/new']
|
||||
])
|
||||
)
|
||||
})
|
||||
})
|
||||
|
|
@ -0,0 +1,132 @@
|
|||
export type KimiSessionIndexIdentity = {
|
||||
changeTimeMs: number
|
||||
mtimeMs: number
|
||||
sizeBytes: number
|
||||
}
|
||||
|
||||
type KimiSessionIndexCacheEntry = {
|
||||
expiresAt: number
|
||||
generation: number
|
||||
identity: KimiSessionIndexIdentity
|
||||
timer: NodeJS.Timeout | null
|
||||
value: Promise<Map<string, string>>
|
||||
}
|
||||
|
||||
export const KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS = 64
|
||||
// Active Vault scans refresh this window; closing the surface releases parsed
|
||||
// index maps soon without making a live Kimi session reread on every scan.
|
||||
export const KIMI_WORK_DIR_CACHE_TTL_MS = 5 * 60_000
|
||||
|
||||
export class KimiSessionIndexCache {
|
||||
private readonly entries = new Map<string, KimiSessionIndexCacheEntry>()
|
||||
private minimumCacheGeneration = 0
|
||||
private nextGeneration = 0
|
||||
|
||||
beginRead(): number {
|
||||
this.nextGeneration += 1
|
||||
return this.nextGeneration
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
for (const entry of this.entries.values()) {
|
||||
if (entry.timer) {
|
||||
clearTimeout(entry.timer)
|
||||
}
|
||||
}
|
||||
this.entries.clear()
|
||||
// Why: a read already awaiting stat/load when its owner clears the cache
|
||||
// may finish later, but must not silently recreate the released entry.
|
||||
this.minimumCacheGeneration = this.nextGeneration + 1
|
||||
}
|
||||
|
||||
delete(indexPath: string, generation = Number.POSITIVE_INFINITY): void {
|
||||
const entry = this.entries.get(indexPath)
|
||||
if (entry && entry.generation <= generation) {
|
||||
this.forget(indexPath, entry)
|
||||
}
|
||||
}
|
||||
|
||||
get(
|
||||
indexPath: string,
|
||||
identity: KimiSessionIndexIdentity,
|
||||
generation: number,
|
||||
load: () => Promise<Map<string, string>>
|
||||
): Promise<Map<string, string>> {
|
||||
if (generation < this.minimumCacheGeneration) {
|
||||
return load()
|
||||
}
|
||||
const cached = this.entries.get(indexPath)
|
||||
const now = Date.now()
|
||||
if (cached && cached.expiresAt > now && identitiesMatch(cached.identity, identity)) {
|
||||
this.remember(indexPath, cached, now)
|
||||
return cached.value
|
||||
}
|
||||
if (cached && cached.generation > generation) {
|
||||
// Why: a slower, older stat must not replace a newer file generation
|
||||
// that another concurrent scan already cached for the same path.
|
||||
return load()
|
||||
}
|
||||
|
||||
const entry: KimiSessionIndexCacheEntry = {
|
||||
expiresAt: now + KIMI_WORK_DIR_CACHE_TTL_MS,
|
||||
generation,
|
||||
identity,
|
||||
timer: null,
|
||||
value: load()
|
||||
}
|
||||
this.remember(indexPath, entry, now)
|
||||
return entry.value
|
||||
}
|
||||
|
||||
has(indexPath: string): boolean {
|
||||
return this.entries.has(indexPath)
|
||||
}
|
||||
|
||||
get size(): number {
|
||||
return this.entries.size
|
||||
}
|
||||
|
||||
private forget(indexPath: string, entry: KimiSessionIndexCacheEntry): void {
|
||||
if (this.entries.get(indexPath) !== entry) {
|
||||
return
|
||||
}
|
||||
if (entry.timer) {
|
||||
clearTimeout(entry.timer)
|
||||
}
|
||||
this.entries.delete(indexPath)
|
||||
}
|
||||
|
||||
private remember(indexPath: string, entry: KimiSessionIndexCacheEntry, now: number): void {
|
||||
const replaced = this.entries.get(indexPath)
|
||||
if (replaced?.timer && replaced !== entry) {
|
||||
clearTimeout(replaced.timer)
|
||||
}
|
||||
if (entry.timer) {
|
||||
clearTimeout(entry.timer)
|
||||
}
|
||||
entry.expiresAt = now + KIMI_WORK_DIR_CACHE_TTL_MS
|
||||
entry.timer = setTimeout(() => this.forget(indexPath, entry), KIMI_WORK_DIR_CACHE_TTL_MS)
|
||||
entry.timer.unref()
|
||||
this.entries.delete(indexPath)
|
||||
this.entries.set(indexPath, entry)
|
||||
|
||||
while (this.entries.size > KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS) {
|
||||
const oldest = this.entries.entries().next().value
|
||||
if (!oldest) {
|
||||
return
|
||||
}
|
||||
this.forget(oldest[0], oldest[1])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function identitiesMatch(
|
||||
left: KimiSessionIndexIdentity,
|
||||
right: KimiSessionIndexIdentity
|
||||
): boolean {
|
||||
return (
|
||||
left.changeTimeMs === right.changeTimeMs &&
|
||||
left.mtimeMs === right.mtimeMs &&
|
||||
left.sizeBytes === right.sizeBytes
|
||||
)
|
||||
}
|
||||
|
|
@ -3,7 +3,12 @@ import { tmpdir } from 'node:os'
|
|||
import { join } from 'node:path'
|
||||
import { afterEach, describe, expect, it } from 'vitest'
|
||||
import { parseKimiSessionFile } from './session-scanner-kimi-parser'
|
||||
import { clearKimiSessionIndexCache } from './session-scanner-kimi-paths'
|
||||
import {
|
||||
clearKimiSessionIndexCache,
|
||||
hasKimiSessionIndexCacheEntryForTests,
|
||||
KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS,
|
||||
readKimiWorkDirBySessionId
|
||||
} from './session-scanner-kimi-paths'
|
||||
import type { FileWithMtime } from './session-scanner-types'
|
||||
|
||||
let tempDirs: string[] = []
|
||||
|
|
@ -218,4 +223,34 @@ describe('parseKimiSessionFile', () => {
|
|||
const session = await parseKimiSessionFile(file, 'darwin')
|
||||
expect(session?.title).toBe('do the thing')
|
||||
})
|
||||
|
||||
it('caps cached session index maps by recent index path', async () => {
|
||||
const root = await mkdtemp(join(tmpdir(), 'orca-kimi-cache-'))
|
||||
tempDirs.push(root)
|
||||
const indexPaths: string[] = []
|
||||
|
||||
for (let index = 0; index <= KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS; index += 1) {
|
||||
const home = join(root, `home-${index}`)
|
||||
await mkdir(home, { recursive: true })
|
||||
const indexPath = join(home, 'session_index.jsonl')
|
||||
await writeFile(
|
||||
indexPath,
|
||||
`${JSON.stringify({ sessionId: `session_${index}`, workDir: `/tmp/kimi-${index}` })}\n`
|
||||
)
|
||||
indexPaths.push(indexPath)
|
||||
}
|
||||
|
||||
for (const indexPath of indexPaths.slice(0, KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS)) {
|
||||
await readKimiWorkDirBySessionId(indexPath)
|
||||
}
|
||||
const refreshedFirst = await readKimiWorkDirBySessionId(indexPaths[0])
|
||||
expect(refreshedFirst.get('session_0')).toBe('/tmp/kimi-0')
|
||||
await readKimiWorkDirBySessionId(indexPaths[KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS])
|
||||
|
||||
expect(hasKimiSessionIndexCacheEntryForTests(indexPaths[0])).toBe(true)
|
||||
expect(hasKimiSessionIndexCacheEntryForTests(indexPaths[1])).toBe(false)
|
||||
expect(
|
||||
hasKimiSessionIndexCacheEntryForTests(indexPaths[KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS])
|
||||
).toBe(true)
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -4,6 +4,13 @@ import { homedir } from 'node:os'
|
|||
import { basename, dirname, join } from 'node:path'
|
||||
import { createInterface } from 'node:readline'
|
||||
import { asRecord, extractString } from './session-scanner-values'
|
||||
import {
|
||||
KimiSessionIndexCache,
|
||||
KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS,
|
||||
KIMI_WORK_DIR_CACHE_TTL_MS
|
||||
} from './session-scanner-kimi-index-cache'
|
||||
|
||||
export { KIMI_WORK_DIR_CACHE_MAX_INDEX_PATHS, KIMI_WORK_DIR_CACHE_TTL_MS }
|
||||
|
||||
// Why: Kimi Code stores sessions under <KIMI_CODE_HOME>/sessions/, mirroring the
|
||||
// CLI's own `KIMI_CODE_HOME ?? ~/.kimi-code` resolution (see kimi-fetcher.ts).
|
||||
|
|
@ -55,38 +62,41 @@ export function kimiPrimaryAgentWirePath(
|
|||
return join(dirname(statePath), 'agents', primaryId, 'wire.jsonl')
|
||||
}
|
||||
|
||||
type WorkDirCacheEntry = {
|
||||
mtimeMs: number
|
||||
map: Promise<Map<string, string>>
|
||||
}
|
||||
|
||||
// Why: every session under one Kimi home shares a single session_index.jsonl.
|
||||
// Re-reading it once per session would be O(n^2); memoize by path + mtime so a
|
||||
// scan reads the index at most once and the cache self-invalidates when Kimi
|
||||
// appends a new session (mtime bump).
|
||||
const workDirCacheByIndexPath = new Map<string, WorkDirCacheEntry>()
|
||||
// Re-reading it once per session would be O(n^2); memoize by path + file
|
||||
// identity so a scan reads the index at most once. Bound and expire entries
|
||||
// because host/WSL/runtime roots can change during one main-process lifetime.
|
||||
const workDirCacheByIndexPath = new KimiSessionIndexCache()
|
||||
|
||||
export function clearKimiSessionIndexCache(): void {
|
||||
workDirCacheByIndexPath.clear()
|
||||
}
|
||||
|
||||
export function hasKimiSessionIndexCacheEntryForTests(indexPath: string): boolean {
|
||||
return workDirCacheByIndexPath.has(indexPath)
|
||||
}
|
||||
|
||||
export async function readKimiWorkDirBySessionId(indexPath: string): Promise<Map<string, string>> {
|
||||
let mtimeMs: number
|
||||
const generation = workDirCacheByIndexPath.beginRead()
|
||||
let identity: Awaited<ReturnType<typeof stat>>
|
||||
try {
|
||||
mtimeMs = (await stat(indexPath)).mtimeMs
|
||||
identity = await stat(indexPath)
|
||||
} catch {
|
||||
// Missing index (e.g. user deleted it): sessions still list, just without cwd.
|
||||
workDirCacheByIndexPath.delete(indexPath, generation)
|
||||
return new Map()
|
||||
}
|
||||
|
||||
const cached = workDirCacheByIndexPath.get(indexPath)
|
||||
if (cached && cached.mtimeMs === mtimeMs) {
|
||||
return cached.map
|
||||
}
|
||||
|
||||
const map = parseKimiSessionIndex(indexPath)
|
||||
workDirCacheByIndexPath.set(indexPath, { mtimeMs, map })
|
||||
return map
|
||||
return workDirCacheByIndexPath.get(
|
||||
indexPath,
|
||||
{
|
||||
changeTimeMs: identity.ctimeMs,
|
||||
mtimeMs: identity.mtimeMs,
|
||||
sizeBytes: identity.size
|
||||
},
|
||||
generation,
|
||||
() => parseKimiSessionIndex(indexPath)
|
||||
)
|
||||
}
|
||||
|
||||
async function parseKimiSessionIndex(indexPath: string): Promise<Map<string, string>> {
|
||||
|
|
|
|||
Loading…
Reference in New Issue