diff --git a/src/lib/ai-proxy.ts b/src/lib/ai-proxy.ts index e20f857..27cb577 100644 --- a/src/lib/ai-proxy.ts +++ b/src/lib/ai-proxy.ts @@ -15,6 +15,7 @@ import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, trackReader, releaseReader, wrapStreamWithCleanup, type FetchWithRetryResult } from "./fetch-utils"; import { getJwt, invalidateJwt } from "./mimo-auth"; import { parseDSML, looksLikeDSML } from "./dsml-parser"; +import { getResponseCache, isDSMLDetectionEnabled, isStreamPassthroughEnabled, ResponseCache } from "./response-cache"; // --- Kimchi API Key (hardcoded untuk deployability di Vercel/CF) -------------- const KIMCHI_API_KEY = "castai_v1_c5b9f4751ccb6c187c7e2f1cb2efbf83ac5d4e22806fd1d7218a0f602fee1777_1d96b89c"; @@ -428,6 +429,24 @@ export async function handleChatCompletion( const wantsStream = req.stream === true; const { url, init } = buildBackendRequest(req, config); + // -- Cache check (non-streaming only) ------------------------------------ + const cache = getResponseCache(); + const cacheKey = !wantsStream && cache + ? ResponseCache.buildKey(req.model, req.messages, wantsStream) + : null; + if (cacheKey && cache) { + const cached = cache.get(cacheKey); + if (cached) { + if (isDevMode()) { + console.log(`[ai-proxy] cache HIT for model=${req.model} key=${cacheKey.slice(0, 12)}`); + } + return new Response(cached.body, { + status: cached.status, + headers: { "Content-Type": "application/json", "Access-Control-Allow-Origin": "*", ...cached.headers }, + }); + } + } + // -- Mimo Free: inject JWT authentication and session affinity -------------- if (config.provider === "mimo-free") { const jwt = await getJwt(); @@ -502,8 +521,11 @@ export async function handleChatCompletion( const contentType = response.headers.get("content-type") ?? ""; const isNativeStream = contentType.includes("text/event-stream"); - if (isNativeStream && (config.provider === "opencode" || config.provider === "mimo-free" || config.provider === "castai")) { - // Passthrough for OpenAI-compatible SSE + if (isNativeStream && isStreamPassthroughEnabled()) { + // Generic passthrough for all native SSE backends (no transform overhead). + // Previously this was restricted to known provider names; now it trusts + // the upstream content-type header, which is more generic and catches + // any new backend that speaks SSE natively. const headers: Record = { "Content-Type": "text/event-stream", "Cache-Control": "no-cache", @@ -517,7 +539,7 @@ export async function handleChatCompletion( ); } - // Transform the stream + // Transform the stream (for non-SSE or passthrough-disabled backends) const transformed = transformStream( response.body!, config, @@ -545,6 +567,15 @@ export async function handleChatCompletion( } const adapted = parseJSONResponse(text, config, req); + // -- Store in cache -------------------------------------------------------- + if (cacheKey && cache) { + const responseBody = JSON.stringify(adapted); + cache.set(cacheKey, responseBody, 200, {}); + if (isDevMode()) { + console.log(`[ai-proxy] cache MISS for model=${req.model} key=${cacheKey.slice(0, 12)} — stored`); + } + } + return new Response(JSON.stringify(adapted), { status: 200, headers: { @@ -586,7 +617,7 @@ function transformStream( // DSML accumulation: buffer text deltas to detect DSML across chunks let dsmlAccumulated = ""; - let dsmlDetecting = true; + let dsmlDetecting = isDSMLDetectionEnabled(); function startKeepalive(controller: ReadableStreamDefaultController) { if (keepaliveTimer) return; diff --git a/src/lib/anthropic-proxy.ts b/src/lib/anthropic-proxy.ts index 95cdb13..dd87451 100644 --- a/src/lib/anthropic-proxy.ts +++ b/src/lib/anthropic-proxy.ts @@ -14,6 +14,7 @@ import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; import { MODEL_ROUTES, type BackendConfig } from "./ai-proxy"; import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, trackReader, releaseReader, wrapStreamWithCleanup, type FetchWithRetryResult } from "./fetch-utils"; import { parseDSML, looksLikeDSML, isCompleteDSML } from "./dsml-parser"; +import { getResponseCache, isDSMLDetectionEnabled, ResponseCache } from "./response-cache"; // --- Types ------------------------------------------------------------------- @@ -692,7 +693,7 @@ async function processStreamChunk( config: BackendConfig, usage: AnthropicResponse["usage"], outputCounter: OutputCounter, - dsmlBuffer?: DSMLStreamBuffer, + dsmlBuffer?: DSMLStreamBuffer | null, ): Promise { const { done, value } = await reader.read(); if (done) { @@ -790,7 +791,7 @@ function transformAnthropicStream( let phase: "init" | "block" | "done" = "init"; const outputCounter: OutputCounter = { chars: 0 }; const usage: AnthropicResponse["usage"] = { input_tokens: 0, output_tokens: 0 }; - const dsmlBuffer = createDSMLStreamBuffer(); + const dsmlBuffer = isDSMLDetectionEnabled() ? createDSMLStreamBuffer() : null; let keepaliveTimer: ReturnType | null = null; const KEEPALIVE_INTERVAL_MS = 15_000; @@ -836,7 +837,7 @@ function transformAnthropicStream( } if (phase === "done") { - const dsmlText = dsmlBuffer.flush(); + const dsmlText = dsmlBuffer?.flush() ?? null; emitDoneEvents(controller, encoder, usage, outputCounter, dsmlText); } } catch (err) { @@ -1059,6 +1060,24 @@ export async function handleAnthropicMessages( // Default anthropic-version to 2023-06-01 (required for prompt caching). const version = anthropicVersion || "2023-06-01"; + // -- Cache check (non-streaming only) ------------------------------------ + const cache = getResponseCache(); + const cacheKey = !wantsStream && cache + ? ResponseCache.buildKey(req.model, req.messages, wantsStream) + : null; + if (cacheKey && cache) { + const cached = cache.get(cacheKey); + if (cached) { + if (isDevMode()) { + console.log(`[anthropic-proxy] cache HIT for model=${req.model} key=${cacheKey.slice(0, 12)}`); + } + return new Response(cached.body, { + status: cached.status, + headers: { "Content-Type": "application/json", "Access-Control-Allow-Origin": "*", ...cached.headers }, + }); + } + } + // Translate Anthropic -> backend const { body: backendBody, headers: extraHeaders } = anthropicToBackend( req, config, backendModel, version, @@ -1148,6 +1167,16 @@ export async function handleAnthropicMessages( } const adapted = backendToAnthropicResponse(parsed, req.model); + + // -- Store in cache -------------------------------------------------------- + if (cacheKey && cache) { + const responseBody = JSON.stringify(adapted); + cache.set(cacheKey, responseBody, 200, {}); + if (isDevMode()) { + console.log(`[anthropic-proxy] cache MISS for model=${req.model} key=${cacheKey.slice(0, 12)} — stored`); + } + } + return new Response(JSON.stringify(adapted), { status: 200, headers: buildJsonHeaders(response), diff --git a/src/lib/fetch-utils.ts b/src/lib/fetch-utils.ts index 7aeb666..effaaee 100644 --- a/src/lib/fetch-utils.ts +++ b/src/lib/fetch-utils.ts @@ -238,9 +238,10 @@ export async function fetchWithRetry( let usedProxy = false; // Strategy: direct first, then proxies as fallback. - // maxAttempts = 1 direct + pool.size proxies, min 3 (all direct if no pool). + // Default: 2 attempts (1 direct + 1 proxy fallback). Override via MAX_RETRIES env. + const MAX_RETRIES = Number(process.env.MAX_RETRIES ?? 2); const poolSize = proxyPool?.size ?? 0; - const maxAttempts = poolSize > 0 ? poolSize + 1 : 3; + const maxAttempts = Math.min(MAX_RETRIES, poolSize > 0 ? poolSize + 1 : MAX_RETRIES); for (let attempt = 0; attempt < maxAttempts; attempt++) { // Exponential backoff between attempts (skip on first attempt) @@ -407,8 +408,9 @@ export async function fetchWithSessionRetry( } // Strategy: direct first, then session-sticky proxies as fallback. - // totalAttempts = 1 direct + sessionPool.size proxies, min 3. - const totalAttempts = maxRetries ?? Math.max(3, sessionPool.size + 1); + // Default: 2 attempts (1 direct + 1 proxy fallback). Override via MAX_RETRIES env. + const MAX_RETRIES = Number(process.env.MAX_RETRIES ?? 2); + const totalAttempts = maxRetries ?? Math.min(MAX_RETRIES, Math.max(2, sessionPool.size + 1)); let lastError: unknown; let lastResponse: Response | undefined; for (let attempt = 0; attempt < totalAttempts; attempt++) { diff --git a/src/lib/response-cache.ts b/src/lib/response-cache.ts new file mode 100644 index 0000000..e414533 --- /dev/null +++ b/src/lib/response-cache.ts @@ -0,0 +1,248 @@ +/** + * LRU Response Cache with TTL. + * + * Designed for LLM proxy responses — reduces bandwidth and upstream + * round-trips for repeated prompts (common during development, retries, + * or shared prefixes). + * + * Strategy: + * - LRU eviction (most recently used survives) + * - Configurable TTL per entry (default 15s — safe for dev, short enough + * that stale responses are unlikely) + * - Cache key = hash of (model + sorted messages + stream flag) + * - Only non-streaming responses are cached (streaming would require + * buffering the entire body, defeating the purpose) + * - Disabled entirely when CACHE_TTL=0 or NODE_ENV=production without + * explicit opt-in + * + * Thread safety: LRU operations happen on a single Map + Doubly Linked + * List, all synchronous — safe within Bun's single-threaded event loop. + */ + +// ─── Cache entry ────────────────────────────────────────────────────────── + +interface CacheEntry { + /** Serialized response body (JSON string). */ + body: string; + /** HTTP status code. */ + status: number; + /** Response headers to forward. */ + headers: Record; + /** When this entry was created (epoch ms). */ + createdAt: number; + /** When this entry expires (epoch ms). */ + expiresAt: number; +} + +// ─── LRU Linked List Node ───────────────────────────────────────────────── + +interface LRUNode { + key: string; + prev: LRUNode | null; + next: LRUNode | null; +} + +// ─── ResponseCache class ────────────────────────────────────────────────── + +export class ResponseCache { + private readonly map = new Map(); + private head: LRUNode | null = null; + private tail: LRUNode | null = null; + private readonly maxSize: number; + private readonly defaultTtlMs: number; + + constructor(opts?: { maxSize?: number; defaultTtlMs?: number }) { + this.maxSize = opts?.maxSize ?? 500; + this.defaultTtlMs = opts?.defaultTtlMs ?? 15_000; // 15 seconds + } + + // ── Public API ────────────────────────────────────────────────────────── + + /** Build a deterministic cache key from an LLM request. */ + static buildKey(model: string, messages: unknown, stream: boolean): string { + // Normalize messages to a stable string representation + const stable = JSON.stringify(messages, stableStringifyReplacer); + const raw = `${model}|${stream}|${stable}`; + return simpleHash(raw); + } + + /** Retrieve a cached response. Returns null if missing or expired. */ + get(key: string): { body: string; status: number; headers: Record } | null { + const entry = this.map.get(key); + if (!entry) return null; + + // Expired — evict and return null + if (Date.now() > entry.expiresAt) { + this.delete(key); + return null; + } + + // Move to front (most recently used) + this.moveToFront(key); + return { body: entry.body, status: entry.status, headers: entry.headers }; + } + + /** Store a response in the cache. */ + set( + key: string, + body: string, + status: number, + headers: Record, + ttlMs?: number, + ): void { + // Enforce max size before inserting + if (this.map.size >= this.maxSize) { + this.evictLRU(); + } + + const now = Date.now(); + const ttl = ttlMs ?? this.defaultTtlMs; + + const entry: CacheEntry = { + body, + status, + headers, + createdAt: now, + expiresAt: now + ttl, + }; + + this.map.set(key, entry); + this.moveToFront(key); + } + + /** Delete a specific key. */ + delete(key: string): void { + this.map.delete(key); + this.removeNode(key); + } + + /** Clear all entries. */ + clear(): void { + this.map.clear(); + this.head = null; + this.tail = null; + } + + /** Current number of entries. */ + get size(): number { + return this.map.size; + } + + /** Sweep expired entries. Call periodically if desired. */ + sweep(): number { + const now = Date.now(); + let removed = 0; + for (const [key, entry] of this.map) { + if (now > entry.expiresAt) { + this.delete(key); + removed++; + } + } + return removed; + } + + // ── LRU internals ─────────────────────────────────────────────────────── + + private moveToFront(key: string): void { + // Remove from current position + this.removeNode(key); + + // Add to front + const node: LRUNode = { key, prev: null, next: this.head }; + if (this.head) { + this.head.prev = node; + } + this.head = node; + if (!this.tail) { + this.tail = node; + } + } + + private removeNode(key: string): void { + // Find the node — linear scan, but bounded by cache size (500). + // For larger caches, maintain a separate Map. + let cur = this.head; + while (cur) { + if (cur.key === key) { + // Unlink + if (cur.prev) cur.prev.next = cur.next; + if (cur.next) cur.next.prev = cur.prev; + if (this.head === cur) this.head = cur.next; + if (this.tail === cur) this.tail = cur.prev; + return; + } + cur = cur.next; + } + } + + private evictLRU(): void { + // Tail is the least recently used + if (!this.tail) return; + const lruKey = this.tail.key; + this.map.delete(lruKey); + this.removeNode(lruKey); + } +} + +// ─── Helpers ────────────────────────────────────────────────────────────── + +/** + * JSON.stringify replacer that sorts object keys for stable hashing. + */ +function stableStringifyReplacer(_key: string, value: unknown): unknown { + if (value !== null && typeof value === "object" && !Array.isArray(value)) { + const sorted: Record = {}; + for (const k of Object.keys(value).sort()) { + sorted[k] = (value as Record)[k]; + } + return sorted; + } + return value; +} + +/** + * Simple, fast non-cryptographic hash (djb2 variant). + * Collisions are theoretically possible but extremely unlikely for + * cache-key use — a collision would serve a wrong cached response, + * which is bounded by TTL. + */ +function simpleHash(input: string): string { + let hash = 5381; + for (let i = 0; i < input.length; i++) { + hash = ((hash << 5) + hash + input.charCodeAt(i)) & 0xffffffff; + } + return hash.toString(36); +} + +// ─── Module-level singleton (lazy) ──────────────────────────────────────── + +let _instance: ResponseCache | null = null; + +/** + * Get or create the shared ResponseCache instance. + * Configured via environment variables: + * CACHE_TTL — TTL in ms (0 = disabled, default 15000) + * CACHE_MAX_SIZE — max entries (default 500) + */ +export function getResponseCache(): ResponseCache | null { + const ttl = Number(process.env.CACHE_TTL ?? 15000); + if (ttl <= 0) return null; // explicitly disabled + + if (!_instance) { + const maxSize = Number(process.env.CACHE_MAX_SIZE ?? 500); + _instance = new ResponseCache({ maxSize, defaultTtlMs: ttl }); + } + return _instance; +} + +/** Check if DSML detection should run (env toggle, default true). */ +export function isDSMLDetectionEnabled(): boolean { + const val = process.env.DSML_DETECTION ?? "true"; + return val === "true" || val === "1"; +} + +/** Check if stream passthrough mode is enabled (env toggle, default true). */ +export function isStreamPassthroughEnabled(): boolean { + const val = process.env.STREAM_PASSTHROUGH ?? "true"; + return val === "true" || val === "1"; +}