From 9cf3dd19041b4e57476cad8a95d17ab7659ec0ba Mon Sep 17 00:00:00 2001 From: Bainily Date: Tue, 24 Mar 2026 19:40:44 +0800 Subject: [PATCH] =?UTF-8?q?feat(proxy):=20improve=20channel=20failover=20s?= =?UTF-8?q?tability=20and=20sticky=20session=20coordination=20/=20?= =?UTF-8?q?=E4=BC=98=E5=8C=96=E4=BB=A3=E7=90=86=E9=80=9A=E9=81=93=E6=95=85?= =?UTF-8?q?=E9=9A=9C=E5=88=87=E6=8D=A2=E4=B8=8E=E7=B2=98=E6=BB=9E=E4=BC=9A?= =?UTF-8?q?=E8=AF=9D=E5=8D=8F=E8=B0=83=20(#268)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 --- src/server/config.ts | 6 + src/server/proxy-core/surfaces/chatSurface.ts | 151 ++++++++- .../surfaces/openAiResponsesSurface.ts | 218 +++++++++---- .../proxy-core/surfaces/sharedSurface.test.ts | 68 ++++ .../proxy-core/surfaces/sharedSurface.ts | 94 +++++- src/server/routes/proxy/endpointFlow.test.ts | 1 - .../services/proxyChannelCoordinator.test.ts | 149 +++++++++ .../services/proxyChannelCoordinator.ts | 295 ++++++++++++++++++ .../services/tokenRouter.selection.test.ts | 48 +++ src/server/services/tokenRouter.ts | 80 +++++ 10 files changed, 1031 insertions(+), 79 deletions(-) create mode 100644 src/server/services/proxyChannelCoordinator.test.ts create mode 100644 src/server/services/proxyChannelCoordinator.ts diff --git a/src/server/config.ts b/src/server/config.ts index 9416ba3..5e5cb7c 100644 --- a/src/server/config.ts +++ b/src/server/config.ts @@ -109,6 +109,12 @@ export function buildConfig(env: NodeJS.ProcessEnv) { 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))), 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))), 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))), diff --git a/src/server/proxy-core/surfaces/chatSurface.ts b/src/server/proxy-core/surfaces/chatSurface.ts index dbffe83..800713f 100644 --- a/src/server/proxy-core/surfaces/chatSurface.ts +++ b/src/server/proxy-core/surfaces/chatSurface.ts @@ -43,6 +43,11 @@ import { summarizeConversationFileInputsInOpenAiBody } from '../capabilities/con import { detectDownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js'; import { canRetryProxyChannel, getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js'; import { + acquireSurfaceChannelLease, + bindSurfaceStickyChannel, + buildSurfaceChannelBusyMessage, + buildSurfaceStickySessionKey, + clearSurfaceStickyChannel, createSurfaceFailureToolkit, createSurfaceDispatchRequest, recordSurfaceSuccess, @@ -126,6 +131,12 @@ export async function handleChatSurfaceRequest( clientContext, downstreamApiKeyId, }); + const stickySessionKey = buildSurfaceStickySessionKey({ + clientContext, + requestedModel, + downstreamPath, + downstreamApiKeyId, + }); const excludeChannelIds: number[] = []; let retryCount = 0; @@ -136,6 +147,7 @@ export async function handleChatSurfaceRequest( downstreamPolicy, excludeChannelIds, retryCount, + stickySessionKey, }); if (!selected) { @@ -250,6 +262,36 @@ export async function handleChatSurfaceRequest( return endpointStrategy.tryRecover(ctx); }; 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 { const endpointResult = await executeEndpointFlow({ @@ -287,6 +329,10 @@ export async function handleChatSurfaceRequest( }); if (!endpointResult.ok) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); const failureOutcome = await failureToolkit.handleUpstreamFailure({ selected, requestedModel, @@ -358,6 +404,10 @@ export async function handleChatSurfaceRequest( ); const latency = Date.now() - startTime; if (streamResult.status === 'failed') { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); await failureToolkit.recordStreamFailure({ selected, requestedModel, @@ -372,6 +422,10 @@ export async function handleChatSurfaceRequest( }); return; } + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); return; } let fallbackData: unknown = null; @@ -387,6 +441,10 @@ export async function handleChatSurfaceRequest( const latency = Date.now() - startTime; const failure = detectProxyFailure({ rawText, usage: parsedUsage }); if (failure) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); const failureOutcome = await failureToolkit.handleDetectedFailure({ selected, requestedModel, @@ -409,6 +467,10 @@ export async function handleChatSurfaceRequest( startSseResponse(); const streamResult = streamSession.consumeUpstreamFinalPayload(fallbackData, fallbackText, reply.raw); if (streamResult.status === 'failed') { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); await failureToolkit.recordStreamFailure({ selected, requestedModel, @@ -424,6 +486,10 @@ export async function handleChatSurfaceRequest( }); return; } + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); } else { startSseResponse(); const upstreamReader = upstream.body?.getReader(); @@ -453,6 +519,10 @@ export async function handleChatSurfaceRequest( const latency = Date.now() - startTime; if (streamResult.status === 'failed') { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); await failureToolkit.recordStreamFailure({ selected, requestedModel, @@ -493,6 +563,10 @@ export async function handleChatSurfaceRequest( errorLabel: '[proxy/chat] failed to record success metrics', }, }); + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); return; } @@ -524,6 +598,10 @@ export async function handleChatSurfaceRequest( const parsedUsage = parseProxyUsage(upstreamData); const failure = detectProxyFailure({ rawText, usage: parsedUsage }); if (failure) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); const failureOutcome = await failureToolkit.handleDetectedFailure({ selected, requestedModel, @@ -562,12 +640,20 @@ export async function handleChatSurfaceRequest( errorLabel: '[proxy/chat] failed to record success metrics', }, }); + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); return reply.send(downstreamResponse); - } catch (err: any) { - const failureOutcome = await failureToolkit.handleExecutionError({ - selected, - requestedModel, + } catch (err: any) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + const failureOutcome = await failureToolkit.handleExecutionError({ + selected, + requestedModel, modelName, errorMessage: err?.message || 'network failure', latencyMs: Date.now() - startTime, @@ -576,10 +662,12 @@ export async function handleChatSurfaceRequest( if (failureOutcome.action === 'retry') { retryCount += 1; continue; + } + return reply.code(failureOutcome.status).send(failureOutcome.payload); + } finally { + channelLease.release(); } - return reply.code(failureOutcome.status).send(failureOutcome.payload); } - } } function deriveCodexSessionCacheKey(input: { @@ -646,6 +734,12 @@ export async function handleClaudeCountTokensSurfaceRequest( clientContext, downstreamApiKeyId, }); + const stickySessionKey = buildSurfaceStickySessionKey({ + clientContext, + requestedModel, + downstreamPath, + downstreamApiKeyId, + }); const excludeChannelIds: number[] = []; let retryCount = 0; @@ -655,6 +749,7 @@ export async function handleClaudeCountTokensSurfaceRequest( downstreamPolicy, excludeChannelIds, retryCount, + stickySessionKey, }); if (!selected) { @@ -692,6 +787,36 @@ export async function handleClaudeCountTokensSurfaceRequest( } const oauth = getOauthInfoFromAccount(selected.account); 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 upstreamRequest = buildClaudeCountTokensUpstreamRequest({ @@ -753,6 +878,10 @@ export async function handleClaudeCountTokensSurfaceRequest( } if (!upstream.ok) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); const failureOutcome = await failureToolkit.handleUpstreamFailure({ selected, requestedModel, @@ -782,8 +911,16 @@ export async function handleClaudeCountTokensSurfaceRequest( retryCount, upstreamPath: upstreamRequest.path, }); + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); return reply.code(upstream.status).type(contentType).send(payload); } catch (error: any) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); const failureOutcome = await failureToolkit.handleExecutionError({ selected, requestedModel, @@ -797,6 +934,8 @@ export async function handleClaudeCountTokensSurfaceRequest( continue; } return reply.code(failureOutcome.status).send(failureOutcome.payload); + } finally { + channelLease.release(); } } } diff --git a/src/server/proxy-core/surfaces/openAiResponsesSurface.ts b/src/server/proxy-core/surfaces/openAiResponsesSurface.ts index a716833..a8f7ba9 100644 --- a/src/server/proxy-core/surfaces/openAiResponsesSurface.ts +++ b/src/server/proxy-core/surfaces/openAiResponsesSurface.ts @@ -40,6 +40,11 @@ import { import { detectDownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js'; import { getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js'; import { + acquireSurfaceChannelLease, + bindSurfaceStickyChannel, + buildSurfaceChannelBusyMessage, + buildSurfaceStickySessionKey, + clearSurfaceStickyChannel, createSurfaceFailureToolkit, createSurfaceDispatchRequest, recordSurfaceSuccess, @@ -178,23 +183,30 @@ export async function handleOpenAiResponsesSurfaceRequest( const downstreamPolicy = getDownstreamRoutingPolicy(request); const downstreamApiKeyId = getProxyAuthContext(request)?.keyId ?? null; const maxRetries = getProxyMaxChannelRetries(); - const failureToolkit = createSurfaceFailureToolkit({ - warningScope: 'responses', - downstreamPath, - maxRetries, - clientContext, - downstreamApiKeyId, - }); - const excludeChannelIds: number[] = []; - let retryCount = 0; + const failureToolkit = createSurfaceFailureToolkit({ + warningScope: 'responses', + downstreamPath, + maxRetries, + clientContext, + downstreamApiKeyId, + }); + const stickySessionKey = buildSurfaceStickySessionKey({ + clientContext, + requestedModel, + downstreamPath, + downstreamApiKeyId, + }); + const excludeChannelIds: number[] = []; + let retryCount = 0; while (retryCount <= maxRetries) { - const selected = await selectSurfaceChannelForAttempt({ - requestedModel, - downstreamPolicy, - excludeChannelIds, - retryCount, - }); + const selected = await selectSurfaceChannelForAttempt({ + requestedModel, + downstreamPolicy, + excludeChannelIds, + retryCount, + stickySessionKey, + }); if (!selected) { await reportProxyAllFailed({ @@ -336,9 +348,39 @@ export async function handleOpenAiResponsesSurfaceRequest( return endpointStrategy.tryRecover(ctx); }; - 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({ siteUrl: selected.site.url, endpointCandidates, @@ -373,10 +415,14 @@ export async function handleOpenAiResponsesSurfaceRequest( }, }); - if (!endpointResult.ok) { - const failureOutcome = await failureToolkit.handleUpstreamFailure({ - selected, - requestedModel, + if (!endpointResult.ok) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + const failureOutcome = await failureToolkit.handleUpstreamFailure({ + selected, + requestedModel, modelName, status: endpointResult.status || 502, errText: endpointResult.errText || 'unknown error', @@ -465,10 +511,14 @@ export async function handleOpenAiResponsesSurfaceRequest( reply.raw, ); const latency = Date.now() - startTime; - if (streamResult.status === 'failed') { - await failureToolkit.recordStreamFailure({ - selected, - requestedModel, + if (streamResult.status === 'failed') { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + await failureToolkit.recordStreamFailure({ + selected, + requestedModel, modelName, errorMessage: streamResult.errorMessage, latencyMs: latency, @@ -479,11 +529,15 @@ export async function handleOpenAiResponsesSurfaceRequest( upstreamPath: successfulUpstreamPath, }); return; - } + } - await finalizeStreamSuccess(parsedUsage, latency); - return; - } + await finalizeStreamSuccess(parsedUsage, latency); + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + return; + } let upstreamData: unknown = rawText; try { upstreamData = JSON.parse(rawText); @@ -497,10 +551,14 @@ export async function handleOpenAiResponsesSurfaceRequest( parsedUsage = parseProxyUsage(upstreamData); const latency = Date.now() - startTime; const failure = detectProxyFailure({ rawText, usage: parsedUsage }); - if (failure) { - const failureOutcome = await failureToolkit.handleDetectedFailure({ - selected, - requestedModel, + if (failure) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + const failureOutcome = await failureToolkit.handleDetectedFailure({ + selected, + requestedModel, modelName, failure, latencyMs: latency, @@ -519,10 +577,14 @@ export async function handleOpenAiResponsesSurfaceRequest( startSseResponse(); const streamResult = streamSession.consumeUpstreamFinalPayload(upstreamData, rawText, reply.raw); - if (streamResult.status === 'failed') { - await failureToolkit.recordStreamFailure({ - selected, - requestedModel, + if (streamResult.status === 'failed') { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + await failureToolkit.recordStreamFailure({ + selected, + requestedModel, modelName, errorMessage: streamResult.errorMessage, latencyMs: latency, @@ -534,11 +596,15 @@ export async function handleOpenAiResponsesSurfaceRequest( runtimeFailureStatus: 502, }); return; - } + } - await finalizeStreamSuccess(parsedUsage, latency); - return; - } + await finalizeStreamSuccess(parsedUsage, latency); + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + return; + } startSseResponse(); @@ -569,10 +635,14 @@ export async function handleOpenAiResponsesSurfaceRequest( rawText += decoder.decode(); const latency = Date.now() - startTime; - if (streamResult.status === 'failed') { - await failureToolkit.recordStreamFailure({ - selected, - requestedModel, + if (streamResult.status === 'failed') { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + await failureToolkit.recordStreamFailure({ + selected, + requestedModel, modelName, errorMessage: streamResult.errorMessage, latencyMs: latency, @@ -589,11 +659,15 @@ export async function handleOpenAiResponsesSurfaceRequest( // Once SSE has been hijacked and bytes may already be on the wire, we // must not attempt to convert stream failures into a fresh HTTP error // response or retry on another channel. Responses stream failures are - // handled in-band by the proxy stream session. + // handled in-band by the proxy stream session. - await finalizeStreamSuccess(parsedUsage, latency); - return; - } + await finalizeStreamSuccess(parsedUsage, latency); + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + return; + } const upstreamContentType = (upstream.headers.get('content-type') || '').toLowerCase(); let rawText = ''; @@ -627,10 +701,14 @@ export async function handleOpenAiResponsesSurfaceRequest( const latency = Date.now() - startTime; const parsedUsage = parseProxyUsage(upstreamData); const failure = detectProxyFailure({ rawText, usage: parsedUsage }); - if (failure) { - const failureOutcome = await failureToolkit.handleDetectedFailure({ - selected, - requestedModel, + if (failure) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + const failureOutcome = await failureToolkit.handleDetectedFailure({ + selected, + requestedModel, modelName, failure, latencyMs: latency, @@ -675,14 +753,22 @@ export async function handleOpenAiResponsesSurfaceRequest( errorLabel: '[responses] post-response bookkeeping failed:', }, }); - } catch (error) { - console.error('[responses] post-response success logging failed:', error); - } - return reply.send(downstreamData); - } catch (err: any) { - const failureOutcome = await failureToolkit.handleExecutionError({ - selected, - requestedModel, + } catch (error) { + console.error('[responses] post-response success logging failed:', error); + } + bindSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + return reply.send(downstreamData); + } catch (err: any) { + clearSurfaceStickyChannel({ + stickySessionKey, + selected, + }); + const failureOutcome = await failureToolkit.handleExecutionError({ + selected, + requestedModel, modelName, errorMessage: err?.message || 'network failure', latencyMs: Date.now() - startTime, @@ -691,8 +777,10 @@ export async function handleOpenAiResponsesSurfaceRequest( if (failureOutcome.action === 'retry') { retryCount += 1; continue; - } - return reply.code(failureOutcome.status).send(failureOutcome.payload); - } - } + } + return reply.code(failureOutcome.status).send(failureOutcome.payload); + } finally { + channelLease.release(); + } + } } diff --git a/src/server/proxy-core/surfaces/sharedSurface.test.ts b/src/server/proxy-core/surfaces/sharedSurface.test.ts index 1947cdf..4fca8f9 100644 --- a/src/server/proxy-core/surfaces/sharedSurface.test.ts +++ b/src/server/proxy-core/surfaces/sharedSurface.test.ts @@ -3,6 +3,7 @@ import { EMPTY_DOWNSTREAM_ROUTING_POLICY } from '../../services/downstreamPolicy const selectChannelMock = vi.fn(); const selectNextChannelMock = vi.fn(); +const selectPreferredChannelMock = vi.fn(); const recordFailureMock = vi.fn(); const refreshModelsAndRebuildRoutesMock = vi.fn(); const composeProxyLogMessageMock = vi.fn(); @@ -20,6 +21,11 @@ const recordSuccessMock = vi.fn(); const resolveProxyUsageWithSelfLogFallbackMock = vi.fn(); const resolveProxyLogBillingMock = 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 consoleErrorMock = vi.spyOn(console, 'error').mockImplementation(() => {}); @@ -27,11 +33,22 @@ vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { selectChannel: (...args: unknown[]) => selectChannelMock(...args), selectNextChannel: (...args: unknown[]) => selectNextChannelMock(...args), + selectPreferredChannel: (...args: unknown[]) => selectPreferredChannelMock(...args), recordFailure: (...args: unknown[]) => recordFailureMock(...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', () => ({ refreshModelsAndRebuildRoutes: (...args: unknown[]) => refreshModelsAndRebuildRoutesMock(...args), })); @@ -95,6 +112,7 @@ describe('selectSurfaceChannelForAttempt', () => { beforeEach(() => { selectChannelMock.mockReset(); selectNextChannelMock.mockReset(); + selectPreferredChannelMock.mockReset(); recordFailureMock.mockReset(); refreshModelsAndRebuildRoutesMock.mockReset(); composeProxyLogMessageMock.mockReset(); @@ -112,6 +130,11 @@ describe('selectSurfaceChannelForAttempt', () => { resolveProxyUsageWithSelfLogFallbackMock.mockReset(); resolveProxyLogBillingMock.mockReset(); refreshOauthAccessTokenSingleflightMock.mockReset(); + getStickyChannelIdMock.mockReset(); + bindStickyChannelMock.mockReset(); + clearStickyChannelMock.mockReset(); + acquireChannelLeaseMock.mockReset(); + buildStickySessionKeyMock.mockReset(); consoleWarnMock.mockClear(); consoleErrorMock.mockClear(); }); @@ -158,6 +181,51 @@ describe('selectSurfaceChannelForAttempt', () => { 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 () => { const selected = { channel: { id: 33 } }; selectChannelMock diff --git a/src/server/proxy-core/surfaces/sharedSurface.ts b/src/server/proxy-core/surfaces/sharedSurface.ts index 53062f3..177955e 100644 --- a/src/server/proxy-core/surfaces/sharedSurface.ts +++ b/src/server/proxy-core/surfaces/sharedSurface.ts @@ -17,6 +17,7 @@ import type { BuiltEndpointRequest } from '../../routes/proxy/endpointFlow.js'; import { buildUpstreamUrl } from '../../routes/proxy/upstreamUrl.js'; import { recordOauthQuotaResetHint } from '../../services/oauth/quota.js'; import { refreshOauthAccessTokenSingleflight } from '../../services/oauth/refreshSingleflight.js'; +import { proxyChannelCoordinator } from '../../services/proxyChannelCoordinator.js'; type SelectedChannel = Awaited>; type SurfaceWarningScope = 'chat' | 'responses'; @@ -103,14 +104,34 @@ export async function selectSurfaceChannelForAttempt(input: { downstreamPolicy: DownstreamRoutingPolicy; excludeChannelIds: number[]; retryCount: number; + stickySessionKey?: string | null; }): Promise { - let selected = input.retryCount === 0 - ? await tokenRouter.selectChannel(input.requestedModel, input.downstreamPolicy) - : await tokenRouter.selectNextChannel( - input.requestedModel, - input.excludeChannelIds, - input.downstreamPolicy, - ); + 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.selectNextChannel( + input.requestedModel, + input.excludeChannelIds, + input.downstreamPolicy, + ); + } if (!selected && input.retryCount === 0) { try { @@ -124,6 +145,65 @@ export async function selectSurfaceChannelForAttempt(input: { 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: { warningScope: string; selected: { diff --git a/src/server/routes/proxy/endpointFlow.test.ts b/src/server/routes/proxy/endpointFlow.test.ts index 42c4399..89b0232 100644 --- a/src/server/routes/proxy/endpointFlow.test.ts +++ b/src/server/routes/proxy/endpointFlow.test.ts @@ -298,7 +298,6 @@ describe('executeEndpointFlow', () => { 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 () => { fetchMock.mockResolvedValueOnce(toUndiciResponse(new Response(JSON.stringify({ ok: true }), { status: 200, diff --git a/src/server/services/proxyChannelCoordinator.test.ts b/src/server/services/proxyChannelCoordinator.test.ts new file mode 100644 index 0000000..0ef4720 --- /dev/null +++ b/src/server/services/proxyChannelCoordinator.test.ts @@ -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(); + } + }); +}); diff --git a/src/server/services/proxyChannelCoordinator.ts b/src/server/services/proxyChannelCoordinator.ts new file mode 100644 index 0000000..debab79 --- /dev/null +++ b/src/server/services/proxyChannelCoordinator.ts @@ -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 | null; +}; + +type ChannelRuntimeState = { + activeLeaseIds: Set; + 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(); +const channelRuntimeStates = new Map(); +let nextLeaseId = 1; + +function shouldUnrefTimer(timer: ReturnType | ReturnType) { + 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(), + 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 { + 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((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 | null = null; + let keepaliveTimer: ReturnType | 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(); diff --git a/src/server/services/tokenRouter.selection.test.ts b/src/server/services/tokenRouter.selection.test.ts index c3f2303..9909e84 100644 --- a/src/server/services/tokenRouter.selection.test.ts +++ b/src/server/services/tokenRouter.selection.test.ts @@ -117,6 +117,54 @@ describe('TokenRouter selection scoring', () => { }).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 () => { config.routingWeights = { baseWeightFactor: 1, diff --git a/src/server/services/tokenRouter.ts b/src/server/services/tokenRouter.ts index 61f1352..8bbf4c2 100644 --- a/src/server/services/tokenRouter.ts +++ b/src/server/services/tokenRouter.ts @@ -1249,6 +1249,27 @@ export class TokenRouter { if (!match) return null; return await this.selectFromMatch(match, requestedModel, downstreamPolicy, excludeChannelIds); } + + async selectPreferredChannel( + requestedModel: string, + preferredChannelId: number, + downstreamPolicy: DownstreamRoutingPolicy = DEFAULT_DOWNSTREAM_POLICY, + excludeChannelIds: number[] = [], + ): Promise { + 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( requestedModel: string, @@ -1843,6 +1864,65 @@ export class TokenRouter { return null; } + private async selectPreferredFromMatch( + match: RouteMatch, + requestedModel: string, + preferredChannelId: number, + excludeChannelIds: number[] = [], + recordSelection = true, + ): Promise { + 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 { let routes = await loadEnabledRoutes();