LaPaoPao/BoosAPI
1
0
Fork 0
forked from boos/BoosAPI
Code Pull Requests 1 Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
src/server/proxy-core/surfaces/chatSurface.ts
+145 -6
View File
@@ -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();
}
}
}
src/server/proxy-core/surfaces/openAiResponsesSurface.ts
+153 -65
View File
@@ -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();
}
}
}
src/server/proxy-core/surfaces/sharedSurface.test.ts
+68
View File
@@ -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
src/server/proxy-core/surfaces/sharedSurface.ts
+87 -7
View File
@@ -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<ReturnType<typeof tokenRouter.selectChannel>>;
type SurfaceWarningScope = 'chat' | 'responses';
@@ -103,14 +104,34 @@ export async function selectSurfaceChannelForAttempt(input: {
downstreamPolicy: DownstreamRoutingPolicy;
excludeChannelIds: number[];
retryCount: number;
stickySessionKey?: string | null;
}): Promise<SelectedChannel> {
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: {