feat(proxy): improve channel failover stability and sticky session coordination / 优化代理通道故障切换与粘滞会话协调 (#268)

* wip: preserve local changes before upstream merge

* Improve proxy routing stability and add stable-first strategy

* Sync explicit-group routing strategy to eligible source routes

* Fix proxy retry and runtime health review issues

* fix: make proxy channel retry attempts configurable

* fix: harden proxy retry recovery and bookkeeping

* fix: address pr 257 review comments

* feat(proxy): add sticky session channel coordination

---------

Co-authored-by: Cita <juricek.chen@gmail.com>
This commit is contained in:
Bainily
2026-03-24 19:40:44 +08:00
committed by GitHub
parent 5b9c83f830
commit 9cf3dd1904
10 changed files with 1031 additions and 79 deletions
+6
View File
@@ -109,6 +109,12 @@ export function buildConfig(env: NodeJS.ProcessEnv) {
routingFallbackUnitCost: Math.max(1e-6, parseNumber(env.ROUTING_FALLBACK_UNIT_COST, 1)), routingFallbackUnitCost: Math.max(1e-6, parseNumber(env.ROUTING_FALLBACK_UNIT_COST, 1)),
tokenRouterCacheTtlMs: Math.max(100, Math.trunc(parseNumber(env.TOKEN_ROUTER_CACHE_TTL_MS, 1_500))), tokenRouterCacheTtlMs: Math.max(100, Math.trunc(parseNumber(env.TOKEN_ROUTER_CACHE_TTL_MS, 1_500))),
proxyMaxChannelAttempts: Math.max(1, Math.trunc(parseNumber(env.PROXY_MAX_CHANNEL_ATTEMPTS, 3))), proxyMaxChannelAttempts: Math.max(1, Math.trunc(parseNumber(env.PROXY_MAX_CHANNEL_ATTEMPTS, 3))),
proxyStickySessionEnabled: parseBoolean(env.PROXY_STICKY_SESSION_ENABLED, true),
proxyStickySessionTtlMs: Math.max(30_000, Math.trunc(parseNumber(env.PROXY_STICKY_SESSION_TTL_MS, 30 * 60 * 1000))),
proxySessionChannelConcurrencyLimit: Math.max(0, Math.trunc(parseNumber(env.PROXY_SESSION_CHANNEL_CONCURRENCY_LIMIT, 2))),
proxySessionChannelQueueWaitMs: Math.max(0, Math.trunc(parseNumber(env.PROXY_SESSION_CHANNEL_QUEUE_WAIT_MS, 1_500))),
proxySessionChannelLeaseTtlMs: Math.max(5_000, Math.trunc(parseNumber(env.PROXY_SESSION_CHANNEL_LEASE_TTL_MS, 90_000))),
proxySessionChannelLeaseKeepaliveMs: Math.max(1_000, Math.trunc(parseNumber(env.PROXY_SESSION_CHANNEL_LEASE_KEEPALIVE_MS, 15_000))),
proxyLogRetentionDays: Math.max(0, Math.trunc(parseNumber(env.PROXY_LOG_RETENTION_DAYS, 30))), proxyLogRetentionDays: Math.max(0, Math.trunc(parseNumber(env.PROXY_LOG_RETENTION_DAYS, 30))),
proxyLogRetentionPruneIntervalMinutes: Math.max(1, Math.trunc(parseNumber(env.PROXY_LOG_RETENTION_PRUNE_INTERVAL_MINUTES, 30))), proxyLogRetentionPruneIntervalMinutes: Math.max(1, Math.trunc(parseNumber(env.PROXY_LOG_RETENTION_PRUNE_INTERVAL_MINUTES, 30))),
proxyFileRetentionDays: Math.max(0, Math.trunc(parseNumber(env.PROXY_FILE_RETENTION_DAYS, 30))), proxyFileRetentionDays: Math.max(0, Math.trunc(parseNumber(env.PROXY_FILE_RETENTION_DAYS, 30))),
@@ -43,6 +43,11 @@ import { summarizeConversationFileInputsInOpenAiBody } from '../capabilities/con
import { detectDownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js'; import { detectDownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js';
import { canRetryProxyChannel, getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js'; import { canRetryProxyChannel, getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js';
import { import {
acquireSurfaceChannelLease,
bindSurfaceStickyChannel,
buildSurfaceChannelBusyMessage,
buildSurfaceStickySessionKey,
clearSurfaceStickyChannel,
createSurfaceFailureToolkit, createSurfaceFailureToolkit,
createSurfaceDispatchRequest, createSurfaceDispatchRequest,
recordSurfaceSuccess, recordSurfaceSuccess,
@@ -126,6 +131,12 @@ export async function handleChatSurfaceRequest(
clientContext, clientContext,
downstreamApiKeyId, downstreamApiKeyId,
}); });
const stickySessionKey = buildSurfaceStickySessionKey({
clientContext,
requestedModel,
downstreamPath,
downstreamApiKeyId,
});
const excludeChannelIds: number[] = []; const excludeChannelIds: number[] = [];
let retryCount = 0; let retryCount = 0;
@@ -136,6 +147,7 @@ export async function handleChatSurfaceRequest(
downstreamPolicy, downstreamPolicy,
excludeChannelIds, excludeChannelIds,
retryCount, retryCount,
stickySessionKey,
}); });
if (!selected) { if (!selected) {
@@ -250,6 +262,36 @@ export async function handleChatSurfaceRequest(
return endpointStrategy.tryRecover(ctx); return endpointStrategy.tryRecover(ctx);
}; };
let startTime = Date.now(); let startTime = Date.now();
const leaseResult = await acquireSurfaceChannelLease({
selected,
});
if (leaseResult.status === 'timeout') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const busyMessage = buildSurfaceChannelBusyMessage(leaseResult.waitMs);
await failureToolkit.log({
selected,
modelRequested: requestedModel,
status: 'failed',
httpStatus: 503,
latencyMs: leaseResult.waitMs,
errorMessage: busyMessage,
retryCount,
});
if (canRetryProxyChannel(retryCount)) {
retryCount += 1;
continue;
}
return reply.code(503).send({
error: {
message: busyMessage,
type: 'server_error',
},
});
}
const channelLease = leaseResult.lease;
try { try {
const endpointResult = await executeEndpointFlow({ const endpointResult = await executeEndpointFlow({
@@ -287,6 +329,10 @@ export async function handleChatSurfaceRequest(
}); });
if (!endpointResult.ok) { if (!endpointResult.ok) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleUpstreamFailure({ const failureOutcome = await failureToolkit.handleUpstreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -358,6 +404,10 @@ export async function handleChatSurfaceRequest(
); );
const latency = Date.now() - startTime; const latency = Date.now() - startTime;
if (streamResult.status === 'failed') { if (streamResult.status === 'failed') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
await failureToolkit.recordStreamFailure({ await failureToolkit.recordStreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -372,6 +422,10 @@ export async function handleChatSurfaceRequest(
}); });
return; return;
} }
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return; return;
} }
let fallbackData: unknown = null; let fallbackData: unknown = null;
@@ -387,6 +441,10 @@ export async function handleChatSurfaceRequest(
const latency = Date.now() - startTime; const latency = Date.now() - startTime;
const failure = detectProxyFailure({ rawText, usage: parsedUsage }); const failure = detectProxyFailure({ rawText, usage: parsedUsage });
if (failure) { if (failure) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleDetectedFailure({ const failureOutcome = await failureToolkit.handleDetectedFailure({
selected, selected,
requestedModel, requestedModel,
@@ -409,6 +467,10 @@ export async function handleChatSurfaceRequest(
startSseResponse(); startSseResponse();
const streamResult = streamSession.consumeUpstreamFinalPayload(fallbackData, fallbackText, reply.raw); const streamResult = streamSession.consumeUpstreamFinalPayload(fallbackData, fallbackText, reply.raw);
if (streamResult.status === 'failed') { if (streamResult.status === 'failed') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
await failureToolkit.recordStreamFailure({ await failureToolkit.recordStreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -424,6 +486,10 @@ export async function handleChatSurfaceRequest(
}); });
return; return;
} }
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
} else { } else {
startSseResponse(); startSseResponse();
const upstreamReader = upstream.body?.getReader(); const upstreamReader = upstream.body?.getReader();
@@ -453,6 +519,10 @@ export async function handleChatSurfaceRequest(
const latency = Date.now() - startTime; const latency = Date.now() - startTime;
if (streamResult.status === 'failed') { if (streamResult.status === 'failed') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
await failureToolkit.recordStreamFailure({ await failureToolkit.recordStreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -493,6 +563,10 @@ export async function handleChatSurfaceRequest(
errorLabel: '[proxy/chat] failed to record success metrics', errorLabel: '[proxy/chat] failed to record success metrics',
}, },
}); });
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return; return;
} }
@@ -524,6 +598,10 @@ export async function handleChatSurfaceRequest(
const parsedUsage = parseProxyUsage(upstreamData); const parsedUsage = parseProxyUsage(upstreamData);
const failure = detectProxyFailure({ rawText, usage: parsedUsage }); const failure = detectProxyFailure({ rawText, usage: parsedUsage });
if (failure) { if (failure) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleDetectedFailure({ const failureOutcome = await failureToolkit.handleDetectedFailure({
selected, selected,
requestedModel, requestedModel,
@@ -562,9 +640,17 @@ export async function handleChatSurfaceRequest(
errorLabel: '[proxy/chat] failed to record success metrics', errorLabel: '[proxy/chat] failed to record success metrics',
}, },
}); });
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return reply.send(downstreamResponse); return reply.send(downstreamResponse);
} catch (err: any) { } catch (err: any) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleExecutionError({ const failureOutcome = await failureToolkit.handleExecutionError({
selected, selected,
requestedModel, requestedModel,
@@ -578,6 +664,8 @@ export async function handleChatSurfaceRequest(
continue; continue;
} }
return reply.code(failureOutcome.status).send(failureOutcome.payload); return reply.code(failureOutcome.status).send(failureOutcome.payload);
} finally {
channelLease.release();
} }
} }
} }
@@ -646,6 +734,12 @@ export async function handleClaudeCountTokensSurfaceRequest(
clientContext, clientContext,
downstreamApiKeyId, downstreamApiKeyId,
}); });
const stickySessionKey = buildSurfaceStickySessionKey({
clientContext,
requestedModel,
downstreamPath,
downstreamApiKeyId,
});
const excludeChannelIds: number[] = []; const excludeChannelIds: number[] = [];
let retryCount = 0; let retryCount = 0;
@@ -655,6 +749,7 @@ export async function handleClaudeCountTokensSurfaceRequest(
downstreamPolicy, downstreamPolicy,
excludeChannelIds, excludeChannelIds,
retryCount, retryCount,
stickySessionKey,
}); });
if (!selected) { if (!selected) {
@@ -692,6 +787,36 @@ export async function handleClaudeCountTokensSurfaceRequest(
} }
const oauth = getOauthInfoFromAccount(selected.account); const oauth = getOauthInfoFromAccount(selected.account);
const startTime = Date.now(); const startTime = Date.now();
const leaseResult = await acquireSurfaceChannelLease({
selected,
});
if (leaseResult.status === 'timeout') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const busyMessage = buildSurfaceChannelBusyMessage(leaseResult.waitMs);
await failureToolkit.log({
selected,
modelRequested: requestedModel,
status: 'failed',
httpStatus: 503,
latencyMs: leaseResult.waitMs,
errorMessage: busyMessage,
retryCount,
});
if (canRetryProxyChannel(retryCount)) {
retryCount += 1;
continue;
}
return reply.code(503).send({
error: {
message: busyMessage,
type: 'server_error',
},
});
}
const channelLease = leaseResult.lease;
const buildRequest = () => { const buildRequest = () => {
const upstreamRequest = buildClaudeCountTokensUpstreamRequest({ const upstreamRequest = buildClaudeCountTokensUpstreamRequest({
@@ -753,6 +878,10 @@ export async function handleClaudeCountTokensSurfaceRequest(
} }
if (!upstream.ok) { if (!upstream.ok) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleUpstreamFailure({ const failureOutcome = await failureToolkit.handleUpstreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -782,8 +911,16 @@ export async function handleClaudeCountTokensSurfaceRequest(
retryCount, retryCount,
upstreamPath: upstreamRequest.path, upstreamPath: upstreamRequest.path,
}); });
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return reply.code(upstream.status).type(contentType).send(payload); return reply.code(upstream.status).type(contentType).send(payload);
} catch (error: any) { } catch (error: any) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleExecutionError({ const failureOutcome = await failureToolkit.handleExecutionError({
selected, selected,
requestedModel, requestedModel,
@@ -797,6 +934,8 @@ export async function handleClaudeCountTokensSurfaceRequest(
continue; continue;
} }
return reply.code(failureOutcome.status).send(failureOutcome.payload); return reply.code(failureOutcome.status).send(failureOutcome.payload);
} finally {
channelLease.release();
} }
} }
} }
@@ -40,6 +40,11 @@ import {
import { detectDownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js'; import { detectDownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js';
import { getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js'; import { getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js';
import { import {
acquireSurfaceChannelLease,
bindSurfaceStickyChannel,
buildSurfaceChannelBusyMessage,
buildSurfaceStickySessionKey,
clearSurfaceStickyChannel,
createSurfaceFailureToolkit, createSurfaceFailureToolkit,
createSurfaceDispatchRequest, createSurfaceDispatchRequest,
recordSurfaceSuccess, recordSurfaceSuccess,
@@ -185,6 +190,12 @@ export async function handleOpenAiResponsesSurfaceRequest(
clientContext, clientContext,
downstreamApiKeyId, downstreamApiKeyId,
}); });
const stickySessionKey = buildSurfaceStickySessionKey({
clientContext,
requestedModel,
downstreamPath,
downstreamApiKeyId,
});
const excludeChannelIds: number[] = []; const excludeChannelIds: number[] = [];
let retryCount = 0; let retryCount = 0;
@@ -194,6 +205,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
downstreamPolicy, downstreamPolicy,
excludeChannelIds, excludeChannelIds,
retryCount, retryCount,
stickySessionKey,
}); });
if (!selected) { if (!selected) {
@@ -337,6 +349,36 @@ export async function handleOpenAiResponsesSurfaceRequest(
}; };
const startTime = Date.now(); const startTime = Date.now();
const leaseResult = await acquireSurfaceChannelLease({
selected,
});
if (leaseResult.status === 'timeout') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const busyMessage = buildSurfaceChannelBusyMessage(leaseResult.waitMs);
await failureToolkit.log({
selected,
modelRequested: requestedModel,
status: 'failed',
httpStatus: 503,
latencyMs: leaseResult.waitMs,
errorMessage: busyMessage,
retryCount,
});
retryCount += 1;
if (retryCount <= maxRetries) {
continue;
}
return reply.code(503).send({
error: {
message: busyMessage,
type: 'server_error',
},
});
}
const channelLease = leaseResult.lease;
try { try {
const endpointResult = await executeEndpointFlow({ const endpointResult = await executeEndpointFlow({
@@ -374,6 +416,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
}); });
if (!endpointResult.ok) { if (!endpointResult.ok) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleUpstreamFailure({ const failureOutcome = await failureToolkit.handleUpstreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -466,6 +512,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
); );
const latency = Date.now() - startTime; const latency = Date.now() - startTime;
if (streamResult.status === 'failed') { if (streamResult.status === 'failed') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
await failureToolkit.recordStreamFailure({ await failureToolkit.recordStreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -482,6 +532,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
} }
await finalizeStreamSuccess(parsedUsage, latency); await finalizeStreamSuccess(parsedUsage, latency);
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return; return;
} }
let upstreamData: unknown = rawText; let upstreamData: unknown = rawText;
@@ -498,6 +552,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
const latency = Date.now() - startTime; const latency = Date.now() - startTime;
const failure = detectProxyFailure({ rawText, usage: parsedUsage }); const failure = detectProxyFailure({ rawText, usage: parsedUsage });
if (failure) { if (failure) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleDetectedFailure({ const failureOutcome = await failureToolkit.handleDetectedFailure({
selected, selected,
requestedModel, requestedModel,
@@ -520,6 +578,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
startSseResponse(); startSseResponse();
const streamResult = streamSession.consumeUpstreamFinalPayload(upstreamData, rawText, reply.raw); const streamResult = streamSession.consumeUpstreamFinalPayload(upstreamData, rawText, reply.raw);
if (streamResult.status === 'failed') { if (streamResult.status === 'failed') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
await failureToolkit.recordStreamFailure({ await failureToolkit.recordStreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -537,6 +599,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
} }
await finalizeStreamSuccess(parsedUsage, latency); await finalizeStreamSuccess(parsedUsage, latency);
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return; return;
} }
@@ -570,6 +636,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
const latency = Date.now() - startTime; const latency = Date.now() - startTime;
if (streamResult.status === 'failed') { if (streamResult.status === 'failed') {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
await failureToolkit.recordStreamFailure({ await failureToolkit.recordStreamFailure({
selected, selected,
requestedModel, requestedModel,
@@ -592,6 +662,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
// handled in-band by the proxy stream session. // handled in-band by the proxy stream session.
await finalizeStreamSuccess(parsedUsage, latency); await finalizeStreamSuccess(parsedUsage, latency);
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return; return;
} }
@@ -628,6 +702,10 @@ export async function handleOpenAiResponsesSurfaceRequest(
const parsedUsage = parseProxyUsage(upstreamData); const parsedUsage = parseProxyUsage(upstreamData);
const failure = detectProxyFailure({ rawText, usage: parsedUsage }); const failure = detectProxyFailure({ rawText, usage: parsedUsage });
if (failure) { if (failure) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleDetectedFailure({ const failureOutcome = await failureToolkit.handleDetectedFailure({
selected, selected,
requestedModel, requestedModel,
@@ -678,8 +756,16 @@ export async function handleOpenAiResponsesSurfaceRequest(
} catch (error) { } catch (error) {
console.error('[responses] post-response success logging failed:', error); console.error('[responses] post-response success logging failed:', error);
} }
bindSurfaceStickyChannel({
stickySessionKey,
selected,
});
return reply.send(downstreamData); return reply.send(downstreamData);
} catch (err: any) { } catch (err: any) {
clearSurfaceStickyChannel({
stickySessionKey,
selected,
});
const failureOutcome = await failureToolkit.handleExecutionError({ const failureOutcome = await failureToolkit.handleExecutionError({
selected, selected,
requestedModel, requestedModel,
@@ -693,6 +779,8 @@ export async function handleOpenAiResponsesSurfaceRequest(
continue; continue;
} }
return reply.code(failureOutcome.status).send(failureOutcome.payload); return reply.code(failureOutcome.status).send(failureOutcome.payload);
} finally {
channelLease.release();
} }
} }
} }
@@ -3,6 +3,7 @@ import { EMPTY_DOWNSTREAM_ROUTING_POLICY } from '../../services/downstreamPolicy
const selectChannelMock = vi.fn(); const selectChannelMock = vi.fn();
const selectNextChannelMock = vi.fn(); const selectNextChannelMock = vi.fn();
const selectPreferredChannelMock = vi.fn();
const recordFailureMock = vi.fn(); const recordFailureMock = vi.fn();
const refreshModelsAndRebuildRoutesMock = vi.fn(); const refreshModelsAndRebuildRoutesMock = vi.fn();
const composeProxyLogMessageMock = vi.fn(); const composeProxyLogMessageMock = vi.fn();
@@ -20,6 +21,11 @@ const recordSuccessMock = vi.fn();
const resolveProxyUsageWithSelfLogFallbackMock = vi.fn(); const resolveProxyUsageWithSelfLogFallbackMock = vi.fn();
const resolveProxyLogBillingMock = vi.fn(); const resolveProxyLogBillingMock = vi.fn();
const refreshOauthAccessTokenSingleflightMock = vi.fn(); const refreshOauthAccessTokenSingleflightMock = vi.fn();
const getStickyChannelIdMock = vi.fn();
const bindStickyChannelMock = vi.fn();
const clearStickyChannelMock = vi.fn();
const acquireChannelLeaseMock = vi.fn();
const buildStickySessionKeyMock = vi.fn();
const consoleWarnMock = vi.spyOn(console, 'warn').mockImplementation(() => {}); const consoleWarnMock = vi.spyOn(console, 'warn').mockImplementation(() => {});
const consoleErrorMock = vi.spyOn(console, 'error').mockImplementation(() => {}); const consoleErrorMock = vi.spyOn(console, 'error').mockImplementation(() => {});
@@ -27,11 +33,22 @@ vi.mock('../../services/tokenRouter.js', () => ({
tokenRouter: { tokenRouter: {
selectChannel: (...args: unknown[]) => selectChannelMock(...args), selectChannel: (...args: unknown[]) => selectChannelMock(...args),
selectNextChannel: (...args: unknown[]) => selectNextChannelMock(...args), selectNextChannel: (...args: unknown[]) => selectNextChannelMock(...args),
selectPreferredChannel: (...args: unknown[]) => selectPreferredChannelMock(...args),
recordFailure: (...args: unknown[]) => recordFailureMock(...args), recordFailure: (...args: unknown[]) => recordFailureMock(...args),
recordSuccess: (...args: unknown[]) => recordSuccessMock(...args), recordSuccess: (...args: unknown[]) => recordSuccessMock(...args),
}, },
})); }));
vi.mock('../../services/proxyChannelCoordinator.js', () => ({
proxyChannelCoordinator: {
getStickyChannelId: (...args: unknown[]) => getStickyChannelIdMock(...args),
bindStickyChannel: (...args: unknown[]) => bindStickyChannelMock(...args),
clearStickyChannel: (...args: unknown[]) => clearStickyChannelMock(...args),
acquireChannelLease: (...args: unknown[]) => acquireChannelLeaseMock(...args),
buildStickySessionKey: (...args: unknown[]) => buildStickySessionKeyMock(...args),
},
}));
vi.mock('../../services/modelService.js', () => ({ vi.mock('../../services/modelService.js', () => ({
refreshModelsAndRebuildRoutes: (...args: unknown[]) => refreshModelsAndRebuildRoutesMock(...args), refreshModelsAndRebuildRoutes: (...args: unknown[]) => refreshModelsAndRebuildRoutesMock(...args),
})); }));
@@ -95,6 +112,7 @@ describe('selectSurfaceChannelForAttempt', () => {
beforeEach(() => { beforeEach(() => {
selectChannelMock.mockReset(); selectChannelMock.mockReset();
selectNextChannelMock.mockReset(); selectNextChannelMock.mockReset();
selectPreferredChannelMock.mockReset();
recordFailureMock.mockReset(); recordFailureMock.mockReset();
refreshModelsAndRebuildRoutesMock.mockReset(); refreshModelsAndRebuildRoutesMock.mockReset();
composeProxyLogMessageMock.mockReset(); composeProxyLogMessageMock.mockReset();
@@ -112,6 +130,11 @@ describe('selectSurfaceChannelForAttempt', () => {
resolveProxyUsageWithSelfLogFallbackMock.mockReset(); resolveProxyUsageWithSelfLogFallbackMock.mockReset();
resolveProxyLogBillingMock.mockReset(); resolveProxyLogBillingMock.mockReset();
refreshOauthAccessTokenSingleflightMock.mockReset(); refreshOauthAccessTokenSingleflightMock.mockReset();
getStickyChannelIdMock.mockReset();
bindStickyChannelMock.mockReset();
clearStickyChannelMock.mockReset();
acquireChannelLeaseMock.mockReset();
buildStickySessionKeyMock.mockReset();
consoleWarnMock.mockClear(); consoleWarnMock.mockClear();
consoleErrorMock.mockClear(); consoleErrorMock.mockClear();
}); });
@@ -158,6 +181,51 @@ describe('selectSurfaceChannelForAttempt', () => {
expect(refreshModelsAndRebuildRoutesMock).not.toHaveBeenCalled(); expect(refreshModelsAndRebuildRoutesMock).not.toHaveBeenCalled();
}); });
it('prefers the sticky session channel on the first attempt when it is still eligible', async () => {
const selected = { channel: { id: 55 } };
getStickyChannelIdMock.mockReturnValueOnce(55);
selectPreferredChannelMock.mockResolvedValueOnce(selected);
const { selectSurfaceChannelForAttempt } = await import('./sharedSurface.js');
const result = await selectSurfaceChannelForAttempt({
requestedModel: 'gpt-5.2',
downstreamPolicy: EMPTY_DOWNSTREAM_ROUTING_POLICY,
excludeChannelIds: [],
retryCount: 0,
stickySessionKey: 'sticky-session',
});
expect(result).toBe(selected);
expect(selectPreferredChannelMock).toHaveBeenCalledWith(
'gpt-5.2',
55,
EMPTY_DOWNSTREAM_ROUTING_POLICY,
[],
);
expect(selectChannelMock).not.toHaveBeenCalled();
expect(clearStickyChannelMock).not.toHaveBeenCalled();
});
it('clears stale sticky bindings and falls back to regular selection when the preferred channel is unavailable', async () => {
const selected = { channel: { id: 22 } };
getStickyChannelIdMock.mockReturnValueOnce(55);
selectPreferredChannelMock.mockResolvedValueOnce(null);
selectChannelMock.mockResolvedValueOnce(selected);
const { selectSurfaceChannelForAttempt } = await import('./sharedSurface.js');
const result = await selectSurfaceChannelForAttempt({
requestedModel: 'gpt-5.2',
downstreamPolicy: EMPTY_DOWNSTREAM_ROUTING_POLICY,
excludeChannelIds: [],
retryCount: 0,
stickySessionKey: 'sticky-session',
});
expect(result).toBe(selected);
expect(clearStickyChannelMock).toHaveBeenCalledWith('sticky-session', 55);
expect(selectChannelMock).toHaveBeenCalledWith('gpt-5.2', EMPTY_DOWNSTREAM_ROUTING_POLICY);
});
it('logs refresh failures and still retries selection once on the first attempt', async () => { it('logs refresh failures and still retries selection once on the first attempt', async () => {
const selected = { channel: { id: 33 } }; const selected = { channel: { id: 33 } };
selectChannelMock selectChannelMock
@@ -17,6 +17,7 @@ import type { BuiltEndpointRequest } from '../../routes/proxy/endpointFlow.js';
import { buildUpstreamUrl } from '../../routes/proxy/upstreamUrl.js'; import { buildUpstreamUrl } from '../../routes/proxy/upstreamUrl.js';
import { recordOauthQuotaResetHint } from '../../services/oauth/quota.js'; import { recordOauthQuotaResetHint } from '../../services/oauth/quota.js';
import { refreshOauthAccessTokenSingleflight } from '../../services/oauth/refreshSingleflight.js'; import { refreshOauthAccessTokenSingleflight } from '../../services/oauth/refreshSingleflight.js';
import { proxyChannelCoordinator } from '../../services/proxyChannelCoordinator.js';
type SelectedChannel = Awaited<ReturnType<typeof tokenRouter.selectChannel>>; type SelectedChannel = Awaited<ReturnType<typeof tokenRouter.selectChannel>>;
type SurfaceWarningScope = 'chat' | 'responses'; type SurfaceWarningScope = 'chat' | 'responses';
@@ -103,14 +104,34 @@ export async function selectSurfaceChannelForAttempt(input: {
downstreamPolicy: DownstreamRoutingPolicy; downstreamPolicy: DownstreamRoutingPolicy;
excludeChannelIds: number[]; excludeChannelIds: number[];
retryCount: number; retryCount: number;
stickySessionKey?: string | null;
}): Promise<SelectedChannel> { }): Promise<SelectedChannel> {
let selected = input.retryCount === 0 let selected: SelectedChannel = null;
if (input.retryCount === 0 && input.stickySessionKey) {
const preferredChannelId = proxyChannelCoordinator.getStickyChannelId(input.stickySessionKey);
if (preferredChannelId && !input.excludeChannelIds.includes(preferredChannelId)) {
selected = await tokenRouter.selectPreferredChannel(
input.requestedModel,
preferredChannelId,
input.downstreamPolicy,
input.excludeChannelIds,
);
if (!selected) {
proxyChannelCoordinator.clearStickyChannel(input.stickySessionKey, preferredChannelId);
}
}
}
if (!selected) {
selected = input.retryCount === 0
? await tokenRouter.selectChannel(input.requestedModel, input.downstreamPolicy) ? await tokenRouter.selectChannel(input.requestedModel, input.downstreamPolicy)
: await tokenRouter.selectNextChannel( : await tokenRouter.selectNextChannel(
input.requestedModel, input.requestedModel,
input.excludeChannelIds, input.excludeChannelIds,
input.downstreamPolicy, input.downstreamPolicy,
); );
}
if (!selected && input.retryCount === 0) { if (!selected && input.retryCount === 0) {
try { try {
@@ -124,6 +145,65 @@ export async function selectSurfaceChannelForAttempt(input: {
return selected; return selected;
} }
export function buildSurfaceStickySessionKey(input: {
clientContext?: DownstreamClientContext | null;
requestedModel: string;
downstreamPath: string;
downstreamApiKeyId?: number | null;
}): string | null {
return proxyChannelCoordinator.buildStickySessionKey({
clientKind: input.clientContext?.clientKind || null,
sessionId: input.clientContext?.sessionId || null,
requestedModel: input.requestedModel,
downstreamPath: input.downstreamPath,
downstreamApiKeyId: input.downstreamApiKeyId,
});
}
export function bindSurfaceStickyChannel(input: {
stickySessionKey?: string | null;
selected: {
channel: { id: number };
account?: { extraConfig?: string | null } | null;
};
}): void {
proxyChannelCoordinator.bindStickyChannel(
input.stickySessionKey,
input.selected.channel.id,
input.selected.account?.extraConfig,
);
}
export function clearSurfaceStickyChannel(input: {
stickySessionKey?: string | null;
selected: {
channel: { id: number };
};
}): void {
proxyChannelCoordinator.clearStickyChannel(
input.stickySessionKey,
input.selected.channel.id,
);
}
export async function acquireSurfaceChannelLease(input: {
selected: {
channel: { id: number };
account?: { extraConfig?: string | null } | null;
};
}) {
return await proxyChannelCoordinator.acquireChannelLease({
channelId: input.selected.channel.id,
accountExtraConfig: input.selected.account?.extraConfig,
});
}
export function buildSurfaceChannelBusyMessage(waitMs: number): string {
return waitMs > 0
? `Channel busy: waited ${waitMs}ms for an available session slot`
: 'Channel busy: no session slot available';
}
export async function writeSurfaceProxyLog(input: { export async function writeSurfaceProxyLog(input: {
warningScope: string; warningScope: string;
selected: { selected: {
@@ -298,7 +298,6 @@ describe('executeEndpointFlow', () => {
expect(fetchMock.mock.calls[0]?.[0]).toBe('https://proxy.internal/base/v1/responses'); expect(fetchMock.mock.calls[0]?.[0]).toBe('https://proxy.internal/base/v1/responses');
}); });
it('normalizes proxyUrl with versioned base paths instead of duplicating path segments', async () => { it('normalizes proxyUrl with versioned base paths instead of duplicating path segments', async () => {
fetchMock.mockResolvedValueOnce(toUndiciResponse(new Response(JSON.stringify({ ok: true }), { fetchMock.mockResolvedValueOnce(toUndiciResponse(new Response(JSON.stringify({ ok: true }), {
status: 200, status: 200,
@@ -0,0 +1,149 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { config } from '../config.js';
import {
proxyChannelCoordinator,
resetProxyChannelCoordinatorState,
} from './proxyChannelCoordinator.js';
describe('proxyChannelCoordinator', () => {
const originalStickyEnabled = config.proxyStickySessionEnabled;
const originalStickyTtlMs = config.proxyStickySessionTtlMs;
const originalConcurrencyLimit = config.proxySessionChannelConcurrencyLimit;
const originalQueueWaitMs = config.proxySessionChannelQueueWaitMs;
const originalLeaseTtlMs = config.proxySessionChannelLeaseTtlMs;
const originalLeaseKeepaliveMs = config.proxySessionChannelLeaseKeepaliveMs;
beforeEach(() => {
vi.useFakeTimers();
config.proxyStickySessionEnabled = true;
config.proxyStickySessionTtlMs = 31_000;
config.proxySessionChannelConcurrencyLimit = 1;
config.proxySessionChannelQueueWaitMs = 200;
config.proxySessionChannelLeaseTtlMs = 100;
config.proxySessionChannelLeaseKeepaliveMs = 30;
resetProxyChannelCoordinatorState();
});
afterEach(() => {
config.proxyStickySessionEnabled = originalStickyEnabled;
config.proxyStickySessionTtlMs = originalStickyTtlMs;
config.proxySessionChannelConcurrencyLimit = originalConcurrencyLimit;
config.proxySessionChannelQueueWaitMs = originalQueueWaitMs;
config.proxySessionChannelLeaseTtlMs = originalLeaseTtlMs;
config.proxySessionChannelLeaseKeepaliveMs = originalLeaseKeepaliveMs;
resetProxyChannelCoordinatorState();
vi.useRealTimers();
});
it('stores sticky bindings for session-scoped channels and expires them by ttl', async () => {
const key = proxyChannelCoordinator.buildStickySessionKey({
clientKind: 'codex',
sessionId: 'turn-123',
requestedModel: 'gpt-5.2',
downstreamPath: '/v1/responses',
downstreamApiKeyId: 9,
});
proxyChannelCoordinator.bindStickyChannel(key, 42, JSON.stringify({ credentialMode: 'session' }));
expect(proxyChannelCoordinator.getStickyChannelId(key)).toBe(42);
await vi.advanceTimersByTimeAsync(31_100);
expect(proxyChannelCoordinator.getStickyChannelId(key)).toBeNull();
});
it('does not store sticky bindings for apikey-only channels', () => {
const key = proxyChannelCoordinator.buildStickySessionKey({
clientKind: 'codex',
sessionId: 'turn-456',
requestedModel: 'gpt-5.2',
downstreamPath: '/v1/responses',
downstreamApiKeyId: 9,
});
proxyChannelCoordinator.bindStickyChannel(key, 42, JSON.stringify({ credentialMode: 'apikey' }));
expect(proxyChannelCoordinator.getStickyChannelId(key)).toBeNull();
});
it('queues requests behind the active lease and grants the next waiter after release', async () => {
const first = await proxyChannelCoordinator.acquireChannelLease({
channelId: 11,
accountExtraConfig: JSON.stringify({ credentialMode: 'session' }),
});
expect(first.status).toBe('acquired');
if (first.status !== 'acquired') return;
let secondSettled = false;
const secondPromise = proxyChannelCoordinator.acquireChannelLease({
channelId: 11,
accountExtraConfig: JSON.stringify({ credentialMode: 'session' }),
}).then((result) => {
secondSettled = true;
return result;
});
await vi.advanceTimersByTimeAsync(50);
expect(secondSettled).toBe(false);
first.lease.release();
await vi.advanceTimersByTimeAsync(0);
const second = await secondPromise;
expect(second.status).toBe('acquired');
if (second.status === 'acquired') {
second.lease.release();
}
});
it('times out queued requests when no slot becomes available', async () => {
const first = await proxyChannelCoordinator.acquireChannelLease({
channelId: 11,
accountExtraConfig: JSON.stringify({ credentialMode: 'session' }),
});
expect(first.status).toBe('acquired');
if (first.status !== 'acquired') return;
const secondPromise = proxyChannelCoordinator.acquireChannelLease({
channelId: 11,
accountExtraConfig: JSON.stringify({ credentialMode: 'session' }),
});
await vi.advanceTimersByTimeAsync(250);
await expect(secondPromise).resolves.toEqual({
status: 'timeout',
waitMs: 200,
});
first.lease.release();
});
it('keeps active leases alive until they are explicitly released', async () => {
const first = await proxyChannelCoordinator.acquireChannelLease({
channelId: 11,
accountExtraConfig: JSON.stringify({ credentialMode: 'session' }),
});
expect(first.status).toBe('acquired');
if (first.status !== 'acquired') return;
let secondSettled = false;
const secondPromise = proxyChannelCoordinator.acquireChannelLease({
channelId: 11,
accountExtraConfig: JSON.stringify({ credentialMode: 'session' }),
}).then((result) => {
secondSettled = true;
return result;
});
await vi.advanceTimersByTimeAsync(180);
expect(first.lease.isActive()).toBe(true);
expect(secondSettled).toBe(false);
first.lease.release();
await vi.advanceTimersByTimeAsync(0);
const second = await secondPromise;
expect(second.status).toBe('acquired');
if (second.status === 'acquired') {
second.lease.release();
}
});
});
@@ -0,0 +1,295 @@
import { config } from '../config.js';
import {
getCredentialModeFromExtraConfig,
hasOauthProvider,
} from './accountExtraConfig.js';
type StickyEntry = {
channelId: number;
expiresAtMs: number;
};
type ActiveLeaseState = {
release: () => void;
};
type ChannelWaiter = {
cancelled: boolean;
resolve: (result: AcquireProxyChannelLeaseResult) => void;
timer: ReturnType<typeof setTimeout> | null;
};
type ChannelRuntimeState = {
activeLeaseIds: Set<number>;
queue: ChannelWaiter[];
};
export type ProxyChannelLease = {
channelId: number;
isActive(): boolean;
release(): void;
touch(): void;
};
export type AcquireProxyChannelLeaseResult =
| { status: 'acquired'; lease: ProxyChannelLease }
| { status: 'timeout'; waitMs: number };
const stickySessionBindings = new Map<string, StickyEntry>();
const channelRuntimeStates = new Map<number, ChannelRuntimeState>();
let nextLeaseId = 1;
function shouldUnrefTimer(timer: ReturnType<typeof setTimeout> | ReturnType<typeof setInterval>) {
if (typeof (timer as { unref?: () => void }).unref === 'function') {
(timer as { unref: () => void }).unref();
}
}
function cleanupExpiredStickyBindings(nowMs = Date.now()): void {
for (const [key, entry] of stickySessionBindings.entries()) {
if (entry.expiresAtMs <= nowMs) {
stickySessionBindings.delete(key);
}
}
}
function isSessionScopedChannel(extraConfig?: string | null): boolean {
return getCredentialModeFromExtraConfig(extraConfig) === 'session'
|| hasOauthProvider(extraConfig);
}
function getStickySessionTtlMs(): number {
return Math.max(30_000, Math.trunc(config.proxyStickySessionTtlMs || 0));
}
function getChannelLeaseTtlMs(): number {
return Math.max(5_000, Math.trunc(config.proxySessionChannelLeaseTtlMs || 0));
}
function getChannelLeaseKeepaliveMs(): number {
return Math.max(1_000, Math.trunc(config.proxySessionChannelLeaseKeepaliveMs || 0));
}
function getChannelQueueWaitMs(): number {
return Math.max(0, Math.trunc(config.proxySessionChannelQueueWaitMs || 0));
}
function getChannelConcurrencyLimit(extraConfig?: string | null): number {
if (!isSessionScopedChannel(extraConfig)) return 0;
return Math.max(0, Math.trunc(config.proxySessionChannelConcurrencyLimit || 0));
}
function getOrCreateChannelRuntimeState(channelId: number): ChannelRuntimeState {
let state = channelRuntimeStates.get(channelId);
if (!state) {
state = {
activeLeaseIds: new Set<number>(),
queue: [],
};
channelRuntimeStates.set(channelId, state);
}
return state;
}
function maybeDeleteChannelRuntimeState(channelId: number): void {
const state = channelRuntimeStates.get(channelId);
if (!state) return;
if (state.activeLeaseIds.size <= 0 && state.queue.every((waiter) => waiter.cancelled)) {
channelRuntimeStates.delete(channelId);
}
}
function createNoopLease(channelId: number): ProxyChannelLease {
return {
channelId,
isActive: () => false,
release: () => {},
touch: () => {},
};
}
class ProxyChannelCoordinator {
buildStickySessionKey(input: {
clientKind?: string | null;
sessionId?: string | null;
requestedModel: string;
downstreamPath: string;
downstreamApiKeyId?: number | null;
}): string | null {
if (!config.proxyStickySessionEnabled) return null;
const sessionId = String(input.sessionId || '').trim();
if (!sessionId) return null;
const requestedModel = String(input.requestedModel || '').trim().toLowerCase();
if (!requestedModel) return null;
const downstreamPath = String(input.downstreamPath || '').trim().toLowerCase() || 'unknown';
const clientKind = String(input.clientKind || 'generic').trim().toLowerCase() || 'generic';
const owner = typeof input.downstreamApiKeyId === 'number' && Number.isFinite(input.downstreamApiKeyId)
? `key:${Math.trunc(input.downstreamApiKeyId)}`
: 'key:anonymous';
return [owner, clientKind, downstreamPath, requestedModel, sessionId].join('|');
}
getStickyChannelId(stickySessionKey?: string | null, nowMs = Date.now()): number | null {
cleanupExpiredStickyBindings(nowMs);
const normalizedKey = String(stickySessionKey || '').trim();
if (!normalizedKey) return null;
const entry = stickySessionBindings.get(normalizedKey);
if (!entry || entry.expiresAtMs <= nowMs) {
stickySessionBindings.delete(normalizedKey);
return null;
}
return entry.channelId;
}
bindStickyChannel(stickySessionKey: string | null | undefined, channelId: number, extraConfig?: string | null): void {
if (!config.proxyStickySessionEnabled) return;
if (!isSessionScopedChannel(extraConfig)) return;
const normalizedKey = String(stickySessionKey || '').trim();
if (!normalizedKey || !Number.isFinite(channelId) || channelId <= 0) return;
cleanupExpiredStickyBindings();
stickySessionBindings.set(normalizedKey, {
channelId: Math.trunc(channelId),
expiresAtMs: Date.now() + getStickySessionTtlMs(),
});
}
clearStickyChannel(stickySessionKey?: string | null, channelId?: number | null): void {
const normalizedKey = String(stickySessionKey || '').trim();
if (!normalizedKey) return;
const existing = stickySessionBindings.get(normalizedKey);
if (!existing) return;
if (typeof channelId === 'number' && Number.isFinite(channelId) && existing.channelId !== Math.trunc(channelId)) {
return;
}
stickySessionBindings.delete(normalizedKey);
}
async acquireChannelLease(input: {
channelId: number;
accountExtraConfig?: string | null;
}): Promise<AcquireProxyChannelLeaseResult> {
const channelId = Math.trunc(input.channelId || 0);
if (channelId <= 0) {
return {
status: 'acquired',
lease: createNoopLease(0),
};
}
const concurrencyLimit = getChannelConcurrencyLimit(input.accountExtraConfig);
if (concurrencyLimit <= 0) {
return {
status: 'acquired',
lease: createNoopLease(channelId),
};
}
const state = getOrCreateChannelRuntimeState(channelId);
if (state.activeLeaseIds.size < concurrencyLimit) {
return {
status: 'acquired',
lease: this.createTrackedLease(channelId, state),
};
}
const waitMs = getChannelQueueWaitMs();
if (waitMs <= 0) {
return {
status: 'timeout',
waitMs: 0,
};
}
return await new Promise<AcquireProxyChannelLeaseResult>((resolve) => {
const waiter: ChannelWaiter = {
cancelled: false,
resolve,
timer: null,
};
waiter.timer = setTimeout(() => {
waiter.cancelled = true;
waiter.timer = null;
maybeDeleteChannelRuntimeState(channelId);
resolve({
status: 'timeout',
waitMs,
});
}, waitMs);
shouldUnrefTimer(waiter.timer);
state.queue.push(waiter);
});
}
private createTrackedLease(channelId: number, state: ChannelRuntimeState): ProxyChannelLease {
const leaseId = nextLeaseId++;
state.activeLeaseIds.add(leaseId);
let released = false;
let expiryTimer: ReturnType<typeof setTimeout> | null = null;
let keepaliveTimer: ReturnType<typeof setInterval> | null = null;
const release = () => {
if (released) return;
released = true;
if (expiryTimer) clearTimeout(expiryTimer);
if (keepaliveTimer) clearInterval(keepaliveTimer);
state.activeLeaseIds.delete(leaseId);
this.drainQueue(channelId);
maybeDeleteChannelRuntimeState(channelId);
};
const touch = () => {
if (released) return;
if (expiryTimer) clearTimeout(expiryTimer);
expiryTimer = setTimeout(() => {
release();
}, getChannelLeaseTtlMs());
shouldUnrefTimer(expiryTimer);
};
touch();
const keepaliveMs = getChannelLeaseKeepaliveMs();
if (keepaliveMs > 0) {
keepaliveTimer = setInterval(() => {
touch();
}, keepaliveMs);
shouldUnrefTimer(keepaliveTimer);
}
return {
channelId,
isActive: () => !released,
release,
touch,
};
}
private drainQueue(channelId: number): void {
const state = channelRuntimeStates.get(channelId);
if (!state) return;
const concurrencyLimit = Math.max(0, Math.trunc(config.proxySessionChannelConcurrencyLimit || 0));
while (state.activeLeaseIds.size < concurrencyLimit && state.queue.length > 0) {
const waiter = state.queue.shift();
if (!waiter || waiter.cancelled) continue;
if (waiter.timer) clearTimeout(waiter.timer);
waiter.timer = null;
waiter.resolve({
status: 'acquired',
lease: this.createTrackedLease(channelId, state),
});
}
}
}
export function resetProxyChannelCoordinatorState(): void {
stickySessionBindings.clear();
channelRuntimeStates.clear();
nextLeaseId = 1;
}
export function isProxyChannelSessionScoped(extraConfig?: string | null): boolean {
return isSessionScopedChannel(extraConfig);
}
export const proxyChannelCoordinator = new ProxyChannelCoordinator();
@@ -117,6 +117,54 @@ describe('TokenRouter selection scoring', () => {
}).returning().get(); }).returning().get();
} }
it('reuses a preferred channel only while it remains healthy', async () => {
config.routingWeights = {
baseWeightFactor: 1,
valueScoreFactor: 0,
costWeight: 0,
balanceWeight: 0,
usageWeight: 0,
};
const route = await createRoute('gpt-5.2');
const site = await createSite('sticky-site');
const account = await createAccount(site.id, 'sticky-user');
const tokenA = await createToken(account.id, 'sticky-a');
const tokenB = await createToken(account.id, 'sticky-b');
const preferredChannel = await db.insert(schema.routeChannels).values({
routeId: route.id,
accountId: account.id,
tokenId: tokenA.id,
priority: 0,
weight: 10,
enabled: true,
failCount: 0,
}).returning().get();
await db.insert(schema.routeChannels).values({
routeId: route.id,
accountId: account.id,
tokenId: tokenB.id,
priority: 0,
weight: 10,
enabled: true,
failCount: 0,
}).run();
const router = new TokenRouter();
const selected = await router.selectPreferredChannel('gpt-5.2', preferredChannel.id);
expect(selected?.channel.id).toBe(preferredChannel.id);
await db.update(schema.routeChannels).set({
failCount: 4,
lastFailAt: new Date().toISOString(),
}).where(eq(schema.routeChannels.id, preferredChannel.id)).run();
invalidateTokenRouterCache();
await expect(router.selectPreferredChannel('gpt-5.2', preferredChannel.id)).resolves.toBeNull();
});
it('normalizes probability across channels on the same site', async () => { it('normalizes probability across channels on the same site', async () => {
config.routingWeights = { config.routingWeights = {
baseWeightFactor: 1, baseWeightFactor: 1,
+80
View File
@@ -1250,6 +1250,27 @@ export class TokenRouter {
return await this.selectFromMatch(match, requestedModel, downstreamPolicy, excludeChannelIds); return await this.selectFromMatch(match, requestedModel, downstreamPolicy, excludeChannelIds);
} }
async selectPreferredChannel(
requestedModel: string,
preferredChannelId: number,
downstreamPolicy: DownstreamRoutingPolicy = DEFAULT_DOWNSTREAM_POLICY,
excludeChannelIds: number[] = [],
): Promise<SelectedChannel | null> {
if (!isModelAllowedByDownstreamPolicy(requestedModel, downstreamPolicy)) return null;
const normalizedPreferredChannelId = Math.trunc(preferredChannelId || 0);
if (normalizedPreferredChannelId <= 0) return null;
await ensureSiteRuntimeHealthStateLoaded();
const match = await this.findRoute(requestedModel, downstreamPolicy);
if (!match) return null;
return await this.selectPreferredFromMatch(
match,
requestedModel,
normalizedPreferredChannelId,
excludeChannelIds,
);
}
async explainSelection( async explainSelection(
requestedModel: string, requestedModel: string,
excludeChannelIds: number[] = [], excludeChannelIds: number[] = [],
@@ -1843,6 +1864,65 @@ export class TokenRouter {
return null; return null;
} }
private async selectPreferredFromMatch(
match: RouteMatch,
requestedModel: string,
preferredChannelId: number,
excludeChannelIds: number[] = [],
recordSelection = true,
): Promise<SelectedChannel | null> {
const mappedModel = resolveMappedModel(requestedModel, match.route.modelMapping);
const requestedByDisplayName = isRouteDisplayNameMatch(requestedModel, match.route.displayName);
const bypassSourceModelCheck = requestedByDisplayName;
const routeStrategy = resolveRouteStrategy(match.route);
const runtimeModelResolver = requestedByDisplayName
? ((candidate: RouteChannelCandidate) => normalizeChannelSourceModel(candidate.channel.sourceModel) || mappedModel)
: mappedModel;
const nowIso = new Date().toISOString();
const nowMs = Date.now();
const available = match.channels.filter((candidate) => (
this.getCandidateEligibilityReasons(candidate, {
requestedModel,
bypassSourceModelCheck,
excludeChannelIds,
nowIso,
}).length === 0
));
const preferred = available.find((candidate) => candidate.channel.id === preferredChannelId);
if (!preferred) return null;
const breakerFiltered = filterSiteRuntimeBrokenCandidatesByModel([preferred], runtimeModelResolver, nowMs);
if (breakerFiltered.candidates.length <= 0) return null;
const selected = breakerFiltered.candidates.find((candidate) => candidate.channel.id === preferredChannelId);
if (!selected) return null;
if (routeStrategy !== 'round_robin' && isChannelRecentlyFailed(selected.channel, nowMs)) {
return null;
}
const tokenValue = this.resolveChannelTokenValue(selected);
if (!tokenValue) return null;
if (recordSelection && (routeStrategy === 'round_robin' || routeStrategy === 'stable_first')) {
await this.recordChannelSelection(selected.channel.id);
}
const actualModel = resolveActualModelForSelectedChannel(
requestedModel,
match.route,
mappedModel,
selected.channel.sourceModel,
);
return {
...selected,
tokenValue,
tokenName: selected.token?.name || 'default',
actualModel,
};
}
private async findRoute(model: string, downstreamPolicy: DownstreamRoutingPolicy): Promise<RouteMatch | null> { private async findRoute(model: string, downstreamPolicy: DownstreamRoutingPolicy): Promise<RouteMatch | null> {
let routes = await loadEnabledRoutes(); let routes = await loadEnabledRoutes();