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>
This commit is contained in:
MythEclipse
2026-06-16 23:50:47 +07:00
co-authored by Claude
parent 6c945dcbf4
commit 9bd6acddc5
5 changed files with 551 additions and 131 deletions
+106 -99
View File
@@ -21,7 +21,6 @@ import {
filterRequestHeaders, filterRequestHeaders,
buildRelayRequest, buildRelayRequest,
createRelayResponse, createRelayResponse,
classifyFetchError,
createErrorResponse, createErrorResponse,
createCorsPreflightResponse, createCorsPreflightResponse,
getCorsHeaders, getCorsHeaders,
@@ -30,7 +29,7 @@ import {
import { checkBodySize } from "./middleware/body-limiter"; import { checkBodySize } from "./middleware/body-limiter";
import { createRateLimiter } from "./middleware/rate-limiter"; import { createRateLimiter } from "./middleware/rate-limiter";
import { logRelayEvent } from "./middleware/logger"; 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 { handleChatCompletion, listModels } from "./lib/ai-proxy";
import { handleAnthropicMessages } from "./lib/anthropic-proxy"; import { handleAnthropicMessages } from "./lib/anthropic-proxy";
import { fetchWithRetry, closeAllActiveReaders, isDevMode } from "./lib/fetch-utils"; import { fetchWithRetry, closeAllActiveReaders, isDevMode } from "./lib/fetch-utils";
@@ -81,6 +80,12 @@ proxyPool.tryLoad(
process.env.PROXY_FILE || process.env.PROXY_LIST || "./proxy.txt", 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 ----------------------------------------------- // --- WebSocket relay data type -----------------------------------------------
interface WSRelayData { interface WSRelayData {
@@ -112,85 +117,85 @@ function handleHealth(): Response {
/** Simple embedded HTML documentation page. */ /** Simple embedded HTML documentation page. */
function handleDocs(): Response { function handleDocs(): Response {
const html = `<!DOCTYPE html> const html = `<!DOCTYPE html>
<html lang="en"> <html lang="en">
<head> <head>
<meta charset="UTF-8"> <meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0"> <meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Edge Proxy Relay — Docs</title> <title>Edge Proxy Relay — Docs</title>
<style> <style>
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; } *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
body { font-family: system-ui, -apple-system, sans-serif; line-height: 1.6; color: #e1e4e8; background: #0d1117; padding: 2rem; } body { font-family: system-ui, -apple-system, sans-serif; line-height: 1.6; color: #e1e4e8; background: #0d1117; padding: 2rem; }
main { max-width: 800px; margin: 0 auto; } main { max-width: 800px; margin: 0 auto; }
h1 { font-size: 2rem; margin-bottom: 0.5rem; color: #58a6ff; } h1 { font-size: 2rem; margin-bottom: 0.5rem; color: #58a6ff; }
h2 { font-size: 1.25rem; margin: 2rem 0 0.75rem; color: #c9d1d9; border-bottom: 1px solid #30363d; padding-bottom: 0.25rem; } h2 { font-size: 1.25rem; margin: 2rem 0 0.75rem; color: #c9d1d9; border-bottom: 1px solid #30363d; padding-bottom: 0.25rem; }
p, li { color: #8b949e; } p, li { color: #8b949e; }
code { background: #161b22; padding: 0.2em 0.4em; border-radius: 4px; font-size: 0.9em; color: #f0f6fc; } code { background: #161b22; padding: 0.2em 0.4em; border-radius: 4px; font-size: 0.9em; color: #f0f6fc; }
pre { background: #161b22; padding: 1rem; border-radius: 6px; overflow-x: auto; margin: 0.75rem 0; border: 1px solid #30363d; } pre { background: #161b22; padding: 1rem; border-radius: 6px; overflow-x: auto; margin: 0.75rem 0; border: 1px solid #30363d; }
pre code { background: none; padding: 0; } pre code { background: none; padding: 0; }
table { width: 100%; border-collapse: collapse; margin: 0.75rem 0; } table { width: 100%; border-collapse: collapse; margin: 0.75rem 0; }
th, td { text-align: left; padding: 0.5rem 0.75rem; border: 1px solid #30363d; } th, td { text-align: left; padding: 0.5rem 0.75rem; border: 1px solid #30363d; }
th { background: #161b22; color: #c9d1d9; } th { background: #161b22; color: #c9d1d9; }
ul { padding-left: 1.5rem; margin: 0.5rem 0; } ul { padding-left: 1.5rem; margin: 0.5rem 0; }
.endpoint { background: #161b22; border: 1px solid #30363d; border-radius: 6px; padding: 1rem; margin: 1rem 0; } .endpoint { background: #161b22; border: 1px solid #30363d; border-radius: 6px; padding: 1rem; margin: 1rem 0; }
.endpoint h3 { color: #58a6ff; font-family: monospace; margin-bottom: 0.5rem; } .endpoint h3 { color: #58a6ff; font-family: monospace; margin-bottom: 0.5rem; }
.status { color: #3fb950; } .status { color: #3fb950; }
a { color: #58a6ff; } a { color: #58a6ff; }
</style> </style>
</head> </head>
<body> <body>
<main> <main>
<h1>Edge Proxy Relay</h1> <h1>Edge Proxy Relay</h1>
<p>Forward HTTP and WebSocket requests to any target server via the <code>x-relay-target</code> header.</p> <p>Forward HTTP and WebSocket requests to any target server via the <code>x-relay-target</code> header.</p>
<h2>Endpoints</h2> <h2>Endpoints</h2>
<div class="endpoint"> <div class="endpoint">
<h3>GET /health</h3> <h3>GET /health</h3>
<p>Health check. Returns <span class="status">200 OK</span> with server status, uptime, and version.</p> <p>Health check. Returns <span class="status">200 OK</span> with server status, uptime, and version.</p>
</div> </div>
<div class="endpoint"> <div class="endpoint">
<h3>GET /docs</h3> <h3>GET /docs</h3>
<p>This page.</p> <p>This page.</p>
</div> </div>
<div class="endpoint"> <div class="endpoint">
<h3>Any Path (Catch-all Relay)</h3> <h3>Any Path (Catch-all Relay)</h3>
<p>Send a request with the <code>x-relay-target</code> header and this proxy forwards it.</p> <p>Send a request with the <code>x-relay-target</code> header and this proxy forwards it.</p>
</div> </div>
<h2>Usage — HTTP Relay</h2> <h2>Usage — HTTP Relay</h2>
<pre><code>curl -s \\ <pre><code>curl -s \\
-H "x-relay-target: https://httpbin.org" \\ -H "x-relay-target: https://httpbin.org" \\
-H "x-relay-path: /get" \\ -H "x-relay-path: /get" \\
"https://your-proxy.example/any/path"</code></pre> "https://your-proxy.example/any/path"</code></pre>
<table> <table>
<tr><th>Header</th><th>Required</th><th>Description</th></tr> <tr><th>Header</th><th>Required</th><th>Description</th></tr>
<tr><td><code>x-relay-target</code></td><td>Yes</td><td>Base URL of the upstream (http:// or https://)</td></tr> <tr><td><code>x-relay-target</code></td><td>Yes</td><td>Base URL of the upstream (http:// or https://)</td></tr>
<tr><td><code>x-relay-path</code></td><td>No</td><td>Path to append (default: <code>/</code>)</td></tr> <tr><td><code>x-relay-path</code></td><td>No</td><td>Path to append (default: <code>/</code>)</td></tr>
</table> </table>
<h2>Usage — WebSocket Relay</h2> <h2>Usage — WebSocket Relay</h2>
<pre><code>const ws = new WebSocket("wss://your-proxy.example/relay", { <pre><code>const ws = new WebSocket("wss://your-proxy.example/relay", {
headers: { "x-relay-target": "wss://echo-websocket.example" }, headers: { "x-relay-target": "wss://echo-websocket.example" },
}); });
ws.onopen = () => ws.send("Hello via relay!"); ws.onopen = () => ws.send("Hello via relay!");
ws.onmessage = (e) => console.log("Got:", e.data);</code></pre> ws.onmessage = (e) => console.log("Got:", e.data);</code></pre>
<h2>Status Codes</h2> <h2>Status Codes</h2>
<table> <table>
<tr><th>Code</th><th>Meaning</th></tr> <tr><th>Code</th><th>Meaning</th></tr>
<tr><td>204</td><td>CORS preflight success (OPTIONS)</td></tr> <tr><td>204</td><td>CORS preflight success (OPTIONS)</td></tr>
<tr><td>400</td><td>Missing <code>x-relay-target</code> header</td></tr> <tr><td>400</td><td>Missing <code>x-relay-target</code> header</td></tr>
<tr><td>403</td><td>Target blocked (SSRF protection / not allowed)</td></tr> <tr><td>403</td><td>Target blocked (SSRF protection / not allowed)</td></tr>
<tr><td>413</td><td>Request body exceeds size limit</td></tr> <tr><td>413</td><td>Request body exceeds size limit</td></tr>
<tr><td>429</td><td>Rate limit exceeded</td></tr> <tr><td>429</td><td>Rate limit exceeded</td></tr>
<tr><td>502</td><td>Upstream network / DNS error</td></tr> <tr><td>502</td><td>Upstream network / DNS error</td></tr>
<tr><td>504</td><td>Upstream timeout</td></tr> <tr><td>504</td><td>Upstream timeout</td></tr>
</table> </table>
</main> </main>
</body> </body>
</html>`; </html>`;
return new Response(html, { return new Response(html, {
status: 200, status: 200,
@@ -204,29 +209,29 @@ ws.onmessage = (e) => console.log("Got:", e.data);</code></pre>
/** Minimal status page shown at the root `/`. */ /** Minimal status page shown at the root `/`. */
function handleIndex(): Response { function handleIndex(): Response {
const html = `<!DOCTYPE html> const html = `<!DOCTYPE html>
<html lang="en"> <html lang="en">
<head> <head>
<meta charset="UTF-8"> <meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0"> <meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Edge Proxy Relay</title> <title>Edge Proxy Relay</title>
<style> <style>
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; } *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
body { font-family: system-ui, -apple-system, sans-serif; line-height: 1.6; color: #e1e4e8; background: #0d1117; display: flex; align-items: center; justify-content: center; min-height: 100vh; } body { font-family: system-ui, -apple-system, sans-serif; line-height: 1.6; color: #e1e4e8; background: #0d1117; display: flex; align-items: center; justify-content: center; min-height: 100vh; }
main { text-align: center; } main { text-align: center; }
h1 { font-size: 2rem; color: #58a6ff; margin-bottom: 0.5rem; } h1 { font-size: 2rem; color: #58a6ff; margin-bottom: 0.5rem; }
p { color: #8b949e; } p { color: #8b949e; }
a { color: #58a6ff; } a { color: #58a6ff; }
.status { color: #3fb950; } .status { color: #3fb950; }
</style> </style>
</head> </head>
<body> <body>
<main> <main>
<h1>Edge Proxy Relay</h1> <h1>Edge Proxy Relay</h1>
<p class="status">Server is running</p> <p class="status">Server is running</p>
<p><a href="/health">/health</a> &middot; <a href="/docs">/docs</a></p> <p><a href="/health">/health</a> &middot; <a href="/docs">/docs</a></p>
</main> </main>
</body> </body>
</html>`; </html>`;
return new Response(html, { return new Response(html, {
status: 200, status: 200,
@@ -484,7 +489,8 @@ const server: Server<WSRelayData> = Bun.serve<WSRelayData>({
if (authErr) return authErr; if (authErr) return authErr;
try { try {
const body = await req.json(); const body = await req.json();
return handleChatCompletion(body, proxyPool); const sessionId = crypto.randomUUID();
return handleChatCompletion(body, proxyPool, sessionPool, sessionId);
} catch { } catch {
return new Response( return new Response(
JSON.stringify({ error: { message: "Invalid JSON body", type: "invalid_request_error" } }), 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; if (authErr) return authErr;
try { try {
const body = await req.json(); const body = await req.json();
return handleAnthropicMessages(body, proxyPool); const sessionId = crypto.randomUUID();
return handleAnthropicMessages(body, proxyPool, sessionPool, sessionId);
} catch { } catch {
return new Response( return new Response(
JSON.stringify({ JSON.stringify({
+90 -21
View File
@@ -12,10 +12,8 @@
* Streaming (SSE) is supported for all backends. * 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 ------------------------------------------------------------------- // --- Types -------------------------------------------------------------------
@@ -315,10 +313,30 @@ function openAIError(status: number, message: string, type: string): Response {
/** /**
* Handle an OpenAI-compatible chat completions request. * 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( export async function handleChatCompletion(
body: unknown, body: unknown,
proxyPool?: ProxyPool, 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> { ): Promise<Response> {
// -- Input validation ------------------------------------------------------- // -- Input validation -------------------------------------------------------
const validationError = validateChatRequest(body); const validationError = validateChatRequest(body);
@@ -340,13 +358,11 @@ export async function handleChatCompletion(
const wantsStream = req.stream === true; const wantsStream = req.stream === true;
const { url, init } = buildBackendRequest(req, config); 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) { if (result.errorClassification) {
return new Response( return new Response(
@@ -389,7 +405,10 @@ export async function handleChatCompletion(
"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Origin": "*",
"X-Accel-Buffering": "no", "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 // Transform the stream
@@ -398,20 +417,26 @@ export async function handleChatCompletion(
config, config,
req, req,
); );
return new Response(transformed, { return new Response(
status: 200, wrapStreamMaybe(transformed, sessionPool, sessionId),
headers: { {
"Content-Type": "text/event-stream", status: 200,
"Cache-Control": "no-cache", headers: {
Connection: "keep-alive", "Content-Type": "text/event-stream",
"Access-Control-Allow-Origin": "*", "Cache-Control": "no-cache",
"X-Accel-Buffering": "no", Connection: "keep-alive",
"Access-Control-Allow-Origin": "*",
"X-Accel-Buffering": "no",
},
}, },
}); );
} }
// -- Handle non-streaming response ------------------------------------------ // -- Handle non-streaming response ------------------------------------------
const text = await response.text(); const text = await response.text();
if (sessionPool && sessionId) {
sessionPool.release(sessionId);
}
const adapted = parseJSONResponse(text, config, req); const adapted = parseJSONResponse(text, config, req);
return new Response(JSON.stringify(adapted), { 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);
},
});
}
+82 -10
View File
@@ -10,9 +10,9 @@
* - Backend SSE stream -> Anthropic SSE events * - 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 { 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 { SSELineBuffer } from "./fetch-utils";
import { isDevMode } 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. * 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( export async function handleAnthropicMessages(
body: unknown, body: unknown,
proxyPool?: ProxyPool, 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> { ): Promise<Response> {
// -- Input validation ------------------------------------------------------- // -- Input validation -------------------------------------------------------
const validationError = validateAnthropicRequest(body); const validationError = validateAnthropicRequest(body);
@@ -483,15 +503,16 @@ export async function handleAnthropicMessages(
const url = config.url; 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 (result.errorClassification) {
if (sessionPool && sessionId) {
sessionPool.release(sessionId);
}
return new Response( return new Response(
JSON.stringify({ JSON.stringify({
type: "error", type: "error",
@@ -516,16 +537,20 @@ export async function handleAnthropicMessages(
if (!response.ok) { if (!response.ok) {
const status = response.status; const status = response.status;
const genericMsg = status >= 500 ? "Upstream server error" : "Upstream rejected request"; const genericMsg = status >= 500 ? "Upstream server error" : "Upstream rejected request";
if (sessionPool && sessionId) {
sessionPool.release(sessionId);
}
return anthropicError(status, genericMsg, "upstream_error"); return anthropicError(status, genericMsg, "upstream_error");
} }
// -- Handle streaming ------------------------------------------------------- // -- Handle streaming -------------------------------------------------------
if (wantsStream) { if (wantsStream) {
const transformed = transformAnthropicStream( let transformed = transformAnthropicStream(
response.body!, response.body!,
req.model, req.model,
config, config,
); );
transformed = wrapAnthropicStreamMaybe(transformed, sessionPool, sessionId);
return new Response(transformed, { return new Response(transformed, {
status: 200, status: 200,
headers: { headers: {
@@ -540,6 +565,9 @@ export async function handleAnthropicMessages(
// -- Handle non-streaming --------------------------------------------------- // -- Handle non-streaming ---------------------------------------------------
const text = await response.text(); const text = await response.text();
if (sessionPool && sessionId) {
sessionPool.release(sessionId);
}
if (text.trimStart().startsWith("data: ")) { if (text.trimStart().startsWith("data: ")) {
const accumulated = accumulateSSEText(text); 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);
},
});
}
+76 -1
View File
@@ -5,7 +5,7 @@
* error sanitization, and graceful shutdown tracking — all in one place. * 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) ─────────────────────── // ─── Active stream tracking (for graceful shutdown) ───────────────────────
@@ -207,3 +207,78 @@ function classifyFetchErrorSafe(error: unknown): {
return { code: "NETWORK_ERROR", status: 502, message: "Upstream unreachable" }; 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) };
}
+197
View File
@@ -179,3 +179,200 @@ export class ProxyPool {
this.failureThreshold = n; 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;
}
}