From f1bc97da63f0bdf00eec8bf5f46b4b62d8593c3d Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Thu, 28 May 2026 12:23:35 +0200 Subject: [PATCH 1/2] feat: Add `BRAINTRUST_CACHE_LOCATION` env var to control caching location --- js/src/logger.ts | 42 ++-- js/src/prompt-cache/cache-config.ts | 99 +++++++++ js/src/prompt-cache/cache-mode.test.ts | 268 ++++++++++++++++++++++++ js/src/prompt-cache/parameters-cache.ts | 16 +- js/src/prompt-cache/prompt-cache.ts | 26 +-- 5 files changed, 405 insertions(+), 46 deletions(-) create mode 100644 js/src/prompt-cache/cache-config.ts create mode 100644 js/src/prompt-cache/cache-mode.test.ts diff --git a/js/src/logger.ts b/js/src/logger.ts index 79a841926..1c6c6ab6b 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -148,8 +148,7 @@ import { devNullWritableStream, } from "./functions/stream"; import iso, { IsoAsyncLocalStorage } from "./isomorph"; -import { canUseDiskCache, DiskCache } from "./prompt-cache/disk-cache"; -import { LRUCache } from "./prompt-cache/lru-cache"; +import { createCacheLayers } from "./prompt-cache/cache-config"; import { PromptCache } from "./prompt-cache/prompt-cache"; import { ParametersCache } from "./prompt-cache/parameters-cache"; import { @@ -711,34 +710,25 @@ export class BraintrustState { this.resetLoginInfo(); - const memoryCache = new LRUCache({ - max: Number(iso.getEnv("BRAINTRUST_PROMPT_CACHE_MEMORY_MAX")) ?? 1 << 10, + const { memoryCache, diskCache } = createCacheLayers({ + memoryMaxEnvVar: "BRAINTRUST_PROMPT_CACHE_MEMORY_MAX", + diskCacheDirEnvVar: "BRAINTRUST_PROMPT_CACHE_DIR", + diskMaxEnvVar: "BRAINTRUST_PROMPT_CACHE_DISK_MAX", + getDefaultDiskCacheDir: () => + `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/prompt_cache`, }); - const diskCache = canUseDiskCache() - ? new DiskCache({ - cacheDir: - iso.getEnv("BRAINTRUST_PROMPT_CACHE_DIR") ?? - `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/prompt_cache`, - max: - Number(iso.getEnv("BRAINTRUST_PROMPT_CACHE_DISK_MAX")) ?? 1 << 20, - }) - : undefined; this.promptCache = new PromptCache({ memoryCache, diskCache }); - const parametersMemoryCache = new LRUCache({ - max: - Number(iso.getEnv("BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX")) ?? 1 << 10, + const { + memoryCache: parametersMemoryCache, + diskCache: parametersDiskCache, + } = createCacheLayers({ + memoryMaxEnvVar: "BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX", + diskCacheDirEnvVar: "BRAINTRUST_PARAMETERS_CACHE_DIR", + diskMaxEnvVar: "BRAINTRUST_PARAMETERS_CACHE_DISK_MAX", + getDefaultDiskCacheDir: () => + `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/parameters_cache`, }); - const parametersDiskCache = canUseDiskCache() - ? new DiskCache({ - cacheDir: - iso.getEnv("BRAINTRUST_PARAMETERS_CACHE_DIR") ?? - `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/parameters_cache`, - max: - Number(iso.getEnv("BRAINTRUST_PARAMETERS_CACHE_DISK_MAX")) ?? - 1 << 20, - }) - : undefined; this.parametersCache = new ParametersCache({ memoryCache: parametersMemoryCache, diskCache: parametersDiskCache, diff --git a/js/src/prompt-cache/cache-config.ts b/js/src/prompt-cache/cache-config.ts new file mode 100644 index 000000000..4d4a34d9a --- /dev/null +++ b/js/src/prompt-cache/cache-config.ts @@ -0,0 +1,99 @@ +import { debugLogger } from "../debug-logger"; +import iso from "../isomorph"; +import { canUseDiskCache, DiskCache } from "./disk-cache"; +import { LRUCache } from "./lru-cache"; + +type CacheMode = "mixed" | "memory" | "disk" | "none"; + +const CACHE_LOCATION_ENV_VAR = "BRAINTRUST_CACHE_LOCATION"; +// Cache max values are entry counts, not byte sizes. +const DEFAULT_CACHE_MEMORY_MAX = 1 << 10; // 2^10 = 1024 entries. +const DEFAULT_CACHE_DISK_MAX = 1 << 20; // 2^20 = 1,048,576 entries. +let warnedInvalidCacheModeEnvValue = false; +let warnedUnavailableDiskCacheMode = false; + +function warnInvalidCacheMode(value: string) { + if (warnedInvalidCacheModeEnvValue) { + return; + } + warnedInvalidCacheModeEnvValue = true; + debugLogger.warn( + `Invalid ${CACHE_LOCATION_ENV_VAR} value "${value}". Expected "mixed", "memory", "disk", or "none". Falling back to "mixed".`, + ); +} + +function warnUnavailableDiskCache() { + if (warnedUnavailableDiskCacheMode) { + return; + } + warnedUnavailableDiskCacheMode = true; + debugLogger.warn( + `Disk cache is not supported on this platform, so ${CACHE_LOCATION_ENV_VAR}="disk" disables prompt and parameters caching.`, + ); +} + +function parseCacheMode(): CacheMode { + const value = iso.getEnv(CACHE_LOCATION_ENV_VAR); + const normalized = value?.trim().toLowerCase(); + if (!normalized) { + return "mixed"; + } + if ( + normalized === "mixed" || + normalized === "memory" || + normalized === "disk" || + normalized === "none" + ) { + return normalized; + } + warnInvalidCacheMode(value ?? ""); + return "mixed"; +} + +function parsePositiveIntegerEnv(envVar: string, defaultValue: number): number { + const value = Number(iso.getEnv(envVar)); + return Number.isInteger(value) && value > 0 ? value : defaultValue; +} + +export function createCacheLayers({ + memoryMaxEnvVar, + diskCacheDirEnvVar, + diskMaxEnvVar, + getDefaultDiskCacheDir, +}: { + memoryMaxEnvVar: string; + diskCacheDirEnvVar: string; + diskMaxEnvVar: string; + getDefaultDiskCacheDir: () => string; +}): { + memoryCache?: LRUCache; + diskCache?: DiskCache; +} { + const mode = parseCacheMode(); + const memoryCache = + mode === "mixed" || mode === "memory" + ? new LRUCache({ + max: parsePositiveIntegerEnv( + memoryMaxEnvVar, + DEFAULT_CACHE_MEMORY_MAX, + ), + }) + : undefined; + + let diskCache: DiskCache | undefined; + if (mode === "mixed" || mode === "disk") { + if (canUseDiskCache()) { + diskCache = new DiskCache({ + cacheDir: iso.getEnv(diskCacheDirEnvVar) ?? getDefaultDiskCacheDir(), + max: parsePositiveIntegerEnv(diskMaxEnvVar, DEFAULT_CACHE_DISK_MAX), + }); + } else if (mode === "disk") { + warnUnavailableDiskCache(); + } + } + + if (diskCache) { + return { memoryCache, diskCache }; + } + return { memoryCache }; +} diff --git a/js/src/prompt-cache/cache-mode.test.ts b/js/src/prompt-cache/cache-mode.test.ts new file mode 100644 index 000000000..0a585b0c2 --- /dev/null +++ b/js/src/prompt-cache/cache-mode.test.ts @@ -0,0 +1,268 @@ +import * as fs from "fs/promises"; +import * as path from "path"; +import { tmpdir } from "os"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import iso from "../isomorph"; +import { configureNode } from "../node/config"; +import { BraintrustState, Prompt, RemoteEvalParameters } from "../logger"; + +configureNode(); + +const CACHE_ENV_VARS = [ + "BRAINTRUST_CACHE_LOCATION", + "BRAINTRUST_DEBUG_LOG_LEVEL", + "BRAINTRUST_PROMPT_CACHE_DIR", + "BRAINTRUST_PROMPT_CACHE_MEMORY_MAX", + "BRAINTRUST_PROMPT_CACHE_DISK_MAX", + "BRAINTRUST_PARAMETERS_CACHE_DIR", + "BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX", + "BRAINTRUST_PARAMETERS_CACHE_DISK_MAX", +]; + +const promptKey = { + projectId: "11111111-1111-4111-8111-111111111111", + slug: "saved-prompt", + version: "v1", +}; + +const parametersKey = { + projectId: "22222222-2222-4222-8222-222222222222", + slug: "saved-parameters", + version: "v1", +}; + +const testPrompt = new Prompt( + { + id: "33333333-3333-4333-8333-333333333333", + _xact_id: "v1", + project_id: promptKey.projectId, + name: "Saved prompt", + slug: promptKey.slug, + }, + {}, + false, +); + +const testParameters = new RemoteEvalParameters({ + id: "44444444-4444-4444-8444-444444444444", + _xact_id: "v1", + project_id: parametersKey.projectId, + name: "Saved parameters", + slug: parametersKey.slug, + description: null, + function_type: "parameters", + function_data: { + type: "parameters", + data: { prefix: "hello" }, + __schema: { + type: "object", + properties: { + prefix: { type: "string" }, + }, + }, + }, +}); + +describe("prompt and parameters cache modes", () => { + const originalEnv = new Map(); + const originalGzip = iso.gzip; + const testDirs: string[] = []; + + beforeEach(() => { + for (const envVar of CACHE_ENV_VARS) { + originalEnv.set(envVar, process.env[envVar]); + delete process.env[envVar]; + } + iso.gzip = originalGzip; + vi.restoreAllMocks(); + }); + + afterEach(async () => { + for (const envVar of CACHE_ENV_VARS) { + const value = originalEnv.get(envVar); + if (value === undefined) { + delete process.env[envVar]; + } else { + process.env[envVar] = value; + } + } + originalEnv.clear(); + iso.gzip = originalGzip; + vi.restoreAllMocks(); + + await Promise.all( + testDirs.map((dir) => fs.rm(dir, { recursive: true, force: true })), + ); + testDirs.length = 0; + }); + + async function makeTempDir() { + const dir = await fs.mkdtemp(path.join(tmpdir(), "braintrust-cache-mode-")); + testDirs.push(dir); + return dir; + } + + async function configureCacheDirs() { + const promptDir = await makeTempDir(); + const parametersDir = await makeTempDir(); + process.env.BRAINTRUST_PROMPT_CACHE_DIR = promptDir; + process.env.BRAINTRUST_PARAMETERS_CACHE_DIR = parametersDir; + return { promptDir, parametersDir }; + } + + it("mixed mode uses memory and disk caches", async () => { + process.env.BRAINTRUST_CACHE_LOCATION = " MiXeD "; + process.env.BRAINTRUST_PROMPT_CACHE_MEMORY_MAX = "1"; + process.env.BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX = "1"; + await configureCacheDirs(); + + const state = new BraintrustState({}); + await state.promptCache.set(promptKey, testPrompt); + await state.promptCache.set({ ...promptKey, slug: "other" }, testPrompt); + await state.parametersCache.set(parametersKey, testParameters); + await state.parametersCache.set( + { ...parametersKey, slug: "other" }, + testParameters, + ); + + const nextState = new BraintrustState({}); + expect(await nextState.promptCache.get(promptKey)).toEqual(testPrompt); + expect(await nextState.parametersCache.get(parametersKey)).toEqual( + testParameters, + ); + }); + + it("memory mode does not write to disk", async () => { + process.env.BRAINTRUST_CACHE_LOCATION = "memory"; + const { promptDir, parametersDir } = await configureCacheDirs(); + + const state = new BraintrustState({}); + await state.promptCache.set(promptKey, testPrompt); + await state.parametersCache.set(parametersKey, testParameters); + + expect(await state.promptCache.get(promptKey)).toEqual(testPrompt); + expect(await state.parametersCache.get(parametersKey)).toEqual( + testParameters, + ); + expect(await fs.readdir(promptDir)).toEqual([]); + expect(await fs.readdir(parametersDir)).toEqual([]); + + const nextState = new BraintrustState({}); + expect(await nextState.promptCache.get(promptKey)).toBeUndefined(); + expect(await nextState.parametersCache.get(parametersKey)).toBeUndefined(); + }); + + it("disk mode does not warm memory", async () => { + process.env.BRAINTRUST_CACHE_LOCATION = "disk"; + const { promptDir, parametersDir } = await configureCacheDirs(); + + const state = new BraintrustState({}); + await state.promptCache.set(promptKey, testPrompt); + await state.parametersCache.set(parametersKey, testParameters); + + const nextState = new BraintrustState({}); + expect(await nextState.promptCache.get(promptKey)).toEqual(testPrompt); + expect(await nextState.parametersCache.get(parametersKey)).toEqual( + testParameters, + ); + + await fs.rm(promptDir, { recursive: true, force: true }); + await fs.rm(parametersDir, { recursive: true, force: true }); + + expect(await nextState.promptCache.get(promptKey)).toBeUndefined(); + expect(await nextState.parametersCache.get(parametersKey)).toBeUndefined(); + }); + + it("none mode disables all cache reads and writes", async () => { + process.env.BRAINTRUST_CACHE_LOCATION = "none"; + const { promptDir, parametersDir } = await configureCacheDirs(); + + const state = new BraintrustState({}); + await state.promptCache.set(promptKey, testPrompt); + await state.parametersCache.set(parametersKey, testParameters); + + expect(await state.promptCache.get(promptKey)).toBeUndefined(); + expect(await state.parametersCache.get(parametersKey)).toBeUndefined(); + expect(await fs.readdir(promptDir)).toEqual([]); + expect(await fs.readdir(parametersDir)).toEqual([]); + }); + + it("invalid cache mode warns once and falls back to mixed", async () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + process.env.BRAINTRUST_CACHE_LOCATION = "invalid"; + process.env.BRAINTRUST_DEBUG_LOG_LEVEL = "warn"; + await configureCacheDirs(); + + const state = new BraintrustState({}); + await state.promptCache.set(promptKey, testPrompt); + await state.parametersCache.set(parametersKey, testParameters); + + const nextState = new BraintrustState({}); + expect(await nextState.promptCache.get(promptKey)).toEqual(testPrompt); + expect(await nextState.parametersCache.get(parametersKey)).toEqual( + testParameters, + ); + + expect(warnSpy).toHaveBeenCalledTimes(1); + expect(warnSpy).toHaveBeenCalledWith( + "[braintrust]", + 'Invalid BRAINTRUST_CACHE_LOCATION value "invalid". Expected "mixed", "memory", "disk", or "none". Falling back to "mixed".', + ); + }); + + it("unset or invalid memory max defaults to 1024 entries", async () => { + process.env.BRAINTRUST_CACHE_LOCATION = "memory"; + process.env.BRAINTRUST_PROMPT_CACHE_MEMORY_MAX = "invalid"; + + const state = new BraintrustState({}); + for (let i = 0; i < 1025; i++) { + await state.promptCache.set( + { ...promptKey, slug: `prompt-${i}` }, + testPrompt, + ); + await state.parametersCache.set( + { ...parametersKey, slug: `parameters-${i}` }, + testParameters, + ); + } + + expect( + await state.promptCache.get({ ...promptKey, slug: "prompt-0" }), + ).toBeUndefined(); + expect( + await state.parametersCache.get({ + ...parametersKey, + slug: "parameters-0", + }), + ).toBeUndefined(); + expect( + await state.promptCache.get({ ...promptKey, slug: "prompt-1024" }), + ).toEqual(testPrompt); + expect( + await state.parametersCache.get({ + ...parametersKey, + slug: "parameters-1024", + }), + ).toEqual(testParameters); + }); + + it("disk mode disables caching and warns once when disk cache is unavailable", async () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + process.env.BRAINTRUST_CACHE_LOCATION = "disk"; + process.env.BRAINTRUST_DEBUG_LOG_LEVEL = "warn"; + iso.gzip = undefined; + + const state = new BraintrustState({}); + await state.promptCache.set(promptKey, testPrompt); + expect(await state.promptCache.get(promptKey)).toBeUndefined(); + await state.parametersCache.set(parametersKey, testParameters); + expect(await state.parametersCache.get(parametersKey)).toBeUndefined(); + + new BraintrustState({}); + expect(warnSpy).toHaveBeenCalledTimes(1); + expect(warnSpy).toHaveBeenCalledWith( + "[braintrust]", + 'Disk cache is not supported on this platform, so BRAINTRUST_CACHE_LOCATION="disk" disables prompt and parameters caching.', + ); + }); +}); diff --git a/js/src/prompt-cache/parameters-cache.ts b/js/src/prompt-cache/parameters-cache.ts index ff1b8d6dc..75fad987a 100644 --- a/js/src/prompt-cache/parameters-cache.ts +++ b/js/src/prompt-cache/parameters-cache.ts @@ -26,11 +26,11 @@ function createCacheKey(key: ParametersKey): string { } export class ParametersCache { - private readonly memoryCache: LRUCache; + private readonly memoryCache?: LRUCache; private readonly diskCache?: DiskCache; constructor(options: { - memoryCache: LRUCache; + memoryCache?: LRUCache; diskCache?: DiskCache; }) { this.memoryCache = options.memoryCache; @@ -40,9 +40,11 @@ export class ParametersCache { async get(key: ParametersKey): Promise { const cacheKey = createCacheKey(key); - const memoryParams = this.memoryCache.get(cacheKey); - if (memoryParams !== undefined) { - return memoryParams; + if (this.memoryCache) { + const memoryParams = this.memoryCache.get(cacheKey); + if (memoryParams !== undefined) { + return memoryParams; + } } if (this.diskCache) { @@ -50,7 +52,7 @@ export class ParametersCache { if (!diskParams) { return undefined; } - this.memoryCache.set(cacheKey, diskParams); + this.memoryCache?.set(cacheKey, diskParams); return diskParams; } @@ -60,7 +62,7 @@ export class ParametersCache { async set(key: ParametersKey, value: RemoteEvalParameters): Promise { const cacheKey = createCacheKey(key); - this.memoryCache.set(cacheKey, value); + this.memoryCache?.set(cacheKey, value); if (this.diskCache) { await this.diskCache.set(cacheKey, value); diff --git a/js/src/prompt-cache/prompt-cache.ts b/js/src/prompt-cache/prompt-cache.ts index 433204f38..e106a32cf 100644 --- a/js/src/prompt-cache/prompt-cache.ts +++ b/js/src/prompt-cache/prompt-cache.ts @@ -57,18 +57,16 @@ function createCacheKey(key: PromptKey): string { } /** - * A two-layer cache for Braintrust prompts with both in-memory and filesystem storage. + * A configurable cache for Braintrust prompts with optional in-memory and filesystem storage. * - * This cache implements either a one or two-layer caching strategy: - * 1. A fast in-memory LRU cache for frequently accessed prompts. - * 2. An optional persistent filesystem-based cache that serves as a backing store. + * This cache can use either layer independently, both layers together, or no layers. */ export class PromptCache { - private readonly memoryCache: LRUCache; + private readonly memoryCache?: LRUCache; private readonly diskCache?: DiskCache; constructor(options: { - memoryCache: LRUCache; + memoryCache?: LRUCache; diskCache?: DiskCache; }) { this.memoryCache = options.memoryCache; @@ -83,9 +81,11 @@ export class PromptCache { const cacheKey = createCacheKey(key); // First check memory cache. - const memoryPrompt = this.memoryCache.get(cacheKey); - if (memoryPrompt !== undefined) { - return memoryPrompt; + if (this.memoryCache) { + const memoryPrompt = this.memoryCache.get(cacheKey); + if (memoryPrompt !== undefined) { + return memoryPrompt; + } } // If not in memory and disk cache exists, check disk cache. @@ -94,8 +94,8 @@ export class PromptCache { if (!diskPrompt) { return undefined; } - // Store in memory cache. - this.memoryCache.set(cacheKey, diskPrompt); + // Store in memory cache if available. + this.memoryCache?.set(cacheKey, diskPrompt); return diskPrompt; } @@ -113,8 +113,8 @@ export class PromptCache { async set(key: PromptKey, value: Prompt): Promise { const cacheKey = createCacheKey(key); - // Update memory cache. - this.memoryCache.set(cacheKey, value); + // Update memory cache if available. + this.memoryCache?.set(cacheKey, value); // Update disk cache if available. if (this.diskCache) { From 1fe8937ce04eb276583a5a92350075ac5ece3d9a Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Thu, 28 May 2026 12:25:51 +0200 Subject: [PATCH 2/2] cs --- .changeset/itchy-falcons-flow.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/itchy-falcons-flow.md diff --git a/.changeset/itchy-falcons-flow.md b/.changeset/itchy-falcons-flow.md new file mode 100644 index 000000000..172a03882 --- /dev/null +++ b/.changeset/itchy-falcons-flow.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Add `BRAINTRUST_CACHE_LOCATION` env var to control caching location