diff --git a/src/index.ts b/src/index.ts index ad2fe94..891a84d 100644 --- a/src/index.ts +++ b/src/index.ts @@ -21,7 +21,6 @@ import { filterRequestHeaders, buildRelayRequest, createRelayResponse, - classifyFetchError, createErrorResponse, createCorsPreflightResponse, getCorsHeaders, @@ -30,7 +29,7 @@ import { import { checkBodySize } from "./middleware/body-limiter"; import { createRateLimiter } from "./middleware/rate-limiter"; import { logRelayEvent } from "./middleware/logger"; -import { ProxyPool } from "./lib/proxy-pool"; +import { ProxyPool, SessionProxyPool } from "./lib/proxy-pool"; import { handleChatCompletion, listModels } from "./lib/ai-proxy"; import { handleAnthropicMessages } from "./lib/anthropic-proxy"; import { fetchWithRetry, closeAllActiveReaders, isDevMode } from "./lib/fetch-utils"; @@ -81,6 +80,12 @@ proxyPool.tryLoad( process.env.PROXY_FILE || process.env.PROXY_LIST || "./proxy.txt", ); +// Session-aware proxy pool wrapping the base pool. +// Sessions are per-request and short-lived — each streaming request gets +// a random session ID so the same proxy is reused for the entire stream. +const sessionPool = new SessionProxyPool(proxyPool); +sessionPool.setFailureThreshold(3); + // --- WebSocket relay data type ----------------------------------------------- interface WSRelayData { @@ -112,85 +117,85 @@ function handleHealth(): Response { /** Simple embedded HTML documentation page. */ function handleDocs(): Response { const html = ` - - - - - Edge Proxy Relay — Docs - - - -
-

Edge Proxy Relay

-

Forward HTTP and WebSocket requests to any target server via the x-relay-target header.

+ + + + + Edge Proxy Relay — Docs + + + +
+

Edge Proxy Relay

+

Forward HTTP and WebSocket requests to any target server via the x-relay-target header.

-

Endpoints

+

Endpoints

-
-

GET /health

-

Health check. Returns 200 OK with server status, uptime, and version.

-
+
+

GET /health

+

Health check. Returns 200 OK with server status, uptime, and version.

+
-
-

GET /docs

-

This page.

-
+
+

GET /docs

+

This page.

+
-
-

Any Path (Catch-all Relay)

-

Send a request with the x-relay-target header and this proxy forwards it.

-
+
+

Any Path (Catch-all Relay)

+

Send a request with the x-relay-target header and this proxy forwards it.

+
-

Usage — HTTP Relay

- 
curl -s \\
-  -H "x-relay-target: https://httpbin.org" \\
-  -H "x-relay-path: /get" \\
-  "https://your-proxy.example/any/path"
- - - - -
HeaderRequiredDescription
x-relay-targetYesBase URL of the upstream (http:// or https://)
x-relay-pathNoPath to append (default: /)
+

Usage — HTTP Relay

+ 
curl -s \\
+	  -H "x-relay-target: https://httpbin.org" \\
+	  -H "x-relay-path: /get" \\
+	  "https://your-proxy.example/any/path"
+ + + + +
HeaderRequiredDescription
x-relay-targetYesBase URL of the upstream (http:// or https://)
x-relay-pathNoPath to append (default: /)
-

Usage — WebSocket Relay

- 
const ws = new WebSocket("wss://your-proxy.example/relay", {
-  headers: { "x-relay-target": "wss://echo-websocket.example" },
-});
-ws.onopen = () => ws.send("Hello via relay!");
-ws.onmessage = (e) => console.log("Got:", e.data);
+

Usage — WebSocket Relay

+ 
const ws = new WebSocket("wss://your-proxy.example/relay", {
+	  headers: { "x-relay-target": "wss://echo-websocket.example" },
+	});
+	ws.onopen = () => ws.send("Hello via relay!");
+	ws.onmessage = (e) => console.log("Got:", e.data);
-

Status Codes

- - - - - - - - - -
CodeMeaning
204CORS preflight success (OPTIONS)
400Missing x-relay-target header
403Target blocked (SSRF protection / not allowed)
413Request body exceeds size limit
429Rate limit exceeded
502Upstream network / DNS error
504Upstream timeout
-
- -`; +

Status Codes

+ + + + + + + + + +
CodeMeaning
204CORS preflight success (OPTIONS)
400Missing x-relay-target header
403Target blocked (SSRF protection / not allowed)
413Request body exceeds size limit
429Rate limit exceeded
502Upstream network / DNS error
504Upstream timeout
+
+ + `; return new Response(html, { status: 200, @@ -204,29 +209,29 @@ ws.onmessage = (e) => console.log("Got:", e.data); /** Minimal status page shown at the root `/`. */ function handleIndex(): Response { const html = ` - - - - - Edge Proxy Relay - - - -
-

Edge Proxy Relay

-

Server is running

-

/health · /docs

-
- -`; + + + + + Edge Proxy Relay + + + +
+

Edge Proxy Relay

+

Server is running

+

/health · /docs

+
+ + `; return new Response(html, { status: 200, @@ -484,7 +489,8 @@ const server: Server = Bun.serve({ if (authErr) return authErr; try { const body = await req.json(); - return handleChatCompletion(body, proxyPool); + const sessionId = crypto.randomUUID(); + return handleChatCompletion(body, proxyPool, sessionPool, sessionId); } catch { return new Response( JSON.stringify({ error: { message: "Invalid JSON body", type: "invalid_request_error" } }), @@ -505,7 +511,8 @@ const server: Server = Bun.serve({ if (authErr) return authErr; try { const body = await req.json(); - return handleAnthropicMessages(body, proxyPool); + const sessionId = crypto.randomUUID(); + return handleAnthropicMessages(body, proxyPool, sessionPool, sessionId); } catch { return new Response( JSON.stringify({ diff --git a/src/lib/ai-proxy.ts b/src/lib/ai-proxy.ts index aa87a52..5cab053 100644 --- a/src/lib/ai-proxy.ts +++ b/src/lib/ai-proxy.ts @@ -12,10 +12,8 @@ * Streaming (SSE) is supported for all backends. */ -import type { ProxyPool } from "./proxy-pool"; -import { fetchWithRetry } from "./fetch-utils"; -import { SSELineBuffer } from "./fetch-utils"; -import { isDevMode } from "./fetch-utils"; +import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; +import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, type FetchWithRetryResult } from "./fetch-utils"; // --- Types ------------------------------------------------------------------- @@ -315,10 +313,30 @@ function openAIError(status: number, message: string, type: string): Response { /** * Handle an OpenAI-compatible chat completions request. + * + * Two calling conventions: + * 1. Standard: (body, proxyPool?) + * 2. Session-aware: (body, proxyPool?, sessionPool, sessionId) + * + * When both `sessionPool` and `sessionId` are present the request uses + * session-sticky proxy allocation via `fetchWithSessionRetry`; otherwise + * the existing `fetchWithRetry` path is used (backward-compatible). */ export async function handleChatCompletion( body: unknown, proxyPool?: ProxyPool, +): Promise; +export async function handleChatCompletion( + body: unknown, + proxyPool?: ProxyPool, + sessionPool?: SessionProxyPool, + sessionId?: string, +): Promise; +export async function handleChatCompletion( + body: unknown, + proxyPool?: ProxyPool, + sessionPool?: SessionProxyPool, + sessionId?: string, ): Promise { // -- Input validation ------------------------------------------------------- const validationError = validateChatRequest(body); @@ -340,13 +358,11 @@ export async function handleChatCompletion( const wantsStream = req.stream === true; const { url, init } = buildBackendRequest(req, config); - // -- Execute (direct -> proxy fallback) with shared retry ------------------- - const result = await fetchWithRetry( - url, - init, - proxyPool, - `openai:${req.model}`, - ); + // -- Execute with session-aware or standard retry -------------------------- + const result: FetchWithRetryResult = + sessionPool && sessionId + ? await fetchWithSessionRetry(url, init, sessionPool, sessionId, `openai:${req.model}`) + : await fetchWithRetry(url, init, proxyPool, `openai:${req.model}`); if (result.errorClassification) { return new Response( @@ -389,7 +405,10 @@ export async function handleChatCompletion( "Access-Control-Allow-Origin": "*", "X-Accel-Buffering": "no", }; - return new Response(response.body, { status: 200, headers }); + return new Response( + wrapStreamMaybe(response.body!, sessionPool, sessionId), + { status: 200, headers }, + ); } // Transform the stream @@ -398,20 +417,26 @@ export async function handleChatCompletion( config, req, ); - return new Response(transformed, { - status: 200, - headers: { - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - Connection: "keep-alive", - "Access-Control-Allow-Origin": "*", - "X-Accel-Buffering": "no", + return new Response( + wrapStreamMaybe(transformed, sessionPool, sessionId), + { + status: 200, + headers: { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + "Access-Control-Allow-Origin": "*", + "X-Accel-Buffering": "no", + }, }, - }); + ); } // -- Handle non-streaming response ------------------------------------------ const text = await response.text(); + if (sessionPool && sessionId) { + sessionPool.release(sessionId); + } const adapted = parseJSONResponse(text, config, req); return new Response(JSON.stringify(adapted), { @@ -505,3 +530,47 @@ function transformStream( }, }); } + +// --- Stream cleanup wrapper ---------------------------------------------------- + +/** + * If a session is active, wrap the stream so the session is released on end/error. + * Otherwise pass through the stream unchanged. + */ +function wrapStreamMaybe( + body: ReadableStream, + sessionPool?: SessionProxyPool, + sessionId?: string, +): ReadableStream { + if (!sessionPool || !sessionId) return body; + return wrapStreamWithCleanup(body, () => sessionPool.release(sessionId)); +} + +/** + * Wraps a ReadableStream and calls `cleanup` when the stream ends, errors, + * or is cancelled by the consumer. + */ +function wrapStreamWithCleanup(body: ReadableStream, cleanup: () => void): ReadableStream { + const reader = body.getReader(); + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + cleanup(); + controller.close(); + return; + } + controller.enqueue(value); + } catch (err) { + cleanup(); + controller.error(err); + } + }, + cancel(reason) { + cleanup(); + reader.cancel(reason); + }, + }); +} diff --git a/src/lib/anthropic-proxy.ts b/src/lib/anthropic-proxy.ts index 0619181..f2d169d 100644 --- a/src/lib/anthropic-proxy.ts +++ b/src/lib/anthropic-proxy.ts @@ -10,9 +10,9 @@ * - Backend SSE stream -> Anthropic SSE events */ -import type { ProxyPool } from "./proxy-pool"; +import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; import { MODEL_ROUTES, type BackendConfig } from "./ai-proxy"; -import { fetchWithRetry } from "./fetch-utils"; +import { fetchWithRetry, fetchWithSessionRetry, type FetchWithRetryResult } from "./fetch-utils"; import { SSELineBuffer } from "./fetch-utils"; import { isDevMode } from "./fetch-utils"; @@ -449,10 +449,30 @@ function anthropicError(status: number, message: string, type: string): Response /** * Handle an Anthropic-compatible messages request. + * + * Two calling conventions: + * 1. Standard: (body, proxyPool?) + * 2. Session-aware: (body, proxyPool?, sessionPool, sessionId) + * + * When both `sessionPool` and `sessionId` are present the request uses + * session-sticky proxy allocation via `fetchWithSessionRetry`; otherwise + * the existing `fetchWithRetry` path is used (backward-compatible). */ export async function handleAnthropicMessages( body: unknown, proxyPool?: ProxyPool, +): Promise; +export async function handleAnthropicMessages( + body: unknown, + proxyPool?: ProxyPool, + sessionPool?: SessionProxyPool, + sessionId?: string, +): Promise; +export async function handleAnthropicMessages( + body: unknown, + proxyPool?: ProxyPool, + sessionPool?: SessionProxyPool, + sessionId?: string, ): Promise { // -- Input validation ------------------------------------------------------- const validationError = validateAnthropicRequest(body); @@ -483,15 +503,16 @@ export async function handleAnthropicMessages( const url = config.url; - // -- Execute (direct -> proxy fallback) with shared retry ------------------- - const result = await fetchWithRetry( - url, - init, - proxyPool, - `anthropic:${req.model}`, - ); + // -- Execute with session-aware or standard retry -------------------------- + const result: FetchWithRetryResult = + sessionPool && sessionId + ? await fetchWithSessionRetry(url, init, sessionPool, sessionId, `anthropic:${req.model}`) + : await fetchWithRetry(url, init, proxyPool, `anthropic:${req.model}`); if (result.errorClassification) { + if (sessionPool && sessionId) { + sessionPool.release(sessionId); + } return new Response( JSON.stringify({ type: "error", @@ -516,16 +537,20 @@ export async function handleAnthropicMessages( if (!response.ok) { const status = response.status; const genericMsg = status >= 500 ? "Upstream server error" : "Upstream rejected request"; + if (sessionPool && sessionId) { + sessionPool.release(sessionId); + } return anthropicError(status, genericMsg, "upstream_error"); } // -- Handle streaming ------------------------------------------------------- if (wantsStream) { - const transformed = transformAnthropicStream( + let transformed = transformAnthropicStream( response.body!, req.model, config, ); + transformed = wrapAnthropicStreamMaybe(transformed, sessionPool, sessionId); return new Response(transformed, { status: 200, headers: { @@ -540,6 +565,9 @@ export async function handleAnthropicMessages( // -- Handle non-streaming --------------------------------------------------- const text = await response.text(); + if (sessionPool && sessionId) { + sessionPool.release(sessionId); + } if (text.trimStart().startsWith("data: ")) { const accumulated = accumulateSSEText(text); @@ -583,3 +611,47 @@ export async function handleAnthropicMessages( }, }); } + +// --- Stream cleanup wrapper ---------------------------------------------------- + +/** + * If a session is active, wrap the stream so the session is released on end/error. + * Otherwise pass through the stream unchanged. + */ +function wrapAnthropicStreamMaybe( + body: ReadableStream, + sessionPool?: SessionProxyPool, + sessionId?: string, +): ReadableStream { + if (!sessionPool || !sessionId) return body; + return wrapAnthropicStreamWithCleanup(body, () => sessionPool.release(sessionId)); +} + +/** + * Wraps a ReadableStream and calls `cleanup` when the stream ends, errors, + * or is cancelled by the consumer. + */ +function wrapAnthropicStreamWithCleanup(body: ReadableStream, cleanup: () => void): ReadableStream { + const reader = body.getReader(); + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + cleanup(); + controller.close(); + return; + } + controller.enqueue(value); + } catch (err) { + cleanup(); + controller.error(err); + } + }, + cancel(reason) { + cleanup(); + reader.cancel(reason); + }, + }); +} diff --git a/src/lib/fetch-utils.ts b/src/lib/fetch-utils.ts index 0db4b28..bbd5ff8 100644 --- a/src/lib/fetch-utils.ts +++ b/src/lib/fetch-utils.ts @@ -5,7 +5,7 @@ * error sanitization, and graceful shutdown tracking — all in one place. */ -import type { ProxyPool } from "./proxy-pool"; +import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; // ─── Active stream tracking (for graceful shutdown) ─────────────────────── @@ -207,3 +207,78 @@ function classifyFetchErrorSafe(error: unknown): { return { code: "NETWORK_ERROR", status: 502, message: "Upstream unreachable" }; } + +// ─── Fetch with session-based proxy retry ──────────────────────────────── + +/** + * Execute an upstream `fetch` using a session-sticky proxy with retry. + * + * Strategy: session-sticky proxy (via SessionProxyPool) on attempt 1; on + * failure the session's proxy is marked failed — which auto-rotates if the + * failure threshold is exceeded — and the request retries with the new proxy. + * + * SSE streams: the initial request is retried normally. Once the response body + * starts streaming, mid-stream errors are **not** retried; the session is + * released and the error is returned to the caller. + */ +export async function fetchWithSessionRetry( + url: string, + init: RequestInit & { proxy?: string }, + sessionPool: SessionProxyPool | undefined, + sessionId: string, + context?: string, + maxRetries = 3, +): Promise { + // Fallback when no session pool is available + if (!sessionPool) { + try { + const response = await fetch(url, init); + return { response }; + } catch (err) { + return { errorClassification: classifyFetchErrorSafe(err) }; + } + } + + let lastError: unknown; + + for (let attempt = 0; attempt < maxRetries; attempt++) { + const proxyUrl = sessionPool.getProxyUrl(sessionId); + if (proxyUrl) { + init.proxy = proxyUrl; + } + + try { + const response = await fetch(url, init); + + if (response.ok) { + sessionPool.markSuccess(sessionId); + return { response }; + } + + // Non-2xx — mark session failed and retry + lastError = new Error(`Upstream returned ${response.status}`); + const rotated = sessionPool.markFailed(sessionId); + + const ctx = context ? `[${context}] ` : ""; + console.warn( + `${ctx}fetchWithSessionRetry attempt ${attempt + 1}/${maxRetries} ` + + `failed with ${response.status}${rotated ? " (rotated proxy)" : ""}`, + ); + } catch (err) { + lastError = err; + const rotated = sessionPool.markFailed(sessionId); + + const ctx = context ? `[${context}] ` : ""; + const errMsg = err instanceof Error ? err.message : String(err); + console.warn( + `${ctx}fetchWithSessionRetry attempt ${attempt + 1}/${maxRetries} ` + + `failed: ${errMsg}${rotated ? " (rotated proxy)" : ""}`, + ); + } + } + + // All attempts exhausted — release session and classify last error + sessionPool.release(sessionId); + const err = lastError ?? new Error("All session proxy attempts failed"); + return { errorClassification: classifyFetchErrorSafe(err) }; +} diff --git a/src/lib/proxy-pool.ts b/src/lib/proxy-pool.ts index ec7d9b4..82d7a1b 100644 --- a/src/lib/proxy-pool.ts +++ b/src/lib/proxy-pool.ts @@ -179,3 +179,200 @@ export class ProxyPool { this.failureThreshold = n; } } + +// --- SessionProxyPool --------------------------------------------------------- + +interface SessionInfo { + proxyIndex: number; + failures: number; +} + +/** + * Session-based sticky proxy allocation on top of ProxyPool. + * + * Each session gets one sticky proxy until: + * - The session is released (cleanup) + * - The proxy exceeds the failure threshold (auto-rotate to next avail) + * - The session explicitly calls release() + * + * New sessions are assigned to the least-loaded proxy (fewest active sessions). + */ +export class SessionProxyPool { + private pool: ProxyPool; + private sessions = new Map(); + /** proxyIndex -> set of session IDs currently using it */ + private proxyUsage = new Map>(); + private failureThreshold: number; + + /** + * @param poolOrPath Existing ProxyPool or file path to load from. + */ + constructor(poolOrPath?: ProxyPool | string) { + if (poolOrPath instanceof ProxyPool) { + this.pool = poolOrPath; + } else { + this.pool = new ProxyPool(); + if (poolOrPath) this.pool.load(poolOrPath); + } + this.failureThreshold = 3; + } + + /** Number of available proxies in the underlying pool. */ + get size(): number { + return this.pool.size; + } + + /** Number of active sessions. */ + get activeSessions(): number { + return this.sessions.size; + } + + // -- Session management ------------------------------------------------------- + + /** + * Assign a sticky proxy to a session. Returns the proxy URL, or null if empty. + * + * If the session already has a proxy, returns the same one (resume). + * Otherwise picks the least-loaded proxy. + */ + acquire(sessionId: string): string | null { + if (this.pool.size === 0) return null; + + const existing = this.sessions.get(sessionId); + if (existing !== undefined) { + return this.formatProxyUrlAtIndex(existing.proxyIndex); + } + + const index = this.pickLeastUsedIndex(); + if (index === -1) return null; + + this.sessions.set(sessionId, { proxyIndex: index, failures: 0 }); + + let usedBy = this.proxyUsage.get(index); + if (!usedBy) { + usedBy = new Set(); + this.proxyUsage.set(index, usedBy); + } + usedBy.add(sessionId); + + return this.formatProxyUrlAtIndex(index); + } + + /** Return the current proxy URL for a session (no rotation), or null. */ + getProxyUrl(sessionId: string): string | null { + const info = this.sessions.get(sessionId); + if (!info) return null; + return this.formatProxyUrlAtIndex(info.proxyIndex); + } + + /** Remove a session from all tracking. */ + release(sessionId: string): void { + const info = this.sessions.get(sessionId); + if (!info) return; + + const usedBy = this.proxyUsage.get(info.proxyIndex); + if (usedBy) { + usedBy.delete(sessionId); + if (usedBy.size === 0) this.proxyUsage.delete(info.proxyIndex); + } + + this.sessions.delete(sessionId); + } + + /** + * Increment failure count for this session's proxy. + * + * If failures >= threshold, auto-rotate to a different proxy. + * Also marks the old proxy as failed in the underlying ProxyPool. + * + * @returns true if the session was rotated to a new proxy. + */ + markFailed(sessionId: string): boolean { + const info = this.sessions.get(sessionId); + if (!info) return false; + + info.failures += 1; + if (info.failures < this.failureThreshold) return false; + + const oldIndex = info.proxyIndex; + + // Mark in underlying pool (public API only works on currentIndex) + const savedIdx = (this.pool as any).currentIndex as number; + (this.pool as any).currentIndex = oldIndex; + this.pool.markFailed(this.failureThreshold); + (this.pool as any).currentIndex = savedIdx; + + // Remove session from old proxy usage tracking + const usedBy = this.proxyUsage.get(oldIndex); + if (usedBy) { + usedBy.delete(sessionId); + if (usedBy.size === 0) this.proxyUsage.delete(oldIndex); + } + + // Pick next available proxy + const newIndex = this.pickLeastUsedIndex(); + if (newIndex === -1 || newIndex === oldIndex) { + // Single-proxy pool or none available — reset failures, stay put + this.sessions.set(sessionId, { proxyIndex: oldIndex, failures: 0 }); + return false; + } + + this.sessions.set(sessionId, { proxyIndex: newIndex, failures: 0 }); + + let newUsedBy = this.proxyUsage.get(newIndex); + if (!newUsedBy) { + newUsedBy = new Set(); + this.proxyUsage.set(newIndex, newUsedBy); + } + newUsedBy.add(sessionId); + + return true; + } + + /** Reset failure count for this session's proxy. */ + markSuccess(sessionId: string): void { + const info = this.sessions.get(sessionId); + if (!info) return; + info.failures = 0; + } + + // -- Internals ---------------------------------------------------------------- + + /** Get the ProxyEntry at a given index. Forward reference to local type. */ + private poolEntryAtIndex(index: number): ProxyEntry | null { + return (this.pool as any).proxies[index] ?? null; + } + + /** Build proxy URL string by index. */ + private formatProxyUrlAtIndex(index: number): string | null { + const entry = this.poolEntryAtIndex(index); + if (!entry) return null; + const auth = entry.username + ? `${encodeURIComponent(entry.username)}:${encodeURIComponent(entry.password)}@` + : ""; + return `http://${auth}${entry.host}:${entry.port}`; + } + + /** Return the index of the proxy with the fewest active sessions, or -1. */ + private pickLeastUsedIndex(): number { + if (this.pool.size === 0) return -1; + + let bestIndex = 0; + let bestCount = Infinity; + + for (let i = 0; i < this.pool.size; i++) { + const count = this.proxyUsage.get(i)?.size ?? 0; + if (count < bestCount) { + bestCount = count; + bestIndex = i; + } + } + + return bestIndex; + } + + /** Override the failure threshold (default 3). */ + setFailureThreshold(n: number): void { + this.failureThreshold = n; + } +}