From a613caced3c43e80cf0c94b5d75758bc07b12f47 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 15 Apr 2026 09:41:42 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=8A=20fix:=20In-Memory=20Endpoint=20To?= =?UTF-8?q?ken=20Config=20Cache=20Isolation=20(#12673)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: endpoint token config not using shared cache in same process (initializing clients) * refactor: Update default max context tokens for agent initialization - Introduced a constant `DEFAULT_MAX_CONTEXT_TOKENS` set to 32000. - Updated the `initializeAgent` function to use this constant instead of hardcoded values for maximum context tokens, improving maintainability and clarity. * refactor: shared caching mechanism for token configuration - Introduced a memoized in-memory cache for Keyv instances to ensure shared access across the same namespace, improving cache efficiency. - Updated the `standardCache` function to utilize the new in-memory cache for the TOKEN_CONFIG namespace. - Refactored the `initializeCustom` function to use the `tokenConfigCache` for better cache management. - Removed redundant tokenCache parameter from `fetchModels` to streamline the function signature. * fix: match TOKEN_CONFIG TTL and add memoization tests Pass Time.THIRTY_MINUTES to tokenConfigCache() to match the TTL used by getLogStores.js, preventing load-order-dependent expiry behavior. Add 7 automated tests covering in-memory memoization: referential identity, cross-call-site data sharing, namespace isolation, first-caller TTL semantics, fallbackStore bypass, and tokenConfigCache parity with direct standardCache access. * fix: export DEFAULT_MAX_CONTEXT_TOKENS, address review nits - Export the constant so tests (and future consumers) reference it directly instead of hardcoding the numeric value. - Add independent TTL assertion for tokenConfigCache (R-1 nit). - Add tokenConfigCache mock to custom/initialize.spec.ts. --- .../src/agents/__tests__/initialize.test.ts | 6 +- packages/api/src/agents/initialize.ts | 6 +- ...tandardCache.in_memory_memoization.spec.ts | 113 ++++++++++++++++++ packages/api/src/cache/cacheFactory.ts | 31 ++++- .../src/endpoints/custom/initialize.spec.ts | 1 + .../api/src/endpoints/custom/initialize.ts | 5 +- packages/api/src/endpoints/models.ts | 5 +- 7 files changed, 152 insertions(+), 15 deletions(-) create mode 100644 packages/api/src/cache/__tests__/cacheFactory/standardCache.in_memory_memoization.spec.ts diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index fe1493628a..be7377954a 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -3,6 +3,7 @@ import { EModelEndpoint } from 'librechat-data-provider'; import type { Agent } from 'librechat-data-provider'; import type { ServerRequest, InitializeResultBase, EndpointTokenConfig } from '~/types'; import type { InitializeAgentDbMethods } from '../initialize'; +import { DEFAULT_MAX_CONTEXT_TOKENS } from '../initialize'; // Mock logger jest.mock('winston', () => ({ @@ -308,11 +309,8 @@ describe('initializeAgent — maxContextTokens', () => { db, ); - // 0 is not used as-is; the formula kicks in. - // optionalChainWithEmptyCheck(0, 200000, 18000) returns 0 (not null/undefined), - // then Number(0) || 18000 = 18000 (the fallback default). expect(result.maxContextTokens).not.toBe(0); - const expected = Math.round((18000 - maxOutputTokens) * 0.95); + const expected = Math.round((DEFAULT_MAX_CONTEXT_TOKENS - maxOutputTokens) * 0.95); expect(result.maxContextTokens).toBe(expected); }); diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 1ebd20bb01..5105b21305 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -68,6 +68,8 @@ export type InitializedAgent = Agent & { maxToolResultChars?: number; }; +export const DEFAULT_MAX_CONTEXT_TOKENS = 32000; + /** * Parameters for initializing an agent * Matches the CJS signature from api/server/services/Endpoints/agents/agent.js @@ -357,7 +359,7 @@ export async function initializeAgent( providerEndpointMap[overrideProvider as keyof typeof providerEndpointMap], options.endpointTokenConfig, ), - 18000, + DEFAULT_MAX_CONTEXT_TOKENS, ); if ( @@ -414,7 +416,7 @@ export async function initializeAgent( agent.additional_instructions = artifactsPromptResult ?? undefined; } - const agentMaxContextNum = Number(agentMaxContextTokens) || 18000; + const agentMaxContextNum = Number(agentMaxContextTokens) || DEFAULT_MAX_CONTEXT_TOKENS; const maxOutputTokensNum = Number(maxOutputTokens) || 0; const baseContextTokens = Math.max(0, agentMaxContextNum - maxOutputTokensNum); diff --git a/packages/api/src/cache/__tests__/cacheFactory/standardCache.in_memory_memoization.spec.ts b/packages/api/src/cache/__tests__/cacheFactory/standardCache.in_memory_memoization.spec.ts new file mode 100644 index 0000000000..16b016f9bd --- /dev/null +++ b/packages/api/src/cache/__tests__/cacheFactory/standardCache.in_memory_memoization.spec.ts @@ -0,0 +1,113 @@ +import { CacheKeys, Time } from 'librechat-data-provider'; + +jest.mock('@keyv/redis', () => ({ + default: jest.fn(), +})); + +jest.mock('../../redisClients', () => ({ + keyvRedisClient: null, + ioredisClient: null, +})); + +jest.mock('../../redisUtils', () => ({ + batchDeleteKeys: jest.fn(), + scanKeys: jest.fn(), +})); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { + error: jest.fn(), + warn: jest.fn(), + debug: jest.fn(), + }, +})); + +describe('standardCache - in-memory memoization', () => { + afterEach(() => { + jest.resetModules(); + }); + + async function loadFactory() { + jest.doMock('../../cacheConfig', () => ({ + cacheConfig: { + FORCED_IN_MEMORY_CACHE_NAMESPACES: [], + REDIS_KEY_PREFIX: '', + GLOBAL_PREFIX_SEPARATOR: '>>', + }, + })); + return import('../../cacheFactory'); + } + + it('returns the same instance for repeated calls with the same namespace', async () => { + const { standardCache } = await loadFactory(); + const a = standardCache('test-ns'); + const b = standardCache('test-ns'); + expect(a).toBe(b); + }); + + it('returns different instances for different namespaces', async () => { + const { standardCache } = await loadFactory(); + const a = standardCache('ns-one'); + const b = standardCache('ns-two'); + expect(a).not.toBe(b); + }); + + it('shares data across separate standardCache calls for the same namespace', async () => { + const { standardCache } = await loadFactory(); + const writer = standardCache(CacheKeys.TOKEN_CONFIG); + await writer.set('model-a', { context: 128000 }); + + const reader = standardCache(CacheKeys.TOKEN_CONFIG); + expect(await reader.get('model-a')).toEqual({ context: 128000 }); + }); + + it('does not leak data between different namespaces', async () => { + const { standardCache } = await loadFactory(); + const cacheA = standardCache('ns-a'); + const cacheB = standardCache('ns-b'); + + await cacheA.set('key', 'value-a'); + await cacheB.set('key', 'value-b'); + + expect(await cacheA.get('key')).toBe('value-a'); + expect(await cacheB.get('key')).toBe('value-b'); + }); + + it('first caller TTL wins for a given namespace', async () => { + const { standardCache } = await loadFactory(); + const first = standardCache('ttl-ns', 500); + const second = standardCache('ttl-ns', 99999); + expect(first).toBe(second); + + type KeyvWithOpts = typeof first & { opts: { ttl?: number } }; + expect((first as KeyvWithOpts).opts.ttl).toBe(500); + }); + + it('does not memoize when a custom fallbackStore is provided', async () => { + const { standardCache } = await loadFactory(); + const storeA = new Map(); + const storeB = new Map(); + const a = standardCache('fb-ns', undefined, storeA); + const b = standardCache('fb-ns', undefined, storeB); + expect(a).not.toBe(b); + }); + + it('tokenConfigCache shares data with standardCache(TOKEN_CONFIG)', async () => { + const { standardCache, tokenConfigCache } = await loadFactory(); + const direct = standardCache(CacheKeys.TOKEN_CONFIG, Time.THIRTY_MINUTES); + await direct.set('openrouter', { 'gpt-4': { context: 128000, prompt: 5, completion: 15 } }); + + const convenience = tokenConfigCache(); + expect(await convenience.get('openrouter')).toEqual({ + 'gpt-4': { context: 128000, prompt: 5, completion: 15 }, + }); + expect(convenience).toBe(direct); + }); + + it('tokenConfigCache creates the instance with THIRTY_MINUTES TTL', async () => { + const { tokenConfigCache } = await loadFactory(); + const cache = tokenConfigCache(); + type KeyvWithOpts = typeof cache & { opts: { ttl?: number } }; + expect((cache as KeyvWithOpts).opts.ttl).toBe(Time.THIRTY_MINUTES); + }); +}); diff --git a/packages/api/src/cache/cacheFactory.ts b/packages/api/src/cache/cacheFactory.ts index 2d7817c2ad..5c5db865db 100644 --- a/packages/api/src/cache/cacheFactory.ts +++ b/packages/api/src/cache/cacheFactory.ts @@ -9,18 +9,33 @@ const KeyvRedis = require('@keyv/redis').default as typeof import('@keyv/redis') import { Keyv } from 'keyv'; import createMemoryStore from 'memorystore'; import { RedisStore } from 'rate-limit-redis'; -import { Time } from 'librechat-data-provider'; import { logger } from '@librechat/data-schemas'; import session, { MemoryStore } from 'express-session'; +import { Time, CacheKeys } from 'librechat-data-provider'; import { RedisStore as ConnectRedis } from 'connect-redis'; import type { SendCommandFn } from 'rate-limit-redis'; import { keyvRedisClient, ioredisClient } from './redisClients'; +import { batchDeleteKeys, scanKeys } from './redisUtils'; import { cacheConfig } from './cacheConfig'; import { violationFile } from './keyvFiles'; -import { batchDeleteKeys, scanKeys } from './redisUtils'; + +/** + * Memoized in-memory Keyv instances keyed by namespace. + * Without Redis, each `new Keyv()` gets its own internal Map, so callers that + * write in one call-site and read in another would see an empty store. + * Memoizing ensures a single shared Map per namespace across the entire bundle. + * + * Only applies to the plain in-memory path (no Redis, no custom fallbackStore). + */ +const inMemoryCacheMap = new Map(); /** * Creates a cache instance using Redis or a fallback store. Suitable for general caching needs. + * + * **In-memory mode** (no Redis, no custom fallbackStore): instances are memoized by + * namespace so that every call-site shares the same underlying `Map`. The first + * caller's TTL wins for a given namespace. + * * @param namespace - The cache namespace. * @param ttl - Time to live for cache entries. * @param fallbackStore - Optional fallback store if Redis is not used. @@ -73,9 +88,19 @@ export const standardCache = (namespace: string, ttl?: number, fallbackStore?: o if (fallbackStore) { return new Keyv({ store: fallbackStore, namespace, ttl }); } - return new Keyv({ namespace, ttl }); + const existing = inMemoryCacheMap.get(namespace); + if (existing) { + return existing; + } + const cache = new Keyv({ namespace, ttl }); + inMemoryCacheMap.set(namespace, cache); + return cache; }; +/** Convenience accessor for the TOKEN_CONFIG cache namespace. */ +export const tokenConfigCache = (): Keyv => + standardCache(CacheKeys.TOKEN_CONFIG, Time.THIRTY_MINUTES); + /** * Creates a cache instance for storing violation data. * Uses a file-based fallback store if Redis is not enabled. diff --git a/packages/api/src/endpoints/custom/initialize.spec.ts b/packages/api/src/endpoints/custom/initialize.spec.ts index eddd7cb515..0accf3171a 100644 --- a/packages/api/src/endpoints/custom/initialize.spec.ts +++ b/packages/api/src/endpoints/custom/initialize.spec.ts @@ -20,6 +20,7 @@ jest.mock('~/endpoints/models', () => ({ jest.mock('~/cache', () => ({ standardCache: jest.fn(() => ({ get: jest.fn().mockResolvedValue(null) })), + tokenConfigCache: jest.fn(() => ({ get: jest.fn().mockResolvedValue(null) })), })); jest.mock('~/utils', () => ({ diff --git a/packages/api/src/endpoints/custom/initialize.ts b/packages/api/src/endpoints/custom/initialize.ts index ea0d2dbf5d..18719c4f02 100644 --- a/packages/api/src/endpoints/custom/initialize.ts +++ b/packages/api/src/endpoints/custom/initialize.ts @@ -1,5 +1,4 @@ import { - CacheKeys, ErrorTypes, envVarRegex, FetchTokenConfig, @@ -13,7 +12,7 @@ import { isUserProvided, checkUserKeyExpiry } from '~/utils'; import { getCustomEndpointConfig } from '~/app/config'; import { fetchModels } from '~/endpoints/models'; import { validateEndpointURL } from '~/auth'; -import { standardCache } from '~/cache'; +import { tokenConfigCache } from '~/cache'; const { PROXY } = process.env; @@ -136,7 +135,7 @@ export async function initializeCustom({ const userId = req.user?.id ?? ''; - const cache = standardCache(CacheKeys.TOKEN_CONFIG); + const cache = tokenConfigCache(); /** tokenConfig is an optional extended property on custom endpoints */ const hasTokenConfig = (endpointConfig as Record).tokenConfig != null; const tokenKey = diff --git a/packages/api/src/endpoints/models.ts b/packages/api/src/endpoints/models.ts index cc91c1a4fa..fd296e59ec 100644 --- a/packages/api/src/endpoints/models.ts +++ b/packages/api/src/endpoints/models.ts @@ -19,7 +19,7 @@ import { logAxiosError, inputSchema, } from '~/utils'; -import { standardCache } from '~/cache'; +import { standardCache, tokenConfigCache } from '~/cache'; export interface FetchModelsParams { /** User ID for API requests */ @@ -195,8 +195,7 @@ export async function fetchModels({ const validationResult = inputSchema.safeParse(input); if (validationResult.success && createTokenConfig) { const endpointTokenConfig = processModelData(input); - const cache = standardCache(CacheKeys.TOKEN_CONFIG); - await cache.set(tokenKey ?? name, endpointTokenConfig); + await tokenConfigCache().set(tokenKey ?? name, endpointTokenConfig); } models = input.data.map((item: { id: string }) => item.id); } catch (error) {