diff --git a/src/server/modules/proxy/proxy.service.ts b/src/server/modules/proxy/proxy.service.ts index cecc5a2..4e135ba 100644 --- a/src/server/modules/proxy/proxy.service.ts +++ b/src/server/modules/proxy/proxy.service.ts @@ -35,6 +35,7 @@ import { getMaxOutputTokens, getThinkingBudget } from '../../../lib/geminiNexus/ import { resolveRequestUserAgent } from './request-user-agent'; import { UpstreamRequestError } from './clients/upstream-error'; import { trafficLogger } from '../../../utils/traffic-logger'; +import { TtlCache } from '../../../utils/ttl-cache'; import { encode } from 'gpt-tokenizer'; export interface ProxyRequestContext { @@ -52,14 +53,17 @@ export interface ProxyConfigOverride { timeout?: number; } +const RESPONSE_CACHE_TTL_MS = 60_000; +const RESPONSE_CACHE_MAX_ENTRIES = 100; + @Injectable() export class ProxyService { private readonly logger = new Logger(ProxyService.name); - private readonly responseCache = new Map< - string, - { response: OpenAIChatResponse; expiry: number } - >(); - private readonly CACHE_TTL_MS = 60000; // 1 minute + /** Bounded: each key is a full request body, so an unlimited cache grows with every unique prompt. */ + private readonly responseCache = new TtlCache( + RESPONSE_CACHE_TTL_MS, + RESPONSE_CACHE_MAX_ENTRIES, + ); constructor( @Inject(TokenManagerService) private readonly tokenManager: TokenManagerService, @@ -2268,19 +2272,11 @@ export class ProxyService { } private getCachedResponse(key: string): OpenAIChatResponse | undefined { - const entry = this.responseCache.get(key); - if (!entry) { - return undefined; - } - if (Date.now() > entry.expiry) { - this.responseCache.delete(key); - return undefined; - } - return entry.response; + return this.responseCache.get(key); } private setCachedResponse(key: string, response: OpenAIChatResponse): void { - this.responseCache.set(key, { response, expiry: Date.now() + this.CACHE_TTL_MS }); + this.responseCache.set(key, response); } private extractAnthropicSessionKey(request: AnthropicChatRequest): string | undefined { diff --git a/src/tests/unit/ttl-cache.test.ts b/src/tests/unit/ttl-cache.test.ts new file mode 100644 index 0000000..52ba02c --- /dev/null +++ b/src/tests/unit/ttl-cache.test.ts @@ -0,0 +1,66 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { TtlCache } from '../../utils/ttl-cache'; + +describe('TtlCache', () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00Z')); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('returns stored values until the ttl elapses', () => { + const cache = new TtlCache(1000, 10); + cache.set('a', 'value'); + expect(cache.get('a')).toBe('value'); + + vi.advanceTimersByTime(999); + expect(cache.get('a')).toBe('value'); + + vi.advanceTimersByTime(1); + expect(cache.get('a')).toBeUndefined(); + expect(cache.size).toBe(0); + }); + + it('returns undefined for unknown keys', () => { + expect(new TtlCache(1000, 10).get('missing')).toBeUndefined(); + }); + + it('drops expired entries on write even if they are never read again', () => { + const cache = new TtlCache(1000, 10); + cache.set('old-1', 1); + cache.set('old-2', 2); + vi.advanceTimersByTime(1500); + + cache.set('fresh', 3); + expect(cache.size).toBe(1); + expect(cache.get('fresh')).toBe(3); + }); + + it('evicts the oldest entries when the limit is exceeded', () => { + const cache = new TtlCache(60_000, 2); + cache.set('a', 1); + cache.set('b', 2); + cache.set('c', 3); + + expect(cache.size).toBe(2); + expect(cache.get('a')).toBeUndefined(); + expect(cache.get('b')).toBe(2); + expect(cache.get('c')).toBe(3); + }); + + it('refreshes ttl and recency when a key is overwritten', () => { + const cache = new TtlCache(1000, 2); + cache.set('a', 1); + cache.set('b', 2); + vi.advanceTimersByTime(600); + cache.set('a', 10); + cache.set('c', 3); + + expect(cache.get('b')).toBeUndefined(); + vi.advanceTimersByTime(600); + expect(cache.get('a')).toBe(10); + }); +}); diff --git a/src/utils/ttl-cache.ts b/src/utils/ttl-cache.ts new file mode 100644 index 0000000..0a91a7d --- /dev/null +++ b/src/utils/ttl-cache.ts @@ -0,0 +1,60 @@ +/** + * Small in-memory cache with a fixed time-to-live and a hard entry limit. + * + * Entries live in insertion order and share one TTL, so expiry order matches insertion order: + * pruning can stop at the first live entry. The size cap evicts the oldest entries, which keeps + * memory bounded even when callers never read a key again (a plain Map only drops expired + * entries on read). + */ +export class TtlCache { + private readonly entries = new Map(); + + constructor( + private readonly ttlMs: number, + private readonly maxEntries: number, + ) {} + + get(key: string): V | undefined { + const entry = this.entries.get(key); + if (!entry) { + return undefined; + } + if (Date.now() >= entry.expiresAt) { + this.entries.delete(key); + return undefined; + } + return entry.value; + } + + set(key: string, value: V): void { + const now = Date.now(); + // Re-inserting moves the key to the newest position so insertion order stays expiry order. + this.entries.delete(key); + this.pruneExpired(now); + this.entries.set(key, { value, expiresAt: now + this.ttlMs }); + this.evictOverflow(); + } + + get size(): number { + return this.entries.size; + } + + private pruneExpired(now: number): void { + for (const [key, entry] of this.entries) { + if (entry.expiresAt > now) { + return; + } + this.entries.delete(key); + } + } + + private evictOverflow(): void { + while (this.entries.size > this.maxEntries) { + const oldestKey = this.entries.keys().next().value; + if (oldestKey === undefined) { + return; + } + this.entries.delete(oldestKey); + } + } +}