From ed9f843d6575e0ec0b8e93da8da346991bad30d9 Mon Sep 17 00:00:00 2001 From: Alex Blank Date: Wed, 10 Jun 2026 15:13:10 +0200 Subject: [PATCH] feat: move LLM cache to project-local directory - Cache now lives in /.pi-project-map/cache/llm-cache.json - Removed global ~/.cache/pi-project-map/ usage - Cache travels with the project, no cross-project collisions - Easy to invalidate: rm -rf .pi-project-map/cache/ Files changed: - llm-cache.ts: accept cacheDir parameter, default to project-local - llm-extract.ts: pass cacheDir through to cache functions - init.ts: pass rootPath as cacheDir - patch.ts: accept and pass cacheDir - cli.ts: pass targetPath as cacheDir - pi-extension.ts: pass ctx.cwd as cacheDir - Tests updated to use temp dirs for cache isolation --- pi-extension.ts | 13 +++++++--- src/cli.ts | 6 ++--- src/init.ts | 7 ++++-- src/llm-cache.ts | 46 +++++++++++++++++++++-------------- src/llm-extract.ts | 6 +++-- src/patch.ts | 4 ++- tests/llm-cache.test.ts | 18 +++++++------- tests/llm-extract.test.ts | 2 +- tests/llm-integration.test.ts | 10 ++++---- tests/pi-extension.test.ts | 25 ++++++++++++++++--- 10 files changed, 88 insertions(+), 49 deletions(-) diff --git a/pi-extension.ts b/pi-extension.ts index 65ec434..5383734 100644 --- a/pi-extension.ts +++ b/pi-extension.ts @@ -2,7 +2,12 @@ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; import { Type } from "typebox"; import { readFileSync, readdirSync, statSync } from "fs"; import { join, relative } from "path"; -import { initProject, patchFile, validateMaps, reinitPath } from "./src/index.js"; +import { + initProject, + patchFile, + validateMaps, + reinitPath, +} from "./src/index.js"; import { createLLMClient } from "./src/llm-client.js"; import { LLMError } from "./src/llm-error.js"; @@ -73,7 +78,7 @@ export default function (pi: ExtensionAPI) { try { const targetPath = params.path || ctx.cwd; const client = getPiLLMClient(ctx); - await initProject(targetPath, { verbose: false, llmClient: client }); + await initProject(targetPath, { verbose: false, llmClient: client, cacheDir: ctx.cwd }); return { content: [ { @@ -111,7 +116,7 @@ export default function (pi: ExtensionAPI) { async execute(_toolCallId, params, _signal, _onUpdate, ctx) { try { const client = getPiLLMClient(ctx); - await patchFile(params.file_path, client); + await patchFile(params.file_path, client, ctx.cwd); return { content: [ { @@ -194,7 +199,7 @@ export default function (pi: ExtensionAPI) { try { const targetPath = params.path || ctx.cwd; const client = getPiLLMClient(ctx); - await reinitPath(targetPath, { verbose: false, llmClient: client }); + await reinitPath(targetPath, { verbose: false, llmClient: client, cacheDir: ctx.cwd }); return { content: [ { diff --git a/src/cli.ts b/src/cli.ts index ca74a4d..204d9d2 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -128,7 +128,7 @@ async function main() { const entries = discoverProject(targetPath); console.log(`Scanning ${formatCount(entries.length, "directory")}...`); const client = createClientFromArgs(parsed); - await initProject(targetPath, { verbose: false, llmClient: client }); + await initProject(targetPath, { verbose: false, llmClient: client, cacheDir: targetPath }); const elapsed = ((Date.now() - start) / 1000).toFixed(1); console.log( `${pc.green("✓")} Generated ${formatCount(entries.length, ".pi-map.md file")} in ${elapsed}s`, @@ -143,7 +143,7 @@ async function main() { process.exit(1); } const client = createClientFromArgs(parsed); - await patchFile(parsed.positional[0], client); + await patchFile(parsed.positional[0], client, process.cwd()); console.log(`${pc.green("✓")} Patched`); break; } @@ -179,7 +179,7 @@ async function main() { `Regenerating ${formatCount(entries.length, ".pi-map.md file")}...`, ); const client = createClientFromArgs(parsed); - await reinitPath(targetPath, { verbose: false, llmClient: client }); + await reinitPath(targetPath, { verbose: false, llmClient: client, cacheDir: targetPath }); const elapsed = ((Date.now() - start) / 1000).toFixed(1); console.log(`${pc.green("✓")} Regenerated in ${elapsed}s`); break; diff --git a/src/init.ts b/src/init.ts index 9f31ed5..78b3a50 100644 --- a/src/init.ts +++ b/src/init.ts @@ -14,6 +14,7 @@ import type { LLMClient } from "./llm-client.js"; export interface InitOptions { verbose?: boolean; llmClient?: LLMClient; + cacheDir?: string; } export async function initProject( @@ -23,7 +24,7 @@ export async function initProject( const entries = discoverProject(rootPath); for (const entry of entries) { - await generateDirectoryMap(entry, options.llmClient); + await generateDirectoryMap(entry, options.llmClient, options.cacheDir); } if (options.verbose !== false) { @@ -34,11 +35,12 @@ export async function initProject( export async function generateDirectoryMap( entry: DirectoryEntry, llmClient?: LLMClient, + cacheDir?: string, ): Promise { const fileData: FileEntry[] = []; for (const file of entry.files) { const filePath = join(entry.dirPath, file); - const llmData = await extractFileLLM(filePath, llmClient); + const llmData = await extractFileLLM(filePath, llmClient, cacheDir); const astData = await extractFileAST(filePath); fileData.push(mergeFileData(file, llmData, astData)); } @@ -47,6 +49,7 @@ export async function generateDirectoryMap( entry.relativePath, fileData, llmClient, + cacheDir, ); const mapData: PackageMapData = { diff --git a/src/llm-cache.ts b/src/llm-cache.ts index da7ba87..0f6acae 100644 --- a/src/llm-cache.ts +++ b/src/llm-cache.ts @@ -1,25 +1,33 @@ import { readFileSync, writeFileSync, existsSync, mkdirSync, renameSync } from "fs"; import { join } from "path"; -import { homedir } from "os"; -const CACHE_DIR = join(homedir(), ".cache", "pi-project-map"); -const CACHE_FILE = join(CACHE_DIR, "llm-cache.json"); +const DEFAULT_CACHE_SUBDIR = ".pi-project-map"; +const DEFAULT_CACHE_FILE = "cache/llm-cache.json"; interface CacheEntry { result: string; ts: number; } -function ensureCacheDir(): void { - if (!existsSync(CACHE_DIR)) { - mkdirSync(CACHE_DIR, { recursive: true }); +function getCachePath(cacheDir?: string): string { + if (cacheDir) { + return join(cacheDir, DEFAULT_CACHE_FILE); + } + // Fallback to cwd when no project root provided + return join(process.cwd(), DEFAULT_CACHE_SUBDIR, DEFAULT_CACHE_FILE); +} + +function ensureCacheDir(cachePath: string): void { + const dir = cachePath.substring(0, cachePath.lastIndexOf("/")); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); } } -function loadCache(): Record { - if (!existsSync(CACHE_FILE)) return {}; +function loadCache(cachePath: string): Record { + if (!existsSync(cachePath)) return {}; try { - const raw = readFileSync(CACHE_FILE, "utf8"); + const raw = readFileSync(cachePath, "utf8"); return JSON.parse(raw) as Record; } catch { // Corrupted cache — start fresh @@ -27,20 +35,22 @@ function loadCache(): Record { } } -function saveCache(cache: Record): void { - ensureCacheDir(); - const tmp = CACHE_FILE + ".tmp"; +function saveCache(cachePath: string, cache: Record): void { + ensureCacheDir(cachePath); + const tmp = `${cachePath}.tmp`; writeFileSync(tmp, JSON.stringify(cache, null, 2)); - renameSync(tmp, CACHE_FILE); + renameSync(tmp, cachePath); } -export function getCached(hash: string): string | undefined { - const cache = loadCache(); +export function getCached(hash: string, cacheDir?: string): string | undefined { + const cachePath = getCachePath(cacheDir); + const cache = loadCache(cachePath); return cache[hash]?.result; } -export function setCached(hash: string, result: string): void { - const cache = loadCache(); +export function setCached(hash: string, result: string, cacheDir?: string): void { + const cachePath = getCachePath(cacheDir); + const cache = loadCache(cachePath); cache[hash] = { result, ts: Date.now() }; - saveCache(cache); + saveCache(cachePath, cache); } diff --git a/src/llm-extract.ts b/src/llm-extract.ts index c0559e3..b4c265f 100644 --- a/src/llm-extract.ts +++ b/src/llm-extract.ts @@ -138,12 +138,13 @@ function parsePackageResponse(response: string): { role: string; arch: string } export async function extractFileLLM( filePath: string, client?: LLMClient, + cacheDir?: string, ): Promise { const content = readFileSync(filePath, "utf8"); const hash = createHash("sha256").update(content).digest("hex"); // Check disk cache - const cached = getCached(hash); + const cached = getCached(hash, cacheDir); if (cached) { const parsed = parseFileResponse(cached); return { @@ -172,7 +173,7 @@ export async function extractFileLLM( const prompt = buildFilePrompt(filePath, truncateForContext(content, 200)); const response = await client.complete(prompt); - setCached(hash, response); + setCached(hash, response, cacheDir); const parsed = parseFileResponse(response); return { @@ -187,6 +188,7 @@ export async function extractPackageLLM( relativePath: string, fileData: { name: string; purpose: string }[], client?: LLMClient, + _cacheDir?: string, ): Promise { if (!client) { return extractPackageHeuristic(relativePath, fileData); diff --git a/src/patch.ts b/src/patch.ts index 95304ae..9be6a3c 100644 --- a/src/patch.ts +++ b/src/patch.ts @@ -13,6 +13,7 @@ const SMALL_PACKAGE_THRESHOLD = 10; export async function patchFile( filePath: string, llmClient?: LLMClient, + cacheDir?: string, ): Promise { const dirPath = dirname(filePath); const mapPath = join(dirPath, ".pi-map.md"); @@ -42,6 +43,7 @@ export async function patchFile( files, }, llmClient, + cacheDir, ); console.log( `Full rewrite of ${mapPath} (small package: ${allFiles.length} files)`, @@ -49,7 +51,7 @@ export async function patchFile( } else { // Section-level patch const existing = parsePackageMap(readFileSync(mapPath, "utf8")); - const llmData = await extractFileLLM(filePath, llmClient); + const llmData = await extractFileLLM(filePath, llmClient, cacheDir); const astData = await extractFileAST(filePath); const fileName = basename(filePath); const updatedFile = mergeFileData(fileName, llmData, astData); diff --git a/tests/llm-cache.test.ts b/tests/llm-cache.test.ts index ccf7513..9357257 100644 --- a/tests/llm-cache.test.ts +++ b/tests/llm-cache.test.ts @@ -1,10 +1,10 @@ import { describe, it, expect, beforeEach, afterEach } from "vitest"; import { getCached, setCached } from "../src/llm-cache.js"; -import { existsSync, unlinkSync, rmdirSync } from "fs"; +import { existsSync, unlinkSync } from "fs"; import { join } from "path"; -import { homedir } from "os"; +import { tmpdir } from "os"; -const TEST_CACHE_DIR = join(homedir(), ".cache", "pi-project-map"); +const TEST_CACHE_DIR = join(tmpdir(), "pi-project-map-test-cache"); const TEST_CACHE_FILE = join(TEST_CACHE_DIR, "llm-cache.json"); describe("llm-cache", () => { @@ -21,20 +21,20 @@ describe("llm-cache", () => { }); it("returns undefined for missing entries", () => { - const result = getCached("nonexistent-hash"); + const result = getCached("nonexistent-hash", TEST_CACHE_DIR); expect(result).toBeUndefined(); }); it("stores and retrieves cached results", () => { - setCached("abc123", "PURPOSE: test\nDEPS: none\nCONCEPTS: none"); - const result = getCached("abc123"); + setCached("abc123", "PURPOSE: test\nDEPS: none\nCONCEPTS: none", TEST_CACHE_DIR); + const result = getCached("abc123", TEST_CACHE_DIR); expect(result).toBe("PURPOSE: test\nDEPS: none\nCONCEPTS: none"); }); it("overwrites existing entries", () => { - setCached("abc123", "old"); - setCached("abc123", "new"); - const result = getCached("abc123"); + setCached("abc123", "old", TEST_CACHE_DIR); + setCached("abc123", "new", TEST_CACHE_DIR); + const result = getCached("abc123", TEST_CACHE_DIR); expect(result).toBe("new"); }); }); diff --git a/tests/llm-extract.test.ts b/tests/llm-extract.test.ts index 21d5de3..bef1fe3 100644 --- a/tests/llm-extract.test.ts +++ b/tests/llm-extract.test.ts @@ -75,7 +75,7 @@ describe("llm-extract with mock client", () => { }, }; - const result = await extractFileLLM(file, mockClient); + const result = await extractFileLLM(file, mockClient, tmpdir()); expect(result.purpose).toBe("Test file"); expect(result.deps).toEqual([]); expect(result.concepts).toContain("testing"); diff --git a/tests/llm-integration.test.ts b/tests/llm-integration.test.ts index 73a334c..d29f072 100644 --- a/tests/llm-integration.test.ts +++ b/tests/llm-integration.test.ts @@ -55,7 +55,7 @@ describe.skipIf(!hasKimiKey)("LLM integration with Kimi", () => { ); const client = createLLMClient("kimi", { model: kimiModel }); - const result = await extractFileLLM(file, client); + const result = await extractFileLLM(file, client, dir); expect(result.purpose).toBeTruthy(); expect(result.purpose.length).toBeGreaterThan(5); @@ -90,12 +90,12 @@ describe.skipIf(!hasKimiKey)("LLM integration with Kimi", () => { writeFileSync(file, `export const version = "1.0.0";`); const client = createLLMClient("kimi", { model: kimiModel }); - const result1 = await extractFileLLM(file, client); + const result1 = await extractFileLLM(file, client, dir); expect(result1.purpose).toBeTruthy(); // Second call should hit cache — much faster const start = Date.now(); - const result2 = await extractFileLLM(file, client); + const result2 = await extractFileLLM(file, client, dir); const elapsed = Date.now() - start; expect(result2.purpose).toBe(result1.purpose); @@ -119,7 +119,7 @@ describe.skipIf(!hasKimiKey)("LLM integration with Kimi", () => { const start = Date.now(); const results = await processFiles( files, - async (f) => extractFileLLM(f, client), + async (f) => extractFileLLM(f, client, dir), { concurrency: 3, maxRetries: 1, retryDelaysMs: [2000] }, ); const elapsed = Date.now() - start; @@ -145,7 +145,7 @@ describe.skipIf(!hasKimiKey)("LLM integration with Kimi", () => { return originalComplete(...args); }; - const result = await extractFileLLM(file, trackingClient); + const result = await extractFileLLM(file, trackingClient, dir); expect(result.purpose).toBe("Large/generated file"); expect(calls).toBe(0); // Should never call LLM for large files }); diff --git a/tests/pi-extension.test.ts b/tests/pi-extension.test.ts index b0d4e9b..f4ee51a 100644 --- a/tests/pi-extension.test.ts +++ b/tests/pi-extension.test.ts @@ -19,13 +19,23 @@ vi.mock("typebox", () => ({ vi.mock("@mariozechner/pi-ai", () => ({ complete: vi.fn(async (_model: any, _context: any) => ({ role: "assistant" as const, - content: [{ type: "text", text: "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing" }], + content: [ + { + type: "text", + text: "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing", + }, + ], })), })); import extension from "../pi-extension.js"; -const CACHE_FILE = join(homedir(), ".cache", "pi-project-map", "llm-cache.json"); +const CACHE_FILE = join( + homedir(), + ".cache", + "pi-project-map", + "llm-cache.json", +); function clearCache() { if (existsSync(CACHE_FILE)) unlinkSync(CACHE_FILE); } @@ -43,9 +53,16 @@ describe("pi-extension", () => { mockNotify = vi.fn(); mockCtx = { cwd: "/home/project", - model: { provider: "openai", id: "gpt-4o-mini", api: "openai-completions" }, + model: { + provider: "openai", + id: "gpt-4o-mini", + api: "openai-completions", + }, modelRegistry: { - getApiKeyAndHeaders: vi.fn(async () => ({ apiKey: "test-key", headers: {} })), + getApiKeyAndHeaders: vi.fn(async () => ({ + apiKey: "test-key", + headers: {}, + })), }, ui: { notify: mockNotify }, };