From 5696bd60b0e5378276bce5bcceac6262d1b9df85 Mon Sep 17 00:00:00 2001 From: Devyash Saini Date: Sat, 24 Jan 2026 13:43:46 +0530 Subject: [PATCH 1/2] refactor(cache): modify api key cache and implement tag cache to use a generic cache store Signed-off-by: Devyash Saini --- src/__tests__/unit/utils/apiKeyCache.test.ts | 2 +- src/utils/apiKeyCache.ts | 163 +++---------------- src/utils/cacheStore.ts | 121 ++++++++++++++ src/utils/tagCache.ts | 6 + src/zod/event.ts | 83 ++++++---- 5 files changed, 198 insertions(+), 177 deletions(-) create mode 100644 src/utils/cacheStore.ts create mode 100644 src/utils/tagCache.ts diff --git a/src/__tests__/unit/utils/apiKeyCache.test.ts b/src/__tests__/unit/utils/apiKeyCache.test.ts index ae9724d..5dd5a11 100644 --- a/src/__tests__/unit/utils/apiKeyCache.test.ts +++ b/src/__tests__/unit/utils/apiKeyCache.test.ts @@ -47,7 +47,7 @@ describe("apiKeyCache", () => { expiresAt, }); - // Advance time by > 5 minutes TTL (cacheTTLMinutes = 5) + // Advance time by > 5 minutes TTL const later = base + 6 * 60 * 1000; (Date as any).now = () => later; diff --git a/src/utils/apiKeyCache.ts b/src/utils/apiKeyCache.ts index 4cce98b..b3ecfc0 100644 --- a/src/utils/apiKeyCache.ts +++ b/src/utils/apiKeyCache.ts @@ -1,152 +1,27 @@ -// USING LOCAL CACHING JUST FOR NOW -import { logger } from "../errors/logger"; +import { Cache } from "./cacheStore"; interface CachedAPIKey { id: string; expiresAt: string; - cachedAt: number; - lastAccessed: number; } -class APIKeyCache { - private cache: Map; // Map - private cacheTTL: number; // in milliseconds - private maxSize: number; - - constructor(cacheTTLMinutes: number = 5, maxSize: number = 1000) { - this.cache = new Map(); - this.cacheTTL = cacheTTLMinutes * 60 * 1000; - this.maxSize = maxSize; - - // Run cleanup every minute to remove expired cache entries - setInterval(() => this.cleanup(), 60 * 1000); - } - - /** - * Get API key data from cache if it exists and is not expired - * @param hash - The HMAC-SHA256 hash of the API key - */ - get(hash: string): CachedAPIKey | null { - const cached = this.cache.get(hash); - - if (!cached) { - return null; - } - - const now = Date.now(); - - // Check if cache entry has expired - if (now - cached.cachedAt > this.cacheTTL) { - this.cache.delete(hash); - return null; - } - - // Check if API key itself has expired - const keyExpiresAt = new Date(cached.expiresAt).getTime(); - if (now > keyExpiresAt) { - this.cache.delete(hash); - return null; - } - - // Update last accessed time for LRU - cached.lastAccessed = now; - - return cached; - } - - /** - * Set API key data in cache - * @param hash - The HMAC-SHA256 hash of the API key - */ - set(hash: string, data: { id: string; expiresAt: string }): void { - // Check if cache is full and evict LRU entry if needed - if (this.cache.size >= this.maxSize && !this.cache.has(hash)) { - this.evictLRU(); - } - - const now = Date.now(); - this.cache.set(hash, { - id: data.id, - expiresAt: data.expiresAt, - cachedAt: now, - lastAccessed: now, - }); - } - - /** - * Remove an API key from cache - * @param hash - The HMAC-SHA256 hash of the API key - */ - delete(hash: string): void { - this.cache.delete(hash); - } - - /** - * Clear all cache entries - */ - clear(): void { - this.cache.clear(); - } - - /** - * Evict least recently used entry from cache - */ - private evictLRU(): void { - let oldestKey: string | null = null; - let oldestTime = Infinity; - - for (const [key, value] of this.cache.entries()) { - if (value.lastAccessed < oldestTime) { - oldestTime = value.lastAccessed; - oldestKey = key; - } - } - - if (oldestKey) { - this.cache.delete(oldestKey); - logger.logDebug( - `Evicted LRU cache entry (cache full at ${this.maxSize})`, - {}, - ); - } - } - - /** - * Clean up expired cache entries - */ - private cleanup(): void { - const now = Date.now(); - const keysToDelete: string[] = []; - - for (const [key, value] of this.cache.entries()) { - // Remove if cache TTL expired or API key expired - const keyExpiresAt = new Date(value.expiresAt).getTime(); - if (now - value.cachedAt > this.cacheTTL || now > keyExpiresAt) { - keysToDelete.push(key); - } - } - - keysToDelete.forEach((key) => this.cache.delete(key)); - - if (keysToDelete.length > 0) { - logger.logDebug( - `Cleaned up ${keysToDelete.length} expired cache entries`, - {}, - ); - } - } - - /** - * Get cache statistics - */ - getStats(): { size: number; maxSize: number; ttlMinutes: number } { +const store = Cache.getStore("api-keys", { + max: 1000, + ttlMs: 5 * 60 * 1000, + validate: (value) => Date.now() <= new Date(value.expiresAt).getTime(), +}); + +export const apiKeyCache = { + get: (hash: string) => store.get(hash) ?? null, + set: (hash: string, data: CachedAPIKey) => store.set(hash, data), + delete: (hash: string) => store.delete(hash), + clear: () => store.clear(), + getStats: () => { // for testing and debugging purposes + const stats = store.getStats(); return { - size: this.cache.size, - maxSize: this.maxSize, - ttlMinutes: this.cacheTTL / (60 * 1000), + size: stats.size, + maxSize: stats.max, + ttlMinutes: stats.ttlMs / (60 * 1000), }; - } -} - -// Export singleton instance -export const apiKeyCache = new APIKeyCache(5, 1000); // 5 minutes TTL, max 1000 entries + }, +}; diff --git a/src/utils/cacheStore.ts b/src/utils/cacheStore.ts new file mode 100644 index 0000000..46f0533 --- /dev/null +++ b/src/utils/cacheStore.ts @@ -0,0 +1,121 @@ +export interface CacheConfig { + max: number; + ttlMs: number; + validate?: (value: V) => boolean; +} + +export const DEFAULT_CACHE_CONFIG: CacheConfig = { + max: 500, + ttlMs: 10 * 60 * 1000, +}; + +export class CacheStore { + private readonly max: number; + private readonly ttlMs: number; + private readonly validate?: (value: V) => boolean; + private readonly store = new Map(); + + constructor(config: CacheConfig) { + this.max = config.max; + this.ttlMs = config.ttlMs; + this.validate = config.validate; + } + + get(key: K): V | undefined { + const entry = this.store.get(key); + if (!entry) return undefined; + + if (this.isExpired(entry.expiresAt) || !this.isValid(entry.value)) { + this.store.delete(key); + return undefined; + } + + this.refreshKey(key, entry); + return entry.value; + } + + has(key: K): boolean { + const entry = this.store.get(key); + if (!entry) return false; + + if (this.isExpired(entry.expiresAt) || !this.isValid(entry.value)) { + this.store.delete(key); + return false; + } + + return true; + } + + set(key: K, value: V): void { + const expiresAt = Date.now() + this.ttlMs; + + if (this.store.has(key)) { + this.store.delete(key); + } else if (this.store.size >= this.max) { + this.evictOldest(); + } + + this.store.set(key, { value, expiresAt }); + } + + delete(key: K): boolean { + return this.store.delete(key); + } + + clear(): void { + this.store.clear(); + } + + size(): number { + return this.store.size; + } + + getStats(): { size: number; max: number; ttlMs: number } { + return { + size: this.store.size, + max: this.max, + ttlMs: this.ttlMs, + }; + } + + private isExpired(expiresAt: number): boolean { + return Date.now() > expiresAt; + } + + private isValid(value: V): boolean { + return this.validate ? this.validate(value) : true; + } + + private refreshKey(key: K, entry: { value: V; expiresAt: number }) { + this.store.delete(key); + this.store.set(key, entry); + } + + private evictOldest() { + const oldestKey = this.store.keys().next().value; + if (oldestKey !== undefined) { + this.store.delete(oldestKey); + } + } +} + +export class Cache { + private static stores = new Map>(); + + static getStore( + name: string, + config?: Partial>, + ): CacheStore { + const existing = Cache.stores.get(name); + if (existing) return existing as CacheStore; + + const merged: CacheConfig = { + ...(DEFAULT_CACHE_CONFIG as CacheConfig), + ...config, + }; + + const store = new CacheStore(merged); + Cache.stores.set(name, store as CacheStore); + return store; + } +} diff --git a/src/utils/tagCache.ts b/src/utils/tagCache.ts new file mode 100644 index 0000000..004a898 --- /dev/null +++ b/src/utils/tagCache.ts @@ -0,0 +1,6 @@ +import { Cache } from "./cacheStore"; + +export const tagCache = Cache.getStore("tags", { + max: 500, + ttlMs: 10 * 60 * 1000, +}); diff --git a/src/zod/event.ts b/src/zod/event.ts index 88f56df..5da3b44 100644 --- a/src/zod/event.ts +++ b/src/zod/event.ts @@ -4,6 +4,7 @@ import { getPostgresDB } from "../storage/db/postgres/db"; import { tagsTable } from "../storage/db/postgres/schema"; import { eq } from "drizzle-orm"; import { EventError } from "../errors/event"; +import { tagCache } from "../utils/tagCache"; const BaseEvent = z.object({ type: z.number(), // overwritten later by discriminators @@ -39,6 +40,11 @@ const SDKCallEvent = BaseEvent.extend({ .transform(async (v) => { // If a tag is provided, fetch the integer value for the tag and store it into debitAmount if (v.debit.case === "tag") { + const cachedAmount = tagCache.get(v.debit.value); + if (cachedAmount !== undefined) { + return { sdkCallType: v.sdkCallType, debitAmount: cachedAmount }; + } + const db = getPostgresDB(); try { const [tagRow] = await db @@ -53,6 +59,7 @@ const SDKCallEvent = BaseEvent.extend({ ); } + tagCache.set(v.debit.value, tagRow.amount); return { sdkCallType: v.sdkCallType, debitAmount: tagRow.amount }; } catch (e) { if (e instanceof EventError) { @@ -110,25 +117,31 @@ const AITokenUsageEvent = BaseEvent.extend({ // Process input debit let inputDebitAmount: number; if (v.inputDebit.case === "inputTag") { - try { - const [tagRow] = await db - .select() - .from(tagsTable) - .where(eq(tagsTable.tag, v.inputDebit.value)) - .limit(1); + const cachedAmount = tagCache.get(v.inputDebit.value); + if (cachedAmount !== undefined) { + inputDebitAmount = cachedAmount; + } else { + try { + const [tagRow] = await db + .select() + .from(tagsTable) + .where(eq(tagsTable.tag, v.inputDebit.value)) + .limit(1); - if (!tagRow) { - throw EventError.validationFailed( - `Input tag not found: ${v.inputDebit.value}`, - ); - } + if (!tagRow) { + throw EventError.validationFailed( + `Input tag not found: ${v.inputDebit.value}`, + ); + } - inputDebitAmount = tagRow.amount; - } catch (e) { - if (e instanceof EventError) { - throw e; + tagCache.set(v.inputDebit.value, tagRow.amount); + inputDebitAmount = tagRow.amount; + } catch (e) { + if (e instanceof EventError) { + throw e; + } + throw EventError.unknown(e as Error); } - throw EventError.unknown(e as Error); } } else { inputDebitAmount = Math.floor(v.inputDebit.value * 100); @@ -137,25 +150,31 @@ const AITokenUsageEvent = BaseEvent.extend({ // Process output debit let outputDebitAmount: number; if (v.outputDebit.case === "outputTag") { - try { - const [tagRow] = await db - .select() - .from(tagsTable) - .where(eq(tagsTable.tag, v.outputDebit.value)) - .limit(1); + const cachedAmount = tagCache.get(v.outputDebit.value); + if (cachedAmount !== undefined) { + outputDebitAmount = cachedAmount; + } else { + try { + const [tagRow] = await db + .select() + .from(tagsTable) + .where(eq(tagsTable.tag, v.outputDebit.value)) + .limit(1); - if (!tagRow) { - throw EventError.validationFailed( - `Output tag not found: ${v.outputDebit.value}`, - ); - } + if (!tagRow) { + throw EventError.validationFailed( + `Output tag not found: ${v.outputDebit.value}`, + ); + } - outputDebitAmount = tagRow.amount; - } catch (e) { - if (e instanceof EventError) { - throw e; + tagCache.set(v.outputDebit.value, tagRow.amount); + outputDebitAmount = tagRow.amount; + } catch (e) { + if (e instanceof EventError) { + throw e; + } + throw EventError.unknown(e as Error); } - throw EventError.unknown(e as Error); } } else { outputDebitAmount = Math.floor(v.outputDebit.value * 100); From 0cf04daf1f17ce186d1c0159cf296a69411ec164 Mon Sep 17 00:00:00 2001 From: Devyash Saini Date: Sat, 24 Jan 2026 13:46:59 +0530 Subject: [PATCH 2/2] refactor: formatting issues Signed-off-by: Devyash Saini --- .github/workflows/commitlint.yml | 14 +++++++------- src/utils/apiKeyCache.ts | 3 ++- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/.github/workflows/commitlint.yml b/.github/workflows/commitlint.yml index f87ba1c..a91b30f 100644 --- a/.github/workflows/commitlint.yml +++ b/.github/workflows/commitlint.yml @@ -1,7 +1,7 @@ -name: commitlint - -on: pull_request - -jobs: - commitlint: - uses: ScrawnDotDev/.github/.github/workflows/commitlint.yml@master +name: commitlint + +on: pull_request + +jobs: + commitlint: + uses: ScrawnDotDev/.github/.github/workflows/commitlint.yml@master diff --git a/src/utils/apiKeyCache.ts b/src/utils/apiKeyCache.ts index b3ecfc0..ba6ce91 100644 --- a/src/utils/apiKeyCache.ts +++ b/src/utils/apiKeyCache.ts @@ -16,7 +16,8 @@ export const apiKeyCache = { set: (hash: string, data: CachedAPIKey) => store.set(hash, data), delete: (hash: string) => store.delete(hash), clear: () => store.clear(), - getStats: () => { // for testing and debugging purposes + getStats: () => { + // for testing and debugging purposes const stats = store.getStats(); return { size: stats.size,