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>
2026-03-24 19:40:44 +08:00
parent 5b9c83f830
commit 9cf3dd1904
10 changed files with 1031 additions and 79 deletions
@@ -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();
  }

+  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,
@@ -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<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(
    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<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> {
    let routes = await loadEnabledRoutes();