feat: implement session-based sticky proxy pool with auto-failover

Add SessionProxyPool for per-session sticky proxy allocation with
load-balanced least-used selection and auto-rotation on failure.
Introduce fetchWithSessionRetry for transparent retry with proxy
rotation. Wire into AI proxy handlers (OpenAI + Anthropic) with
stream lifecycle cleanup.

Co-Authored-By: Claude <noreply@anthropic.com>

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 {
@@ -484,7 +489,8 @@ const server: Server<WSRelayData> = Bun.serve<WSRelayData>({
 			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<WSRelayData> = Bun.serve<WSRelayData>({
 			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({

src/lib/ai-proxy.ts

@@ -12,10 +12,8 @@
  * Streaming (SSE) is supported for all backends.
  */
 
-import type { ProxyPool } from "./proxy-pool";
+import type { ProxyPool, SessionProxyPool } from "./proxy-pool";
-import { fetchWithRetry } from "./fetch-utils";
+import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, type FetchWithRetryResult } from "./fetch-utils";
-import { SSELineBuffer } from "./fetch-utils";
-import { isDevMode } 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<Response>;
+export async function handleChatCompletion(
+	body: unknown,
+	proxyPool?: ProxyPool,
+	sessionPool?: SessionProxyPool,
+	sessionId?: string,
+): Promise<Response>;
+export async function handleChatCompletion(
+	body: unknown,
+	proxyPool?: ProxyPool,
+	sessionPool?: SessionProxyPool,
+	sessionId?: string,
 ): Promise<Response> {
 	// -- 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 -------------------
+	// -- Execute with session-aware or standard retry --------------------------
-	const result = await fetchWithRetry(
+	const result: FetchWithRetryResult =
-		url,
+		sessionPool && sessionId
-		init,
+			? await fetchWithSessionRetry(url, init, sessionPool, sessionId, `openai:${req.model}`)
-		proxyPool,
+			: await fetchWithRetry(url, init, proxyPool, `openai:${req.model}`);
-		`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,7 +417,9 @@ export async function handleChatCompletion(
 			config,
 			req,
 		);
-		return new Response(transformed, {
+		return new Response(
+			wrapStreamMaybe(transformed, sessionPool, sessionId),
+			{
 				status: 200,
 				headers: {
 					"Content-Type": "text/event-stream",
@@ -407,11 +428,15 @@ export async function handleChatCompletion(
 					"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);
+		},
+	});
+}

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<Response>;
+export async function handleAnthropicMessages(
+	body: unknown,
+	proxyPool?: ProxyPool,
+	sessionPool?: SessionProxyPool,
+	sessionId?: string,
+): Promise<Response>;
+export async function handleAnthropicMessages(
+	body: unknown,
+	proxyPool?: ProxyPool,
+	sessionPool?: SessionProxyPool,
+	sessionId?: string,
 ): Promise<Response> {
 	// -- 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 -------------------
+	// -- Execute with session-aware or standard retry --------------------------
-	const result = await fetchWithRetry(
+	const result: FetchWithRetryResult =
-		url,
+		sessionPool && sessionId
-		init,
+			? await fetchWithSessionRetry(url, init, sessionPool, sessionId, `anthropic:${req.model}`)
-		proxyPool,
+			: await fetchWithRetry(url, init, proxyPool, `anthropic:${req.model}`);
-		`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);
+		},
+	});
+}

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<FetchWithRetryResult> {
+  // 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) };
+}

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<string, SessionInfo>();
+	/** proxyIndex -> set of session IDs currently using it */
+	private proxyUsage = new Map<number, Set<string>>();
+	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;
+	}
+}