perf(core): Add content-addressed token-count disk cache

intent(token-count-cache): warm-run repacks of the same repo spend ~600ms in BPE tokenization for ~1000 files; persist token counts across CLI invocations so the metrics phase becomes ~free on the second run
decision(cache-shape): single shared JSON under $TMPDIR/repomix-cache/ keyed by `${encoding}:${byteLength}:${md5_16}` — content-addressed entries dedupe across repos (vendored copies, shared boilerplate) and any change to encoding/length/digest auto-invalidates without explicit accounting
decision(wrapper-cache): reuse the same cache for the output-wrapper token count — wrapper string is byte-stable across runs whenever file set, headers, instructions, and template are unchanged, and a hit replaces a ~30ms worker round-trip with MD5+Map.get
decision(save-strategy): await save at end of pack() rather than fire-and-forget so newly produced entries are not lost when a fast-exiting CLI tears down before the write flushes
decision(atomicity): write to ${cacheFile}.${pid}.tmp then rename onto destination so concurrent invocations or SIGINT mid-write cannot leave torn JSON that nukes the cache for the next run
decision(eviction): FIFO on insertion order at MAX_CACHE_ENTRIES=100_000 — true LRU would dirty the map on every read and force a write on every warm run, which costs more than it saves
decision(disable-switch): env-var only (REPOMIX_TOKEN_CACHE=0, REPOMIX_TOKEN_CACHE_PATH=…) — no config schema entry needed for a transparent perf cache with graceful degradation
rejected(per-repo-cache-file): would lose cross-repo dedupe and require placing state under .git or a path-hash directory; single shared file with FIFO eviction is simpler and still bounded
rejected(true-LRU): every getCached would require Map.delete+set to refresh order, dirtying state on read and forcing a save on warm runs even when no new entries were produced
constraint(test-isolation): tests share $TMPDIR with the developer's real cache; vitestSetup defaults REPOMIX_TOKEN_CACHE=0 so the suite neither reads nor writes the host cache, and tokenCountCache.test.ts overrides REPOMIX_TOKEN_CACHE_PATH per test to a fresh tmpdir
constraint(double-md5): cache key is computed once during miss-detection and carried alongside the file into the worker-result map so cold-cache runs do not re-hash content (~10ms saved on 1000 files)
learned(fs-permissions): cache contains digests of user file contents; the cache directory is created with mode 0o700 and the file with mode 0o600 so it is not world-readable on shared hosts
This commit is contained in:
Kazuki Yamada
2026-05-10 15:00:42 +09:00
parent e27d8be1c4
commit 343e6c8c8e
7 changed files with 437 additions and 18 deletions
+53 -14
View File
@@ -4,6 +4,7 @@ import type { RepomixProgressCallback } from '../../shared/types.js';
import type { ProcessedFile } from '../file/fileTypes.js';
import { type MetricsTaskRunner, runBatchTokenCount } from './metricsWorkerRunner.js';
import type { TokenEncoding } from './TokenCounter.js';
import { contentCacheKey, getCached, setCached } from './tokenCountCache.js';
import type { FileMetrics } from './workers/types.js';
// Batch size for grouping files into worker tasks to reduce IPC overhead.
@@ -33,31 +34,64 @@ export const calculateFileMetrics = async (
const startTime = process.hrtime.bigint();
logger.trace(`Starting file metrics calculation for ${filesToProcess.length} files using worker pool`);
// Split files into batches to reduce IPC round-trips
const batches: ProcessedFile[][] = [];
for (let i = 0; i < filesToProcess.length; i += METRICS_BATCH_SIZE) {
batches.push(filesToProcess.slice(i, i + METRICS_BATCH_SIZE));
// Resolve cache hits before dispatching to workers. MD5 hashing each file
// costs ~0.01 ms, far less than a worker round-trip. The key computed here
// is carried forward to the miss path so we never hash the same content
// twice.
type UncachedEntry = { file: ProcessedFile; key: string };
const cachedResults: FileMetrics[] = [];
const uncachedFiles: UncachedEntry[] = [];
for (const file of filesToProcess) {
const key = contentCacheKey(tokenCounterEncoding, file.content);
const cached = getCached(key);
if (cached !== undefined) {
cachedResults.push({ path: file.path, charCount: file.content.length, tokenCount: cached });
} else {
uncachedFiles.push({ file, key });
}
}
logger.trace(`Split ${filesToProcess.length} files into ${batches.length} batches for token counting`);
const cacheHits = cachedResults.length;
const cacheMisses = uncachedFiles.length;
logger.trace(`Token count cache: ${cacheHits} hits, ${cacheMisses} misses`);
let completedItems = 0;
if (cacheMisses === 0) {
// All files were in cache — reconstruct results in original order
const resultMap = new Map(cachedResults.map((r) => [r.path, r]));
const allResults = filesToProcess.map((file) => resultMap.get(file.path) as FileMetrics);
const duration = Number(process.hrtime.bigint() - startTime) / 1e6;
logger.trace(`File metrics calculation completed in ${duration.toFixed(2)}ms (all from cache)`);
progressCallback(`Calculating metrics... (${allResults.length}/${filesToProcess.length})`);
return allResults;
}
// Split uncached files into batches to reduce IPC round-trips
const batches: UncachedEntry[][] = [];
for (let i = 0; i < uncachedFiles.length; i += METRICS_BATCH_SIZE) {
batches.push(uncachedFiles.slice(i, i + METRICS_BATCH_SIZE));
}
logger.trace(`Split ${uncachedFiles.length} uncached files into ${batches.length} batches for token counting`);
let completedItems = cacheHits;
const batchResults = await Promise.all(
batches.map(async (batch) => {
const tokenCounts = await runBatchTokenCount(deps.taskRunner, {
items: batch.map((file) => ({ content: file.content, path: file.path })),
items: batch.map(({ file }) => ({ content: file.content, path: file.path })),
encoding: tokenCounterEncoding,
});
const results: FileMetrics[] = batch.map((file, index) => ({
path: file.path,
charCount: file.content.length,
tokenCount: tokenCounts[index],
}));
const results: FileMetrics[] = batch.map(({ file, key }, index) => {
const tokenCount = tokenCounts[index];
// Reuse the key computed during miss-detection to avoid re-hashing.
setCached(key, tokenCount);
return { path: file.path, charCount: file.content.length, tokenCount };
});
completedItems += batch.length;
const lastFile = batch[batch.length - 1];
const lastFile = batch[batch.length - 1].file;
progressCallback(
`Calculating metrics... (${completedItems}/${filesToProcess.length}) ${pc.dim(lastFile.path)}`,
);
@@ -67,7 +101,12 @@ export const calculateFileMetrics = async (
}),
);
const allResults = batchResults.flat();
// Merge cached and worker results back in original file order.
const workerResultMap = new Map(batchResults.flat().map((r) => [r.path, r]));
const cachedResultMap = new Map(cachedResults.map((r) => [r.path, r]));
const allResults = filesToProcess.map((file) => {
return (cachedResultMap.get(file.path) ?? workerResultMap.get(file.path)) as FileMetrics;
});
const endTime = process.hrtime.bigint();
const duration = Number(endTime - startTime) / 1e6;
+17 -4
View File
@@ -12,6 +12,7 @@ import { calculateGitLogMetrics } from './calculateGitLogMetrics.js';
import { calculateOutputMetrics } from './calculateOutputMetrics.js';
import { type MetricsTaskRunner, runTokenCount } from './metricsWorkerRunner.js';
import type { TokenEncoding } from './TokenCounter.js';
import { contentCacheKey, getCached, setCached } from './tokenCountCache.js';
import type { MetricsWorkerResult, MetricsWorkerTask } from './workers/calculateMetricsWorker.js';
export interface CalculateMetricsResult {
@@ -159,10 +160,22 @@ export const calculateMetrics = async (
? (async () => {
// Dispatch wrapper tokenization immediately — a worker may already be
// idle while file metrics batches still occupy the other workers.
const wrapperTokensPromise = runTokenCount(taskRunner, {
content: outputWrapper,
encoding: config.tokenCount.encoding,
});
// The wrapper string is byte-stable across runs whenever the file
// set, headers, instructions, and template format are unchanged, so
// we reuse the same content-addressed disk cache as per-file token
// counts. Any change to the wrapper automatically misses.
const wrapperCacheKey = contentCacheKey(config.tokenCount.encoding, outputWrapper);
const cachedWrapperTokens = getCached(wrapperCacheKey);
const wrapperTokensPromise =
cachedWrapperTokens !== undefined
? Promise.resolve(cachedWrapperTokens)
: runTokenCount(taskRunner, {
content: outputWrapper,
encoding: config.tokenCount.encoding,
}).then((tokens) => {
setCached(wrapperCacheKey, tokens);
return tokens;
});
const [allFileMetrics, wrapperTokens] = await Promise.all([fileMetricsPromise, wrapperTokensPromise]);
const fileTokensSum = allFileMetrics.reduce((sum, f) => sum + f.tokenCount, 0);
logger.trace(
+159
View File
@@ -0,0 +1,159 @@
import { createHash } from 'node:crypto';
import fs from 'node:fs/promises';
import os from 'node:os';
import path from 'node:path';
import { logger } from '../../shared/logger.js';
import type { TokenEncoding } from './tokenEncodings.js';
// Cache schema version. Bump when the on-disk format changes incompatibly so
// stale caches are discarded silently.
const CACHE_VERSION = 1;
// Hard cap on the number of entries persisted to disk. At ~32 bytes per JSON
// entry, 100k entries ≈ 3 MB. Eviction is FIFO on insertion order — when the
// cap is exceeded the oldest entries are dropped at save time.
export const MAX_CACHE_ENTRIES = 100_000;
const CACHE_DIR_NAME = 'repomix-cache';
const CACHE_FILE_NAME = 'token-counts.json';
interface CacheData {
version: number;
// key: `${encoding}:${byteLength}:${md5_16}`, value: tokenCount
entries: Record<string, number>;
}
interface CacheState {
loaded: boolean;
dirty: boolean;
entries: Map<string, number>;
}
const createState = (): CacheState => ({
loaded: false,
dirty: false,
entries: new Map(),
});
let state = createState();
/**
* Returns the absolute path to the on-disk cache file.
* `REPOMIX_TOKEN_CACHE_PATH` overrides the default location for tests and
* explicit user configuration.
*/
export const getCacheFilePath = (): string => {
const override = process.env.REPOMIX_TOKEN_CACHE_PATH;
if (override) return override;
return path.join(os.tmpdir(), CACHE_DIR_NAME, CACHE_FILE_NAME);
};
/**
* Returns true when the cache is disabled via `REPOMIX_TOKEN_CACHE=0`.
* Any other value (including unset) leaves it enabled.
*/
export const isCacheDisabled = (): boolean => {
return process.env.REPOMIX_TOKEN_CACHE === '0';
};
/**
* Load the on-disk cache into memory. Errors (missing file, corrupt JSON,
* version mismatch) degrade silently to an empty cache so first runs and
* deleted caches keep working.
*/
export const loadTokenCountCache = async (): Promise<void> => {
if (state.loaded) return;
state.loaded = true;
if (isCacheDisabled()) {
logger.trace('Token count cache disabled via REPOMIX_TOKEN_CACHE=0');
return;
}
const cacheFile = getCacheFilePath();
try {
const raw = await fs.readFile(cacheFile, 'utf8');
const data = JSON.parse(raw) as CacheData;
if (data?.version !== CACHE_VERSION || !data.entries) {
logger.trace('Token count cache version mismatch — discarding');
return;
}
for (const [key, value] of Object.entries(data.entries)) {
if (typeof value === 'number') {
state.entries.set(key, value);
}
}
logger.trace(`Loaded ${state.entries.size} token count cache entries from ${cacheFile}`);
} catch {
logger.trace('Token count cache not found or unreadable — starting fresh');
}
};
/**
* Persist the in-memory cache to disk. Writes to a temporary sibling and
* renames over the destination so concurrent invocations and interrupts
* cannot leave a torn JSON file. Caller should await this so newly produced
* entries are not lost on process exit.
*/
export const saveTokenCountCache = async (): Promise<void> => {
if (!state.dirty || state.entries.size === 0) return;
if (isCacheDisabled()) return;
const cacheFile = getCacheFilePath();
const cacheDir = path.dirname(cacheFile);
try {
// Restrict directory permissions so the cache (which contains digests of
// user file contents) is not world-readable on shared hosts.
await fs.mkdir(cacheDir, { recursive: true, mode: 0o700 });
// FIFO eviction: Map iteration order is insertion order, so the oldest
// entries appear first. When over the cap, drop the head of the list.
let entriesToSave = state.entries;
if (state.entries.size > MAX_CACHE_ENTRIES) {
const arr = [...state.entries.entries()];
entriesToSave = new Map(arr.slice(arr.length - MAX_CACHE_ENTRIES));
}
const data: CacheData = {
version: CACHE_VERSION,
entries: Object.fromEntries(entriesToSave),
};
const tmpFile = `${cacheFile}.${process.pid}.tmp`;
await fs.writeFile(tmpFile, JSON.stringify(data), { mode: 0o600 });
await fs.rename(tmpFile, cacheFile);
state.dirty = false;
logger.trace(`Saved ${entriesToSave.size} token count cache entries to ${cacheFile}`);
} catch (error) {
logger.trace('Failed to save token count cache:', error);
}
};
/**
* Build a cache key for content under a specific token encoding.
*
* Format: `${encoding}:${byteLength}:${md5_16}`. Including the byte length
* makes the key tolerant to MD5 collisions on differently-sized inputs and
* keeps the digest portion short (16 hex chars / 64 bits) for compact JSON.
*/
export const contentCacheKey = (encoding: TokenEncoding, content: string): string => {
const byteLength = Buffer.byteLength(content);
const digest = createHash('md5').update(content).digest('hex').slice(0, 16);
return `${encoding}:${byteLength}:${digest}`;
};
export const getCached = (key: string): number | undefined => {
return state.entries.get(key);
};
export const setCached = (key: string, tokenCount: number): void => {
state.entries.set(key, tokenCount);
state.dirty = true;
};
/**
* Test-only: drop all in-memory state so each test starts with a clean slate.
*/
export const __resetTokenCountCacheForTests = (): void => {
state = createState();
};
+16
View File
@@ -12,6 +12,7 @@ import type { ProcessedFile } from './file/fileTypes.js';
import { getGitDiffs } from './git/gitDiffHandle.js';
import { getGitLogs } from './git/gitLogHandle.js';
import { calculateMetrics, createMetricsTaskRunner } from './metrics/calculateMetrics.js';
import { loadTokenCountCache, saveTokenCountCache } from './metrics/tokenCountCache.js';
import { prefetchSortData, sortOutputFiles } from './output/outputSort.js';
import { produceOutput } from './packager/produceOutput.js';
import type { SuspiciousFileResult } from './security/securityCheck.js';
@@ -79,6 +80,12 @@ export const pack = async (
logMemoryUsage('Pack - Start');
// Kick off the token-count cache load in the background so it is ready by
// the time `calculateFileMetrics` reads it. The load itself is small (a few
// hundred KB of JSON at most), but starting it here lets it overlap with
// file search and collection rather than blocking the metrics phase.
const tokenCacheLoadPromise = loadTokenCountCache();
// Pre-fetch git file-change counts for sortOutputFiles while search and
// collection are in flight, so the later sortOutputFiles call is a cache hit.
const sortDataPromise = deps.prefetchSortData(config).catch((error) => {
@@ -208,6 +215,10 @@ export const pack = async (
// Ensure warm-up task completes before metrics calculation
await metricsWarmupPromise;
// Ensure the token-count cache is loaded before calculateFileMetrics reads
// from it. The load was started at the very beginning of pack() and is
// typically already resolved; this await is a safety net for fast machines.
await tokenCacheLoadPromise;
// Generate and write output, overlapping with metrics calculation.
// File and git metrics don't depend on the output, so they start immediately
@@ -255,6 +266,11 @@ export const pack = async (
skippedFiles: allSkippedFiles,
};
// Persist the token-count cache for future runs. Awaited so newly produced
// entries are not lost if the CLI exits immediately after pack(). The save
// is atomic (writeFile-to-tmp + rename) and silently swallows errors.
await saveTokenCountCache();
logMemoryUsage('Pack - End');
return result;
+183
View File
@@ -0,0 +1,183 @@
import fs from 'node:fs/promises';
import os from 'node:os';
import path from 'node:path';
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
import {
__resetTokenCountCacheForTests,
contentCacheKey,
getCached,
isCacheDisabled,
loadTokenCountCache,
MAX_CACHE_ENTRIES,
saveTokenCountCache,
setCached,
} from '../../../src/core/metrics/tokenCountCache.js';
describe('tokenCountCache', () => {
let tmpDir: string;
let cacheFile: string;
const originalDisableEnv = process.env.REPOMIX_TOKEN_CACHE;
const originalPathEnv = process.env.REPOMIX_TOKEN_CACHE_PATH;
beforeEach(async () => {
tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'repomix-cache-test-'));
cacheFile = path.join(tmpDir, 'token-counts.json');
process.env.REPOMIX_TOKEN_CACHE_PATH = cacheFile;
delete process.env.REPOMIX_TOKEN_CACHE;
__resetTokenCountCacheForTests();
});
afterEach(async () => {
if (originalDisableEnv === undefined) {
delete process.env.REPOMIX_TOKEN_CACHE;
} else {
process.env.REPOMIX_TOKEN_CACHE = originalDisableEnv;
}
if (originalPathEnv === undefined) {
delete process.env.REPOMIX_TOKEN_CACHE_PATH;
} else {
process.env.REPOMIX_TOKEN_CACHE_PATH = originalPathEnv;
}
__resetTokenCountCacheForTests();
await fs.rm(tmpDir, { recursive: true, force: true });
});
describe('contentCacheKey', () => {
it('produces stable keys for identical content/encoding pairs', () => {
const a = contentCacheKey('o200k_base', 'hello world');
const b = contentCacheKey('o200k_base', 'hello world');
expect(a).toBe(b);
});
it('differs across encodings', () => {
const a = contentCacheKey('o200k_base', 'hello');
const b = contentCacheKey('cl100k_base', 'hello');
expect(a).not.toBe(b);
});
it('differs across content', () => {
const a = contentCacheKey('o200k_base', 'hello');
const b = contentCacheKey('o200k_base', 'world');
expect(a).not.toBe(b);
});
it('embeds the byte length so same-digest different-length inputs cannot collide', () => {
const key = contentCacheKey('o200k_base', 'abc');
const [encoding, byteLength] = key.split(':');
expect(encoding).toBe('o200k_base');
expect(byteLength).toBe('3');
});
});
describe('load + save round-trip', () => {
it('persists entries to disk and reloads them in a fresh process', async () => {
const key = contentCacheKey('o200k_base', 'hello');
setCached(key, 42);
await saveTokenCountCache();
__resetTokenCountCacheForTests();
await loadTokenCountCache();
expect(getCached(key)).toBe(42);
});
it('skips writing when nothing has been added since the last save', async () => {
// Nothing dirty → no file should exist
await saveTokenCountCache();
await expect(fs.stat(cacheFile)).rejects.toThrow();
});
it('starts fresh when the cache file is missing', async () => {
await loadTokenCountCache();
expect(getCached('whatever:0:0000000000000000')).toBeUndefined();
});
it('starts fresh when the cache file is corrupt JSON', async () => {
await fs.mkdir(path.dirname(cacheFile), { recursive: true });
await fs.writeFile(cacheFile, '{not valid json');
await loadTokenCountCache();
const key = contentCacheKey('o200k_base', 'x');
expect(getCached(key)).toBeUndefined();
// Subsequent writes should still succeed
setCached(key, 7);
await saveTokenCountCache();
__resetTokenCountCacheForTests();
await loadTokenCountCache();
expect(getCached(key)).toBe(7);
});
it('discards entries when the on-disk version does not match', async () => {
await fs.mkdir(path.dirname(cacheFile), { recursive: true });
await fs.writeFile(cacheFile, JSON.stringify({ version: 999, entries: { 'o200k_base:1:abc': 5 } }));
await loadTokenCountCache();
expect(getCached('o200k_base:1:abc')).toBeUndefined();
});
it('is atomic: the destination file is replaced via rename, not overwritten in place', async () => {
const key = contentCacheKey('o200k_base', 'first');
setCached(key, 1);
await saveTokenCountCache();
const firstStat = await fs.stat(cacheFile);
const key2 = contentCacheKey('o200k_base', 'second');
setCached(key2, 2);
await saveTokenCountCache();
const secondStat = await fs.stat(cacheFile);
// Atomic rename gives the file a new inode on most filesystems.
// The data, at minimum, must have been replaced cleanly.
const raw = await fs.readFile(cacheFile, 'utf8');
const data = JSON.parse(raw);
expect(data.entries[key]).toBe(1);
expect(data.entries[key2]).toBe(2);
expect(secondStat.size).toBeGreaterThanOrEqual(firstStat.size);
});
});
describe('FIFO eviction', () => {
it('drops the oldest entries once the cap is exceeded on save', async () => {
// Use a stub cap by directly inserting MAX + 5 entries
for (let i = 0; i < MAX_CACHE_ENTRIES + 5; i++) {
setCached(`o200k_base:1:${i.toString(16).padStart(16, '0')}`, i);
}
await saveTokenCountCache();
__resetTokenCountCacheForTests();
await loadTokenCountCache();
// Oldest 5 dropped; newest preserved
expect(getCached(`o200k_base:1:${(0).toString(16).padStart(16, '0')}`)).toBeUndefined();
expect(getCached(`o200k_base:1:${(4).toString(16).padStart(16, '0')}`)).toBeUndefined();
expect(getCached(`o200k_base:1:${(5).toString(16).padStart(16, '0')}`)).toBe(5);
expect(getCached(`o200k_base:1:${(MAX_CACHE_ENTRIES + 4).toString(16).padStart(16, '0')}`)).toBe(
MAX_CACHE_ENTRIES + 4,
);
});
});
describe('disable switch', () => {
it('reports disabled when REPOMIX_TOKEN_CACHE=0', () => {
process.env.REPOMIX_TOKEN_CACHE = '0';
expect(isCacheDisabled()).toBe(true);
});
it('treats unset and other values as enabled', () => {
delete process.env.REPOMIX_TOKEN_CACHE;
expect(isCacheDisabled()).toBe(false);
process.env.REPOMIX_TOKEN_CACHE = '1';
expect(isCacheDisabled()).toBe(false);
});
it('skips load and save when disabled', async () => {
process.env.REPOMIX_TOKEN_CACHE = '0';
const key = contentCacheKey('o200k_base', 'hello');
setCached(key, 99);
await saveTokenCountCache();
// No file should have been written
await expect(fs.stat(cacheFile)).rejects.toThrow();
__resetTokenCountCacheForTests();
await loadTokenCountCache();
expect(getCached(key)).toBeUndefined();
});
});
});
+8
View File
@@ -0,0 +1,8 @@
// Disable the token-count disk cache by default for the entire test suite so
// that (a) test runs do not read or write the developer's real cache file in
// $TMPDIR and (b) tests asserting on worker dispatch behavior are not skewed
// by entries left behind by a previous run. Tests that exercise the cache
// directly explicitly clear this variable in their own setup.
if (process.env.REPOMIX_TOKEN_CACHE === undefined) {
process.env.REPOMIX_TOKEN_CACHE = '0';
}
+1
View File
@@ -5,6 +5,7 @@ export default defineConfig({
globals: true,
environment: 'node',
include: ['tests/**/*.test.ts'],
setupFiles: ['tests/testing/vitestSetup.ts'],
coverage: {
include: ['src/**/*'],
exclude: ['src/index.ts'],