From a5aeccff5194190b9618e5175b741be1f116e653 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 14:05:22 -0400 Subject: [PATCH 01/23] feat(proxy): add bounded streaming discovery --- .../v1/_lib/proxy/discovery-coordinator.ts | 255 +++++++++ src/app/v1/_lib/proxy/discovery-validity.ts | 156 ++++++ src/app/v1/_lib/proxy/forwarder.ts | 530 +++++++++++++++++- src/app/v1/_lib/proxy/provider-selector.ts | 54 +- src/app/v1/_lib/proxy/response-handler.ts | 83 ++- src/app/v1/_lib/proxy/session.ts | 14 + src/app/v1/_lib/proxy/stream-finalization.ts | 6 + src/lib/config/system-settings-cache.ts | 7 + src/types/system-config.ts | 17 + .../unit/proxy/discovery-coordinator.test.ts | 63 +++ tests/unit/proxy/discovery-validity.test.ts | 63 +++ 11 files changed, 1235 insertions(+), 13 deletions(-) create mode 100644 src/app/v1/_lib/proxy/discovery-coordinator.ts create mode 100644 src/app/v1/_lib/proxy/discovery-validity.ts create mode 100644 tests/unit/proxy/discovery-coordinator.test.ts create mode 100644 tests/unit/proxy/discovery-validity.test.ts diff --git a/src/app/v1/_lib/proxy/discovery-coordinator.ts b/src/app/v1/_lib/proxy/discovery-coordinator.ts new file mode 100644 index 000000000..44d9de404 --- /dev/null +++ b/src/app/v1/_lib/proxy/discovery-coordinator.ts @@ -0,0 +1,255 @@ +/** + * Pure state machine for bounded provider discovery. + * + * The coordinator deliberately has no network or timer dependencies. The + * forwarder owns attempts and calls these methods at event boundaries. This + * keeps cancellation and stale-event handling deterministic and testable. + */ + +export type DiscoveryAttemptKind = "normal" | "fallback"; +export type DiscoveryState = + | "STICKY_PROBING" + | "DISCOVERY_RACING" + | "FALLBACK_READY_HELD" + | "FALLBACK_ACTIVE" + | "WINNER_COMMITTED" + | "TERMINAL_FAILED"; + +export type DiscoveryAttempt = { + id: string; + providerId: number; + priority: number; + kind: DiscoveryAttemptKind; + ready: boolean; + pending: boolean; + round: number; + launchOrder: number; +}; + +export type DiscoveryAction = + | { type: "commit_normal"; attemptId: string } + | { type: "promote_fallback"; attemptId: string } + | { type: "cancel"; attemptIds: string[] } + | { type: "launch"; slots: number } + | { type: "none" } + | { type: "terminal_failure" }; + +export type DiscoveryCoordinatorOptions = { + concurrency: number; + maxRounds: number; +}; + +function compareAttempts(a: DiscoveryAttempt, b: DiscoveryAttempt): number { + return a.priority - b.priority || a.launchOrder - b.launchOrder; +} + +export class DiscoveryCoordinator { + readonly concurrency: number; + readonly maxRounds: number; + state: DiscoveryState = "DISCOVERY_RACING"; + round = 1; + private attempts = new Map(); + private requestEpoch = 0; + private roundEpoch = 0; + + constructor(options: DiscoveryCoordinatorOptions) { + this.concurrency = Math.max(1, Math.floor(options.concurrency)); + this.maxRounds = Math.max(1, Math.floor(options.maxRounds)); + } + + get epochs(): { requestEpoch: number; roundEpoch: number } { + return { requestEpoch: this.requestEpoch, roundEpoch: this.roundEpoch }; + } + + beginRound(): { requestEpoch: number; roundEpoch: number; round: number } { + this.roundEpoch += 1; + return { ...this.epochs, round: this.round }; + } + + addAttempt(attempt: DiscoveryAttempt): boolean { + if (this.isTerminal || this.attempts.has(attempt.id)) return false; + this.attempts.set(attempt.id, { ...attempt, round: this.round }); + return true; + } + + removeAttempt(id: string): void { + this.attempts.delete(id); + } + + get isTerminal(): boolean { + return this.state === "WINNER_COMMITTED" || this.state === "TERMINAL_FAILED"; + } + + get activeAttempts(): DiscoveryAttempt[] { + return Array.from(this.attempts.values()).filter((attempt) => attempt.pending); + } + + get snapshot(): DiscoveryAttempt[] { + return Array.from(this.attempts.values()).map((attempt) => ({ ...attempt })); + } + + /** Ignore events from a cancelled request or an old round. */ + acceptsEpoch(requestEpoch: number, roundEpoch: number): boolean { + return requestEpoch === this.requestEpoch && roundEpoch === this.roundEpoch; + } + + markReady( + id: string, + requestEpoch = this.requestEpoch, + roundEpoch = this.roundEpoch + ): DiscoveryAction { + if (!this.acceptsEpoch(requestEpoch, roundEpoch) || this.isTerminal) return { type: "none" }; + const attempt = this.attempts.get(id); + if (!attempt?.pending) return { type: "none" }; + attempt.ready = true; + return this.chooseReadyNormal(); + } + + markFailed( + id: string, + requestEpoch = this.requestEpoch, + roundEpoch = this.roundEpoch + ): DiscoveryAction { + if (!this.acceptsEpoch(requestEpoch, roundEpoch) || this.isTerminal) return { type: "none" }; + const attempt = this.attempts.get(id); + if (!attempt) return { type: "none" }; + attempt.pending = false; + attempt.ready = false; + return this.afterAttemptState(); + } + + /** A normal ready result may win only after priority gating is satisfied. */ + private chooseReadyNormal(ignorePriorityGate = false): DiscoveryAction { + const readyNormal = Array.from(this.attempts.values()) + .filter((attempt) => attempt.pending && attempt.ready && attempt.kind === "normal") + .sort(compareAttempts); + if (readyNormal.length === 0) return { type: "none" }; + const bestPriority = readyNormal[0].priority; + if (!ignorePriorityGate) { + const higherTierPending = Array.from(this.attempts.values()).some( + (attempt) => + attempt.pending && + attempt.kind === "normal" && + !attempt.ready && + attempt.priority < bestPriority + ); + if (higherTierPending) return { type: "none" }; + } + const sameTier = readyNormal.filter((attempt) => attempt.priority === bestPriority); + const winner = sameTier[0]; + this.state = "WINNER_COMMITTED"; + winner.pending = false; + return { + type: "commit_normal", + attemptId: winner.id, + }; + } + + /** + * Close the current SLA window. At a boundary a ready normal always wins; + * otherwise the best still-pending normal becomes the sole fallback. A + * fallback that is merely ready is held until no normal can still win. + */ + onRoundBoundary(requestEpoch = this.requestEpoch, roundEpoch = this.roundEpoch): DiscoveryAction { + if (!this.acceptsEpoch(requestEpoch, roundEpoch) || this.isTerminal) return { type: "none" }; + const readyAction = this.chooseReadyNormal(true); + if (readyAction.type === "commit_normal") return readyAction; + + const currentFallback = Array.from(this.attempts.values()).find( + (attempt) => attempt.pending && attempt.kind === "fallback" + ); + if (currentFallback?.ready) { + currentFallback.pending = false; + this.state = "FALLBACK_ACTIVE"; + return { type: "promote_fallback", attemptId: currentFallback.id }; + } + + const pendingNormal = Array.from(this.attempts.values()) + .filter((attempt) => attempt.pending && attempt.kind === "normal") + .sort(compareAttempts); + if (currentFallback && pendingNormal.length > 0) { + for (const attempt of pendingNormal) attempt.pending = false; + if (this.round < this.maxRounds) { + this.round += 1; + this.roundEpoch += 1; + this.state = "DISCOVERY_RACING"; + return { type: "launch", slots: Math.max(1, this.concurrency - 1) }; + } + return { type: "none" }; + } + if (pendingNormal.length === 0) { + if (currentFallback) { + this.state = "FALLBACK_READY_HELD"; + return { type: "none" }; + } + return this.finishOrLaunch(); + } + + const fallback = pendingNormal[0]; + fallback.kind = "fallback"; + this.state = "FALLBACK_READY_HELD"; + const losers = pendingNormal.slice(1).map((attempt) => attempt.id); + for (const id of losers) this.attempts.get(id)!.pending = false; + + if (currentFallback) { + currentFallback.pending = true; + return { type: "cancel", attemptIds: losers }; + } + return { type: "cancel", attemptIds: losers }; + } + + onDeadline(): DiscoveryAction { + if (this.isTerminal) return { type: "none" }; + const fallback = Array.from(this.attempts.values()).find( + (attempt) => attempt.pending && attempt.kind === "fallback" && attempt.ready + ); + if (fallback) { + fallback.pending = false; + this.state = "FALLBACK_ACTIVE"; + return { type: "promote_fallback", attemptId: fallback.id }; + } + this.state = "TERMINAL_FAILED"; + return { type: "terminal_failure" }; + } + + commitWinner(id: string): DiscoveryAction { + const attempt = this.attempts.get(id); + if (!attempt || this.isTerminal) return { type: "none" }; + attempt.pending = false; + this.state = "WINNER_COMMITTED"; + return { + type: attempt.kind === "fallback" ? "promote_fallback" : "commit_normal", + attemptId: id, + }; + } + + cancelRequest(): DiscoveryAction { + this.requestEpoch += 1; + this.roundEpoch += 1; + const ids = this.activeAttempts.map((attempt) => attempt.id); + for (const attempt of this.attempts.values()) attempt.pending = false; + this.state = "TERMINAL_FAILED"; + return { type: "cancel", attemptIds: ids }; + } + + private afterAttemptState(): DiscoveryAction { + const pending = this.activeAttempts; + if (pending.length === 0) return this.finishOrLaunch(); + const fallback = pending.find((attempt) => attempt.kind === "fallback"); + if (fallback?.ready && pending.every((attempt) => attempt.kind === "fallback")) { + return this.commitWinner(fallback.id); + } + return { type: "none" }; + } + + private finishOrLaunch(): DiscoveryAction { + if (this.round >= this.maxRounds) { + this.state = "TERMINAL_FAILED"; + return { type: "terminal_failure" }; + } + this.round += 1; + this.roundEpoch += 1; + this.state = "DISCOVERY_RACING"; + return { type: "launch", slots: Math.max(1, this.concurrency - 1) }; + } +} diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts new file mode 100644 index 000000000..29ad6a960 --- /dev/null +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -0,0 +1,156 @@ +export type DiscoveryProtocol = + | "anthropic" + | "openai-chat" + | "openai-responses" + | "gemini" + | "unknown"; + +export type DiscoveryValidity = { + ready: boolean; + terminal: boolean; + error: boolean; +}; + +function hasContent(value: unknown): boolean { + if (typeof value === "string") return value.trim().length > 0; + if (!value || typeof value !== "object") return false; + if (Array.isArray(value)) return value.some(hasContent); + const object = value as Record; + return [ + "text", + "content", + "delta", + "output_text", + "thinking", + "tool_use", + "tool_calls", + "functionCall", + "function_call", + "arguments", + "input", + "parts", + ].some((key) => hasContent(object[key])); +} + +function classifyJson(value: unknown, protocol: DiscoveryProtocol): DiscoveryValidity { + if (!value || typeof value !== "object") return { ready: false, terminal: false, error: true }; + const object = value as Record; + if ( + object.error || + object.failed || + object.type === "error" || + object.type === "response.failed" + ) { + return { ready: false, terminal: true, error: true }; + } + if (protocol === "openai-chat") { + const choices = Array.isArray(object.choices) ? object.choices : []; + const ready = choices.some((choice) => { + if (!choice || typeof choice !== "object") return false; + const choiceObject = choice as Record; + const delta = choiceObject.delta; + return hasContent(delta) || hasContent(choiceObject.message); + }); + return { ready, terminal: false, error: false }; + } + if (protocol === "openai-responses") { + if (object.type === "response.completed" || object.type === "response.done") { + return { ready: false, terminal: true, error: false }; + } + return { + ready: + (object.type === "response.output_text.delta" && hasContent(object.delta)) || + (object.type === "response.function_call_arguments.delta" && hasContent(object.delta)) || + (object.type === "response.output_item.added" && hasContent(object.item)), + terminal: false, + error: false, + }; + } + if (protocol === "gemini") { + const candidates = Array.isArray(object.candidates) ? object.candidates : []; + return { + ready: candidates.some((candidate) => hasContent(candidate)), + terminal: false, + error: false, + }; + } + // Anthropic SSE data events: message_start/message_delta are metadata; a + // content_block_delta or tool use is the first deliverable event. + if ( + object.type === "message_start" || + object.type === "message_delta" || + object.type === "ping" + ) { + return { ready: false, terminal: false, error: false }; + } + if (object.type === "message_stop") { + return { ready: false, terminal: true, error: false }; + } + return { + ready: + (object.type === "content_block_delta" && hasContent(object.delta)) || + (object.type === "content_block_start" && hasContent(object.content_block)) || + hasContent(object.content), + terminal: false, + error: false, + }; +} + +export function classifyDiscoveryChunk( + chunk: Uint8Array | string, + protocol: DiscoveryProtocol +): DiscoveryValidity { + const text = typeof chunk === "string" ? chunk : new TextDecoder().decode(chunk); + if (!text.trim() || text.trim().startsWith(":")) + return { ready: false, terminal: false, error: false }; + if (text.includes("[DONE]")) return { ready: false, terminal: true, error: false }; + + const lines = text.split(/\r?\n/); + let sawTerminal = false; + let sawError = false; + let sawReady = false; + for (const line of lines) { + const candidate = line.startsWith("data:") ? line.slice(5).trim() : line.trim(); + if (!candidate || candidate.startsWith(":")) continue; + try { + const result = classifyJson(JSON.parse(candidate), protocol); + sawTerminal ||= result.terminal; + sawError ||= result.error; + sawReady ||= result.ready; + } catch { + // A raw JSON response may arrive in a single chunk. Plain text is not + // a protocol-safe winner; keep waiting for a parseable event. + } + } + return { ready: sawReady && !sawError, terminal: sawTerminal, error: sawError }; +} + +export class DiscoveryValidityParser { + private buffered = ""; + private readonly decoder = new TextDecoder(); + private _ready = false; + private _terminal = false; + private _error = false; + + constructor(readonly protocol: DiscoveryProtocol) {} + + push(chunk: Uint8Array | string): DiscoveryValidity { + this.buffered += + typeof chunk === "string" ? chunk : this.decoder.decode(chunk, { stream: true }); + const result = classifyDiscoveryChunk(this.buffered, this.protocol); + this._ready ||= result.ready; + this._terminal ||= result.terminal; + this._error ||= result.error; + return { ready: this._ready && !this._error, terminal: this._terminal, error: this._error }; + } + + get ready(): boolean { + return this._ready && !this._error; + } + get terminal(): boolean { + return this._terminal; + } + get error(): boolean { + return this._error; + } +} diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 632807ef2..05cff7803 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -63,6 +63,8 @@ import { buildProxyUrl } from "../url"; import { rectifyBillingHeader } from "./billing-header-rectifier"; import { bindClientAbortListener } from "./client-abort-listener"; import { deriveClientSafeUpstreamErrorMessage } from "./client-error-message"; +import { DiscoveryCoordinator } from "./discovery-coordinator"; +import { type DiscoveryProtocol, DiscoveryValidityParser } from "./discovery-validity"; import { isStandardProxyEndpointPath } from "./endpoint-family-catalog"; import { resolveEndpointPolicy, shouldEnforceStrictEndpointPoolPolicy } from "./endpoint-policy"; import { @@ -1197,6 +1199,12 @@ export class ProxyForwarder { throw new Error("代理上下文缺少供应商或鉴权信息"); } + if (await ProxyForwarder.shouldUseStreamingDiscovery(session)) { + const discoveryPromise = ProxyForwarder.sendStreamingWithDiscovery(session); + void discoveryPromise.catch(() => undefined); + return await discoveryPromise; + } + if (ProxyForwarder.shouldUseStreamingHedge(session)) { const hedgePromise = ProxyForwarder.sendStreamingWithHedge(session); void hedgePromise.catch(() => undefined); @@ -2392,7 +2400,8 @@ export class ProxyForwarder { baseUrl: string, endpointAudit?: { endpointId: number | null; endpointUrl: string }, attemptNumber?: number, - deferDetailSnapshotPersistence: boolean = false + deferDetailSnapshotPersistence: boolean = false, + externalAbortSignal?: AbortSignal ): Promise { if (!provider) { throw new Error("Provider is required"); @@ -3029,12 +3038,16 @@ export class ProxyForwarder { const clientSignal = session.clientAbortSignal; if (clientSignal) abortTransportFrom(clientSignal); }); + const cleanupExternalTransportSignal = bindClientAbortListener(externalAbortSignal, () => { + if (externalAbortSignal) abortTransportFrom(externalAbortSignal); + }); const cleanupCombinedSignal = () => { cleanupResponseTransportSignal(); cleanupClientTransportSignal(); + cleanupExternalTransportSignal(); }; logger.debug("ProxyForwarder: Combined abort signals", { - signalCount: session.clientAbortSignal ? 2 : 1, + signalCount: (session.clientAbortSignal ? 1 : 0) + (externalAbortSignal ? 1 : 0) + 1, }); const init: UndiciFetchOptions = { @@ -3798,6 +3811,13 @@ export class ProxyForwarder { private static shouldUseStreamingHedge(session: ProxySession): boolean { const endpointPolicy = ProxyForwarder.getEndpointPolicy(session); + const routing = session as ProxySession & { + routingMode?: string; + disableStreamingHedge?: boolean; + }; + if (routing.routingMode === "lease_conflict_single" || routing.disableStreamingHedge === true) { + return false; + } return ( (endpointPolicy?.allowRetry ?? true) && (endpointPolicy?.allowProviderSwitch ?? true) && @@ -3806,6 +3826,47 @@ export class ProxyForwarder { ); } + private static async shouldUseStreamingDiscovery(session: ProxySession): Promise { + const settings = await getCachedSystemSettings(); + if (SessionManager.getVersionedBindingCapabilityState() !== "available") { + return false; + } + const endpointPolicy = ProxyForwarder.getEndpointPolicy(session); + const protocol = ProxyForwarder.discoveryProtocol(session); + const message = session.request.message as Record; + const routing = session as ProxySession & { + routingMode?: string; + disableStreamingHedge?: boolean; + }; + return ( + settings.discoveryEnabled === true && + endpointPolicy.allowRetry && + endpointPolicy.allowProviderSwitch && + message.stream === true && + !endpointPolicy.bypassForwarderPreprocessing && + protocol !== "unknown" && + routing.routingMode !== "lease_conflict_single" && + routing.disableStreamingHedge !== true && + !session.isRawCrossProviderFallbackEnabled() + ); + } + + private static discoveryProtocol(session: ProxySession): DiscoveryProtocol { + switch (session.originalFormat) { + case "claude": + return "anthropic"; + case "openai": + return "openai-chat"; + case "response": + return "openai-responses"; + case "gemini": + case "gemini-cli": + return "gemini"; + default: + return "unknown"; + } + } + private static getEndpointPolicy(session: ProxySession) { const policySession = session as unknown as { getEndpointPolicy?: (() => ReturnType) | undefined; @@ -4766,6 +4827,471 @@ export class ProxyForwarder { } } + /** + * Bounded Discovery path. It intentionally lives beside legacy Hedge so the + * existing loser billing and retry semantics remain unchanged while the + * feature is rolled out behind discoveryEnabled. + */ + private static async sendStreamingWithDiscovery(session: ProxySession): Promise { + const initialProvider = session.provider; + if (!initialProvider) throw new Error("代理上下文缺少供应商"); + const settings = await getCachedSystemSettings(); + const concurrency = Math.max(1, Math.floor(settings.discoveryConcurrency ?? 2)); + const maxRounds = Math.max(1, Math.floor(settings.maxDiscoveryRounds ?? 2)); + const discoverySlaMs = Math.max(1, settings.discoverySlaMs ?? 10_000); + const stickySlaMs = Math.max(discoverySlaMs, settings.stickySlaMs ?? 20_000); + const totalTimeoutMs = Math.max(stickySlaMs, settings.racingTotalTimeoutMs ?? 60_000); + const protocol = ProxyForwarder.discoveryProtocol(session); + const coordinator = new DiscoveryCoordinator({ concurrency, maxRounds }); + const bindingKeyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + // Provider selection normally populates this snapshot for a reused Sticky. + // For a cold start, initialize it once so finalization can use generation + // CAS without another read during winner commit or timeout cleanup. + let bindingSnapshot = session.getSessionBindingSnapshot(); + if (!bindingSnapshot && session.sessionId && bindingKeyId != null) { + const binding = await SessionManager.getSessionBindingSnapshot( + session.sessionId, + bindingKeyId + ); + if (binding.status === "ok") { + bindingSnapshot = binding.snapshot; + session.setSessionBindingSnapshot(binding.snapshot); + } + } + const attempts = new Map< + string, + StreamingHedgeAttempt & { + id: string; + kind: "normal" | "fallback"; + controller: AbortController; + parser: DiscoveryValidityParser; + chunks: Uint8Array[]; + pending: boolean; + ready: boolean; + round: number; + } + >(); + const launched = new Set(); + let sequence = 0; + let currentRound = 1; + let winner: (typeof attempts extends Map ? V : never) | null = null; + let committed = false; + let settled = false; + let noMoreCandidates = false; + let lastError: Error | null = null; + let lastErrorCategory: ErrorCategory | null = null; + let totalTimer: NodeJS.Timeout | null = null; + let roundTimer: NodeJS.Timeout | null = null; + let stickyTimer: NodeJS.Timeout | null = null; + let resolveResult: ((result: { response?: Response; error?: Error }) => void) | null = null; + const resultPromise = new Promise<{ response?: Response; error?: Error }>((resolve) => { + resolveResult = resolve; + }); + + const releaseProviderRef = (providerId: number) => { + if (!session.sessionId) return; + const consumer = (session as { consumeProviderSessionRef?: (id: number) => boolean }) + .consumeProviderSessionRef; + if (consumer?.call(session, providerId)) { + void RateLimitService.releaseProviderSession(providerId, session.sessionId); + } + }; + + const cancelAttempt = (attempt: (typeof winner & { id: string }) | null, reason: string) => { + if (!attempt?.pending) return; + attempt.pending = false; + try { + attempt.controller.abort(new Error(reason)); + } catch { + /* abort is best effort */ + } + void attempt.reader?.cancel(reason).catch(() => undefined); + try { + attempt.releaseAgent?.(); + } catch { + /* release is idempotent */ + } + releaseProviderRef(attempt.provider.id); + }; + + const cancelLosers = (keep: typeof winner = null) => { + for (const attempt of attempts.values()) { + if (attempt !== keep) cancelAttempt(attempt, "discovery_loser"); + } + }; + + const settleFailure = async (error: Error) => { + if (settled) return; + settled = true; + if (totalTimer) clearTimeout(totalTimer); + if (roundTimer) clearTimeout(roundTimer); + if (stickyTimer) clearTimeout(stickyTimer); + cancelLosers(); + const attempted = new Set(launched); + await ProxyForwarder.clearSessionProviderBindings(session, attempted); + resolveResult?.({ error }); + }; + + const commit = async (attempt: typeof winner) => { + if (!attempt || committed || settled || !attempt.response || !attempt.reader) return; + committed = true; + winner = attempt; + attempt.pending = false; + if (totalTimer) clearTimeout(totalTimer); + if (roundTimer) clearTimeout(roundTimer); + if (stickyTimer) clearTimeout(stickyTimer); + cancelLosers(attempt); + session.setProvider(attempt.provider); + if (attempt.session !== session) + ProxyForwarder.syncWinningAttemptSession(session, attempt.session); + + setDeferredStreamingFinalization(session, { + providerId: attempt.provider.id, + providerName: attempt.provider.name, + providerPriority: attempt.provider.priority || 0, + attemptNumber: attempt.sequence, + totalProvidersAttempted: launched.size, + isFirstAttempt: attempt.provider.id === initialProvider.id, + isFailoverSuccess: attempt.provider.id !== initialProvider.id, + endpointId: attempt.endpointAudit.endpointId, + endpointUrl: attempt.endpointAudit.endpointUrl, + upstreamStatusCode: attempt.response.status, + isHedgeWinner: false, + billHedgeLosers: false, + bindingIntent: + attempt.kind === "fallback" + ? "none" + : bindingSnapshot?.providerId == null + ? "create" + : "renew", + bindingSnapshot, + requiresCompletionMarker: attempt.kind !== "fallback", + }); + const prefix = + attempt.chunks.length === 1 + ? attempt.chunks[0] + : (() => { + const size = attempt.chunks.reduce((sum, chunk) => sum + chunk.byteLength, 0); + const output = new Uint8Array(size); + let offset = 0; + for (const chunk of attempt.chunks) { + output.set(chunk, offset); + offset += chunk.byteLength; + } + return output; + })(); + resolveResult?.({ + response: new Response( + ProxyForwarder.buildBufferedFirstChunkStream(prefix, attempt.reader), + { + status: attempt.response.status, + statusText: attempt.response.statusText, + headers: attempt.response.headers, + } + ), + }); + }; + + const chooseCandidate = async (): Promise => { + const candidates = await ProxyProviderResolver.pickDiscoveryProviders( + session, + 1, + Array.from(launched) + ); + if (!candidates[0]) { + noMoreCandidates = true; + return null; + } + return candidates[0]; + }; + + const launch = async (provider: Provider, kind: "normal" | "fallback"): Promise => { + if (settled || committed || launched.has(provider.id)) return; + launched.add(provider.id); + if (provider.id !== initialProvider.id && session.sessionId) { + const limit = provider.limitConcurrentSessions || 0; + const check = await RateLimitService.checkAndTrackProviderSession( + provider.id, + session.sessionId, + limit + ); + if (!check.allowed) { + launched.delete(provider.id); + throw new ProxyError(check.reason || "Provider concurrent limit reached", 503); + } + if (check.referenced) session.recordProviderSessionRef(provider.id); + } + const endpoint = await ProxyForwarder.resolveStreamingHedgeEndpoint(session, provider); + const attemptSession = + provider.id === initialProvider.id + ? session + : ProxyForwarder.createStreamingShadowSession(session, provider); + attemptSession.setProvider(provider); + const controller = new AbortController(); + const id = `${provider.id}:${sequence + 1}`; + const attempt = { + id, + kind, + controller, + parser: new DiscoveryValidityParser(protocol), + chunks: [], + pending: true, + ready: false, + round: currentRound, + provider, + session: attemptSession, + baseUrl: endpoint.baseUrl, + endpointAudit: { endpointId: endpoint.endpointId, endpointUrl: endpoint.endpointUrl }, + modelRedirect: undefined, + responseController: null, + clearResponseTimeout: null, + firstByteTimeoutMs: 0, + sequence: ++sequence, + requestAttemptCount: 1, + reactiveRectifierRetryState: { + thinkingSignatureRetried: false, + thinkingBudgetRetried: false, + thinkingEffortConflictRetried: false, + geminiFunctionIdRetried: false, + }, + settled: false, + thresholdTriggered: false, + thresholdTimer: null, + reader: null, + response: null, + releaseAgent: null, + agentReleased: false, + billAsLoser: false, + loserBillingStarted: false, + firstChunk: null, + billingSnapshot: null, + } as typeof winner & { + id: string; + kind: "normal" | "fallback"; + controller: AbortController; + parser: DiscoveryValidityParser; + chunks: Uint8Array[]; + pending: boolean; + ready: boolean; + round: number; + }; + attempts.set(id, attempt); + coordinator.addAttempt({ + id, + providerId: provider.id, + priority: provider.priority || 0, + kind, + ready: false, + pending: true, + round: currentRound, + launchOrder: attempt.sequence, + }); + + void ProxyForwarder.doForward( + attempt.session, + { ...provider, firstByteTimeoutStreamingMs: 0 }, + endpoint.baseUrl, + attempt.endpointAudit, + attempt.requestAttemptCount, + true, + controller.signal + ) + .then(async (response) => { + const runtime = attempt.session as ProxySessionWithAttemptRuntime; + attempt.responseController = runtime.responseController ?? null; + attempt.clearResponseTimeout = runtime.clearResponseTimeout ?? null; + attempt.releaseAgent = runtime.releaseAgent ?? null; + attempt.clearResponseTimeout?.(); + attempt.response = response; + if (!response.body) + throw new EmptyResponseError(provider.id, provider.name, "empty_body"); + attempt.reader = response.body.getReader(); + while (!committed && !settled && attempt.pending) { + const item = await attempt.reader.read(); + if (item.done) throw new EmptyResponseError(provider.id, provider.name, "empty_body"); + if (!item.value || item.value.byteLength === 0) continue; + attempt.chunks.push(item.value); + const validity = attempt.parser.push(item.value); + if (validity.error || validity.terminal) + throw new ProxyError("Invalid upstream discovery response", 502); + if (!validity.ready) continue; + attempt.ready = true; + const normalPendingHigher = Array.from(attempts.values()).some( + (other) => + other.pending && + other.kind === "normal" && + (other.provider.priority || 0) < (provider.priority || 0) + ); + if (normalPendingHigher) continue; + const action = coordinator.markReady(id); + if (action.type === "commit_normal" || action.type === "promote_fallback") + await commit(attempt); + return; + } + }) + .catch(async (error) => { + if (committed || settled || !attempt.pending) return; + attempt.pending = false; + coordinator.markFailed(id); + lastError = error instanceof Error ? error : new Error(String(error)); + lastErrorCategory = await categorizeErrorAsync(lastError); + session.addProviderToChain(provider, { + ...attempt.endpointAudit, + reason: "retry_failed", + attemptNumber: attempt.sequence, + statusCode: lastError instanceof ProxyError ? lastError.statusCode : undefined, + errorMessage: lastError.message, + }); + if ( + lastErrorCategory === ErrorCategory.PROVIDER_ERROR && + !(lastError instanceof ProxyError && lastError.statusCode === 404) + ) { + await recordFailure(provider.id, lastError).catch(() => undefined); + } + attempt.releaseAgent?.(); + releaseProviderRef(provider.id); + const replacement = await chooseCandidate(); + if (replacement && !committed && !settled) await launch(replacement, "normal"); + if ( + Array.from(attempts.values()).every((candidate) => !candidate.pending) && + noMoreCandidates + ) { + await settleFailure( + ProxyForwarder.resolveHedgeTerminalError(lastError, lastErrorCategory) + ); + } + }); + }; + + const launchNextRound = async () => { + if (settled || committed) return; + currentRound += 1; + if (currentRound > maxRounds) return; + coordinator.beginRound(); + const candidate = await chooseCandidate(); + if (candidate) { + try { + await launch(candidate, "normal"); + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + noMoreCandidates = true; + } + } + if (!committed && !settled) { + roundTimer = setTimeout(() => void onBoundary(), discoverySlaMs); + } + }; + + const onBoundary = async () => { + if (settled || committed) return; + const fallback = Array.from(attempts.values()).find( + (attempt) => attempt.pending && attempt.kind === "fallback" + ); + const pendingNormals = Array.from(attempts.values()) + .filter((attempt) => attempt.pending && attempt.kind === "normal") + .sort( + (a, b) => + (a.provider.priority || 0) - (b.provider.priority || 0) || a.sequence - b.sequence + ); + const readyNormal = pendingNormals.find((attempt) => attempt.ready); + if (readyNormal) { + await commit(readyNormal); + return; + } + // The fallback is held during the SLA window, but at the round + // boundary it may take over when no normal result is ready. + if (fallback?.ready) { + await commit(fallback); + return; + } + if (pendingNormals[0]) { + if (fallback) { + for (const loser of pendingNormals) cancelAttempt(loser, "discovery_round_boundary"); + if (currentRound < maxRounds) await launchNextRound(); + return; + } + pendingNormals[0].kind = "fallback"; + for (const loser of pendingNormals.slice(1)) + cancelAttempt(loser, "discovery_round_boundary"); + if (currentRound < maxRounds) await launchNextRound(); + return; + } + if (currentRound < maxRounds) await launchNextRound(); + else await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); + }; + + const cleanupAbort = bindClientAbortListener(session.clientAbortSignal, () => { + if (settled || committed) return; + void settleFailure(new ProxyError("Request aborted by client", 499, undefined, true)); + }); + + totalTimer = setTimeout(() => { + if (settled || committed) return; + const fallback = Array.from(attempts.values()).find( + (attempt) => attempt.pending && attempt.kind === "fallback" && attempt.ready + ); + if (fallback) void commit(fallback); + else void settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); + }, totalTimeoutMs); + + try { + const hasSticky = session.shouldReuseProvider() && !!session.sessionId; + try { + await launch(initialProvider, "normal"); + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + const replacement = await chooseCandidate(); + if (replacement) await launch(replacement, "normal"); + else await settleFailure(ProxyForwarder.resolveHedgeTerminalError(lastError, null)); + } + const initial = hasSticky ? concurrency - 1 : Math.max(0, concurrency - 1); + if (hasSticky) { + stickyTimer = setTimeout(() => { + const sticky = Array.from(attempts.values()).find( + (attempt) => attempt.pending && attempt.provider.id === initialProvider.id + ); + if (sticky) { + sticky.kind = "fallback"; + if (bindingSnapshot && bindingSnapshot.providerId === initialProvider.id) { + void SessionManager.clearVersionedSessionProvider( + bindingSnapshot, + initialProvider.id, + Math.ceil((settings.stickyTimeoutCooldownMs ?? 300_000) / 1000) + ).catch((error) => + logger.debug("[Discovery] Failed to clear timed-out Sticky", { error }) + ); + } + if (currentRound < maxRounds) void launchNextRound(); + else void onBoundary(); + } + }, stickySlaMs); + } else { + const candidates = await ProxyProviderResolver.pickDiscoveryProviders( + session, + initial, + Array.from(launched) + ); + for (const provider of candidates) { + try { + await launch(provider, "normal"); + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + noMoreCandidates = true; + } + } + roundTimer = setTimeout(() => void onBoundary(), discoverySlaMs); + } + const result = await resultPromise; + if (result.error) throw result.error; + return result.response as Response; + } finally { + cleanupAbort(); + if (totalTimer) clearTimeout(totalTimer); + if (roundTimer) clearTimeout(roundTimer); + if (stickyTimer) clearTimeout(stickyTimer); + } + } + private static async resolveStreamingHedgeEndpoint( session: ProxySession, provider: Provider diff --git a/src/app/v1/_lib/proxy/provider-selector.ts b/src/app/v1/_lib/proxy/provider-selector.ts index 10c397000..ee290e1aa 100644 --- a/src/app/v1/_lib/proxy/provider-selector.ts +++ b/src/app/v1/_lib/proxy/provider-selector.ts @@ -457,6 +457,43 @@ export class ProxyProviderResolver { return provider; } + /** + * Select a bounded Discovery batch using the exact same filters, priority + * and weighted selection as the normal selector. The method is intentionally + * additive: legacy initial selection/fallback keeps its existing behavior. + */ + static async pickDiscoveryProviders( + session: ProxySession, + count: number, + excludeIds: number[] = [] + ): Promise { + const selected: Provider[] = []; + const excluded = new Set(excludeIds); + const limit = Math.max(0, Math.floor(count)); + const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + while (selected.length < limit) { + const provider = await ProxyProviderResolver.pickRandomProviderWithExclusion( + session, + Array.from(excluded) + ); + if (!provider || excluded.has(provider.id)) break; + if (session.sessionId && keyId != null) { + const cooldown = await SessionManager.isSessionProviderCoolingDown( + session.sessionId, + keyId, + provider.id + ); + if (cooldown.status === "ok" && cooldown.coolingDown) { + excluded.add(provider.id); + continue; + } + } + selected.push(provider); + excluded.add(provider.id); + } + return selected; + } + /** * 查找可复用的供应商（基于 session） */ @@ -465,9 +502,22 @@ export class ProxyProviderResolver { return null; } - // 从 Redis 读取该 session 绑定的 provider + // Read the binding once and retain its generation for Discovery timeout + // cleanup/finalization. Re-reading here would allow an older request to + // clear a newer binding (ABA). const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; - const providerId = await SessionManager.getSessionProvider(session.sessionId, keyId); + let providerId: number | null = null; + if (keyId != null) { + const binding = await SessionManager.getSessionBindingSnapshot(session.sessionId, keyId); + if (binding.status === "ok") { + session.setSessionBindingSnapshot(binding.snapshot); + providerId = binding.snapshot.providerId; + } else if (binding.legacyFallbackAllowed) { + providerId = await SessionManager.getSessionProvider(session.sessionId, keyId); + } + } else { + providerId = await SessionManager.getSessionProvider(session.sessionId, keyId); + } if (!providerId) { logger.debug("ProviderSelector: Session has no bound provider", { sessionId: session.sessionId, diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index acd5eb26f..bc36a1a39 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -1150,6 +1150,15 @@ function finalizeDeferredStreamingFinalizationIfNeeded( const clearSessionBinding = async () => { if (!session.sessionId) return; const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + if (meta?.bindingIntent === "none") return; + if (meta?.bindingSnapshot && keyId != null) { + await SessionManager.clearVersionedSessionProvider( + meta.bindingSnapshot, + providerIdForPersistence, + 0 + ); + return; + } await SessionManager.clearSessionProvider(session.sessionId, providerIdForPersistence, keyId); }; @@ -1165,6 +1174,10 @@ function finalizeDeferredStreamingFinalizationIfNeeded( const detected = shouldDetectFake200 ? detectUpstreamErrorFromSseOrJsonText(allContent) : ({ isError: false } as const); + const completionMarkerMissing = + meta?.requiresCompletionMarker === true && + streamEndedNormally && + !hasStreamCompletionMarker(allContent); let clientAbortGateUsage: FinalizeDeferredStreamingResult["clientAbortGateUsage"]; const clientAbortCompleteSuccess = (() => { if (!clientAborted || upstreamStatusCode < 200 || upstreamStatusCode >= 300) { @@ -1209,6 +1222,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( effectiveStatusCode = 502; } errorMessage = detected.detail ? `${detected.code}: ${detected.detail}` : detected.code; + } else if (completionMarkerMissing) { + effectiveStatusCode = 502; + errorMessage = "STREAM_COMPLETION_MARKER_MISSING"; } else if (clientAbortCompleteSuccess) { effectiveStatusCode = upstreamStatusCode; errorMessage = null; @@ -1239,6 +1255,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( const shouldClearSessionBindingOnFailure = ((clientAborted || !streamEndedNormally) && !clientAbortCompleteSuccess) || detected.isError || + completionMarkerMissing || (upstreamStatusCode >= 400 && errorMessage !== null); // 未启用延迟结算 / provider 缺失： @@ -1326,6 +1343,43 @@ function finalizeDeferredStreamingFinalizationIfNeeded( }; } + if (completionMarkerMissing) { + session.addProviderToChain(providerForChain, { + endpointId: meta.endpointId, + endpointUrl: meta.endpointUrl, + reason: "retry_failed", + attemptNumber: meta.attemptNumber, + statusCode: effectiveStatusCode, + errorMessage: errorMessage ?? undefined, + }); + + const commitSideEffects = async () => { + await clearSessionBinding(); + if (session.getEndpointPolicy().allowCircuitBreakerAccounting) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(errorMessage ?? "STREAM_ABORTED")); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record missing stream completion marker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } + } + }; + + return { + effectiveStatusCode, + errorMessage, + providerIdForPersistence, + isHedgeWinner, + billHedgeLosers, + clientAbortGateUsage, + commitSideEffects, + }; + } + if (detected.isError) { logger.warn("[ResponseHandler] SSE completed but body indicates error (fake 200)", { providerId: meta.providerId, @@ -1487,15 +1541,26 @@ function finalizeDeferredStreamingFinalizationIfNeeded( } // Hedge winner: commitWinner() already performed session binding and chain logging. - if (!meta.isHedgeWinner && session.sessionId) { - const result = await SessionManager.updateSessionBindingSmart( - session.sessionId, - meta.providerId, - meta.providerPriority, - meta.isFirstAttempt, - meta.isFailoverSuccess, - session.authState?.key?.id ?? session.messageContext?.key?.id ?? null - ); + if (meta.bindingIntent !== "none" && !meta.isHedgeWinner && session.sessionId) { + const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + const result = + meta.bindingSnapshot && keyId != null + ? await SessionManager.compareAndSetSessionProvider( + meta.bindingSnapshot, + meta.providerId + ).then((cas) => ({ + updated: cas.status === "ok", + reason: cas.status === "ok" ? "discovery_generation_cas" : cas.reason, + details: cas.status, + })) + : await SessionManager.updateSessionBindingSmart( + session.sessionId, + meta.providerId, + meta.providerPriority, + meta.isFirstAttempt, + meta.isFailoverSuccess, + keyId + ); if (result.updated) { logger.info("[ResponseHandler] Session binding updated (stream finalized)", { diff --git a/src/app/v1/_lib/proxy/session.ts b/src/app/v1/_lib/proxy/session.ts index 10d6b4740..18ff8bcae 100644 --- a/src/app/v1/_lib/proxy/session.ts +++ b/src/app/v1/_lib/proxy/session.ts @@ -1,6 +1,7 @@ import type { Context } from "hono"; import { logger } from "@/lib/logger"; import { writeLiveChain } from "@/lib/redis/live-chain-store"; +import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import { clientRequestsContext1m as clientRequestsContext1mHelper } from "@/lib/special-attributes"; import { type ResolvedPricing, @@ -206,6 +207,11 @@ export class ProxySession { // 失败切换 provider 时只能释放这里记录过的引用，避免 hedge/fallback 释放未 acquire 的 Redis 计数。 private providerSessionRefs = new Set(); + // Snapshot captured during provider selection. Discovery reuses this exact + // generation for timeout cleanup/finalization instead of performing a + // second read that could race with another request's binding update. + private sessionBindingSnapshot: SessionBindingSnapshot | null = null; + private constructor(init: { startTime: number; method: string; @@ -349,6 +355,14 @@ export class ProxySession { } } + setSessionBindingSnapshot(snapshot: SessionBindingSnapshot | null): void { + this.sessionBindingSnapshot = snapshot; + } + + getSessionBindingSnapshot(): SessionBindingSnapshot | null { + return this.sessionBindingSnapshot; + } + recordProviderSessionRef(providerId: number): void { if (!this.providerSessionRefs) { this.providerSessionRefs = new Set(); diff --git a/src/app/v1/_lib/proxy/stream-finalization.ts b/src/app/v1/_lib/proxy/stream-finalization.ts index 0f989bb3b..497207f1f 100644 --- a/src/app/v1/_lib/proxy/stream-finalization.ts +++ b/src/app/v1/_lib/proxy/stream-finalization.ts @@ -1,3 +1,4 @@ +import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import type { ProxySession } from "./session"; /** @@ -35,6 +36,11 @@ export type DeferredStreamingFinalization = { * coexists with asynchronously accumulated loser costs without clobbering. */ billHedgeLosers?: boolean; + /** Discovery delays binding until the stream has a valid completion marker. */ + bindingIntent?: "create" | "renew" | "none"; + bindingSnapshot?: SessionBindingSnapshot | null; + /** Discovery winners must satisfy the protocol completion marker before binding. */ + requiresCompletionMarker?: boolean; }; const deferredMeta = new WeakMap(); diff --git a/src/lib/config/system-settings-cache.ts b/src/lib/config/system-settings-cache.ts index 736b55238..6649603f5 100644 --- a/src/lib/config/system-settings-cache.ts +++ b/src/lib/config/system-settings-cache.ts @@ -166,6 +166,13 @@ export async function getCachedSystemSettings(): Promise { publicStatusWindowHours: DEFAULT_SETTINGS.publicStatusWindowHours, publicStatusAggregationIntervalMinutes: DEFAULT_SETTINGS.publicStatusAggregationIntervalMinutes, + discoveryEnabled: false, + discoveryConcurrency: 2, + maxDiscoveryRounds: 2, + discoverySlaMs: 10_000, + stickySlaMs: 20_000, + racingTotalTimeoutMs: 60_000, + stickyTimeoutCooldownMs: 300_000, quotaDbRefreshIntervalSeconds: 10, quotaLeasePercent5h: 0.05, quotaLeasePercentDaily: 0.05, diff --git a/src/types/system-config.ts b/src/types/system-config.ts index 3be265c11..16ef36a98 100644 --- a/src/types/system-config.ts +++ b/src/types/system-config.ts @@ -146,6 +146,15 @@ export interface SystemSettings { publicStatusWindowHours: number; publicStatusAggregationIntervalMinutes: number; + /** Bounded streaming Discovery (PR2; persisted/configured in PR3). */ + discoveryEnabled?: boolean; + discoveryConcurrency?: number; + maxDiscoveryRounds?: number; + discoverySlaMs?: number; + stickySlaMs?: number; + racingTotalTimeoutMs?: number; + stickyTimeoutCooldownMs?: number; + createdAt: Date; updatedAt: Date; } @@ -170,6 +179,14 @@ export interface UpdateSystemSettingsInput { // 供应商竞速输家计费（可选） billHedgeLosers?: boolean; + discoveryEnabled?: boolean; + discoveryConcurrency?: number; + maxDiscoveryRounds?: number; + discoverySlaMs?: number; + stickySlaMs?: number; + racingTotalTimeoutMs?: number; + stickyTimeoutCooldownMs?: number; + // 系统时区配置（可选） timezone?: string | null; diff --git a/tests/unit/proxy/discovery-coordinator.test.ts b/tests/unit/proxy/discovery-coordinator.test.ts new file mode 100644 index 000000000..70b09b6d0 --- /dev/null +++ b/tests/unit/proxy/discovery-coordinator.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from "vitest"; +import { DiscoveryCoordinator } from "@/app/v1/_lib/proxy/discovery-coordinator"; + +const attempt = (id: string, priority: number, kind: "normal" | "fallback" = "normal") => ({ + id, + providerId: Number(id.replace(/\D/g, "")) || 1, + priority, + kind, + ready: false, + pending: true, + round: 1, + launchOrder: Number(id.replace(/\D/g, "")) || 1, +}); + +describe("DiscoveryCoordinator", () => { + it("commits the highest priority ready normal attempt", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("a", 10)); + coordinator.addAttempt(attempt("b", 1)); + expect(coordinator.markReady("a")).toEqual({ type: "none" }); + expect(coordinator.markReady("b")).toEqual({ type: "commit_normal", attemptId: "b" }); + expect(coordinator.state).toBe("WINNER_COMMITTED"); + }); + + it("promotes one pending normal to fallback at a round boundary", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("a", 1)); + coordinator.addAttempt(attempt("b", 2)); + const action = coordinator.onRoundBoundary(); + expect(action.type).toBe("cancel"); + expect(coordinator.snapshot.find((item) => item.id === "a")?.kind).toBe("fallback"); + expect(coordinator.snapshot.filter((item) => item.pending)).toHaveLength(1); + }); + + it("ignores callbacks from an old request epoch", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("a", 1)); + const epoch = coordinator.epochs; + coordinator.cancelRequest(); + expect(coordinator.markReady("a", epoch.requestEpoch, epoch.roundEpoch)).toEqual({ + type: "none", + }); + }); + + it("promotes a ready fallback at the round boundary when no normal is ready", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("a", 1, "fallback")); + coordinator.addAttempt(attempt("b", 1, "normal")); + expect(coordinator.markReady("a")).toEqual({ type: "none" }); + expect(coordinator.onRoundBoundary()).toEqual({ type: "promote_fallback", attemptId: "a" }); + expect(coordinator.snapshot.find((item) => item.id === "a")?.kind).toBe("fallback"); + }); + + it("chooses the best ready normal at a boundary", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 3, maxRounds: 1 }); + coordinator.addAttempt(attempt("a", 10)); + coordinator.addAttempt(attempt("b", 1)); + coordinator.addAttempt(attempt("c", 5)); + coordinator.markReady("a"); + coordinator.markReady("c"); + expect(coordinator.onRoundBoundary()).toEqual({ type: "commit_normal", attemptId: "c" }); + }); +}); diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts new file mode 100644 index 000000000..4d5e6cae2 --- /dev/null +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from "vitest"; +import { + DiscoveryValidityParser, + classifyDiscoveryChunk, +} from "@/app/v1/_lib/proxy/discovery-validity"; + +describe("discovery validity", () => { + it("does not treat Anthropic metadata as a winner", () => { + expect(classifyDiscoveryChunk('data: {"type":"message_start"}\n\n', "anthropic").ready).toBe( + false + ); + expect( + classifyDiscoveryChunk( + 'data: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n', + "anthropic" + ).ready + ).toBe(true); + }); + + it("accepts OpenAI Chat delta and rejects DONE", () => { + expect( + classifyDiscoveryChunk('data: {"choices":[{"delta":{"content":"hi"}}]}\n', "openai-chat") + .ready + ).toBe(true); + expect(classifyDiscoveryChunk("data: [DONE]\n", "openai-chat").terminal).toBe(true); + }); + + it("rejects errors even when a later chunk contains content", () => { + const parser = new DiscoveryValidityParser("openai-responses"); + expect(parser.push('{"type":"response.failed","error":{"message":"no"}}').error).toBe(true); + expect(parser.push('{"type":"response.output_text.delta","delta":"late"}').ready).toBe(false); + }); + + it("does not promote empty tool or content events", () => { + expect( + classifyDiscoveryChunk( + 'data: {"type":"content_block_start","content_block":{"type":"text","text":""}}\n', + "anthropic" + ).ready + ).toBe(false); + expect( + classifyDiscoveryChunk( + 'data: {"choices":[{"delta":{"tool_calls":[{"function":{}}]}}]}\n', + "openai-chat" + ).ready + ).toBe(false); + expect( + classifyDiscoveryChunk( + 'data: {"type":"response.output_text.delta","delta":" "}\n', + "openai-responses" + ).ready + ).toBe(false); + }); + + it("accepts a non-empty function call delta as deliverable content", () => { + expect( + classifyDiscoveryChunk( + 'data: {"type":"response.function_call_arguments.delta","delta":"{\\"x\\":1}"}\n', + "openai-responses" + ).ready + ).toBe(true); + }); +}); From 31c4324d06606dba7815c34421bfa7cf2302b1b0 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 14:25:58 -0400 Subject: [PATCH 02/23] fix(proxy): keep discovery disabled path compatible --- src/app/v1/_lib/proxy/forwarder.ts | 8 +++++++- tests/integration/proxy-hedge-lifecycle.test.ts | 6 ++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 05cff7803..122e00d71 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -3828,7 +3828,13 @@ export class ProxyForwarder { private static async shouldUseStreamingDiscovery(session: ProxySession): Promise { const settings = await getCachedSystemSettings(); - if (SessionManager.getVersionedBindingCapabilityState() !== "available") { + if (settings.discoveryEnabled !== true) { + return false; + } + if ( + typeof SessionManager.getVersionedBindingCapabilityState !== "function" || + SessionManager.getVersionedBindingCapabilityState() !== "available" + ) { return false; } const endpointPolicy = ProxyForwarder.getEndpointPolicy(session); diff --git a/tests/integration/proxy-hedge-lifecycle.test.ts b/tests/integration/proxy-hedge-lifecycle.test.ts index ec3e9a533..390727090 100644 --- a/tests/integration/proxy-hedge-lifecycle.test.ts +++ b/tests/integration/proxy-hedge-lifecycle.test.ts @@ -518,7 +518,8 @@ describe("proxy hedge transport/lifecycle integration (persistence and control-p outputTokens: 3, providerId: 2, statusCode: 200, - }) + }), + expect.any(Object) ); expect(state.updateMessageRequestDetailsIfUnfinalized).not.toHaveBeenCalled(); @@ -615,7 +616,8 @@ describe("proxy hedge transport/lifecycle integration (persistence and control-p expect(state.durableTerminal).toHaveBeenCalledOnce(); expect(state.durableTerminal).toHaveBeenCalledWith( MESSAGE.id, - expect.objectContaining({ statusCode: 502 }) + expect.objectContaining({ statusCode: 502 }), + expect.any(Object) ); expect(agents.release).toHaveBeenCalledOnce(); expect(agents.pool.getPoolStats().activeRequests).toBe(0); From d4cba72c6206fc01d723e94221a4a94635a92854 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 14:54:02 -0400 Subject: [PATCH 03/23] fix(proxy): close discovery coordinator lifecycle gaps --- .../v1/_lib/proxy/discovery-coordinator.ts | 45 ++++- src/app/v1/_lib/proxy/discovery-validity.ts | 51 +++++- src/app/v1/_lib/proxy/forwarder.ts | 154 +++++++++++++----- .../unit/proxy/discovery-coordinator.test.ts | 22 ++- tests/unit/proxy/discovery-validity.test.ts | 14 ++ 5 files changed, 227 insertions(+), 59 deletions(-) diff --git a/src/app/v1/_lib/proxy/discovery-coordinator.ts b/src/app/v1/_lib/proxy/discovery-coordinator.ts index 44d9de404..4d4462b2c 100644 --- a/src/app/v1/_lib/proxy/discovery-coordinator.ts +++ b/src/app/v1/_lib/proxy/discovery-coordinator.ts @@ -29,8 +29,13 @@ export type DiscoveryAttempt = { export type DiscoveryAction = | { type: "commit_normal"; attemptId: string } | { type: "promote_fallback"; attemptId: string } - | { type: "cancel"; attemptIds: string[] } - | { type: "launch"; slots: number } + | { type: "cancel"; attemptIds: string[]; promoteAttemptId?: string } + | { + type: "launch"; + slots: number; + cancelAttemptIds?: string[]; + promoteAttemptId?: string; + } | { type: "none" } | { type: "terminal_failure" }; @@ -63,6 +68,7 @@ export class DiscoveryCoordinator { beginRound(): { requestEpoch: number; roundEpoch: number; round: number } { this.roundEpoch += 1; + if (!this.isTerminal) this.state = "DISCOVERY_RACING"; return { ...this.epochs, round: this.round }; } @@ -102,6 +108,17 @@ export class DiscoveryCoordinator { const attempt = this.attempts.get(id); if (!attempt?.pending) return { type: "none" }; attempt.ready = true; + if (attempt.kind === "fallback") { + const pendingNormal = Array.from(this.attempts.values()).some( + (candidate) => candidate.pending && candidate.kind === "normal" + ); + if (!pendingNormal) { + attempt.pending = false; + this.state = "FALLBACK_ACTIVE"; + return { type: "promote_fallback", attemptId: attempt.id }; + } + return { type: "none" }; + } return this.chooseReadyNormal(); } @@ -168,14 +185,19 @@ export class DiscoveryCoordinator { .filter((attempt) => attempt.pending && attempt.kind === "normal") .sort(compareAttempts); if (currentFallback && pendingNormal.length > 0) { + const cancelAttemptIds = pendingNormal.map((attempt) => attempt.id); for (const attempt of pendingNormal) attempt.pending = false; if (this.round < this.maxRounds) { this.round += 1; this.roundEpoch += 1; this.state = "DISCOVERY_RACING"; - return { type: "launch", slots: Math.max(1, this.concurrency - 1) }; + return { + type: "launch", + slots: Math.max(1, this.concurrency - 1), + cancelAttemptIds, + }; } - return { type: "none" }; + return { type: "cancel", attemptIds: cancelAttemptIds }; } if (pendingNormal.length === 0) { if (currentFallback) { @@ -191,11 +213,18 @@ export class DiscoveryCoordinator { const losers = pendingNormal.slice(1).map((attempt) => attempt.id); for (const id of losers) this.attempts.get(id)!.pending = false; - if (currentFallback) { - currentFallback.pending = true; - return { type: "cancel", attemptIds: losers }; + if (this.round < this.maxRounds) { + this.round += 1; + this.roundEpoch += 1; + this.state = "DISCOVERY_RACING"; + return { + type: "launch", + slots: Math.max(1, this.concurrency - 1), + cancelAttemptIds: losers, + promoteAttemptId: fallback.id, + }; } - return { type: "cancel", attemptIds: losers }; + return { type: "cancel", attemptIds: losers, promoteAttemptId: fallback.id }; } onDeadline(): DiscoveryAction { diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts index 29ad6a960..d4ece4eff 100644 --- a/src/app/v1/_lib/proxy/discovery-validity.ts +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -137,11 +137,58 @@ export class DiscoveryValidityParser { push(chunk: Uint8Array | string): DiscoveryValidity { this.buffered += typeof chunk === "string" ? chunk : this.decoder.decode(chunk, { stream: true }); - const result = classifyDiscoveryChunk(this.buffered, this.protocol); + + // SSE streams are line framed. Consume each completed line once instead + // of reparsing the complete prefix on every chunk (which is quadratic on + // long streams). Keep only the unfinished line for the next push. + if (this.buffered.includes("\n")) { + const lines = this.buffered.split(/\r?\n/); + this.buffered = lines.pop() ?? ""; + for (const line of lines) this.consumeLine(line); + } + + // Some providers return one raw JSON object without an SSE newline. Parse + // it only when the complete object is available; incomplete JSON remains + // buffered and is not repeatedly scanned as a protocol event. + const tail = this.buffered.trim(); + if (tail) { + const candidate = tail.startsWith("data:") ? tail.slice(5).trim() : tail; + if (candidate === "[DONE]") { + this._terminal = true; + this.buffered = ""; + } else if (candidate.startsWith("{") || candidate.startsWith("[")) { + try { + const value = JSON.parse(candidate) as unknown; + this.consumeValue(value); + this.buffered = ""; + } catch { + // Keep incomplete raw JSON until the next chunk completes it. + } + } + } + + return { ready: this._ready && !this._error, terminal: this._terminal, error: this._error }; + } + + private consumeLine(line: string): void { + const candidate = line.startsWith("data:") ? line.slice(5).trim() : line.trim(); + if (!candidate || candidate.startsWith(":")) return; + if (candidate === "[DONE]") { + this._terminal = true; + return; + } + try { + this.consumeValue(JSON.parse(candidate) as unknown); + } catch { + // Ignore comments and incomplete/non-JSON protocol lines. + } + } + + private consumeValue(value: unknown): void { + const result = classifyJson(value, this.protocol); this._ready ||= result.ready; this._terminal ||= result.terminal; this._error ||= result.error; - return { ready: this._ready && !this._error, terminal: this._terminal, error: this._error }; } get ready(): boolean { diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 122e00d71..18182a2e1 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -63,7 +63,7 @@ import { buildProxyUrl } from "../url"; import { rectifyBillingHeader } from "./billing-header-rectifier"; import { bindClientAbortListener } from "./client-abort-listener"; import { deriveClientSafeUpstreamErrorMessage } from "./client-error-message"; -import { DiscoveryCoordinator } from "./discovery-coordinator"; +import { type DiscoveryAction, DiscoveryCoordinator } from "./discovery-coordinator"; import { type DiscoveryProtocol, DiscoveryValidityParser } from "./discovery-validity"; import { isStandardProxyEndpointPath } from "./endpoint-family-catalog"; import { resolveEndpointPolicy, shouldEnforceStrictEndpointPoolPolicy } from "./endpoint-policy"; @@ -4875,6 +4875,7 @@ export class ProxyForwarder { pending: boolean; ready: boolean; round: number; + readerTransferred: boolean; } >(); const launched = new Set(); @@ -4889,6 +4890,7 @@ export class ProxyForwarder { let totalTimer: NodeJS.Timeout | null = null; let roundTimer: NodeJS.Timeout | null = null; let stickyTimer: NodeJS.Timeout | null = null; + let executeCoordinatorAction: (action: DiscoveryAction) => Promise = async () => {}; let resolveResult: ((result: { response?: Response; error?: Error }) => void) | null = null; const resultPromise = new Promise<{ response?: Response; error?: Error }>((resolve) => { resolveResult = resolve; @@ -4904,6 +4906,7 @@ export class ProxyForwarder { }; const cancelAttempt = (attempt: (typeof winner & { id: string }) | null, reason: string) => { + if (attempt?.readerTransferred) return; if (!attempt?.pending) return; attempt.pending = false; try { @@ -4911,7 +4914,12 @@ export class ProxyForwarder { } catch { /* abort is best effort */ } - void attempt.reader?.cancel(reason).catch(() => undefined); + try { + const cancelPromise = attempt.reader?.cancel(reason); + cancelPromise?.catch(() => undefined); + } catch (error) { + logger.debug("[Discovery] Reader cancel failed", { reason, error }); + } try { attempt.releaseAgent?.(); } catch { @@ -4933,8 +4941,18 @@ export class ProxyForwarder { if (roundTimer) clearTimeout(roundTimer); if (stickyTimer) clearTimeout(stickyTimer); cancelLosers(); - const attempted = new Set(launched); - await ProxyForwarder.clearSessionProviderBindings(session, attempted); + if (bindingSnapshot) { + if (bindingSnapshot.providerId != null) { + await SessionManager.clearVersionedSessionProvider( + bindingSnapshot, + bindingSnapshot.providerId, + 0 + ); + } + } else { + const attempted = new Set(launched); + await ProxyForwarder.clearSessionProviderBindings(session, attempted); + } resolveResult?.({ error }); }; @@ -4943,6 +4961,9 @@ export class ProxyForwarder { committed = true; winner = attempt; attempt.pending = false; + // From this point ResponseHandler owns the reader and agent release. + // No coordinator/timer path may cancel or release this attempt again. + attempt.readerTransferred = true; if (totalTimer) clearTimeout(totalTimer); if (roundTimer) clearTimeout(roundTimer); if (stickyTimer) clearTimeout(stickyTimer); @@ -5044,6 +5065,7 @@ export class ProxyForwarder { pending: true, ready: false, round: currentRound, + readerTransferred: false, provider, session: attemptSession, baseUrl: endpoint.baseUrl, @@ -5080,6 +5102,7 @@ export class ProxyForwarder { pending: boolean; ready: boolean; round: number; + readerTransferred: boolean; }; attempts.set(id, attempt); coordinator.addAttempt({ @@ -5138,7 +5161,7 @@ export class ProxyForwarder { .catch(async (error) => { if (committed || settled || !attempt.pending) return; attempt.pending = false; - coordinator.markFailed(id); + const failureAction = coordinator.markFailed(id); lastError = error instanceof Error ? error : new Error(String(error)); lastErrorCategory = await categorizeErrorAsync(lastError); session.addProviderToChain(provider, { @@ -5156,8 +5179,17 @@ export class ProxyForwarder { } attempt.releaseAgent?.(); releaseProviderRef(provider.id); - const replacement = await chooseCandidate(); - if (replacement && !committed && !settled) await launch(replacement, "normal"); + const actionOwnsNextStep = + failureAction.type === "promote_fallback" || + failureAction.type === "launch" || + failureAction.type === "terminal_failure"; + if (actionOwnsNextStep) { + await executeCoordinatorAction(failureAction); + } + if (!actionOwnsNextStep && !committed && !settled) { + const replacement = await chooseCandidate(); + if (replacement) await launch(replacement, "normal"); + } if ( Array.from(attempts.values()).every((candidate) => !candidate.pending) && noMoreCandidates @@ -5166,6 +5198,12 @@ export class ProxyForwarder { ProxyForwarder.resolveHedgeTerminalError(lastError, lastErrorCategory) ); } + }) + .catch((error) => { + logger.warn("[Discovery] Attempt completion handler failed", { + providerId: provider.id, + error, + }); }); }; @@ -5184,60 +5222,75 @@ export class ProxyForwarder { } } if (!committed && !settled) { - roundTimer = setTimeout(() => void onBoundary(), discoverySlaMs); + roundTimer = setTimeout(() => { + void onBoundary().catch((error) => + logger.warn("[Discovery] Round boundary failed", { error }) + ); + }, discoverySlaMs); } }; - const onBoundary = async () => { + executeCoordinatorAction = async (action) => { if (settled || committed) return; - const fallback = Array.from(attempts.values()).find( - (attempt) => attempt.pending && attempt.kind === "fallback" - ); - const pendingNormals = Array.from(attempts.values()) - .filter((attempt) => attempt.pending && attempt.kind === "normal") - .sort( - (a, b) => - (a.provider.priority || 0) - (b.provider.priority || 0) || a.sequence - b.sequence - ); - const readyNormal = pendingNormals.find((attempt) => attempt.ready); - if (readyNormal) { - await commit(readyNormal); + if (action.type === "cancel" || action.type === "launch") { + const cancelIds = + action.type === "cancel" ? action.attemptIds : (action.cancelAttemptIds ?? []); + for (const id of cancelIds) { + const attempt = attempts.get(id); + // Coordinator marks cancelled attempts non-pending before returning + // the action. Restore the transport-facing state long enough for the + // exactly-once cancellation/release path to run. + if (attempt && !attempt.readerTransferred) attempt.pending = true; + if (attempt) cancelAttempt(attempt, "discovery_round_boundary"); + } + if (action.promoteAttemptId) { + const fallback = attempts.get(action.promoteAttemptId); + if (fallback) fallback.kind = "fallback"; + if (action.type === "cancel" && currentRound < maxRounds) { + await launchNextRound(); + } + } + } + if (action.type === "commit_normal" || action.type === "promote_fallback") { + const attempt = attempts.get(action.attemptId); + if (attempt) await commit(attempt); return; } - // The fallback is held during the SLA window, but at the round - // boundary it may take over when no normal result is ready. - if (fallback?.ready) { - await commit(fallback); + if (action.type === "launch") { + await launchNextRound(); return; } - if (pendingNormals[0]) { - if (fallback) { - for (const loser of pendingNormals) cancelAttempt(loser, "discovery_round_boundary"); - if (currentRound < maxRounds) await launchNextRound(); - return; - } - pendingNormals[0].kind = "fallback"; - for (const loser of pendingNormals.slice(1)) - cancelAttempt(loser, "discovery_round_boundary"); - if (currentRound < maxRounds) await launchNextRound(); + if (action.type === "terminal_failure") { + await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); return; } - if (currentRound < maxRounds) await launchNextRound(); - else await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); + if (action.type === "none") { + const fallbackPending = Array.from(attempts.values()).some( + (attempt) => attempt.pending && attempt.kind === "fallback" + ); + if (fallbackPending && currentRound < maxRounds) { + await launchNextRound(); + } + } + }; + + const onBoundary = async () => { + if (settled || committed) return; + await executeCoordinatorAction(coordinator.onRoundBoundary()); }; const cleanupAbort = bindClientAbortListener(session.clientAbortSignal, () => { if (settled || committed) return; - void settleFailure(new ProxyError("Request aborted by client", 499, undefined, true)); + void settleFailure(new ProxyError("Request aborted by client", 499, undefined, true)).catch( + (error) => logger.warn("[Discovery] Client abort cleanup failed", { error }) + ); }); totalTimer = setTimeout(() => { if (settled || committed) return; - const fallback = Array.from(attempts.values()).find( - (attempt) => attempt.pending && attempt.kind === "fallback" && attempt.ready + void executeCoordinatorAction(coordinator.onDeadline()).catch((error) => + logger.warn("[Discovery] Deadline action failed", { error }) ); - if (fallback) void commit(fallback); - else void settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); }, totalTimeoutMs); try { @@ -5267,8 +5320,15 @@ export class ProxyForwarder { logger.debug("[Discovery] Failed to clear timed-out Sticky", { error }) ); } - if (currentRound < maxRounds) void launchNextRound(); - else void onBoundary(); + if (currentRound < maxRounds) { + void launchNextRound().catch((error) => + logger.warn("[Discovery] Sticky round launch failed", { error }) + ); + } else { + void onBoundary().catch((error) => + logger.warn("[Discovery] Sticky boundary failed", { error }) + ); + } } }, stickySlaMs); } else { @@ -5285,7 +5345,11 @@ export class ProxyForwarder { noMoreCandidates = true; } } - roundTimer = setTimeout(() => void onBoundary(), discoverySlaMs); + roundTimer = setTimeout(() => { + void onBoundary().catch((error) => + logger.warn("[Discovery] Round boundary failed", { error }) + ); + }, discoverySlaMs); } const result = await resultPromise; if (result.error) throw result.error; diff --git a/tests/unit/proxy/discovery-coordinator.test.ts b/tests/unit/proxy/discovery-coordinator.test.ts index 70b09b6d0..35585c473 100644 --- a/tests/unit/proxy/discovery-coordinator.test.ts +++ b/tests/unit/proxy/discovery-coordinator.test.ts @@ -27,7 +27,11 @@ describe("DiscoveryCoordinator", () => { coordinator.addAttempt(attempt("a", 1)); coordinator.addAttempt(attempt("b", 2)); const action = coordinator.onRoundBoundary(); - expect(action.type).toBe("cancel"); + expect(action).toMatchObject({ + type: "launch", + promoteAttemptId: "a", + cancelAttemptIds: ["b"], + }); expect(coordinator.snapshot.find((item) => item.id === "a")?.kind).toBe("fallback"); expect(coordinator.snapshot.filter((item) => item.pending)).toHaveLength(1); }); @@ -45,9 +49,7 @@ describe("DiscoveryCoordinator", () => { it("promotes a ready fallback at the round boundary when no normal is ready", () => { const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); coordinator.addAttempt(attempt("a", 1, "fallback")); - coordinator.addAttempt(attempt("b", 1, "normal")); - expect(coordinator.markReady("a")).toEqual({ type: "none" }); - expect(coordinator.onRoundBoundary()).toEqual({ type: "promote_fallback", attemptId: "a" }); + expect(coordinator.markReady("a")).toEqual({ type: "promote_fallback", attemptId: "a" }); expect(coordinator.snapshot.find((item) => item.id === "a")?.kind).toBe("fallback"); }); @@ -60,4 +62,16 @@ describe("DiscoveryCoordinator", () => { coordinator.markReady("c"); expect(coordinator.onRoundBoundary()).toEqual({ type: "commit_normal", attemptId: "c" }); }); + + it("reports normal attempts cancelled when retaining an existing fallback", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 3 }); + coordinator.addAttempt(attempt("fallback", 1, "fallback")); + coordinator.addAttempt(attempt("normal", 2)); + const action = coordinator.onRoundBoundary(); + expect(action).toEqual({ + type: "launch", + slots: 1, + cancelAttemptIds: ["normal"], + }); + }); }); diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index 4d5e6cae2..83feb4f8e 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -60,4 +60,18 @@ describe("discovery validity", () => { ).ready ).toBe(true); }); + + it("consumes split SSE lines incrementally without waiting for the full stream", () => { + const parser = new DiscoveryValidityParser("openai-chat"); + expect(parser.push('data: {"choices":[{"delta":{"content":"hel')).toEqual({ + ready: false, + terminal: false, + error: false, + }); + expect(parser.push('lo"}}]}\n\n')).toEqual({ + ready: true, + terminal: false, + error: false, + }); + }); }); From 9c7793c86e29b6fc8fc4d56163d3c9c3e23e3d76 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 14:56:20 -0400 Subject: [PATCH 04/23] fix(proxy): guard optional discovery request message --- src/app/v1/_lib/proxy/forwarder.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 18182a2e1..753e7d77f 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -3848,7 +3848,7 @@ export class ProxyForwarder { settings.discoveryEnabled === true && endpointPolicy.allowRetry && endpointPolicy.allowProviderSwitch && - message.stream === true && + message?.stream === true && !endpointPolicy.bypassForwarderPreprocessing && protocol !== "unknown" && routing.routingMode !== "lease_conflict_single" && From ee5557ecaf36aaf0915c58d6626128935de023ee Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 14:58:09 -0400 Subject: [PATCH 05/23] fix(proxy): handle replacement launch failures --- src/app/v1/_lib/proxy/forwarder.ts | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 753e7d77f..a96f93e90 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5300,8 +5300,19 @@ export class ProxyForwarder { } catch (error) { lastError = error instanceof Error ? error : new Error(String(error)); const replacement = await chooseCandidate(); - if (replacement) await launch(replacement, "normal"); - else await settleFailure(ProxyForwarder.resolveHedgeTerminalError(lastError, null)); + if (replacement) { + try { + await launch(replacement, "normal"); + } catch (replacementError) { + lastError = + replacementError instanceof Error + ? replacementError + : new Error(String(replacementError)); + await settleFailure(ProxyForwarder.resolveHedgeTerminalError(lastError, null)); + } + } else { + await settleFailure(ProxyForwarder.resolveHedgeTerminalError(lastError, null)); + } } const initial = hasSticky ? concurrency - 1 : Math.max(0, concurrency - 1); if (hasSticky) { From 43c6733e434963e5500e9c33dfc7b504771977fb Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 15:55:23 -0400 Subject: [PATCH 06/23] fix(proxy): preserve complete discovery candidates --- .../v1/_lib/proxy/discovery-coordinator.ts | 7 +++++ src/app/v1/_lib/proxy/forwarder.ts | 26 ++++++++++++------- .../unit/proxy/discovery-coordinator.test.ts | 13 ++++++++++ tests/unit/proxy/discovery-validity.test.ts | 12 +++++++++ 4 files changed, 49 insertions(+), 9 deletions(-) diff --git a/src/app/v1/_lib/proxy/discovery-coordinator.ts b/src/app/v1/_lib/proxy/discovery-coordinator.ts index 4d4462b2c..a38d0813b 100644 --- a/src/app/v1/_lib/proxy/discovery-coordinator.ts +++ b/src/app/v1/_lib/proxy/discovery-coordinator.ts @@ -264,6 +264,13 @@ export class DiscoveryCoordinator { private afterAttemptState(): DiscoveryAction { const pending = this.activeAttempts; if (pending.length === 0) return this.finishOrLaunch(); + + // A higher-priority attempt may have been the only gate preventing a + // ready lower-priority candidate from winning. Once that attempt fails, + // re-run the normal winner selection before waiting for another boundary. + const readyNormal = this.chooseReadyNormal(); + if (readyNormal.type === "commit_normal") return readyNormal; + const fallback = pending.find((attempt) => attempt.kind === "fallback"); if (fallback?.ready && pending.every((attempt) => attempt.kind === "fallback")) { return this.commitWinner(fallback.id); diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index a96f93e90..7648319c5 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5137,24 +5137,32 @@ export class ProxyForwarder { attempt.reader = response.body.getReader(); while (!committed && !settled && attempt.pending) { const item = await attempt.reader.read(); - if (item.done) throw new EmptyResponseError(provider.id, provider.name, "empty_body"); + if (item.done) { + // A ready candidate may have reached EOF while waiting for a + // higher-priority attempt. Its buffered prefix remains a valid + // response and must stay promotable. + if (attempt.ready) return; + throw new EmptyResponseError(provider.id, provider.name, "empty_body"); + } if (!item.value || item.value.byteLength === 0) continue; attempt.chunks.push(item.value); const validity = attempt.parser.push(item.value); - if (validity.error || validity.terminal) + // A single read can contain both deliverable content and the + // protocol terminator. Terminal is only invalid when no content + // was observed; otherwise the buffered candidate is complete. + if (validity.error || (validity.terminal && !validity.ready)) throw new ProxyError("Invalid upstream discovery response", 502); if (!validity.ready) continue; attempt.ready = true; - const normalPendingHigher = Array.from(attempts.values()).some( - (other) => - other.pending && - other.kind === "normal" && - (other.provider.priority || 0) < (provider.priority || 0) - ); - if (normalPendingHigher) continue; const action = coordinator.markReady(id); if (action.type === "commit_normal" || action.type === "promote_fallback") await commit(attempt); + // Do not issue another reader request after a complete candidate; + // the buffered stream is already sufficient for later promotion. + if (validity.terminal) return; + // The coordinator owns priority gating. A ready lower-priority + // candidate stays held while a higher tier is still pending. + if (action.type === "none") continue; return; } }) diff --git a/tests/unit/proxy/discovery-coordinator.test.ts b/tests/unit/proxy/discovery-coordinator.test.ts index 35585c473..328a02e97 100644 --- a/tests/unit/proxy/discovery-coordinator.test.ts +++ b/tests/unit/proxy/discovery-coordinator.test.ts @@ -74,4 +74,17 @@ describe("DiscoveryCoordinator", () => { cancelAttemptIds: ["normal"], }); }); + + it("retains a lower-priority ready candidate until the higher tier fails", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("high", 1)); + coordinator.addAttempt(attempt("low", 10)); + + expect(coordinator.markReady("low")).toEqual({ type: "none" }); + expect(coordinator.snapshot.find((item) => item.id === "low")).toMatchObject({ + ready: true, + pending: true, + }); + expect(coordinator.markFailed("high")).toEqual({ type: "commit_normal", attemptId: "low" }); + }); }); diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index 83feb4f8e..408fca35a 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -74,4 +74,16 @@ describe("discovery validity", () => { error: false, }); }); + + it("keeps ready when content and the terminal marker arrive in one read", () => { + const parser = new DiscoveryValidityParser("openai-chat"); + + expect( + parser.push('data: {"choices":[{"delta":{"content":"done"}}]}\n\ndata: [DONE]\n\n') + ).toEqual({ + ready: true, + terminal: true, + error: false, + }); + }); }); From 7e6dadf9cf85c5fc17a9bab622dbbdf64a167910 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 16:05:01 -0400 Subject: [PATCH 07/23] fix(proxy): pause held discovery readers --- src/app/v1/_lib/proxy/forwarder.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 7648319c5..97be3effa 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5161,8 +5161,9 @@ export class ProxyForwarder { // the buffered stream is already sufficient for later promotion. if (validity.terminal) return; // The coordinator owns priority gating. A ready lower-priority - // candidate stays held while a higher tier is still pending. - if (action.type === "none") continue; + // candidate stays held while a higher tier is still pending. Stop + // reading so later chunks are not consumed before promotion. + if (action.type === "none") return; return; } }) From 8e0cedba26e9309b069cc73203e550422a3acf92 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 16:13:11 -0400 Subject: [PATCH 08/23] fix(proxy): close discovery attempt cleanup gaps --- .../v1/_lib/proxy/discovery-coordinator.ts | 14 +++++++ src/app/v1/_lib/proxy/discovery-validity.ts | 12 +++++- src/app/v1/_lib/proxy/forwarder.ts | 37 ++++++++++++------- .../unit/proxy/discovery-coordinator.test.ts | 15 ++++++++ tests/unit/proxy/discovery-validity.test.ts | 15 ++++++++ ...forwarder-provider-session-release.test.ts | 16 ++++++++ 6 files changed, 94 insertions(+), 15 deletions(-) diff --git a/src/app/v1/_lib/proxy/discovery-coordinator.ts b/src/app/v1/_lib/proxy/discovery-coordinator.ts index a38d0813b..4a94baf70 100644 --- a/src/app/v1/_lib/proxy/discovery-coordinator.ts +++ b/src/app/v1/_lib/proxy/discovery-coordinator.ts @@ -122,6 +122,20 @@ export class DiscoveryCoordinator { return this.chooseReadyNormal(); } + /** Convert a timed-out Sticky attempt into the request's fallback lane. */ + demoteToFallback( + id: string, + requestEpoch = this.requestEpoch, + roundEpoch = this.roundEpoch + ): boolean { + if (!this.acceptsEpoch(requestEpoch, roundEpoch) || this.isTerminal) return false; + const attempt = this.attempts.get(id); + if (!attempt?.pending) return false; + attempt.kind = "fallback"; + this.state = "FALLBACK_READY_HELD"; + return true; + } + markFailed( id: string, requestEpoch = this.requestEpoch, diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts index d4ece4eff..198d77699 100644 --- a/src/app/v1/_lib/proxy/discovery-validity.ts +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -28,10 +28,20 @@ function hasContent(value: unknown): boolean { "function_call", "arguments", "input", + "partial_json", "parts", ].some((key) => hasContent(object[key])); } +function hasAnthropicContentBlock(value: unknown): boolean { + if (!value || typeof value !== "object" || Array.isArray(value)) return false; + const block = value as Record; + if (typeof block.type !== "string" || block.type.length === 0) return false; + // Text blocks need non-empty text; tool_use/thinking/image blocks are + // deliverable as soon as their typed block starts, even with empty input. + return block.type === "text" ? hasContent(block.text) : true; +} + function classifyJson(value: unknown, protocol: DiscoveryProtocol): DiscoveryValidity { if (!value || typeof value !== "object") return { ready: false, terminal: false, error: true }; const object = value as Record; @@ -89,7 +99,7 @@ function classifyJson(value: unknown, protocol: DiscoveryProtocol): DiscoveryVal return { ready: (object.type === "content_block_delta" && hasContent(object.delta)) || - (object.type === "content_block_start" && hasContent(object.content_block)) || + (object.type === "content_block_start" && hasAnthropicContentBlock(object.content_block)) || hasContent(object.content), terminal: false, error: false, diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 97be3effa..fdf0de21c 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -4897,12 +4897,7 @@ export class ProxyForwarder { }); const releaseProviderRef = (providerId: number) => { - if (!session.sessionId) return; - const consumer = (session as { consumeProviderSessionRef?: (id: number) => boolean }) - .consumeProviderSessionRef; - if (consumer?.call(session, providerId)) { - void RateLimitService.releaseProviderSession(providerId, session.sessionId); - } + ProxyForwarder.releaseProviderSessionRef(session, providerId); }; const cancelAttempt = (attempt: (typeof winner & { id: string }) | null, reason: string) => { @@ -5035,6 +5030,7 @@ export class ProxyForwarder { const launch = async (provider: Provider, kind: "normal" | "fallback"): Promise => { if (settled || committed || launched.has(provider.id)) return; launched.add(provider.id); + let providerSessionRefRecorded = false; if (provider.id !== initialProvider.id && session.sessionId) { const limit = provider.limitConcurrentSessions || 0; const check = await RateLimitService.checkAndTrackProviderSession( @@ -5046,9 +5042,19 @@ export class ProxyForwarder { launched.delete(provider.id); throw new ProxyError(check.reason || "Provider concurrent limit reached", 503); } - if (check.referenced) session.recordProviderSessionRef(provider.id); + if (check.referenced) { + session.recordProviderSessionRef(provider.id); + providerSessionRefRecorded = true; + } + } + let endpoint: Awaited>; + try { + endpoint = await ProxyForwarder.resolveStreamingHedgeEndpoint(session, provider); + } catch (error) { + launched.delete(provider.id); + if (providerSessionRefRecorded) releaseProviderRef(provider.id); + throw error; } - const endpoint = await ProxyForwarder.resolveStreamingHedgeEndpoint(session, provider); const attemptSession = provider.id === initialProvider.id ? session @@ -5330,6 +5336,7 @@ export class ProxyForwarder { (attempt) => attempt.pending && attempt.provider.id === initialProvider.id ); if (sticky) { + if (!coordinator.demoteToFallback(sticky.id)) return; sticky.kind = "fallback"; if (bindingSnapshot && bindingSnapshot.providerId === initialProvider.id) { void SessionManager.clearVersionedSessionProvider( @@ -5641,15 +5648,17 @@ export class ProxyForwarder { return; } + ProxyForwarder.releaseProviderSessionRef(session, providerId); + } + + private static releaseProviderSessionRef(session: ProxySession, providerId: number): boolean { + if (!session.sessionId) return false; const providerSessionRefConsumer = ( - session as { consumeProviderSessionRef?: (providerId: number) => boolean } + session as { consumeProviderSessionRef?: (id: number) => boolean } ).consumeProviderSessionRef; - - if (!providerSessionRefConsumer?.call(session, providerId)) { - return; - } - + if (!providerSessionRefConsumer?.call(session, providerId)) return false; void RateLimitService.releaseProviderSession(providerId, session.sessionId); + return true; } private static buildAllProvidersUnavailableError(finalError?: Error | null): ProxyError { diff --git a/tests/unit/proxy/discovery-coordinator.test.ts b/tests/unit/proxy/discovery-coordinator.test.ts index 328a02e97..3a41a8d5e 100644 --- a/tests/unit/proxy/discovery-coordinator.test.ts +++ b/tests/unit/proxy/discovery-coordinator.test.ts @@ -53,6 +53,21 @@ describe("DiscoveryCoordinator", () => { expect(coordinator.snapshot.find((item) => item.id === "a")?.kind).toBe("fallback"); }); + it("keeps Sticky demotion synchronized with the coordinator", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("sticky", 1)); + + expect(coordinator.demoteToFallback("sticky")).toBe(true); + expect(coordinator.snapshot.find((item) => item.id === "sticky")).toMatchObject({ + kind: "fallback", + pending: true, + }); + expect(coordinator.markReady("sticky")).toEqual({ + type: "promote_fallback", + attemptId: "sticky", + }); + }); + it("chooses the best ready normal at a boundary", () => { const coordinator = new DiscoveryCoordinator({ concurrency: 3, maxRounds: 1 }); coordinator.addAttempt(attempt("a", 10)); diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index 408fca35a..7a6f95fc4 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -86,4 +86,19 @@ describe("discovery validity", () => { error: false, }); }); + + it("accepts Anthropic tool-use starts and partial JSON deltas", () => { + expect( + classifyDiscoveryChunk( + 'data: {"type":"content_block_start","content_block":{"type":"tool_use","id":"tu_1","name":"search","input":{}}}\n', + "anthropic" + ).ready + ).toBe(true); + expect( + classifyDiscoveryChunk( + 'data: {"type":"content_block_delta","delta":{"type":"input_json_delta","partial_json":"{\\"q\\":1}"}}\n', + "anthropic" + ).ready + ).toBe(true); + }); }); diff --git a/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts b/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts index efc8a9028..0fa7ebafa 100644 --- a/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts +++ b/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts @@ -68,6 +68,22 @@ describe("ProxyForwarder provider failure session release", () => { expect(mocks.releaseProviderSession).not.toHaveBeenCalled(); }); + it("endpoint resolution rollback releases only a recorded provider ref", async () => { + const { ProxyForwarder } = await import("@/app/v1/_lib/proxy/forwarder"); + const forwarderInternals = ProxyForwarder as unknown as { + releaseProviderSessionRef: (session: ProxySession, providerId: number) => boolean; + }; + const consumeProviderSessionRef = vi.fn(() => true); + const session = { + sessionId: "sess_endpoint_failure", + consumeProviderSessionRef, + } as unknown as ProxySession; + + expect(forwarderInternals.releaseProviderSessionRef(session, 42)).toBe(true); + expect(consumeProviderSessionRef).toHaveBeenCalledWith(42); + expect(mocks.releaseProviderSession).toHaveBeenCalledWith(42, "sess_endpoint_failure"); + }); + it("重复标记同一供应商时只释放一次，避免 hedge 路径重复 ZREM", async () => { const { ProxyForwarder } = await import("@/app/v1/_lib/proxy/forwarder"); const forwarderInternals = ProxyForwarder as unknown as { From 0f7e1701f6882b3d94cb17a89cd7f6d8ff1a7d27 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 16:27:18 -0400 Subject: [PATCH 09/23] fix(proxy): execute discovery normal winner actions --- src/app/v1/_lib/proxy/forwarder.ts | 1 + .../proxy-forwarder-hedge-first-byte.test.ts | 91 +++++++++++++++++++ 2 files changed, 92 insertions(+) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index fdf0de21c..df00d6009 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5195,6 +5195,7 @@ export class ProxyForwarder { attempt.releaseAgent?.(); releaseProviderRef(provider.id); const actionOwnsNextStep = + failureAction.type === "commit_normal" || failureAction.type === "promote_fallback" || failureAction.type === "launch" || failureAction.type === "terminal_failure"; diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index fe5d7f4a8..1458d825d 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -3,6 +3,7 @@ import { resolveEndpointPolicy } from "@/app/v1/_lib/proxy/endpoint-policy"; const mocks = vi.hoisted(() => ({ pickRandomProviderWithExclusion: vi.fn(), + pickDiscoveryProviders: vi.fn(), recordSuccess: vi.fn(), recordFailure: vi.fn(async () => {}), getCircuitState: vi.fn(() => "closed"), @@ -37,6 +38,7 @@ const mocks = vi.hoisted(() => ({ storeSessionSpecialSettings: vi.fn(async () => {}), storeSessionRequestPhaseSnapshot: vi.fn(async () => {}), storeSessionResponsePhaseSnapshot: vi.fn(async () => {}), + getVersionedBindingCapabilityState: vi.fn(() => "available"), })); vi.mock("@/lib/logger", () => ({ @@ -90,6 +92,7 @@ vi.mock("@/lib/rate-limit/service", () => ({ vi.mock("@/lib/session-manager", () => ({ SessionManager: { + getVersionedBindingCapabilityState: mocks.getVersionedBindingCapabilityState, updateSessionBindingSmart: mocks.updateSessionBindingSmart, updateSessionProvider: mocks.updateSessionProvider, clearSessionProvider: mocks.clearSessionProvider, @@ -103,6 +106,7 @@ vi.mock("@/lib/session-manager", () => ({ vi.mock("@/app/v1/_lib/proxy/provider-selector", () => ({ ProxyProviderResolver: { pickRandomProviderWithExclusion: mocks.pickRandomProviderWithExclusion, + pickDiscoveryProviders: mocks.pickDiscoveryProviders, }, })); @@ -333,6 +337,10 @@ function withThinkingBlocks(session: ProxySession): void { describe("ProxyForwarder - first-byte hedge scheduling", () => { beforeEach(() => { vi.clearAllMocks(); + mocks.getCachedSystemSettings.mockResolvedValue({ + enableThinkingSignatureRectifier: true, + enableThinkingBudgetRectifier: true, + }); mocks.checkAndTrackProviderSession.mockResolvedValue({ allowed: true, count: 1, @@ -2061,6 +2069,89 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { } }); + test("Discovery commits a lower-priority ready stream when the higher tier fails", async () => { + vi.useFakeTimers(); + + try { + const high = createProvider({ id: 1, name: "high", priority: 1 }); + const low = createProvider({ id: 2, name: "low", priority: 10 }); + const session = createSession(); + session.setProvider(high); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 100, + stickySlaMs: 100, + racingTotalTimeoutMs: 500, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([low]); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + + doForward.mockImplementationOnce( + async (_attemptSession, _provider, _baseUrl, _audit, _count, _stream, signal) => { + await new Promise((_resolve, reject) => { + const timer = setTimeout(() => reject(new Error("high tier failed")), 30); + signal?.addEventListener( + "abort", + () => { + clearTimeout(timer); + reject(new Error("high tier aborted")); + }, + { once: true } + ); + }); + } + ); + doForward.mockImplementationOnce( + async (_attemptSession, _provider, _baseUrl, _audit, _count, _stream, signal) => { + const stream = new ReadableStream({ + start(controller) { + const timer = setTimeout(() => { + controller.enqueue( + new TextEncoder().encode( + 'data: {"type":"content_block_delta","delta":{"text":"low"}}\n\n' + ) + ); + controller.close(); + }, 5); + signal?.addEventListener( + "abort", + () => { + clearTimeout(timer); + controller.close(); + }, + { once: true } + ); + }, + }); + return new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + ); + + const responsePromise = ProxyForwarder.send(session); + await vi.advanceTimersByTimeAsync(10); + await vi.advanceTimersByTimeAsync(30); + const response = await responsePromise; + + expect(await response.text()).toContain('"low"'); + expect(session.provider?.id).toBe(low.id); + expect(doForward).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); + test("removes streaming hedge client abort listener after winner response is returned", async () => { const clientAbortController = new AbortController(); const addSpy = vi.spyOn(clientAbortController.signal, "addEventListener"); From 0690278f9e17539b6a611359c924fecff6e60526 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 16:32:17 -0400 Subject: [PATCH 10/23] fix(proxy): keep discovery timers and candidates consistent --- src/app/v1/_lib/proxy/forwarder.ts | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index df00d6009..91d78cdaf 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -3845,7 +3845,6 @@ export class ProxyForwarder { disableStreamingHedge?: boolean; }; return ( - settings.discoveryEnabled === true && endpointPolicy.allowRetry && endpointPolicy.allowProviderSwitch && message?.stream === true && @@ -4890,6 +4889,12 @@ export class ProxyForwarder { let totalTimer: NodeJS.Timeout | null = null; let roundTimer: NodeJS.Timeout | null = null; let stickyTimer: NodeJS.Timeout | null = null; + const clearRoundTimer = () => { + if (roundTimer) { + clearTimeout(roundTimer); + roundTimer = null; + } + }; let executeCoordinatorAction: (action: DiscoveryAction) => Promise = async () => {}; let resolveResult: ((result: { response?: Response; error?: Error }) => void) | null = null; const resultPromise = new Promise<{ response?: Response; error?: Error }>((resolve) => { @@ -5225,6 +5230,7 @@ export class ProxyForwarder { const launchNextRound = async () => { if (settled || committed) return; + clearRoundTimer(); currentRound += 1; if (currentRound > maxRounds) return; coordinator.beginRound(); @@ -5234,10 +5240,10 @@ export class ProxyForwarder { await launch(candidate, "normal"); } catch (error) { lastError = error instanceof Error ? error : new Error(String(error)); - noMoreCandidates = true; } } if (!committed && !settled) { + clearRoundTimer(); roundTimer = setTimeout(() => { void onBoundary().catch((error) => logger.warn("[Discovery] Round boundary failed", { error }) @@ -5365,14 +5371,15 @@ export class ProxyForwarder { initial, Array.from(launched) ); + if (candidates.length === 0) noMoreCandidates = true; for (const provider of candidates) { try { await launch(provider, "normal"); } catch (error) { lastError = error instanceof Error ? error : new Error(String(error)); - noMoreCandidates = true; } } + clearRoundTimer(); roundTimer = setTimeout(() => { void onBoundary().catch((error) => logger.warn("[Discovery] Round boundary failed", { error }) @@ -5385,7 +5392,7 @@ export class ProxyForwarder { } finally { cleanupAbort(); if (totalTimer) clearTimeout(totalTimer); - if (roundTimer) clearTimeout(roundTimer); + clearRoundTimer(); if (stickyTimer) clearTimeout(stickyTimer); } } From 7d49cde3394dc5bd782fc2c0731916f3ed18e7ff Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 16:36:29 -0400 Subject: [PATCH 11/23] fix(discovery): clear binding against snapshot provider --- src/app/v1/_lib/proxy/response-handler.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index bc36a1a39..70d0909b5 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -1154,7 +1154,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( if (meta?.bindingSnapshot && keyId != null) { await SessionManager.clearVersionedSessionProvider( meta.bindingSnapshot, - providerIdForPersistence, + meta.bindingSnapshot.providerId, 0 ); return; From a643ff8e231168f3ac0a39b972c3c9a978588bcb Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 16:59:11 -0400 Subject: [PATCH 12/23] fix(discovery): honor effective group priority --- src/app/v1/_lib/proxy/forwarder.ts | 2 +- src/app/v1/_lib/proxy/provider-selector.ts | 7 +++++++ .../provider-selector-group-priority.test.ts | 16 ++++++++++++++++ .../proxy-forwarder-hedge-first-byte.test.ts | 4 ++++ 4 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 91d78cdaf..120b15dd9 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5119,7 +5119,7 @@ export class ProxyForwarder { coordinator.addAttempt({ id, providerId: provider.id, - priority: provider.priority || 0, + priority: ProxyProviderResolver.resolveEffectivePriorityForSession(provider, session), kind, ready: false, pending: true, diff --git a/src/app/v1/_lib/proxy/provider-selector.ts b/src/app/v1/_lib/proxy/provider-selector.ts index ee290e1aa..d4ce6fdbe 100644 --- a/src/app/v1/_lib/proxy/provider-selector.ts +++ b/src/app/v1/_lib/proxy/provider-selector.ts @@ -1179,6 +1179,13 @@ export class ProxyProviderResolver { return provider.priority ?? 0; } + static resolveEffectivePriorityForSession(provider: Provider, session: ProxySession): number { + return ProxyProviderResolver.resolveEffectivePriority( + provider, + getEffectiveProviderGroup(session) + ); + } + /** * 优先级分层：只选择最高优先级的供应商（支持分组优先级覆盖） */ diff --git a/tests/unit/proxy/provider-selector-group-priority.test.ts b/tests/unit/proxy/provider-selector-group-priority.test.ts index e0c6e6da5..143c851fa 100644 --- a/tests/unit/proxy/provider-selector-group-priority.test.ts +++ b/tests/unit/proxy/provider-selector-group-priority.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +import type { ProxySession } from "@/app/v1/_lib/proxy/session"; import type { Provider } from "@/types/provider"; import { ProxyProviderResolver } from "@/app/v1/_lib/proxy/provider-selector"; @@ -128,6 +129,21 @@ describe("resolveEffectivePriority", () => { // "cli,admin" - only "cli" matches, should return 3 expect(ProxyProviderResolver.resolveEffectivePriority(provider, "cli,admin")).toBe(3); }); + + it("resolves Discovery priority from the authenticated key group", () => { + const provider = makeProvider({ + priority: 10, + groupPriorities: { cli: 1 }, + }); + const session = { + authState: { + key: { providerGroup: "cli" }, + user: { providerGroup: "chat" }, + }, + } as unknown as ProxySession; + + expect(ProxyProviderResolver.resolveEffectivePriorityForSession(provider, session)).toBe(1); + }); }); describe("selectTopPriority with group context", () => { diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index 1458d825d..5654e4745 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -4,6 +4,9 @@ import { resolveEndpointPolicy } from "@/app/v1/_lib/proxy/endpoint-policy"; const mocks = vi.hoisted(() => ({ pickRandomProviderWithExclusion: vi.fn(), pickDiscoveryProviders: vi.fn(), + resolveEffectivePriorityForSession: vi.fn( + (provider: { priority?: number | null }) => provider.priority ?? 0 + ), recordSuccess: vi.fn(), recordFailure: vi.fn(async () => {}), getCircuitState: vi.fn(() => "closed"), @@ -107,6 +110,7 @@ vi.mock("@/app/v1/_lib/proxy/provider-selector", () => ({ ProxyProviderResolver: { pickRandomProviderWithExclusion: mocks.pickRandomProviderWithExclusion, pickDiscoveryProviders: mocks.pickDiscoveryProviders, + resolveEffectivePriorityForSession: mocks.resolveEffectivePriorityForSession, }, })); From ca12624f53c701c1d7f5dd5cb4e5095fd2b3ee50 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 19:09:54 -0400 Subject: [PATCH 13/23] fix(discovery): harden bounded streaming lifecycle --- .env.example | 1 + .../v1/_lib/proxy/discovery-coordinator.ts | 56 +- src/app/v1/_lib/proxy/discovery-validity.ts | 41 +- src/app/v1/_lib/proxy/forwarder.ts | 1074 ++++++++++++---- src/app/v1/_lib/proxy/provider-selector.ts | 4 +- src/app/v1/_lib/proxy/response-handler.ts | 721 ++++++++--- src/app/v1/_lib/proxy/session.ts | 53 +- src/app/v1/_lib/proxy/stream-finalization.ts | 13 + src/lib/config/env.schema.ts | 4 + src/lib/observability/discovery-metrics.ts | 126 ++ .../integration/proxy-hedge-lifecycle.test.ts | 127 +- .../unit/proxy/discovery-coordinator.test.ts | 60 +- tests/unit/proxy/discovery-validity.test.ts | 42 + .../proxy-forwarder-hedge-first-byte.test.ts | 1098 +++++++++++++++++ ...forwarder-provider-session-release.test.ts | 14 + ...esponse-handler-client-abort-drain.test.ts | 257 ++++ ...handler-endpoint-circuit-isolation.test.ts | 761 ++++++++++++ ...gemini-stream-passthrough-timeouts.test.ts | 2 +- 18 files changed, 4051 insertions(+), 403 deletions(-) create mode 100644 src/lib/observability/discovery-metrics.ts diff --git a/.env.example b/.env.example index 19b417d51..653ce134c 100644 --- a/.env.example +++ b/.env.example @@ -97,6 +97,7 @@ ENABLE_API_KEY_REDIS_CACHE="true" # 是否启用 API Key Redis 缓存（ # 降低该值会按签发时间收紧已签发 ADMIN_TOKEN 签名 cookie 的剩余寿命，且不会延长其原始 exp。 AUTH_SESSION_TTL_SECONDS=604800 # Web UI 登录态过期时间（秒，默认 604800 = 7 天，范围 60-31536000） SESSION_TTL=300 # 代理请求上下文缓存时间（秒，默认 300 = 5 分钟；不控制 Web UI 登录态） +DISCOVERY_ROLLOUT_PERCENT=100 # Discovery 运维灰度比例（0-100，按 API Key + Session 稳定分桶） STORE_SESSION_MESSAGES=false # 会话消息存储模式（默认：false） # - false：存储请求/响应体但对 message 内容脱敏 [REDACTED] # - true：原样存储 message 内容（注意隐私和存储空间影响） diff --git a/src/app/v1/_lib/proxy/discovery-coordinator.ts b/src/app/v1/_lib/proxy/discovery-coordinator.ts index 4a94baf70..f1330df92 100644 --- a/src/app/v1/_lib/proxy/discovery-coordinator.ts +++ b/src/app/v1/_lib/proxy/discovery-coordinator.ts @@ -66,7 +66,21 @@ export class DiscoveryCoordinator { return { requestEpoch: this.requestEpoch, roundEpoch: this.roundEpoch }; } + startStickyProbe(): void { + if (!this.isTerminal) this.state = "STICKY_PROBING"; + } + + startDiscoveryAfterSticky(): void { + if (this.state === "STICKY_PROBING" || this.state === "FALLBACK_READY_HELD") { + this.state = "DISCOVERY_RACING"; + } + } + beginRound(): { requestEpoch: number; roundEpoch: number; round: number } { + if (this.isTerminal || this.round >= this.maxRounds) { + return { ...this.epochs, round: this.round }; + } + this.round += 1; this.roundEpoch += 1; if (!this.isTerminal) this.state = "DISCOVERY_RACING"; return { ...this.epochs, round: this.round }; @@ -82,8 +96,22 @@ export class DiscoveryCoordinator { this.attempts.delete(id); } + /** Mark an already-running attempt as the sole fallback for this request. */ + promoteToFallback(id: string): boolean { + if (this.isTerminal) return false; + const attempt = this.attempts.get(id); + if (!attempt?.pending) return false; + attempt.kind = "fallback"; + this.state = "FALLBACK_READY_HELD"; + return true; + } + get isTerminal(): boolean { - return this.state === "WINNER_COMMITTED" || this.state === "TERMINAL_FAILED"; + return ( + this.state === "WINNER_COMMITTED" || + this.state === "FALLBACK_ACTIVE" || + this.state === "TERMINAL_FAILED" + ); } get activeAttempts(): DiscoveryAttempt[] { @@ -143,9 +171,11 @@ export class DiscoveryCoordinator { ): DiscoveryAction { if (!this.acceptsEpoch(requestEpoch, roundEpoch) || this.isTerminal) return { type: "none" }; const attempt = this.attempts.get(id); - if (!attempt) return { type: "none" }; + if (!attempt?.pending) return { type: "none" }; attempt.pending = false; attempt.ready = false; + const readyAction = this.chooseReadyNormal(); + if (readyAction.type !== "none") return readyAction; return this.afterAttemptState(); } @@ -166,8 +196,7 @@ export class DiscoveryCoordinator { ); if (higherTierPending) return { type: "none" }; } - const sameTier = readyNormal.filter((attempt) => attempt.priority === bestPriority); - const winner = sameTier[0]; + const winner = readyNormal[0]; this.state = "WINNER_COMMITTED"; winner.pending = false; return { @@ -202,12 +231,11 @@ export class DiscoveryCoordinator { const cancelAttemptIds = pendingNormal.map((attempt) => attempt.id); for (const attempt of pendingNormal) attempt.pending = false; if (this.round < this.maxRounds) { - this.round += 1; - this.roundEpoch += 1; + this.beginRound(); this.state = "DISCOVERY_RACING"; return { type: "launch", - slots: Math.max(1, this.concurrency - 1), + slots: Math.max(0, this.concurrency - 1), cancelAttemptIds, }; } @@ -228,12 +256,11 @@ export class DiscoveryCoordinator { for (const id of losers) this.attempts.get(id)!.pending = false; if (this.round < this.maxRounds) { - this.round += 1; - this.roundEpoch += 1; + this.beginRound(); this.state = "DISCOVERY_RACING"; return { type: "launch", - slots: Math.max(1, this.concurrency - 1), + slots: Math.max(0, this.concurrency - 1), cancelAttemptIds: losers, promoteAttemptId: fallback.id, }; @@ -243,6 +270,8 @@ export class DiscoveryCoordinator { onDeadline(): DiscoveryAction { if (this.isTerminal) return { type: "none" }; + const readyNormal = this.chooseReadyNormal(true); + if (readyNormal.type === "commit_normal") return readyNormal; const fallback = Array.from(this.attempts.values()).find( (attempt) => attempt.pending && attempt.kind === "fallback" && attempt.ready ); @@ -259,7 +288,7 @@ export class DiscoveryCoordinator { const attempt = this.attempts.get(id); if (!attempt || this.isTerminal) return { type: "none" }; attempt.pending = false; - this.state = "WINNER_COMMITTED"; + this.state = attempt.kind === "fallback" ? "FALLBACK_ACTIVE" : "WINNER_COMMITTED"; return { type: attempt.kind === "fallback" ? "promote_fallback" : "commit_normal", attemptId: id, @@ -297,9 +326,8 @@ export class DiscoveryCoordinator { this.state = "TERMINAL_FAILED"; return { type: "terminal_failure" }; } - this.round += 1; - this.roundEpoch += 1; + this.beginRound(); this.state = "DISCOVERY_RACING"; - return { type: "launch", slots: Math.max(1, this.concurrency - 1) }; + return { type: "launch", slots: this.concurrency }; } } diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts index 198d77699..3325295a7 100644 --- a/src/app/v1/_lib/proxy/discovery-validity.ts +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -9,8 +9,12 @@ export type DiscoveryValidity = { ready: boolean; terminal: boolean; error: boolean; + limitExceeded?: boolean; }; +export const DISCOVERY_PREFIX_MAX_BYTES = 1024 * 1024; +export const DISCOVERY_EVENT_MAX_COUNT = 1024; + function hasContent(value: unknown): boolean { if (typeof value === "string") return value.trim().length > 0; if (!value || typeof value !== "object") return false; @@ -26,9 +30,12 @@ function hasContent(value: unknown): boolean { "tool_calls", "functionCall", "function_call", + "function", "arguments", - "input", "partial_json", + "id", + "name", + "input", "parts", ].some((key) => hasContent(object[key])); } @@ -141,10 +148,21 @@ export class DiscoveryValidityParser { private _ready = false; private _terminal = false; private _error = false; + private _limitExceeded = false; + private bytesSeen = 0; + private eventsSeen = 0; constructor(readonly protocol: DiscoveryProtocol) {} push(chunk: Uint8Array | string): DiscoveryValidity { + this.bytesSeen += + typeof chunk === "string" ? new TextEncoder().encode(chunk).byteLength : chunk.byteLength; + if (!this._ready && this.bytesSeen > DISCOVERY_PREFIX_MAX_BYTES) { + this._error = true; + this._limitExceeded = true; + this.buffered = ""; + return this.result; + } this.buffered += typeof chunk === "string" ? chunk : this.decoder.decode(chunk, { stream: true }); @@ -177,7 +195,7 @@ export class DiscoveryValidityParser { } } - return { ready: this._ready && !this._error, terminal: this._terminal, error: this._error }; + return this.result; } private consumeLine(line: string): void { @@ -195,6 +213,12 @@ export class DiscoveryValidityParser { } private consumeValue(value: unknown): void { + this.eventsSeen += 1; + if (!this._ready && this.eventsSeen > DISCOVERY_EVENT_MAX_COUNT) { + this._error = true; + this._limitExceeded = true; + return; + } const result = classifyJson(value, this.protocol); this._ready ||= result.ready; this._terminal ||= result.terminal; @@ -210,4 +234,17 @@ export class DiscoveryValidityParser { get error(): boolean { return this._error; } + + get limitExceeded(): boolean { + return this._limitExceeded; + } + + private get result(): DiscoveryValidity { + return { + ready: this._ready && !this._error, + terminal: this._terminal, + error: this._error, + ...(this._limitExceeded ? { limitExceeded: true } : {}), + }; + } } diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 120b15dd9..e31a81b17 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -24,12 +24,17 @@ import { PROTECTED_AUTH_HEADER_NAMES } from "@/lib/custom-headers"; import { recordEndpointFailure, recordEndpointSuccess } from "@/lib/endpoint-circuit-breaker"; import { applyGeminiGoogleSearchOverrideWithAudit } from "@/lib/gemini/provider-overrides"; import { logger } from "@/lib/logger"; +import { + DiscoveryRequestMetrics, + recordDiscoveryControlEvent, +} from "@/lib/observability/discovery-metrics"; import { getEndpointFilterStats, getPreferredProviderEndpoints, } from "@/lib/provider-endpoints/endpoint-selector"; import { getGlobalAgentPool, getProxyAgentForProvider } from "@/lib/proxy-agent"; import { RateLimitService } from "@/lib/rate-limit/service"; +import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import { SessionManager } from "@/lib/session-manager"; import { detectUpstreamErrorFromSseOrJsonText, @@ -55,6 +60,7 @@ import { HeaderProcessor, resolveAnthropicAuthHeaders } from "../headers"; import { evaluateResponsesWsEligibility, getResponsesWsSessionId, + isWebsocketClientRequest, } from "../responses-ws/eligibility"; import { RESERVED_INTERNAL_HEADERS } from "../responses-ws/internal-secret"; import { markResponsesWsUnsupported } from "../responses-ws/unsupported-cache"; @@ -190,6 +196,21 @@ function applyProviderCustomHeaders( const RETRY_LIMITS = PROVIDER_LIMITS.MAX_RETRY_ATTEMPTS; const MAX_PROVIDER_SWITCHES = 20; // 保险栓：最多切换 20 次供应商（防止无限循环） +const DISCOVERY_LEASE_HANDOFF_GRACE_SECONDS = 5; + +function isDiscoveryRolloutEligible(keyId: number, sessionId: string, percent: number): boolean { + const normalizedPercent = Math.max(0, Math.min(100, Math.floor(percent))); + if (normalizedPercent === 0) return false; + if (normalizedPercent === 100) return true; + + // FNV-1a provides a deterministic bucket without persisting rollout state. + let hash = 0x811c9dc5; + for (const character of `${keyId}:${sessionId}`) { + hash ^= character.charCodeAt(0); + hash = Math.imul(hash, 0x01000193); + } + return (hash >>> 0) % 100 < normalizedPercent; +} type CacheTtlOption = CacheTtlPreference | null | undefined; @@ -199,6 +220,41 @@ type ProxySessionWithAttemptRuntime = ProxySession & { releaseAgent?: () => void; }; +type DiscoveryCancellationKind = + | "discovery_sla_timeout" + | "discovery_loser" + | "request_deadline" + | "client_abort"; + +class DiscoveryCancellationError extends Error { + readonly kind: DiscoveryCancellationKind; + + constructor(kind: DiscoveryCancellationKind) { + super(kind); + this.name = "DiscoveryCancellationError"; + this.kind = kind; + } +} + +class DiscoveryValidityLimitError extends Error { + constructor() { + super("Discovery response prefix exceeded the validation limit"); + this.name = "DiscoveryValidityLimitError"; + } +} + +type PreparedStreamingDiscovery = { + settings: SystemSettings; + bindingSnapshot: SessionBindingSnapshot; + requestStartedAt: number; + lease: { + sessionId: string; + keyId: number; + ownerToken: string; + ttlSeconds: number; + }; +}; + type StreamingHedgeAttempt = { provider: Provider; session: ProxySession; @@ -1199,8 +1255,18 @@ export class ProxyForwarder { throw new Error("代理上下文缺少供应商或鉴权信息"); } - if (await ProxyForwarder.shouldUseStreamingDiscovery(session)) { - const discoveryPromise = ProxyForwarder.sendStreamingWithDiscovery(session); + const requestStartedAt = Date.now(); + const discoverySettings = await getCachedSystemSettings(); + const preparedDiscovery = await ProxyForwarder.prepareStreamingDiscovery( + session, + discoverySettings, + requestStartedAt + ); + if (preparedDiscovery) { + const discoveryPromise = ProxyForwarder.sendStreamingWithDiscovery( + session, + preparedDiscovery + ); void discoveryPromise.catch(() => undefined); return await discoveryPromise; } @@ -1480,6 +1546,7 @@ export class ProxyForwarder { endpointId: activeEndpoint.endpointId, endpointUrl: endpointAudit.endpointUrl, upstreamStatusCode: response.status, + bindingIntent: session.isSessionBindingAllowed() ? undefined : "none", }); logger.info("ProxyForwarder: Streaming response received, deferring finalization", { @@ -1667,7 +1734,7 @@ export class ProxyForwarder { } // ⭐ 成功后绑定 session 到供应商（智能绑定策略） - if (session.sessionId) { + if (session.sessionId && session.isSessionBindingAllowed()) { // 使用智能绑定策略（故障转移优先 + 稳定性优化） const result = await SessionManager.updateSessionBindingSmart( session.sessionId, @@ -3257,6 +3324,15 @@ export class ProxyForwarder { syscall?: string; // 系统调用：如 'getaddrinfo'、'connect'、'read'、'write' }; + const externalAbortReason = externalAbortSignal?.reason; + if ( + externalAbortSignal?.aborted && + externalAbortReason instanceof DiscoveryCancellationError + ) { + cleanupCombinedSignal(); + throw externalAbortReason; + } + // ⭐ SSL 证书错误检测：标记 Agent 为不健康，下次请求将创建新 Agent const sslErrorCacheKey = proxyConfig?.cacheKey ?? directConnectionCacheKey; const sslErrorDispatcherId = proxyConfig?.dispatcherId ?? directConnectionDispatcherId; @@ -3811,11 +3887,7 @@ export class ProxyForwarder { private static shouldUseStreamingHedge(session: ProxySession): boolean { const endpointPolicy = ProxyForwarder.getEndpointPolicy(session); - const routing = session as ProxySession & { - routingMode?: string; - disableStreamingHedge?: boolean; - }; - if (routing.routingMode === "lease_conflict_single" || routing.disableStreamingHedge === true) { + if (session.isStreamingHedgeDisabled()) { return false; } return ( @@ -3826,34 +3898,99 @@ export class ProxyForwarder { ); } - private static async shouldUseStreamingDiscovery(session: ProxySession): Promise { - const settings = await getCachedSystemSettings(); + private static async prepareStreamingDiscovery( + session: ProxySession, + settings: SystemSettings, + requestStartedAt: number + ): Promise { if (settings.discoveryEnabled !== true) { - return false; - } - if ( - typeof SessionManager.getVersionedBindingCapabilityState !== "function" || - SessionManager.getVersionedBindingCapabilityState() !== "available" - ) { - return false; + return null; } const endpointPolicy = ProxyForwarder.getEndpointPolicy(session); const protocol = ProxyForwarder.discoveryProtocol(session); const message = session.request.message as Record; - const routing = session as ProxySession & { - routingMode?: string; - disableStreamingHedge?: boolean; - }; - return ( - endpointPolicy.allowRetry && - endpointPolicy.allowProviderSwitch && - message?.stream === true && - !endpointPolicy.bypassForwarderPreprocessing && - protocol !== "unknown" && - routing.routingMode !== "lease_conflict_single" && - routing.disableStreamingHedge !== true && - !session.isRawCrossProviderFallbackEnabled() + if ( + !endpointPolicy.allowRetry || + !endpointPolicy.allowProviderSwitch || + message?.stream !== true || + endpointPolicy.bypassForwarderPreprocessing || + protocol === "unknown" || + isWebsocketClientRequest(session.headers) || + session.isStreamingHedgeDisabled() || + session.isRawCrossProviderFallbackEnabled() + ) { + return null; + } + + const sessionId = session.sessionId; + const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + if (!sessionId || keyId == null) { + return null; + } + if (!isDiscoveryRolloutEligible(keyId, sessionId, getEnvConfig().DISCOVERY_ROLLOUT_PERCENT)) { + return null; + } + + const capabilityState = await SessionManager.ensureVersionedBindingCapability(); + if (capabilityState !== "available") { + return null; + } + + let bindingSnapshot = session.getSessionBindingSnapshot(); + if ( + !bindingSnapshot || + bindingSnapshot.sessionId !== sessionId || + bindingSnapshot.keyId !== keyId + ) { + const binding = await SessionManager.getSessionBindingSnapshot(sessionId, keyId); + if (binding.status !== "ok") { + // A foreign or irreconcilable mirror must never be mutated by this + // request. Infrastructure unavailability still falls back to the + // established legacy wrapper. + if (binding.status === "conflict") session.setSessionBindingAllowed(false); + return null; + } + bindingSnapshot = binding.snapshot; + session.setSessionBindingSnapshot(binding.snapshot); + } + + const ttlSeconds = Math.max( + 1, + Math.ceil(Math.max(1, settings.racingTotalTimeoutMs ?? 60_000) / 1000) + ); + const lease = await SessionManager.acquireSessionDiscoveryLease( + sessionId, + keyId, + ttlSeconds + DISCOVERY_LEASE_HANDOFF_GRACE_SECONDS ); + if (lease.status !== "acquired") { + if (lease.status === "conflict") { + session.disableStreamingHedge(); + session.setSessionBindingAllowed(false); + logger.info("[Discovery] Lease conflict; routing request in single-upstream mode", { + sessionId, + keyId, + }); + recordDiscoveryControlEvent("lease_conflict", { + requestId: session.messageContext?.id ?? null, + sessionId, + keyId, + }); + } + return null; + } + + return { + settings, + bindingSnapshot, + requestStartedAt, + lease: { + sessionId, + keyId, + ownerToken: lease.ownerToken, + ttlSeconds, + }, + }; } private static discoveryProtocol(session: ProxySession): DiscoveryProtocol { @@ -4623,7 +4760,7 @@ export class ProxyForwarder { // A non-hedged request is finalized through response-handler. Updating // here as well would perform a duplicate binding read/CAS before the // stream has passed its final validation. - if (session.sessionId && isActualHedgeWin) { + if (session.sessionId && isActualHedgeWin && session.isSessionBindingAllowed()) { void (async () => { const bindingResult = await SessionManager.updateSessionBindingSmart( session.sessionId!, @@ -4672,6 +4809,7 @@ export class ProxyForwarder { upstreamStatusCode: attempt.response.status, isHedgeWinner: isActualHedgeWin, billHedgeLosers, + bindingIntent: session.isSessionBindingAllowed() ? undefined : "none", }); const response = new Response( @@ -4716,7 +4854,9 @@ export class ProxyForwarder { } if (checkResult.referenced) { - session.recordProviderSessionRef(provider.id); + session.recordProviderSessionRef(provider.id, { + retainOnSuccess: checkResult.tracked, + }); } } @@ -4837,32 +4977,36 @@ export class ProxyForwarder { * existing loser billing and retry semantics remain unchanged while the * feature is rolled out behind discoveryEnabled. */ - private static async sendStreamingWithDiscovery(session: ProxySession): Promise { + private static async sendStreamingWithDiscovery( + session: ProxySession, + prepared: PreparedStreamingDiscovery + ): Promise { const initialProvider = session.provider; if (!initialProvider) throw new Error("代理上下文缺少供应商"); - const settings = await getCachedSystemSettings(); - const concurrency = Math.max(1, Math.floor(settings.discoveryConcurrency ?? 2)); + const { settings, lease, requestStartedAt } = prepared; + const concurrency = Math.max(2, Math.floor(settings.discoveryConcurrency ?? 2)); const maxRounds = Math.max(1, Math.floor(settings.maxDiscoveryRounds ?? 2)); const discoverySlaMs = Math.max(1, settings.discoverySlaMs ?? 10_000); - const stickySlaMs = Math.max(discoverySlaMs, settings.stickySlaMs ?? 20_000); - const totalTimeoutMs = Math.max(stickySlaMs, settings.racingTotalTimeoutMs ?? 60_000); + // Respect the configured Sticky budget. The settings validator already + // checks the total pre-winner window; a shorter Sticky SLA is a valid + // deliberate choice and must not be silently expanded at runtime. + const stickySlaMs = Math.max(1, settings.stickySlaMs ?? 20_000); + const totalTimeoutMs = Math.max(1, settings.racingTotalTimeoutMs ?? 60_000); + const racingDeadlineAt = requestStartedAt + totalTimeoutMs; const protocol = ProxyForwarder.discoveryProtocol(session); + const rawCrossProviderFallbackEnabled = session.isRawCrossProviderFallbackEnabled(); const coordinator = new DiscoveryCoordinator({ concurrency, maxRounds }); - const bindingKeyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; - // Provider selection normally populates this snapshot for a reused Sticky. - // For a cold start, initialize it once so finalization can use generation - // CAS without another read during winner commit or timeout cleanup. - let bindingSnapshot = session.getSessionBindingSnapshot(); - if (!bindingSnapshot && session.sessionId && bindingKeyId != null) { - const binding = await SessionManager.getSessionBindingSnapshot( - session.sessionId, - bindingKeyId - ); - if (binding.status === "ok") { - bindingSnapshot = binding.snapshot; - session.setSessionBindingSnapshot(binding.snapshot); - } - } + const discoveryMetrics = new DiscoveryRequestMetrics( + { + requestId: session.messageContext?.id ?? null, + sessionId: lease.sessionId, + keyId: lease.keyId, + }, + requestStartedAt + ); + let bindingSnapshot: SessionBindingSnapshot = prepared.bindingSnapshot; + let bindingWriteAllowed = true; + let leaseTransferred = false; const attempts = new Map< string, StreamingHedgeAttempt & { @@ -4875,6 +5019,11 @@ export class ProxyForwarder { ready: boolean; round: number; readerTransferred: boolean; + readerCancelled: boolean; + providerSessionRefOwned: boolean; + providerSessionRefRetainOnSuccess: boolean; + providerSessionRefReleased: boolean; + cancellationKind: DiscoveryCancellationKind | null; } >(); const launched = new Set(); @@ -4889,75 +5038,172 @@ export class ProxyForwarder { let totalTimer: NodeJS.Timeout | null = null; let roundTimer: NodeJS.Timeout | null = null; let stickyTimer: NodeJS.Timeout | null = null; + let roundLaunchesInProgress = 0; + const roundLaunchIdleWaiters = new Set<() => void>(); + let fallbackPromotionBlocked = false; + const hasSticky = + session.shouldReuseProvider() && + !!session.sessionId && + bindingSnapshot.providerId === initialProvider.id; + let stickyProbeActive = hasSticky; + if (hasSticky) coordinator.startStickyProbe(); + const waitForRoundLaunches = (): Promise => { + if (roundLaunchesInProgress === 0) return Promise.resolve(); + return new Promise((resolve) => roundLaunchIdleWaiters.add(resolve)); + }; + const notifyRoundLaunchIdle = () => { + if (roundLaunchesInProgress !== 0) return; + for (const resolve of roundLaunchIdleWaiters) resolve(); + roundLaunchIdleWaiters.clear(); + }; const clearRoundTimer = () => { if (roundTimer) { clearTimeout(roundTimer); roundTimer = null; } }; - let executeCoordinatorAction: (action: DiscoveryAction) => Promise = async () => {}; + let executeCoordinatorAction: ( + action: DiscoveryAction, + terminalCancellationKind?: DiscoveryCancellationKind + ) => Promise = async () => {}; let resolveResult: ((result: { response?: Response; error?: Error }) => void) | null = null; const resultPromise = new Promise<{ response?: Response; error?: Error }>((resolve) => { resolveResult = resolve; }); - const releaseProviderRef = (providerId: number) => { - ProxyForwarder.releaseProviderSessionRef(session, providerId); + const releaseProviderRef = (attempt: (typeof winner & { id: string }) | null) => { + if (!attempt?.providerSessionRefOwned || attempt.providerSessionRefReleased) return; + attempt.providerSessionRefReleased = true; + ProxyForwarder.releaseProviderSessionRef(session, attempt.provider.id); }; - const cancelAttempt = (attempt: (typeof winner & { id: string }) | null, reason: string) => { + const getAttemptModelRedirect = (attempt: (typeof winner & { id: string }) | null) => { + if (!attempt) return undefined; + if (attempt.modelRedirect !== undefined) return attempt.modelRedirect; + const redirect = attempt.session.getCurrentModelRedirect(attempt.provider.id); + if (redirect) attempt.modelRedirect = structuredClone(redirect); + return attempt.modelRedirect; + }; + + const cleanupAttempt = ( + attempt: (typeof winner & { id: string }) | null, + cancellationKind: DiscoveryCancellationKind | null + ) => { if (attempt?.readerTransferred) return; - if (!attempt?.pending) return; + if (!attempt) return; attempt.pending = false; - try { - attempt.controller.abort(new Error(reason)); - } catch { - /* abort is best effort */ + if (cancellationKind && !attempt.cancellationKind) { + attempt.cancellationKind = cancellationKind; } - try { - const cancelPromise = attempt.reader?.cancel(reason); - cancelPromise?.catch(() => undefined); - } catch (error) { - logger.debug("[Discovery] Reader cancel failed", { reason, error }); + if (!attempt.controller.signal.aborted) { + try { + attempt.controller.abort( + cancellationKind + ? new DiscoveryCancellationError(cancellationKind) + : new Error("discovery_attempt_failed") + ); + } catch { + /* abort is best effort */ + } } - try { - attempt.releaseAgent?.(); - } catch { - /* release is idempotent */ + if (attempt.reader && !attempt.readerCancelled) { + attempt.readerCancelled = true; + try { + const cancelPromise = attempt.reader.cancel( + cancellationKind ?? "discovery_attempt_failed" + ); + cancelPromise.catch((error) => + discoveryMetrics.cancelFailed(attempt.id, attempt.provider.id, error) + ); + } catch (error) { + discoveryMetrics.cancelFailed(attempt.id, attempt.provider.id, error); + logger.debug("[Discovery] Reader cancel failed", { cancellationKind, error }); + } + } + if (attempt.releaseAgent && !attempt.agentReleased) { + attempt.agentReleased = true; + try { + attempt.releaseAgent(); + } catch { + /* release is idempotent */ + } } - releaseProviderRef(attempt.provider.id); + releaseProviderRef(attempt); + attempt.chunks.length = 0; + discoveryMetrics.attemptFinished(attempt.id, { + providerId: attempt.provider.id, + outcome: cancellationKind ? "cancelled" : "failed", + cancellationKind, + }); }; - const cancelLosers = (keep: typeof winner = null) => { + const cancelAttempt = ( + attempt: (typeof winner & { id: string }) | null, + cancellationKind: DiscoveryCancellationKind + ) => cleanupAttempt(attempt, cancellationKind); + + const cancelLosers = ( + keep: typeof winner = null, + cancellationKind: DiscoveryCancellationKind = "discovery_loser" + ) => { for (const attempt of attempts.values()) { - if (attempt !== keep) cancelAttempt(attempt, "discovery_loser"); + if (attempt !== keep) cancelAttempt(attempt, cancellationKind); } }; - const settleFailure = async (error: Error) => { + const settleFailure = async ( + error: Error, + options: { + preserveBinding?: boolean; + cancellationKind?: DiscoveryCancellationKind; + } = {} + ) => { if (settled) return; settled = true; if (totalTimer) clearTimeout(totalTimer); if (roundTimer) clearTimeout(roundTimer); if (stickyTimer) clearTimeout(stickyTimer); - cancelLosers(); - if (bindingSnapshot) { + cancelLosers(null, options.cancellationKind ?? "discovery_loser"); + if (!options.preserveBinding && bindingWriteAllowed && session.isSessionBindingAllowed()) { if (bindingSnapshot.providerId != null) { - await SessionManager.clearVersionedSessionProvider( + void SessionManager.clearVersionedSessionProvider( bindingSnapshot, bindingSnapshot.providerId, 0 - ); + ).catch((bindingError) => { + logger.warn("[Discovery] Terminal binding clear failed", { + sessionId: bindingSnapshot.sessionId, + keyId: bindingSnapshot.keyId, + providerId: bindingSnapshot.providerId, + error: bindingError instanceof Error ? bindingError.message : String(bindingError), + }); + }); } - } else { - const attempted = new Set(launched); - await ProxyForwarder.clearSessionProviderBindings(session, attempted); } + const statusCode = error instanceof ProxyError ? error.statusCode : 503; + discoveryMetrics.finish({ + outcome: + options.cancellationKind === "client_abort" || statusCode === 499 + ? "client_abort" + : options.cancellationKind === "request_deadline" + ? "deadline" + : "failed", + statusCode, + winnerOrigin: "none", + }); resolveResult?.({ error }); }; const commit = async (attempt: typeof winner) => { - if (!attempt || committed || settled || !attempt.response || !attempt.reader) return; + if ( + !attempt || + committed || + settled || + !attempt.ready || + !attempt.response || + !attempt.reader + ) + return; committed = true; winner = attempt; attempt.pending = false; @@ -4968,6 +5214,20 @@ export class ProxyForwarder { if (roundTimer) clearTimeout(roundTimer); if (stickyTimer) clearTimeout(stickyTimer); cancelLosers(attempt); + discoveryMetrics.attemptFinished(attempt.id, { + providerId: attempt.provider.id, + outcome: "winner", + }); + discoveryMetrics.finish({ + outcome: "success", + statusCode: attempt.response.status, + winnerOrigin: attempt.kind, + winnerProviderId: attempt.provider.id, + winnerRound: + hasSticky && attempt.provider.id === initialProvider.id && stickyProbeActive + ? 0 + : attempt.round, + }); session.setProvider(attempt.provider); if (attempt.session !== session) ProxyForwarder.syncWinningAttemptSession(session, attempt.session); @@ -4986,14 +5246,18 @@ export class ProxyForwarder { isHedgeWinner: false, billHedgeLosers: false, bindingIntent: - attempt.kind === "fallback" + attempt.kind === "fallback" || !bindingWriteAllowed || !session.isSessionBindingAllowed() ? "none" : bindingSnapshot?.providerId == null ? "create" : "renew", bindingSnapshot, requiresCompletionMarker: attempt.kind !== "fallback", + discoveryLease: lease, + providerSessionRefOwned: attempt.providerSessionRefOwned, + providerSessionRefRetainOnSuccess: attempt.providerSessionRefRetainOnSuccess, }); + leaseTransferred = true; const prefix = attempt.chunks.length === 1 ? attempt.chunks[0] @@ -5032,11 +5296,94 @@ export class ProxyForwarder { return candidates[0]; }; - const launch = async (provider: Provider, kind: "normal" | "fallback"): Promise => { - if (settled || committed || launched.has(provider.id)) return; + const clearCapturedStickyBinding = async (cooldownTtlSeconds: number): Promise => { + if ( + !bindingWriteAllowed || + !session.isSessionBindingAllowed() || + bindingSnapshot.providerId !== initialProvider.id + ) { + return; + } + try { + const cleared = await SessionManager.clearVersionedSessionProvider( + bindingSnapshot, + initialProvider.id, + cooldownTtlSeconds + ); + if (cleared.status === "ok") { + bindingSnapshot = cleared.snapshot; + session.setSessionBindingSnapshot(cleared.snapshot); + return; + } + bindingWriteAllowed = false; + logger.debug("[Discovery] Failed to clear captured Sticky", { + providerId: initialProvider.id, + reason: cleared.reason, + }); + } catch (error) { + bindingWriteAllowed = false; + logger.debug("[Discovery] Failed to clear captured Sticky", { + providerId: initialProvider.id, + error, + }); + } + }; + + const scheduleRoundBoundary = (delayMs: number) => { + clearRoundTimer(); + const epoch = coordinator.epochs; + const remainingMs = Math.max(0, racingDeadlineAt - Date.now()); + roundTimer = setTimeout( + () => { + if (Date.now() >= racingDeadlineAt) { + void executeCoordinatorAction(coordinator.onDeadline(), "request_deadline").catch( + (error) => logger.warn("[Discovery] Deadline action failed", { error }) + ); + return; + } + void executeCoordinatorAction( + coordinator.onRoundBoundary(epoch.requestEpoch, epoch.roundEpoch) + ).catch((error) => logger.warn("[Discovery] Round boundary failed", { error })); + }, + Math.min(delayMs, remainingMs) + ); + }; + + const launch = async ( + provider: Provider, + kind: "normal" | "fallback", + options?: { + attemptSession?: ProxySession; + requestAttemptCount?: number; + retryState?: ReactiveRectifierRetryState; + providerSessionRefTransfer?: { + owned: boolean; + retainOnSuccess: boolean; + }; + } + ): Promise => { + const transferredProviderSessionRef = options?.providerSessionRefTransfer?.owned === true; + let providerSessionRefTracked = transferredProviderSessionRef; + let providerSessionRefRetainOnSuccess = + options?.providerSessionRefTransfer?.retainOnSuccess === true; + const rollbackLaunch = () => { + if (providerSessionRefTracked) { + ProxyForwarder.releaseProviderSessionRef(session, provider.id); + providerSessionRefTracked = false; + providerSessionRefRetainOnSuccess = false; + } + }; + if (settled || committed || launched.has(provider.id)) { + rollbackLaunch(); + return; + } launched.add(provider.id); - let providerSessionRefRecorded = false; - if (provider.id !== initialProvider.id && session.sessionId) { + if (!transferredProviderSessionRef && provider.id === initialProvider.id) { + providerSessionRefTracked = session.hasProviderSessionRef(provider.id); + providerSessionRefRetainOnSuccess = + providerSessionRefTracked && session.shouldRetainProviderSessionRefOnSuccess(provider.id); + } + if (!providerSessionRefTracked && session.sessionId) { const limit = provider.limitConcurrentSessions || 0; const check = await RateLimitService.checkAndTrackProviderSession( provider.id, @@ -5044,27 +5391,47 @@ export class ProxyForwarder { limit ); if (!check.allowed) { - launched.delete(provider.id); throw new ProxyError(check.reason || "Provider concurrent limit reached", 503); } if (check.referenced) { - session.recordProviderSessionRef(provider.id); - providerSessionRefRecorded = true; + session.recordProviderSessionRef(provider.id, { + retainOnSuccess: check.tracked, + }); + providerSessionRefTracked = true; + providerSessionRefRetainOnSuccess = check.tracked; } } + if (settled || committed) { + rollbackLaunch(); + return; + } let endpoint: Awaited>; try { endpoint = await ProxyForwarder.resolveStreamingHedgeEndpoint(session, provider); } catch (error) { - launched.delete(provider.id); - if (providerSessionRefRecorded) releaseProviderRef(provider.id); + rollbackLaunch(); throw error; } - const attemptSession = - provider.id === initialProvider.id - ? session - : ProxyForwarder.createStreamingShadowSession(session, provider); - attemptSession.setProvider(provider); + if (settled || committed) { + rollbackLaunch(); + return; + } + let attemptSession: ProxySession; + try { + attemptSession = + options?.attemptSession ?? + (provider.id === initialProvider.id + ? session + : ProxyForwarder.createStreamingShadowSession(session, provider)); + attemptSession.setProvider(provider); + } catch (error) { + rollbackLaunch(); + throw error; + } + if (settled || committed) { + rollbackLaunch(); + return; + } const controller = new AbortController(); const id = `${provider.id}:${sequence + 1}`; const attempt = { @@ -5077,6 +5444,11 @@ export class ProxyForwarder { ready: false, round: currentRound, readerTransferred: false, + readerCancelled: false, + providerSessionRefOwned: providerSessionRefTracked, + providerSessionRefRetainOnSuccess, + providerSessionRefReleased: false, + cancellationKind: null, provider, session: attemptSession, baseUrl: endpoint.baseUrl, @@ -5086,8 +5458,8 @@ export class ProxyForwarder { clearResponseTimeout: null, firstByteTimeoutMs: 0, sequence: ++sequence, - requestAttemptCount: 1, - reactiveRectifierRetryState: { + requestAttemptCount: options?.requestAttemptCount ?? 1, + reactiveRectifierRetryState: options?.retryState ?? { thinkingSignatureRetried: false, thinkingBudgetRetried: false, thinkingEffortConflictRetried: false, @@ -5114,9 +5486,13 @@ export class ProxyForwarder { ready: boolean; round: number; readerTransferred: boolean; + readerCancelled: boolean; + providerSessionRefOwned: boolean; + providerSessionRefRetainOnSuccess: boolean; + providerSessionRefReleased: boolean; + cancellationKind: DiscoveryCancellationKind | null; }; - attempts.set(id, attempt); - coordinator.addAttempt({ + const registered = coordinator.addAttempt({ id, providerId: provider.id, priority: ProxyProviderResolver.resolveEffectivePriorityForSession(provider, session), @@ -5126,6 +5502,20 @@ export class ProxyForwarder { round: currentRound, launchOrder: attempt.sequence, }); + if (!registered || settled || committed) { + rollbackLaunch(); + return; + } + attempts.set(id, attempt); + discoveryMetrics.attemptStarted({ + attemptId: id, + providerId: provider.id, + round: + stickyProbeActive && provider.id === initialProvider.id && kind === "normal" + ? 0 + : currentRound, + kind, + }); void ProxyForwarder.doForward( attempt.session, @@ -5143,11 +5533,17 @@ export class ProxyForwarder { attempt.releaseAgent = runtime.releaseAgent ?? null; attempt.clearResponseTimeout?.(); attempt.response = response; + if (!attempt.pending || committed || settled) { + if (response.body && !attempt.reader) attempt.reader = response.body.getReader(); + cleanupAttempt(attempt, attempt.cancellationKind); + return; + } if (!response.body) throw new EmptyResponseError(provider.id, provider.name, "empty_body"); attempt.reader = response.body.getReader(); while (!committed && !settled && attempt.pending) { const item = await attempt.reader.read(); + if (attempt.readerTransferred || committed || settled || !attempt.pending) return; if (item.done) { // A ready candidate may have reached EOF while waiting for a // higher-priority attempt. Its buffered prefix remains a valid @@ -5161,10 +5557,30 @@ export class ProxyForwarder { // A single read can contain both deliverable content and the // protocol terminator. Terminal is only invalid when no content // was observed; otherwise the buffered candidate is complete. + if (validity.limitExceeded) { + discoveryMetrics.event("parser_limit", { + attemptId: id, + providerId: provider.id, + round: attempt.round, + }); + throw new DiscoveryValidityLimitError(); + } if (validity.error || (validity.terminal && !validity.ready)) throw new ProxyError("Invalid upstream discovery response", 502); if (!validity.ready) continue; attempt.ready = true; + if ( + attempt.kind === "fallback" && + (fallbackPromotionBlocked || roundLaunchesInProgress > 0) + ) { + // The next wave has been reserved but its normal attempts may + // still be awaiting selection/endpoint setup. Keep the fallback + // ready-held until those slots are registered or exhausted. + return; + } + // Record readiness even when the priority gate holds this attempt. + // The coordinator can then promote the buffered stream if the + // higher-priority attempt fails or the round closes. const action = coordinator.markReady(id); if (action.type === "commit_normal" || action.type === "promote_fallback") await commit(attempt); @@ -5180,25 +5596,183 @@ export class ProxyForwarder { }) .catch(async (error) => { if (committed || settled || !attempt.pending) return; - attempt.pending = false; - const failureAction = coordinator.markFailed(id); lastError = error instanceof Error ? error : new Error(String(error)); lastErrorCategory = await categorizeErrorAsync(lastError); + const errorMessage = + lastError instanceof ProxyError + ? lastError.getDetailedErrorMessage() + : lastError.message; + + if (attempt.endpointAudit.endpointId != null) { + const isTimeoutError = lastError instanceof ProxyError && lastError.statusCode === 524; + if (isTimeoutError || lastErrorCategory === ErrorCategory.SYSTEM_ERROR) { + await recordEndpointFailure(attempt.endpointAudit.endpointId, lastError).catch( + () => undefined + ); + } + } + + if (lastErrorCategory === ErrorCategory.CLIENT_ABORT) { + attempt.pending = false; + coordinator.cancelRequest(); + session.addProviderToChain(provider, { + ...attempt.endpointAudit, + reason: "client_abort", + attemptNumber: attempt.sequence, + errorMessage: "Client aborted request", + }); + cleanupAttempt(attempt, "client_abort"); + await settleFailure( + lastError instanceof ProxyError + ? lastError + : new ProxyError("Request aborted by client", 499, undefined, true), + { preserveBinding: true, cancellationKind: "client_abort" } + ); + return; + } + + if (lastErrorCategory === ErrorCategory.LOCAL_OVERLOAD) { + const admission = findDbPoolAdmissionError(lastError); + const safeAdmissionMessage = admission?.message ?? "Database pool admission exceeded"; + session.addProviderToChain(provider, { + ...attempt.endpointAudit, + reason: "system_error", + attemptNumber: attempt.sequence, + errorMessage: safeAdmissionMessage, + errorDetails: { + system: { + errorType: "DbPoolAdmissionError", + errorName: "DbPoolAdmissionError", + errorMessage: safeAdmissionMessage, + errorCode: admission?.code, + }, + request: buildRequestDetails(session), + }, + }); + cleanupAttempt(attempt, null); + await settleFailure(lastError, { preserveBinding: true }); + return; + } + + // A failure can race the async selection/endpoint setup of the wave + // that is meant to replace it. Wait until those reserved slots have + // either registered or rolled back before the coordinator decides + // whether another round is needed. + await waitForRoundLaunches(); + if (committed || settled || !attempt.pending) return; + + // Preserve the existing provider-local rectifier contract before + // classifying a 400 as terminal. The rectifier mutates the shadow + // request session, so retry the same attempt session rather than + // creating a fresh unrectified shadow from the parent session. + const rectifier = await tryApplyReactiveRectifier({ + provider, + requestSession: attempt.session, + persistSession: session, + errorMessage, + attemptNumber: attempt.requestAttemptCount, + retryAttemptNumber: attempt.requestAttemptCount + 1, + retryState: attempt.reactiveRectifierRetryState, + }); + if (rectifier.matched && rectifier.applied) { + const providerSessionRefTransfer = { + owned: attempt.providerSessionRefOwned && !attempt.providerSessionRefReleased, + retainOnSuccess: attempt.providerSessionRefRetainOnSuccess, + }; + // The provider-local retry keeps the same concurrency slot. Move + // ownership to the replacement launch before cleaning the failed + // transport so there is no release/reacquire race window. + if (providerSessionRefTransfer.owned) attempt.providerSessionRefOwned = false; + attempt.pending = false; + coordinator.removeAttempt(id); + cleanupAttempt(attempt, null); + session.addProviderToChain(provider, { + ...buildRetryFailedChainEntry( + provider, + attempt.endpointAudit, + attempt.requestAttemptCount, + lastError, + errorMessage, + rectifier.requestDetailsBeforeRectify, + rawCrossProviderFallbackEnabled + ), + modelRedirect: getAttemptModelRedirect(attempt), + }); + launched.delete(provider.id); + try { + await launch(provider, attempt.kind, { + attemptSession: attempt.session, + requestAttemptCount: attempt.requestAttemptCount + 1, + retryState: attempt.reactiveRectifierRetryState, + providerSessionRefTransfer, + }); + } catch (retryLaunchError) { + lastError = + retryLaunchError instanceof Error + ? retryLaunchError + : new Error(String(retryLaunchError)); + lastErrorCategory = await categorizeErrorAsync(lastError); + if (stickyProbeActive && provider.id === initialProvider.id) { + stickyProbeActive = false; + if (stickyTimer) { + clearTimeout(stickyTimer); + stickyTimer = null; + } + await clearCapturedStickyBinding(0); + coordinator.startDiscoveryAfterSticky(); + await launchNextRound(concurrency, true); + return; + } + await settleFailure( + ProxyForwarder.resolveHedgeTerminalError(lastError, lastErrorCategory) + ); + } + return; + } + + const failedStickyProbe = + stickyProbeActive && attempt.kind === "normal" && provider.id === initialProvider.id; + attempt.pending = false; + const failureAction = failedStickyProbe + ? ({ type: "none" } as const) + : coordinator.markFailed(id); + if (failedStickyProbe) coordinator.removeAttempt(id); session.addProviderToChain(provider, { ...attempt.endpointAudit, reason: "retry_failed", attemptNumber: attempt.sequence, statusCode: lastError instanceof ProxyError ? lastError.statusCode : undefined, - errorMessage: lastError.message, + errorMessage, }); if ( + !(lastError instanceof DiscoveryValidityLimitError) && lastErrorCategory === ErrorCategory.PROVIDER_ERROR && !(lastError instanceof ProxyError && lastError.statusCode === 404) ) { await recordFailure(provider.id, lastError).catch(() => undefined); } - attempt.releaseAgent?.(); - releaseProviderRef(provider.id); + cleanupAttempt(attempt, null); + if (lastErrorCategory === ErrorCategory.NON_RETRYABLE_CLIENT_ERROR) { + // Client/input errors are independent of the selected provider. + // Stop Discovery immediately so the same invalid request is not + // fanned out or masked by a later generic fallback error. + await settleFailure( + ProxyForwarder.resolveHedgeTerminalError(lastError, lastErrorCategory), + { preserveBinding: true } + ); + return; + } + if (failedStickyProbe) { + stickyProbeActive = false; + if (stickyTimer) { + clearTimeout(stickyTimer); + stickyTimer = null; + } + await clearCapturedStickyBinding(0); + coordinator.startDiscoveryAfterSticky(); + await launchNextRound(concurrency, true); + return; + } const actionOwnsNextStep = failureAction.type === "commit_normal" || failureAction.type === "promote_fallback" || @@ -5209,7 +5783,18 @@ export class ProxyForwarder { } if (!actionOwnsNextStep && !committed && !settled) { const replacement = await chooseCandidate(); - if (replacement) await launch(replacement, "normal"); + if (replacement) { + try { + await launch(replacement, "normal"); + } catch (launchError) { + lastError = + launchError instanceof Error ? launchError : new Error(String(launchError)); + // A single launch failure does not prove that the remaining + // candidate pool is exhausted. The current round boundary can + // still advance or retry selection from the remaining pool. + noMoreCandidates = false; + } + } } if ( Array.from(attempts.values()).every((candidate) => !candidate.pending) && @@ -5228,31 +5813,60 @@ export class ProxyForwarder { }); }; - const launchNextRound = async () => { + const launchNextRound = async (slots: number, coordinatorAlreadyAdvanced = false) => { if (settled || committed) return; + roundLaunchesInProgress += 1; clearRoundTimer(); - currentRound += 1; - if (currentRound > maxRounds) return; - coordinator.beginRound(); - const candidate = await chooseCandidate(); - if (candidate) { - try { - await launch(candidate, "normal"); - } catch (error) { - lastError = error instanceof Error ? error : new Error(String(error)); + try { + if (coordinatorAlreadyAdvanced) { + currentRound = coordinator.round; + } else { + const nextRound = coordinator.beginRound(); + currentRound = nextRound.round; } - } - if (!committed && !settled) { - clearRoundTimer(); - roundTimer = setTimeout(() => { - void onBoundary().catch((error) => - logger.warn("[Discovery] Round boundary failed", { error }) + if (currentRound > maxRounds || slots <= 0) return; + const candidates = await ProxyProviderResolver.pickDiscoveryProviders( + session, + slots, + Array.from(launched) + ); + if (candidates.length === 0) { + noMoreCandidates = true; + } + for (const candidate of candidates) { + try { + await launch(candidate, "normal"); + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + // Launch setup failures do not establish pool exhaustion. + noMoreCandidates = false; + } + } + const hasPendingAttempt = Array.from(attempts.values()).some((attempt) => attempt.pending); + if (!hasPendingAttempt) { + await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); + return; + } + if (!committed && !settled) { + scheduleRoundBoundary(discoverySlaMs); + } + } finally { + roundLaunchesInProgress = Math.max(0, roundLaunchesInProgress - 1); + if (roundLaunchesInProgress === 0) { + notifyRoundLaunchIdle(); + fallbackPromotionBlocked = false; + const readyFallback = Array.from(attempts.values()).find( + (attempt) => attempt.pending && attempt.ready && attempt.kind === "fallback" ); - }, discoverySlaMs); + if (readyFallback && !committed && !settled) { + const action = coordinator.markReady(readyFallback.id); + if (action.type === "promote_fallback") await commit(readyFallback); + } + } } }; - executeCoordinatorAction = async (action) => { + executeCoordinatorAction = async (action, terminalCancellationKind) => { if (settled || committed) return; if (action.type === "cancel" || action.type === "launch") { const cancelIds = @@ -5263,15 +5877,18 @@ export class ProxyForwarder { // the action. Restore the transport-facing state long enough for the // exactly-once cancellation/release path to run. if (attempt && !attempt.readerTransferred) attempt.pending = true; - if (attempt) cancelAttempt(attempt, "discovery_round_boundary"); + if (attempt) cancelAttempt(attempt, "discovery_sla_timeout"); } if (action.promoteAttemptId) { const fallback = attempts.get(action.promoteAttemptId); - if (fallback) fallback.kind = "fallback"; - if (action.type === "cancel" && currentRound < maxRounds) { - await launchNextRound(); + if (fallback) { + fallback.kind = "fallback"; + discoveryMetrics.fallbackPromoted(fallback.id, fallback.provider.id, fallback.round); } } + // A final-round fallback may still be waiting for a protocol-valid + // prefix. Keep it alive; markReady will commit it when it becomes safe. + if (action.type === "cancel") return; } if (action.type === "commit_normal" || action.type === "promote_fallback") { const attempt = attempts.get(action.attemptId); @@ -5279,113 +5896,105 @@ export class ProxyForwarder { return; } if (action.type === "launch") { - await launchNextRound(); + await launchNextRound(action.slots, true); return; } if (action.type === "terminal_failure") { - await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); + await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError), { + cancellationKind: terminalCancellationKind, + }); return; } - if (action.type === "none") { - const fallbackPending = Array.from(attempts.values()).some( - (attempt) => attempt.pending && attempt.kind === "fallback" - ); - if (fallbackPending && currentRound < maxRounds) { - await launchNextRound(); - } - } - }; - - const onBoundary = async () => { - if (settled || committed) return; - await executeCoordinatorAction(coordinator.onRoundBoundary()); }; const cleanupAbort = bindClientAbortListener(session.clientAbortSignal, () => { if (settled || committed) return; - void settleFailure(new ProxyError("Request aborted by client", 499, undefined, true)).catch( - (error) => logger.warn("[Discovery] Client abort cleanup failed", { error }) - ); + if (stickyTimer) clearTimeout(stickyTimer); + coordinator.cancelRequest(); + void settleFailure(new ProxyError("Request aborted by client", 499, undefined, true), { + preserveBinding: true, + cancellationKind: "client_abort", + }).catch((error) => logger.warn("[Discovery] Client abort cleanup failed", { error })); }); - totalTimer = setTimeout(() => { - if (settled || committed) return; - void executeCoordinatorAction(coordinator.onDeadline()).catch((error) => - logger.warn("[Discovery] Deadline action failed", { error }) - ); - }, totalTimeoutMs); + totalTimer = setTimeout( + () => { + if (settled || committed) return; + void executeCoordinatorAction(coordinator.onDeadline(), "request_deadline").catch((error) => + logger.warn("[Discovery] Deadline action failed", { error }) + ); + }, + Math.max(0, racingDeadlineAt - Date.now()) + ); - try { - const hasSticky = session.shouldReuseProvider() && !!session.sessionId; + const orchestrate = async () => { + let initialLaunchFailed = false; try { await launch(initialProvider, "normal"); } catch (error) { + initialLaunchFailed = true; + stickyProbeActive = false; lastError = error instanceof Error ? error : new Error(String(error)); - const replacement = await chooseCandidate(); - if (replacement) { - try { - await launch(replacement, "normal"); - } catch (replacementError) { - lastError = - replacementError instanceof Error - ? replacementError - : new Error(String(replacementError)); - await settleFailure(ProxyForwarder.resolveHedgeTerminalError(lastError, null)); - } - } else { - await settleFailure(ProxyForwarder.resolveHedgeTerminalError(lastError, null)); - } } - const initial = hasSticky ? concurrency - 1 : Math.max(0, concurrency - 1); + if (settled || committed) return; + if (hasSticky) { - stickyTimer = setTimeout(() => { - const sticky = Array.from(attempts.values()).find( - (attempt) => attempt.pending && attempt.provider.id === initialProvider.id - ); - if (sticky) { - if (!coordinator.demoteToFallback(sticky.id)) return; - sticky.kind = "fallback"; - if (bindingSnapshot && bindingSnapshot.providerId === initialProvider.id) { - void SessionManager.clearVersionedSessionProvider( - bindingSnapshot, - initialProvider.id, - Math.ceil((settings.stickyTimeoutCooldownMs ?? 300_000) / 1000) - ).catch((error) => - logger.debug("[Discovery] Failed to clear timed-out Sticky", { error }) - ); - } - if (currentRound < maxRounds) { - void launchNextRound().catch((error) => - logger.warn("[Discovery] Sticky round launch failed", { error }) - ); - } else { - void onBoundary().catch((error) => - logger.warn("[Discovery] Sticky boundary failed", { error }) + if (initialLaunchFailed) { + await clearCapturedStickyBinding(0); + coordinator.startDiscoveryAfterSticky(); + await launchNextRound(concurrency, true); + } else { + stickyTimer = setTimeout( + () => { + if (!stickyProbeActive || settled || committed) return; + if (Date.now() >= racingDeadlineAt) { + void executeCoordinatorAction(coordinator.onDeadline(), "request_deadline").catch( + (error) => logger.warn("[Discovery] Deadline action failed", { error }) + ); + return; + } + const sticky = Array.from(attempts.values()).find( + (attempt) => attempt.pending && attempt.provider.id === initialProvider.id ); - } - } - }, stickySlaMs); - } else { - const candidates = await ProxyProviderResolver.pickDiscoveryProviders( - session, - initial, - Array.from(launched) - ); - if (candidates.length === 0) noMoreCandidates = true; - for (const provider of candidates) { - try { - await launch(provider, "normal"); - } catch (error) { - lastError = error instanceof Error ? error : new Error(String(error)); - } - } - clearRoundTimer(); - roundTimer = setTimeout(() => { - void onBoundary().catch((error) => - logger.warn("[Discovery] Round boundary failed", { error }) + if (sticky) { + if (!coordinator.demoteToFallback(sticky.id)) return; + stickyProbeActive = false; + coordinator.startDiscoveryAfterSticky(); + sticky.kind = "fallback"; + discoveryMetrics.fallbackPromoted(sticky.id, sticky.provider.id, 0); + fallbackPromotionBlocked = true; + if (bindingSnapshot && bindingSnapshot.providerId === initialProvider.id) { + void clearCapturedStickyBinding( + Math.ceil((settings.stickyTimeoutCooldownMs ?? 300_000) / 1000) + ).finally(() => { + void launchNextRound(Math.max(0, concurrency - 1), true).catch((error) => + logger.warn("[Discovery] Sticky round launch failed", { error }) + ); + }); + return; + } + void launchNextRound(Math.max(0, concurrency - 1), true).catch((error) => + logger.warn("[Discovery] Sticky round launch failed", { error }) + ); + } + }, + Math.min(stickySlaMs, Math.max(0, racingDeadlineAt - Date.now())) ); - }, discoverySlaMs); + } + } else { + const activeNormals = Array.from(attempts.values()).filter( + (attempt) => attempt.pending && attempt.kind === "normal" + ).length; + await launchNextRound(Math.max(0, concurrency - activeNormals), true); } + }; + + void orchestrate().catch(async (error) => { + const normalized = error instanceof Error ? error : new Error(String(error)); + await settleFailure(ProxyForwarder.resolveHedgeTerminalError(normalized, null)); + }); + + try { const result = await resultPromise; if (result.error) throw result.error; return result.response as Response; @@ -5394,6 +6003,29 @@ export class ProxyForwarder { if (totalTimer) clearTimeout(totalTimer); clearRoundTimer(); if (stickyTimer) clearTimeout(stickyTimer); + if (!leaseTransferred) { + void SessionManager.releaseSessionDiscoveryLease( + lease.sessionId, + lease.keyId, + lease.ownerToken + ) + .then((released) => { + if (released.status !== "released") { + logger.debug("[Discovery] Lease release skipped", { + sessionId: lease.sessionId, + keyId: lease.keyId, + status: released.status, + }); + } + }) + .catch((releaseError) => { + logger.warn("[Discovery] Lease release failed", { + sessionId: lease.sessionId, + keyId: lease.keyId, + error: releaseError instanceof Error ? releaseError.message : String(releaseError), + }); + }); + } } } @@ -5627,7 +6259,7 @@ export class ProxyForwarder { session: ProxySession, expectedProviderId: number | null ): Promise { - if (!session.sessionId) return; + if (!session.sessionId || !session.isSessionBindingAllowed()) return; const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; await SessionManager.clearSessionProvider(session.sessionId, expectedProviderId, keyId); } @@ -5636,7 +6268,7 @@ export class ProxyForwarder { session: ProxySession, expectedProviderIds: Iterable ): Promise { - if (!session.sessionId) return; + if (!session.sessionId || !session.isSessionBindingAllowed()) return; const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; await SessionManager.clearSessionProviders(session.sessionId, expectedProviderIds, keyId); } @@ -5665,7 +6297,13 @@ export class ProxyForwarder { session as { consumeProviderSessionRef?: (id: number) => boolean } ).consumeProviderSessionRef; if (!providerSessionRefConsumer?.call(session, providerId)) return false; - void RateLimitService.releaseProviderSession(providerId, session.sessionId); + void RateLimitService.releaseProviderSession(providerId, session.sessionId).catch((error) => { + logger.warn("ProxyForwarder: Failed to release Provider session reference", { + providerId, + sessionId: session.sessionId, + error: error instanceof Error ? error.message : String(error), + }); + }); return true; } diff --git a/src/app/v1/_lib/proxy/provider-selector.ts b/src/app/v1/_lib/proxy/provider-selector.ts index d4ce6fdbe..e66e18272 100644 --- a/src/app/v1/_lib/proxy/provider-selector.ts +++ b/src/app/v1/_lib/proxy/provider-selector.ts @@ -296,7 +296,9 @@ export class ProxyProviderResolver { // === 成功 === if (checkResult.referenced) { - session.recordProviderSessionRef(session.provider.id); + session.recordProviderSessionRef(session.provider.id, { + retainOnSuccess: checkResult.tracked, + }); } logger.debug("ProviderSelector: Session tracked atomically", { diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index 70d0909b5..a9ca3a7de 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -9,10 +9,12 @@ import { getEnvConfig } from "@/lib/config/env.schema"; import { getCachedSystemSettings } from "@/lib/config/system-settings-cache"; import { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; import { logger } from "@/lib/logger"; +import { recordDiscoveryControlEvent } from "@/lib/observability/discovery-metrics"; import { requestCloudPriceTableSync } from "@/lib/price-sync/cloud-price-updater"; import { ProxyStatusTracker } from "@/lib/proxy-status-tracker"; import { RateLimitService } from "@/lib/rate-limit"; import { deleteLiveChain } from "@/lib/redis/live-chain-store"; +import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import { SessionManager } from "@/lib/session-manager"; import { SessionTracker } from "@/lib/session-tracker"; import { CODEX_1M_CONTEXT_TOKEN_THRESHOLD } from "@/lib/special-attributes"; @@ -58,6 +60,7 @@ import { isClientAbortError, isTransportError } from "./errors"; import type { ProxySession } from "./session"; import { consumeDeferredStreamingFinalization, + type DeferredStreamingFinalization, peekDeferredStreamingFinalization, } from "./stream-finalization"; @@ -95,8 +98,157 @@ function resolveStreamTaskStaleTimeoutMs(): number { const STREAM_FINALIZATION_MAX_MS = 120_000; const STREAM_FAILURE_PERSISTENCE_MAX_MS = 5_000; +const DISCOVERY_LEASE_RELEASE_MAX_MS = 5_000; const NON_STREAM_TERMINAL_PERSISTENCE_ERROR = Symbol("non_stream_terminal_persistence_error"); +type DiscoveryLeaseLifecycle = { + active: boolean; + ensureOwned: () => Promise; + release: () => Promise; +}; + +function isSessionBindingMutationAllowed(session: ProxySession): boolean { + const checker = (session as ProxySession & { isSessionBindingAllowed?: () => boolean }) + .isSessionBindingAllowed; + return checker?.call(session) !== false; +} + +function startDiscoveryLeaseLifecycle(session: ProxySession): DiscoveryLeaseLifecycle { + const lease = peekDeferredStreamingFinalization(session)?.discoveryLease; + if (!lease) { + return { + active: false, + ensureOwned: async () => true, + release: async () => undefined, + }; + } + + let renewalTimer: ReturnType | null = null; + let renewalInFlight: Promise | null = null; + let releasePromise: Promise | null = null; + let ownershipState: "unknown" | "owned" | "lost" = "unknown"; + + const stopRenewal = () => { + if (renewalTimer) { + clearInterval(renewalTimer); + renewalTimer = null; + } + }; + + const renew = (): Promise => { + if (releasePromise || ownershipState === "lost") return Promise.resolve(false); + if (renewalInFlight) return renewalInFlight; + + const operation = (async () => { + const result = await SessionManager.renewSessionDiscoveryLease( + lease.sessionId, + lease.keyId, + lease.ownerToken, + lease.ttlSeconds + ); + if (result.status !== "renewed") { + ownershipState = "lost"; + stopRenewal(); + logger.warn("[ResponseHandler] Discovery lease renewal stopped", { + sessionId: lease.sessionId, + keyId: lease.keyId, + status: result.status, + reason: "reason" in result ? result.reason : undefined, + }); + return false; + } + ownershipState = "owned"; + return true; + })() + .catch((error) => { + ownershipState = "lost"; + stopRenewal(); + logger.warn("[ResponseHandler] Discovery lease renewal failed", { + sessionId: lease.sessionId, + keyId: lease.keyId, + error: error instanceof Error ? error.message : String(error), + }); + return false; + }) + .finally(() => { + if (renewalInFlight === operation) renewalInFlight = null; + }); + renewalInFlight = operation; + return operation; + }; + + // Ownership transfers from the Forwarder to the finalizer without delaying + // the downstream response. Renew immediately so an expired/lost token is + // observed before any terminal Session binding mutation is attempted. + const handoffRenewal = renew(); + const renewalIntervalMs = Math.max(250, Math.floor((lease.ttlSeconds * 1000) / 3)); + renewalTimer = setInterval(() => { + void renew(); + }, renewalIntervalMs); + renewalTimer.unref?.(); + + return { + active: true, + ensureOwned: async () => { + if (!(await handoffRenewal) || releasePromise || ownershipState !== "owned") return false; + // Revalidate with the owner token at the mutation boundary. This runs in + // post-terminal side effects, so it cannot add latency to downstream TTFB. + return renew(); + }, + release: () => { + if (releasePromise) return releasePromise; + stopRenewal(); + releasePromise = (async () => { + try { + const result = await raceWithTimeout( + SessionManager.releaseSessionDiscoveryLease( + lease.sessionId, + lease.keyId, + lease.ownerToken + ), + DISCOVERY_LEASE_RELEASE_MAX_MS, + "discovery_lease_release_timeout" + ); + if (result.status !== "released") { + logger.debug("[ResponseHandler] Discovery lease release skipped", { + sessionId: lease.sessionId, + keyId: lease.keyId, + status: result.status, + reason: "reason" in result ? result.reason : undefined, + }); + } + } catch (error) { + logger.warn("[ResponseHandler] Discovery lease release failed", { + sessionId: lease.sessionId, + keyId: lease.keyId, + error: error instanceof Error ? error.message : String(error), + }); + } + })(); + return releasePromise; + }, + }; +} + +async function releaseOwnedProviderSessionRef( + session: ProxySession, + meta: DeferredStreamingFinalization | null, + retainAsBaseline: boolean +): Promise { + if (retainAsBaseline || meta?.providerSessionRefOwned !== true || !session.sessionId) return; + if (!session.consumeProviderSessionRef(meta.providerId)) return; + + try { + await RateLimitService.releaseProviderSession(meta.providerId, session.sessionId); + } catch (error) { + logger.warn("[ResponseHandler] Failed to release Discovery Provider session reference", { + sessionId: session.sessionId, + providerId: meta.providerId, + error: error instanceof Error ? error.message : String(error), + }); + } +} + type MessageRequestTerminalDetails = Parameters[1]; type NonStreamTerminalPersistenceError = Error & { [NON_STREAM_TERMINAL_PERSISTENCE_ERROR]: true; @@ -974,8 +1126,31 @@ function hasPositiveBillableTokens(usage: UsageMetrics | null): boolean { return tokens > 0; } -const FINISH_REASON_MARKER = /"finish_reason"\s*:\s*"[a-z_]+"/; -const GEMINI_FINISH_REASON_MARKER = /"finishReason"\s*:\s*"[A-Z_]+"/; +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function hasOpenAIChatCompletionMarker(data: unknown): boolean { + if (!isRecord(data) || !Array.isArray(data.choices)) return false; + return data.choices.some( + (choice) => + isRecord(choice) && + typeof choice.finish_reason === "string" && + choice.finish_reason.trim().length > 0 + ); +} + +function hasGeminiCompletionMarker(data: unknown, format: ProxySession["originalFormat"]): boolean { + if (!isRecord(data)) return false; + const payload = format === "gemini-cli" && isRecord(data.response) ? data.response : data; + if (!Array.isArray(payload.candidates)) return false; + return payload.candidates.some( + (candidate) => + isRecord(candidate) && + typeof candidate.finishReason === "string" && + candidate.finishReason.trim().length > 0 + ); +} /** * 判断流式响应文本中是否存在“与格式匹配的终止完成标记”，用以区分 @@ -984,16 +1159,38 @@ const GEMINI_FINISH_REASON_MARKER = /"finishReason"\s*:\s*"[A-Z_]+"/; * 仅 usage>0 不足以证明完成：Anthropic 在首个 `message_start` 即带 usage、 * Gemini 在中间事件即带 usageMetadata，截断流同样会出现正向 token。 */ -function hasStreamCompletionMarker(text: string): boolean { - if ( - text.includes("response.completed") || // OpenAI Responses / Codex - text.includes("message_stop") || // Anthropic Messages - text.includes("[DONE]") // OpenAI Chat Completions - ) { - return true; +function hasStreamCompletionMarker(text: string, format: ProxySession["originalFormat"]): boolean { + const events = parseSSEData(text); + + switch (format) { + case "response": + return events.some( + (event) => + event.event === "response.completed" && + isRecord(event.data) && + event.data.type === "response.completed" && + isRecord(event.data.response) + ); + case "claude": + return events.some( + (event) => + event.event === "message_stop" && + isRecord(event.data) && + event.data.type === "message_stop" + ); + case "openai": + return events.some( + (event) => + event.event === "message" && + ((typeof event.data === "string" && event.data.trim() === "[DONE]") || + hasOpenAIChatCompletionMarker(event.data)) + ); + case "gemini": + case "gemini-cli": + return events.some( + (event) => event.event === "message" && hasGeminiCompletionMarker(event.data, format) + ); } - // OpenAI chat / Gemini：非空 finish reason 标记最终块。 - return FINISH_REASON_MARKER.test(text) || GEMINI_FINISH_REASON_MARKER.test(text); } export async function resolveBillableUsageMetricsForCost( @@ -1117,6 +1314,12 @@ type FinalizeDeferredStreamingResult = { }; /** Circuit and Session side effects, committed after durable terminal details. */ commitSideEffects?: () => Promise; + /** Attempt-scoped ref cleanup; idempotent via ProxySession ownership consumption. */ + finalizeAttemptResources?: () => Promise; + /** Whether terminal helpers may create auxiliary Sticky bindings (for example Codex cache keys). */ + allowAuxiliarySessionBinding: boolean; + /** Discovery auxiliary bindings must wait for, and depend on, the primary generation CAS. */ + confirmAuxiliarySessionBinding: () => Promise; }; /** @@ -1142,28 +1345,131 @@ function finalizeDeferredStreamingFinalizationIfNeeded( upstreamStatusCode: number, streamEndedNormally: boolean, clientAborted: boolean, + discoveryLeaseLifecycle: DiscoveryLeaseLifecycle, abortReason?: string ): FinalizeDeferredStreamingResult { const meta = consumeDeferredStreamingFinalization(session); const provider = session.provider; const providerIdForPersistence = meta?.providerId ?? provider?.id ?? null; const clearSessionBinding = async () => { - if (!session.sessionId) return; + if (!session.sessionId || !isSessionBindingMutationAllowed(session)) return; const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; - if (meta?.bindingIntent === "none") return; - if (meta?.bindingSnapshot && keyId != null) { - await SessionManager.clearVersionedSessionProvider( + if (meta?.bindingIntent === "none" || meta?.bindingIntent === "create") return; + if (meta?.bindingIntent === "renew") { + // A client disconnect is not evidence that the Sticky Provider failed. + // Discovery renewals may only clear the exact binding snapshot that was + // observed before the request. + if (clientAborted) return; + if ( + !meta.bindingSnapshot || + keyId == null || + meta.bindingSnapshot.keyId !== keyId || + meta.bindingSnapshot.sessionId !== session.sessionId || + meta.bindingSnapshot.providerId !== meta.providerId + ) { + logger.debug("[ResponseHandler] Discovery binding clear skipped", { + sessionId: session.sessionId, + keyId, + expectedProviderId: meta.providerId, + reason: "missing_or_mismatched_snapshot", + }); + return; + } + if (!(await discoveryLeaseLifecycle.ensureOwned())) { + logger.warn( + "[ResponseHandler] Discovery binding clear skipped after lease ownership loss", + { + sessionId: meta.bindingSnapshot.sessionId, + keyId: meta.bindingSnapshot.keyId, + expectedProviderId: meta.bindingSnapshot.providerId, + } + ); + return; + } + const cleared = await SessionManager.clearVersionedSessionProvider( meta.bindingSnapshot, - meta.bindingSnapshot.providerId, + meta.providerId, 0 ); + if (cleared.status !== "ok") { + logger.debug("[ResponseHandler] Discovery binding clear skipped", { + sessionId: meta.bindingSnapshot.sessionId, + keyId: meta.bindingSnapshot.keyId, + expectedProviderId: meta.bindingSnapshot.providerId, + reason: cleared.reason, + }); + } return; } + + // Legacy deferred finalization has no explicit binding intent and keeps its + // pre-Discovery behavior. await SessionManager.clearSessionProvider(session.sessionId, providerIdForPersistence, keyId); }; + let retainProviderSessionRef = false; + const finalizeProviderSessionRef = () => + releaseOwnedProviderSessionRef(session, meta, retainProviderSessionRef); + + const compareAndSetDiscoveryBinding = async ( + snapshot: SessionBindingSnapshot, + providerId: number, + keyId: number + ) => { + if ( + !snapshot || + snapshot.keyId !== keyId || + (session.sessionId !== null && snapshot.sessionId !== session.sessionId) + ) + return { updated: false, reason: "missing_snapshot", details: "missing_snapshot" }; + + if (!(await discoveryLeaseLifecycle.ensureOwned())) { + return { + updated: false, + reason: "discovery_lease_not_owned", + details: "lease_lost_or_unavailable", + }; + } + + const cas = await SessionManager.compareAndSetSessionProvider(snapshot, providerId); + if (cas.status === "conflict") { + recordDiscoveryControlEvent("binding_cas_conflict", { + requestId: session.messageContext?.id ?? null, + sessionId: snapshot.sessionId, + keyId, + providerId, + reason: cas.reason, + }); + } + + return { + updated: cas.status === "ok", + reason: cas.status === "ok" ? "discovery_generation_cas" : cas.reason, + details: cas.status, + }; + }; + const isHedgeWinner = meta?.isHedgeWinner === true; const billHedgeLosers = meta?.billHedgeLosers === true; + const allowAuxiliarySessionBinding = + isSessionBindingMutationAllowed(session) && + (meta?.bindingIntent === undefined || (meta.bindingIntent !== "none" && !clientAborted)); + const hasDiscoveryBindingIntent = + meta?.bindingIntent === "create" || meta?.bindingIntent === "renew"; + let resolvePrimaryDiscoveryBinding: ((updated: boolean) => void) | null = null; + const primaryDiscoveryBinding = hasDiscoveryBindingIntent + ? new Promise((resolve) => { + resolvePrimaryDiscoveryBinding = resolve; + }) + : Promise.resolve(allowAuxiliarySessionBinding); + let primaryDiscoveryBindingSettled = false; + const settlePrimaryDiscoveryBinding = (updated: boolean) => { + if (primaryDiscoveryBindingSettled) return; + primaryDiscoveryBindingSettled = true; + resolvePrimaryDiscoveryBinding?.(updated); + }; + const confirmAuxiliarySessionBinding = async () => + allowAuxiliarySessionBinding && (await primaryDiscoveryBinding); // 仅在“上游 HTTP=200 且流自然结束”时做“假 200”检测： // - 非 200：HTTP 已经表明失败（无需额外启发式） @@ -1177,7 +1483,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( const completionMarkerMissing = meta?.requiresCompletionMarker === true && streamEndedNormally && - !hasStreamCompletionMarker(allContent); + !hasStreamCompletionMarker(allContent, session.originalFormat); let clientAbortGateUsage: FinalizeDeferredStreamingResult["clientAbortGateUsage"]; const clientAbortCompleteSuccess = (() => { if (!clientAborted || upstreamStatusCode < 200 || upstreamStatusCode >= 300) { @@ -1196,7 +1502,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( // completion marker is present, proving the upstream finished before the // client stopped reading. Otherwise keep the pre-PR safe default (499, // unbilled). - if (!hasStreamCompletionMarker(allContent)) { + if (!hasStreamCompletionMarker(allContent, session.originalFormat)) { return false; } @@ -1262,6 +1568,16 @@ function finalizeDeferredStreamingFinalizationIfNeeded( // - 只返回“内部状态码 + 错误原因”，由调用方写入统计； // - 不在这里更新熔断/绑定（meta 缺失意味着 Forwarder 没有启用延迟结算；provider 缺失意味着无法归因）。 if (!meta || !provider) { + const commitSideEffects = + shouldClearSessionBindingOnFailure || meta?.providerSessionRefOwned === true + ? async () => { + try { + if (shouldClearSessionBindingOnFailure) await clearSessionBinding(); + } finally { + await finalizeProviderSessionRef(); + } + } + : undefined; return { effectiveStatusCode, errorMessage, @@ -1269,7 +1585,10 @@ function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, - commitSideEffects: shouldClearSessionBindingOnFailure ? clearSessionBinding : undefined, + commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1313,22 +1632,26 @@ function finalizeDeferredStreamingFinalizationIfNeeded( }); const commitSideEffects = async () => { - await clearSessionBinding(); + try { + await clearSessionBinding(); - if (!clientAborted && session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(meta.providerId, new Error(errorMessage ?? "STREAM_ABORTED")); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record streaming failure in circuit breaker", { - providerId: meta.providerId, - sessionId: session.sessionId ?? null, - error: cbError, - }); - } + if (!clientAborted && session.getEndpointPolicy().allowCircuitBreakerAccounting) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(errorMessage ?? "STREAM_ABORTED")); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record streaming failure in circuit breaker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } - // Stream aborts are key-level errors. The endpoint delivered HTTP 200, - // so only the Provider circuit is updated here. + // Stream aborts are key-level errors. The endpoint delivered HTTP 200, + // so only the Provider circuit is updated here. + } + } finally { + await finalizeProviderSessionRef(); } }; @@ -1340,6 +1663,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1354,18 +1680,22 @@ function finalizeDeferredStreamingFinalizationIfNeeded( }); const commitSideEffects = async () => { - await clearSessionBinding(); - if (session.getEndpointPolicy().allowCircuitBreakerAccounting) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(meta.providerId, new Error(errorMessage ?? "STREAM_ABORTED")); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record missing stream completion marker", { - providerId: meta.providerId, - sessionId: session.sessionId ?? null, - error: cbError, - }); + try { + await clearSessionBinding(); + if (session.getEndpointPolicy().allowCircuitBreakerAccounting) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(errorMessage ?? "STREAM_ABORTED")); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record missing stream completion marker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } } + } finally { + await finalizeProviderSessionRef(); } }; @@ -1377,6 +1707,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1413,23 +1746,27 @@ function finalizeDeferredStreamingFinalizationIfNeeded( }); const commitSideEffects = async () => { - await clearSessionBinding(); + try { + await clearSessionBinding(); - // 404 is RESOURCE_NOT_FOUND and must not penalize the Provider circuit. - if ( - effectiveStatusCode !== 404 && - session.getEndpointPolicy().allowCircuitBreakerAccounting - ) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(meta.providerId, new Error(detected.code)); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record fake-200 error in circuit breaker", { - providerId: meta.providerId, - sessionId: session.sessionId ?? null, - error: cbError, - }); + // 404 is RESOURCE_NOT_FOUND and must not penalize the Provider circuit. + if ( + effectiveStatusCode !== 404 && + session.getEndpointPolicy().allowCircuitBreakerAccounting + ) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(detected.code)); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record fake-200 error in circuit breaker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } } + } finally { + await finalizeProviderSessionRef(); } }; @@ -1441,6 +1778,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1471,22 +1811,26 @@ function finalizeDeferredStreamingFinalizationIfNeeded( }); const commitSideEffects = async () => { - await clearSessionBinding(); + try { + await clearSessionBinding(); - if ( - effectiveStatusCode !== 404 && - session.getEndpointPolicy().allowCircuitBreakerAccounting - ) { - try { - const { recordFailure } = await import("@/lib/circuit-breaker"); - await recordFailure(meta.providerId, new Error(errorMessage)); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record non-200 error in circuit breaker", { - providerId: meta.providerId, - sessionId: session.sessionId ?? null, - error: cbError, - }); + if ( + effectiveStatusCode !== 404 && + session.getEndpointPolicy().allowCircuitBreakerAccounting + ) { + try { + const { recordFailure } = await import("@/lib/circuit-breaker"); + await recordFailure(meta.providerId, new Error(errorMessage)); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record non-200 error in circuit breaker", { + providerId: meta.providerId, + sessionId: session.sessionId ?? null, + error: cbError, + }); + } } + } finally { + await finalizeProviderSessionRef(); } }; @@ -1498,6 +1842,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1517,42 +1864,47 @@ function finalizeDeferredStreamingFinalizationIfNeeded( } const commitSideEffects = async () => { - if (meta.endpointId != null) { + let primaryDiscoveryBindingUpdated = false; + try { + if (meta.endpointId != null) { + try { + const { recordEndpointSuccess } = await import("@/lib/endpoint-circuit-breaker"); + await recordEndpointSuccess(meta.endpointId); + } catch (endpointError) { + logger.warn("[ResponseHandler] Failed to record endpoint success (stream finalized)", { + endpointId: meta.endpointId, + providerId: meta.providerId, + error: endpointError, + }); + } + } + try { - const { recordEndpointSuccess } = await import("@/lib/endpoint-circuit-breaker"); - await recordEndpointSuccess(meta.endpointId); - } catch (endpointError) { - logger.warn("[ResponseHandler] Failed to record endpoint success (stream finalized)", { - endpointId: meta.endpointId, + const { recordSuccess } = await import("@/lib/circuit-breaker"); + await recordSuccess(meta.providerId); + } catch (cbError) { + logger.warn("[ResponseHandler] Failed to record streaming success in circuit breaker", { providerId: meta.providerId, - error: endpointError, + error: cbError, }); } - } - try { - const { recordSuccess } = await import("@/lib/circuit-breaker"); - await recordSuccess(meta.providerId); - } catch (cbError) { - logger.warn("[ResponseHandler] Failed to record streaming success in circuit breaker", { - providerId: meta.providerId, - error: cbError, - }); - } - - // Hedge winner: commitWinner() already performed session binding and chain logging. - if (meta.bindingIntent !== "none" && !meta.isHedgeWinner && session.sessionId) { - const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; - const result = - meta.bindingSnapshot && keyId != null - ? await SessionManager.compareAndSetSessionProvider( - meta.bindingSnapshot, - meta.providerId - ).then((cas) => ({ - updated: cas.status === "ok", - reason: cas.status === "ok" ? "discovery_generation_cas" : cas.reason, - details: cas.status, - })) + // A client abort may still be billable when a completion marker was + // already buffered, but it must never create or renew Sticky state. + if ( + meta.bindingIntent !== "none" && + !meta.isHedgeWinner && + !clientAborted && + session.sessionId && + isSessionBindingMutationAllowed(session) + ) { + const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + const isDiscoveryBinding = + meta.bindingIntent === "create" || meta.bindingIntent === "renew"; + const result = isDiscoveryBinding + ? meta.bindingSnapshot && keyId != null + ? await compareAndSetDiscoveryBinding(meta.bindingSnapshot, meta.providerId, keyId) + : { updated: false, reason: "missing_snapshot", details: "missing_snapshot" } : await SessionManager.updateSessionBindingSmart( session.sessionId, meta.providerId, @@ -1562,48 +1914,56 @@ function finalizeDeferredStreamingFinalizationIfNeeded( keyId ); - if (result.updated) { - logger.info("[ResponseHandler] Session binding updated (stream finalized)", { - sessionId: session.sessionId, - providerId: meta.providerId, - providerName: meta.providerName, - priority: meta.providerPriority, - reason: result.reason, - details: result.details, - attemptNumber: meta.attemptNumber, - totalProvidersAttempted: meta.totalProvidersAttempted, - }); - } else { - logger.debug("[ResponseHandler] Session binding not updated (stream finalized)", { - sessionId: session.sessionId, - providerId: meta.providerId, - providerName: meta.providerName, - priority: meta.providerPriority, - reason: result.reason, - details: result.details, - }); - } + primaryDiscoveryBindingUpdated = isDiscoveryBinding && result.updated; + retainProviderSessionRef = + primaryDiscoveryBindingUpdated && meta.providerSessionRefRetainOnSuccess === true; - if (session.shouldTrackSessionObservability()) { - void SessionManager.updateSessionProvider(session.sessionId, { - providerId: meta.providerId, - providerName: meta.providerName, - }).catch((err) => { - logger.error( - "[ResponseHandler] Failed to update session provider info (stream finalized)", - { error: err } - ); - }); + if (result.updated) { + logger.info("[ResponseHandler] Session binding updated (stream finalized)", { + sessionId: session.sessionId, + providerId: meta.providerId, + providerName: meta.providerName, + priority: meta.providerPriority, + reason: result.reason, + details: result.details, + attemptNumber: meta.attemptNumber, + totalProvidersAttempted: meta.totalProvidersAttempted, + }); + } else { + logger.debug("[ResponseHandler] Session binding not updated (stream finalized)", { + sessionId: session.sessionId, + providerId: meta.providerId, + providerName: meta.providerName, + priority: meta.providerPriority, + reason: result.reason, + details: result.details, + }); + } + + if (session.shouldTrackSessionObservability()) { + void SessionManager.updateSessionProvider(session.sessionId, { + providerId: meta.providerId, + providerName: meta.providerName, + }).catch((err) => { + logger.error( + "[ResponseHandler] Failed to update session provider info (stream finalized)", + { error: err } + ); + }); + } } - } - logger.info("[ResponseHandler] Streaming request finalized as success", { - providerId: meta.providerId, - providerName: meta.providerName, - attemptNumber: meta.attemptNumber, - totalProvidersAttempted: meta.totalProvidersAttempted, - statusCode: meta.upstreamStatusCode, - }); + logger.info("[ResponseHandler] Streaming request finalized as success", { + providerId: meta.providerId, + providerName: meta.providerName, + attemptNumber: meta.attemptNumber, + totalProvidersAttempted: meta.totalProvidersAttempted, + statusCode: meta.upstreamStatusCode, + }); + } finally { + settlePrimaryDiscoveryBinding(primaryDiscoveryBindingUpdated); + await finalizeProviderSessionRef(); + } }; return { @@ -1614,6 +1974,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1869,7 +2232,7 @@ export class ProxyResponseHandler { } const postTerminalSideEffects: Array<() => Promise> = []; - if (session.sessionId) { + if (session.sessionId && isSessionBindingMutationAllowed(session)) { const sessionId = session.sessionId; postTerminalSideEffects.push(async () => { const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; @@ -2044,7 +2407,9 @@ export class ProxyResponseHandler { const sessionId = session.sessionId; postTerminalSideEffects.push(async () => { const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; - await SessionManager.clearSessionProvider(sessionId, provider.id, keyId); + if (isSessionBindingMutationAllowed(session)) { + await SessionManager.clearSessionProvider(sessionId, provider.id, keyId); + } const sessionUsagePayload: SessionUsageUpdate = { status: @@ -2178,7 +2543,8 @@ export class ProxyResponseHandler { statusCode >= 200 && statusCode < 300 && session.sessionId && - provider.id + provider.id && + isSessionBindingMutationAllowed(session) ) { try { const responseData = JSON.parse(responseText) as Record; @@ -2594,10 +2960,16 @@ export class ProxyResponseHandler { private static async handleStream(session: ProxySession, response: Response): Promise { const messageContext = session.messageContext; const provider = session.provider; + const discoveryLeaseLifecycle = startDiscoveryLeaseLifecycle(session); if (!messageContext || !provider || !response.body) { discardBeforeResponseBodySnapshot(session); releaseSessionAgent(session); + const deferredMeta = peekDeferredStreamingFinalization(session); + void (async () => { + await releaseOwnedProviderSessionRef(session, deferredMeta, false); + await discoveryLeaseLifecycle.release(); + })(); return response; } @@ -2695,16 +3067,26 @@ export class ProxyResponseHandler { let transportReleased = false; let commitSideEffectsScheduled = false; let latestCommitSideEffects: (() => Promise) | undefined; + let latestFinalizeAttemptResources: (() => Promise) | undefined; const scheduleCommitSideEffects = (effect: (() => Promise) | undefined) => { - if (!effect || commitSideEffectsScheduled) return; + if ( + (!effect && !latestFinalizeAttemptResources && !discoveryLeaseLifecycle.active) || + commitSideEffectsScheduled + ) + return; commitSideEffectsScheduled = true; + const finalizeAttemptResources = latestFinalizeAttemptResources; return schedulePostTerminalSideEffects({ taskId, providerId: provider.id, sessionId: session.sessionId, commit: async (signal) => { - if (signal.aborted) return; - await effect(); + try { + if (!signal.aborted) await effect?.(); + } finally { + await finalizeAttemptResources?.(); + await discoveryLeaseLifecycle.release(); + } }, }); }; @@ -2900,9 +3282,11 @@ export class ProxyResponseHandler { statusCode, streamEndedNormally, clientAborted, + discoveryLeaseLifecycle, abortReason ); latestCommitSideEffects = finalized.commitSideEffects; + latestFinalizeAttemptResources = finalized.finalizeAttemptResources; const finalizedUsage = await finalizeRequestStats( session, allContent, @@ -2970,9 +3354,11 @@ export class ProxyResponseHandler { statusCode, false, clientAborted, + discoveryLeaseLifecycle, abortReason ); latestCommitSideEffects = finalized.commitSideEffects; + latestFinalizeAttemptResources = finalized.finalizeAttemptResources; await finalizeRequestStats( session, @@ -3013,6 +3399,12 @@ export class ProxyResponseHandler { } } finally { releaseTransportResources(); + if (!commitSideEffectsScheduled) { + void (async () => { + await latestFinalizeAttemptResources?.(); + await discoveryLeaseLifecycle.release(); + })(); + } } }; @@ -3297,15 +3689,30 @@ export class ProxyResponseHandler { let terminalDetailsPersisted = false; let streamCommitSideEffectsScheduled = false; let latestStreamCommitSideEffects: Array<() => Promise> = []; + let latestStreamFinalizeAttemptResources: (() => Promise) | undefined; const scheduleStreamCommitSideEffects = () => { - if (latestStreamCommitSideEffects.length === 0 || streamCommitSideEffectsScheduled) return; + if ( + (latestStreamCommitSideEffects.length === 0 && + !latestStreamFinalizeAttemptResources && + !discoveryLeaseLifecycle.active) || + streamCommitSideEffectsScheduled + ) + return; streamCommitSideEffectsScheduled = true; const committedEffects = [...latestStreamCommitSideEffects]; + const finalizeAttemptResources = latestStreamFinalizeAttemptResources; return schedulePostTerminalSideEffects({ taskId, providerId: provider.id, sessionId: session.sessionId, - commit: (signal) => runPostTerminalSideEffects(committedEffects, signal), + commit: async (signal) => { + try { + await runPostTerminalSideEffects(committedEffects, signal); + } finally { + await finalizeAttemptResources?.(); + await discoveryLeaseLifecycle.release(); + } + }, }); }; let streamFailurePersistencePromise: Promise | null = null; @@ -3356,11 +3763,13 @@ export class ProxyResponseHandler { statusCode, streamEndedNormally, clientAborted, + discoveryLeaseLifecycle, abortReason ); latestStreamCommitSideEffects = finalized.commitSideEffects ? [finalized.commitSideEffects] : []; + latestStreamFinalizeAttemptResources = finalized.finalizeAttemptResources; const effectiveStatusCode = finalized.effectiveStatusCode; const streamErrorMessage = finalized.errorMessage; const providerIdForPersistence = finalized.providerIdForPersistence; @@ -3457,7 +3866,8 @@ export class ProxyResponseHandler { effectiveStatusCode >= 200 && effectiveStatusCode < 300 && session.sessionId && - provider.id + provider.id && + finalized.allowAuxiliarySessionBinding ) { try { const sseEvents = parseSSEData(allContent); @@ -3645,6 +4055,7 @@ export class ProxyResponseHandler { if (codexCacheBinding) { const { sessionId, promptCacheKey, providerId, keyId } = codexCacheBinding; postTerminalSideEffects.push(async () => { + if (!(await finalized.confirmAuxiliarySessionBinding())) return; try { await SessionManager.updateSessionWithCodexCacheKey( sessionId, @@ -4067,6 +4478,12 @@ export class ProxyResponseHandler { clearClientAbortDrainTimer(); clearIdleTimer(); // 清除静默期计时器（防止泄漏） releaseSessionAgent(session); + if (!streamCommitSideEffectsScheduled) { + void (async () => { + await latestStreamFinalizeAttemptResources?.(); + await discoveryLeaseLifecycle.release(); + })(); + } } }; diff --git a/src/app/v1/_lib/proxy/session.ts b/src/app/v1/_lib/proxy/session.ts index 18ff8bcae..0ef7901c1 100644 --- a/src/app/v1/_lib/proxy/session.ts +++ b/src/app/v1/_lib/proxy/session.ts @@ -121,6 +121,11 @@ export class ProxySession { // Session ID（用于会话粘性和并发限流） sessionId: string | null; + // Discovery lease conflicts must stay on a single upstream and must not + // mutate a binding owned by the in-flight discovery request. + private streamingHedgeDisabled = false; + private sessionBindingAllowed = true; + // 客户端 IP（由 ProxyAuthenticator 按系统设置的 ip_extraction_config 解析后写入） clientIp: string | null = null; @@ -203,9 +208,10 @@ export class ProxySession { */ private providersSnapshot: Provider[] | null = null; - // 本请求已通过 Provider 并发检查获得的引用。 - // 失败切换 provider 时只能释放这里记录过的引用，避免 hedge/fallback 释放未 acquire 的 Redis 计数。 - private providerSessionRefs = new Set(); + // 本请求已通过 Provider 并发检查获得的引用。tracked=true 表示这次 + // acquire 同时创建了 Provider Session 基线；Sticky CAS 成功时只有 + // 该引用可以保留，已有基线上的普通 attempt 引用必须在终态释放。 + private providerSessionRefs = new Map>(); // Snapshot captured during provider selection. Discovery reuses this exact // generation for timeout cleanup/finalization instead of performing a @@ -363,25 +369,50 @@ export class ProxySession { return this.sessionBindingSnapshot; } - recordProviderSessionRef(providerId: number): void { + recordProviderSessionRef(providerId: number, options: { retainOnSuccess?: boolean } = {}): void { if (!this.providerSessionRefs) { - this.providerSessionRefs = new Set(); + this.providerSessionRefs = new Map>(); } if (Number.isInteger(providerId) && providerId > 0) { - this.providerSessionRefs.add(providerId); + const refs = this.providerSessionRefs.get(providerId) ?? []; + refs.push({ retainOnSuccess: options.retainOnSuccess === true }); + this.providerSessionRefs.set(providerId, refs); } } consumeProviderSessionRef(providerId: number): boolean { - if (!this.providerSessionRefs?.has(providerId)) { - return false; - } - - this.providerSessionRefs.delete(providerId); + const refs = this.providerSessionRefs?.get(providerId); + if (!refs || refs.length === 0) return false; + refs.shift(); + if (refs.length === 0) this.providerSessionRefs.delete(providerId); return true; } + hasProviderSessionRef(providerId: number): boolean { + return (this.providerSessionRefs?.get(providerId)?.length ?? 0) > 0; + } + + shouldRetainProviderSessionRefOnSuccess(providerId: number): boolean { + return this.providerSessionRefs?.get(providerId)?.[0]?.retainOnSuccess === true; + } + + disableStreamingHedge(): void { + this.streamingHedgeDisabled = true; + } + + isStreamingHedgeDisabled(): boolean { + return this.streamingHedgeDisabled === true; + } + + setSessionBindingAllowed(allowed: boolean): void { + this.sessionBindingAllowed = allowed; + } + + isSessionBindingAllowed(): boolean { + return this.sessionBindingAllowed !== false; + } + setCacheTtlResolved(ttl: CacheTtlResolved | null): void { this.cacheTtlResolved = ttl; } diff --git a/src/app/v1/_lib/proxy/stream-finalization.ts b/src/app/v1/_lib/proxy/stream-finalization.ts index 497207f1f..1ba8f7ada 100644 --- a/src/app/v1/_lib/proxy/stream-finalization.ts +++ b/src/app/v1/_lib/proxy/stream-finalization.ts @@ -1,6 +1,13 @@ import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import type { ProxySession } from "./session"; +export type DeferredStreamingDiscoveryLease = { + sessionId: string; + keyId: number; + ownerToken: string; + ttlSeconds: number; +}; + /** * 流式响应（SSE）在“收到响应头”时无法确定成功与否： * - 上游可能返回 HTTP 200，但 body 是错误 JSON（假 200） @@ -41,6 +48,12 @@ export type DeferredStreamingFinalization = { bindingSnapshot?: SessionBindingSnapshot | null; /** Discovery winners must satisfy the protocol completion marker before binding. */ requiresCompletionMarker?: boolean; + /** Lease already acquired by Forwarder and owned until terminal side effects finish. */ + discoveryLease?: DeferredStreamingDiscoveryLease; + /** Whether this attempt owns a Provider concurrent-session reference. */ + providerSessionRefOwned?: boolean; + /** CAS success converts this attempt ref into the binding baseline when true. */ + providerSessionRefRetainOnSuccess?: boolean; }; const deferredMeta = new WeakMap(); diff --git a/src/lib/config/env.schema.ts b/src/lib/config/env.schema.ts index 80d56a986..29c66acca 100644 --- a/src/lib/config/env.schema.ts +++ b/src/lib/config/env.schema.ts @@ -183,6 +183,10 @@ export const EnvSchema = z.object({ // 超时后主动断开该输家连接，仅用已收到的内容尝试计费（通常计不出 -> 跳过）。 HEDGE_LOSER_DRAIN_TIMEOUT_MS: z.coerce.number().int().min(1000).default(120_000), + // Operational canary for the Discovery scheduler. The database feature + // switch remains authoritative; this percentage only narrows eligibility. + DISCOVERY_ROLLOUT_PERCENT: z.coerce.number().int().min(0).max(100).default(100), + DASHBOARD_LOGS_POLL_INTERVAL_MS: z.coerce.number().int().min(250).max(60000).default(5000), // Langfuse Observability (optional, auto-enabled when keys are set) diff --git a/src/lib/observability/discovery-metrics.ts b/src/lib/observability/discovery-metrics.ts new file mode 100644 index 000000000..16cd93135 --- /dev/null +++ b/src/lib/observability/discovery-metrics.ts @@ -0,0 +1,126 @@ +import { logger } from "@/lib/logger"; + +export type DiscoveryLifecycleEvent = + | "request_started" + | "attempt_started" + | "attempt_finished" + | "fallback_promoted" + | "parser_limit" + | "cancel_failed" + | "lease_conflict" + | "binding_cas_conflict" + | "request_finished"; + +export type DiscoveryWinnerOrigin = "normal" | "fallback" | "none"; + +type DiscoveryMetricIdentity = { + requestId: number | string | null; + sessionId: string; + keyId: number; +}; + +export function recordDiscoveryControlEvent( + event: "lease_conflict" | "binding_cas_conflict", + context: DiscoveryMetricIdentity & Record +): void { + logger.info("[DiscoveryMetric] Control event", { event, ...context }); +} + +export class DiscoveryRequestMetrics { + private readonly attemptStartedAt = new Map(); + private readonly fallbackAttempts = new Set(); + private attempts = 0; + private active = 0; + private maxActive = 0; + private maxRound = 0; + private providerMs = 0; + private fallbackPromotions = 0; + private cancelFailures = 0; + private finished = false; + + constructor( + private readonly identity: DiscoveryMetricIdentity, + private readonly startedAt: number + ) { + this.event("request_started"); + } + + event(event: DiscoveryLifecycleEvent, context: Record = {}): void { + logger.debug("[DiscoveryMetric] Lifecycle event", { + event, + ...this.identity, + elapsedMs: Math.max(0, Date.now() - this.startedAt), + ...context, + }); + } + + attemptStarted(options: { + attemptId: string; + providerId: number; + round: number; + kind: DiscoveryWinnerOrigin; + }): void { + if (this.attemptStartedAt.has(options.attemptId)) return; + this.attemptStartedAt.set(options.attemptId, Date.now()); + this.attempts += 1; + this.active += 1; + this.maxActive = Math.max(this.maxActive, this.active); + this.maxRound = Math.max(this.maxRound, options.round); + this.event("attempt_started", options); + } + + attemptFinished( + attemptId: string, + context: { providerId: number; outcome: string; cancellationKind?: string | null } + ): void { + const startedAt = this.attemptStartedAt.get(attemptId); + if (startedAt == null) return; + this.attemptStartedAt.delete(attemptId); + this.active = Math.max(0, this.active - 1); + const durationMs = Math.max(0, Date.now() - startedAt); + this.providerMs += durationMs; + this.event("attempt_finished", { attemptId, durationMs, ...context }); + } + + fallbackPromoted(attemptId: string, providerId: number, round: number): void { + if (this.fallbackAttempts.has(attemptId)) return; + this.fallbackAttempts.add(attemptId); + this.fallbackPromotions += 1; + this.maxRound = Math.max(this.maxRound, round); + this.event("fallback_promoted", { attemptId, providerId, round }); + } + + cancelFailed(attemptId: string, providerId: number, error: unknown): void { + this.cancelFailures += 1; + this.event("cancel_failed", { + attemptId, + providerId, + error: error instanceof Error ? error.message : String(error), + }); + } + + finish(context: { + outcome: "success" | "failed" | "client_abort" | "deadline"; + statusCode: number; + winnerOrigin?: DiscoveryWinnerOrigin; + winnerProviderId?: number | null; + winnerRound?: number | null; + }): void { + if (this.finished) return; + this.finished = true; + const elapsedMs = Math.max(0, Date.now() - this.startedAt); + logger.info("[DiscoveryMetric] Request aggregate", { + event: "request_finished", + ...this.identity, + ...context, + elapsedMs, + ttfbMs: context.outcome === "success" ? elapsedMs : null, + attemptsPerRequest: this.attempts, + maxActiveAttempts: this.maxActive, + rounds: this.maxRound, + providerMs: this.providerMs, + fallbackPromotions: this.fallbackPromotions, + cancelFailures: this.cancelFailures, + }); + } +} diff --git a/tests/integration/proxy-hedge-lifecycle.test.ts b/tests/integration/proxy-hedge-lifecycle.test.ts index 390727090..0f2b6b3b0 100644 --- a/tests/integration/proxy-hedge-lifecycle.test.ts +++ b/tests/integration/proxy-hedge-lifecycle.test.ts @@ -16,10 +16,45 @@ const state = vi.hoisted(() => { return { addLoserCost: vi.fn(), billHedgeLosers: false, - durableTerminal: vi.fn(async () => {}), + discoveryEnabled: false, + acquireDiscoveryLease: vi.fn(async () => ({ + status: "acquired", + ownerToken: "integration-lease", + legacyFallbackAllowed: false, + })), + releaseDiscoveryLease: vi.fn(async () => ({ + status: "released", + legacyFallbackAllowed: false, + })), + renewDiscoveryLease: vi.fn(async () => ({ + status: "renewed", + legacyFallbackAllowed: false, + })), + compareAndSetBinding: vi.fn(async () => ({ + status: "ok", + source: "updated", + legacyFallbackAllowed: false, + snapshot: { + sessionId: "integration-discovery", + keyId: 22, + providerId: 2, + generation: "g2", + }, + })), + durableTerminal: vi.fn( + async ( + _id: number, + details: unknown, + options?: { onCommitted?: (details: unknown) => void | Promise } + ) => { + await options?.onCommitted?.(details); + return true; + } + ), http2Error: ((): Error | null => null)(), loserBilled: Promise.withResolvers(), pickAlternative: vi.fn(), + pickDiscovery: vi.fn(), providers: Array.from([]), recordFailure: vi.fn(async () => {}), settleLeaseBudgets: vi.fn(async () => {}), @@ -42,10 +77,17 @@ vi.mock("@/lib/config", async (importOriginal) => { ...actual, getCachedSystemSettings: async () => ({ billHedgeLosers: state.billHedgeLosers, + discoveryConcurrency: 2, + discoveryEnabled: state.discoveryEnabled, + discoverySlaMs: 100, enableBillingHeaderRectifier: false, enableClaudeMetadataUserIdInjection: false, enableThinkingBudgetRectifier: false, enableThinkingSignatureRectifier: false, + maxDiscoveryRounds: 1, + racingTotalTimeoutMs: 500, + stickySlaMs: 100, + stickyTimeoutCooldownMs: 300_000, }), isHttp2Enabled: async () => { if (state.http2Error) throw state.http2Error; @@ -57,8 +99,41 @@ vi.mock("@/lib/config/system-settings-cache", () => ({ getCachedSystemSettings: async () => ({ billNonSuccessfulRequests: false }), })); vi.mock("@/app/v1/_lib/proxy/provider-selector", () => ({ - ProxyProviderResolver: { pickRandomProviderWithExclusion: state.pickAlternative }, + ProxyProviderResolver: { + pickDiscoveryProviders: state.pickDiscovery, + pickRandomProviderWithExclusion: state.pickAlternative, + resolveEffectivePriorityForSession: (provider: Provider) => provider.priority ?? 0, + }, })); +vi.mock("@/lib/session-manager", async (importOriginal) => { + const actual = await importOriginal(); + class TestSessionManager extends actual.SessionManager { + static override async ensureVersionedBindingCapability() { + return "available" as const; + } + static override async getSessionBindingSnapshot(sessionId: string, keyId: number) { + return { + status: "ok" as const, + source: "existing" as const, + legacyFallbackAllowed: false as const, + snapshot: { sessionId, keyId, providerId: null, generation: "g1" }, + }; + } + static override async acquireSessionDiscoveryLease() { + return state.acquireDiscoveryLease(); + } + static override async renewSessionDiscoveryLease() { + return state.renewDiscoveryLease(); + } + static override async releaseSessionDiscoveryLease() { + return state.releaseDiscoveryLease(); + } + static override async compareAndSetSessionProvider() { + return state.compareAndSetBinding(); + } + } + return { ...actual, SessionManager: TestSessionManager }; +}); vi.mock("@/lib/provider-endpoints/endpoint-selector", () => ({ getEndpointFilterStats: vi.fn(async () => null), getPreferredProviderEndpoints: vi.fn(async () => []), @@ -154,7 +229,10 @@ vi.mock("@/lib/session-tracker", () => ({ vi.mock("@/lib/proxy-status-tracker", () => ({ ProxyStatusTracker: { getInstance: () => ({ endRequest: vi.fn() }) }, })); -vi.mock("@/lib/redis/live-chain-store", () => ({ deleteLiveChain: vi.fn(async () => {}) })); +vi.mock("@/lib/redis/live-chain-store", () => ({ + deleteLiveChain: vi.fn(async () => {}), + writeLiveChain: vi.fn(async () => {}), +})); const CREATED_AT = new Date(0); const USER = { @@ -386,6 +464,7 @@ beforeEach(async () => { await resetGlobalAgentPool(); vi.clearAllMocks(); state.billHedgeLosers = false; + state.discoveryEnabled = false; state.http2Error = null; state.loserBilled = Promise.withResolvers(); state.providers.length = 0; @@ -394,6 +473,10 @@ beforeEach(async () => { state.pickAlternative.mockImplementation(async (_session: unknown, excludedIds: number[]) => { return state.providers.find((provider) => !excludedIds.includes(provider.id)) ?? null; }); + state.pickDiscovery.mockImplementation( + async (_session: unknown, count: number, excludedIds: number[]) => + state.providers.filter((provider) => !excludedIds.includes(provider.id)).slice(0, count) + ); }); afterEach(async () => { @@ -403,6 +486,44 @@ afterEach(async () => { }); describe("proxy hedge transport/lifecycle integration (persistence and control-plane seams mocked)", () => { + it("runs a leased Discovery race over real loopback transports and cancels the loser", async () => { + const [loser, winner] = await Promise.all([startUpstream(), startUpstream()]); + const client = new AbortController(); + try { + state.discoveryEnabled = true; + const initialProvider = createProvider(1, loser.baseUrl, 0); + const winningProvider = createProvider(2, winner.baseUrl, 0); + winningProvider.priority = initialProvider.priority; + state.providers.push(winningProvider); + const session = await createSession(initialProvider, "/v1/messages", client.signal); + session.sessionId = "integration-discovery"; + const agents = watchAgentReleases(2); + + const forwarded = ProxyForwarder.send(session); + await Promise.all([loser.response, winner.response]); + await winner.send( + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"type":"text_delta","text":"winner"}}\n\n' + + 'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ); + + const downstream = await ProxyResponseHandler.dispatch(session, await forwarded); + await expect(downstream.text()).resolves.toContain("winner"); + await settleTasks(); + await loser.terminated; + await agents.released; + + expect(loser.abortCount()).toBe(1); + expect(winner.abortCount()).toBe(0); + expect(state.acquireDiscoveryLease).toHaveBeenCalledTimes(1); + expect(state.compareAndSetBinding).toHaveBeenCalledTimes(1); + expect(state.releaseDiscoveryLease).toHaveBeenCalledTimes(1); + expect(agents.pool.getPoolStats().activeRequests).toBe(0); + } finally { + client.abort(new Error("fixture cleanup")); + await Promise.all([loser.close(), winner.close()]); + } + }); + it("fences loser timers after winner settlement and releases each launched transport once", async () => { const [slow, winner, fenced] = await Promise.all([ startUpstream(), diff --git a/tests/unit/proxy/discovery-coordinator.test.ts b/tests/unit/proxy/discovery-coordinator.test.ts index 3a41a8d5e..bdbb7ff20 100644 --- a/tests/unit/proxy/discovery-coordinator.test.ts +++ b/tests/unit/proxy/discovery-coordinator.test.ts @@ -13,6 +13,20 @@ const attempt = (id: string, priority: number, kind: "normal" | "fallback" = "no }); describe("DiscoveryCoordinator", () => { + it("keeps Sticky probing outside the Discovery round counter", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 1 }); + coordinator.startStickyProbe(); + expect(coordinator.state).toBe("STICKY_PROBING"); + expect(coordinator.round).toBe(1); + + coordinator.addAttempt(attempt("sticky", 1)); + expect(coordinator.demoteToFallback("sticky")).toBe(true); + coordinator.startDiscoveryAfterSticky(); + + expect(coordinator.state).toBe("DISCOVERY_RACING"); + expect(coordinator.round).toBe(1); + }); + it("commits the highest priority ready normal attempt", () => { const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); coordinator.addAttempt(attempt("a", 10)); @@ -44,6 +58,7 @@ describe("DiscoveryCoordinator", () => { expect(coordinator.markReady("a", epoch.requestEpoch, epoch.roundEpoch)).toEqual({ type: "none", }); + expect(coordinator.markFailed("a")).toEqual({ type: "none" }); }); it("promotes a ready fallback at the round boundary when no normal is ready", () => { @@ -100,6 +115,49 @@ describe("DiscoveryCoordinator", () => { ready: true, pending: true, }); - expect(coordinator.markFailed("high")).toEqual({ type: "commit_normal", attemptId: "low" }); + expect(coordinator.markFailed("high")).toEqual({ + type: "commit_normal", + attemptId: "low", + }); + }); + + it("treats fallback promotion as terminal", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("fallback", 1, "fallback")); + + expect(coordinator.markReady("fallback")).toEqual({ + type: "promote_fallback", + attemptId: "fallback", + }); + expect(coordinator.state).toBe("FALLBACK_ACTIVE"); + expect(coordinator.isTerminal).toBe(true); + expect(coordinator.markFailed("fallback")).toEqual({ type: "none" }); + }); + + it("keeps coordinator kind in sync when a running Sticky becomes fallback", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("sticky", 1)); + + expect(coordinator.promoteToFallback("sticky")).toBe(true); + expect(coordinator.snapshot.find((item) => item.id === "sticky")?.kind).toBe("fallback"); + }); + + it("opens a full new round when all normal attempts fail", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 3, maxRounds: 2 }); + coordinator.addAttempt(attempt("a", 1)); + coordinator.addAttempt(attempt("b", 2)); + + expect(coordinator.markFailed("a")).toEqual({ type: "none" }); + expect(coordinator.markFailed("b")).toEqual({ type: "launch", slots: 3 }); + expect(coordinator.round).toBe(2); + }); + + it("commits a ready normal candidate at the total deadline", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("high", 1)); + coordinator.addAttempt(attempt("normal", 2)); + coordinator.markReady("normal"); + + expect(coordinator.onDeadline()).toEqual({ type: "commit_normal", attemptId: "normal" }); }); }); diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index 7a6f95fc4..9c5ca8b39 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it } from "vitest"; import { + DISCOVERY_EVENT_MAX_COUNT, + DISCOVERY_PREFIX_MAX_BYTES, DiscoveryValidityParser, classifyDiscoveryChunk, } from "@/app/v1/_lib/proxy/discovery-validity"; @@ -87,6 +89,27 @@ describe("discovery validity", () => { }); }); + it("accepts Anthropic tool-use partial JSON as deliverable content", () => { + const parser = new DiscoveryValidityParser("anthropic"); + + expect( + parser.push('data: {"type":"content_block_delta","delta":{"partial_json":"{\\"x\\":1}"}}\n\n') + ).toMatchObject({ ready: true, error: false }); + expect(parser.push('data: {"type":"message_stop"}\n\n')).toMatchObject({ + ready: true, + terminal: true, + error: false, + }); + }); + + it("accepts nested OpenAI Chat tool-call arguments", () => { + expect( + parserForOpenAIChatToolCall().push( + 'data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"{\\"x\\":1}"}}]}}]}\n\n' + ) + ).toMatchObject({ ready: true, error: false }); + }); + it("accepts Anthropic tool-use starts and partial JSON deltas", () => { expect( classifyDiscoveryChunk( @@ -101,4 +124,23 @@ describe("discovery validity", () => { ).ready ).toBe(true); }); + + it("fails a metadata-only prefix after the byte limit", () => { + const parser = new DiscoveryValidityParser("openai-chat"); + const result = parser.push(`:${"x".repeat(DISCOVERY_PREFIX_MAX_BYTES + 1)}`); + expect(result).toMatchObject({ ready: false, error: true, limitExceeded: true }); + }); + + it("fails metadata-only protocol events after the event limit", () => { + const parser = new DiscoveryValidityParser("anthropic"); + let result = parser.push(""); + for (let index = 0; index <= DISCOVERY_EVENT_MAX_COUNT; index += 1) { + result = parser.push('data: {"type":"ping"}\n'); + } + expect(result).toMatchObject({ ready: false, error: true, limitExceeded: true }); + }); }); + +function parserForOpenAIChatToolCall(): DiscoveryValidityParser { + return new DiscoveryValidityParser("openai-chat"); +} diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index 5654e4745..bd2cafab4 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -42,6 +42,33 @@ const mocks = vi.hoisted(() => ({ storeSessionRequestPhaseSnapshot: vi.fn(async () => {}), storeSessionResponsePhaseSnapshot: vi.fn(async () => {}), getVersionedBindingCapabilityState: vi.fn(() => "available"), + ensureVersionedBindingCapability: vi.fn(async () => "available"), + getSessionBindingSnapshot: vi.fn(async (sessionId: string, keyId: number) => ({ + status: "ok", + legacyFallbackAllowed: false, + source: "existing", + snapshot: { sessionId, keyId, providerId: null, generation: "g-test" }, + })), + acquireSessionDiscoveryLease: vi.fn(async () => ({ + status: "acquired", + ownerToken: "lease-test", + legacyFallbackAllowed: false, + })), + releaseSessionDiscoveryLease: vi.fn(async () => ({ + status: "released", + legacyFallbackAllowed: false, + })), + clearVersionedSessionProvider: vi.fn(async (snapshot: unknown) => ({ + status: "ok", + legacyFallbackAllowed: false, + source: "cleared", + snapshot: { + ...(snapshot as Record), + providerId: null, + generation: "g-cleared", + }, + })), + isWebsocketClientRequest: vi.fn(() => false), })); vi.mock("@/lib/logger", () => ({ @@ -69,6 +96,11 @@ vi.mock("@/lib/provider-endpoints/endpoint-selector", () => ({ getEndpointFilterStats: mocks.getEndpointFilterStats, })); +vi.mock("@/app/v1/_lib/responses-ws/eligibility", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, isWebsocketClientRequest: mocks.isWebsocketClientRequest }; +}); + vi.mock("@/lib/endpoint-circuit-breaker", () => ({ recordEndpointSuccess: mocks.recordEndpointSuccess, recordEndpointFailure: mocks.recordEndpointFailure, @@ -96,6 +128,11 @@ vi.mock("@/lib/rate-limit/service", () => ({ vi.mock("@/lib/session-manager", () => ({ SessionManager: { getVersionedBindingCapabilityState: mocks.getVersionedBindingCapabilityState, + ensureVersionedBindingCapability: mocks.ensureVersionedBindingCapability, + getSessionBindingSnapshot: mocks.getSessionBindingSnapshot, + acquireSessionDiscoveryLease: mocks.acquireSessionDiscoveryLease, + releaseSessionDiscoveryLease: mocks.releaseSessionDiscoveryLease, + clearVersionedSessionProvider: mocks.clearVersionedSessionProvider, updateSessionBindingSmart: mocks.updateSessionBindingSmart, updateSessionProvider: mocks.updateSessionProvider, clearSessionProvider: mocks.clearSessionProvider, @@ -131,8 +168,10 @@ import { import { ProxyForwarder } from "@/app/v1/_lib/proxy/forwarder"; import { ModelRedirector } from "@/app/v1/_lib/proxy/model-redirector"; import { ProxySession } from "@/app/v1/_lib/proxy/session"; +import { DbPoolAdmissionError } from "@/drizzle/admitted-client"; import { logger } from "@/lib/logger"; import type { Provider } from "@/types/provider"; +import type { SystemSettings } from "@/types/system-config"; type AttemptRuntime = { clearResponseTimeout?: () => void; @@ -230,6 +269,8 @@ function createSession(clientAbortSignal: AbortSignal | null = null): ProxySessi provider: null, messageContext: null, sessionId: "sess-hedge", + streamingHedgeDisabled: false, + sessionBindingAllowed: true, requestSequence: 1, originalFormat: "claude", providerType: null, @@ -351,8 +392,124 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { tracked: true, referenced: true, }); + mocks.ensureVersionedBindingCapability.mockResolvedValue("available"); + mocks.getSessionBindingSnapshot.mockImplementation( + async (sessionId: string, keyId: number) => ({ + status: "ok", + legacyFallbackAllowed: false, + source: "existing", + snapshot: { sessionId, keyId, providerId: null, generation: "g-test" }, + }) + ); + mocks.acquireSessionDiscoveryLease.mockResolvedValue({ + status: "acquired", + ownerToken: "lease-test", + legacyFallbackAllowed: false, + }); + mocks.releaseSessionDiscoveryLease.mockResolvedValue({ + status: "released", + legacyFallbackAllowed: false, + }); + mocks.clearVersionedSessionProvider.mockImplementation(async (snapshot: unknown) => ({ + status: "ok", + legacyFallbackAllowed: false, + source: "cleared", + snapshot: { + ...(snapshot as Record), + providerId: null, + generation: "g-cleared", + }, + })); + mocks.categorizeErrorAsync.mockResolvedValue(ProxyErrorCategory.PROVIDER_ERROR); + mocks.isWebsocketClientRequest.mockReturnValue(false); + }); + + test("Discovery actively probes an unknown binding capability before acquiring its lease", async () => { + const provider = createProvider({ id: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 20 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + mocks.getVersionedBindingCapabilityState.mockReturnValueOnce("unknown"); + mocks.ensureVersionedBindingCapability.mockResolvedValueOnce("available"); + + const prepareStreamingDiscovery = ( + ProxyForwarder as unknown as { + prepareStreamingDiscovery: ( + session: ProxySession, + settings: SystemSettings, + requestStartedAt: number + ) => Promise; + } + ).prepareStreamingDiscovery; + const prepared = await prepareStreamingDiscovery( + session, + { + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 50, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + } as SystemSettings, + Date.now() + ); + + expect(prepared).not.toBeNull(); + expect(mocks.ensureVersionedBindingCapability).toHaveBeenCalledTimes(1); + expect(mocks.getSessionBindingSnapshot).toHaveBeenCalledWith("sess-hedge", 20); + expect(mocks.acquireSessionDiscoveryLease).toHaveBeenCalledTimes(1); }); + test.each(["unknown", "unavailable"] as const)( + "Discovery fails closed when the binding capability probe returns %s", + async (capabilityState) => { + const provider = createProvider({ id: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 21 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + mocks.ensureVersionedBindingCapability.mockResolvedValueOnce(capabilityState); + + const prepareStreamingDiscovery = ( + ProxyForwarder as unknown as { + prepareStreamingDiscovery: ( + session: ProxySession, + settings: SystemSettings, + requestStartedAt: number + ) => Promise; + } + ).prepareStreamingDiscovery; + const prepared = await prepareStreamingDiscovery( + session, + { + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 50, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + } as SystemSettings, + Date.now() + ); + + expect(prepared).toBeNull(); + expect(mocks.ensureVersionedBindingCapability).toHaveBeenCalledTimes(1); + expect(mocks.getSessionBindingSnapshot).not.toHaveBeenCalled(); + expect(mocks.acquireSessionDiscoveryLease).not.toHaveBeenCalled(); + } + ); + test("shadow session redirect should not overwrite initial provider redirect and winner should keep its own redirect", () => { const requestedModel = "claude-haiku-4-5-20251001"; const fireworksRedirect = "accounts/fireworks/routers/kimi-k2p5-turbo"; @@ -2080,6 +2237,12 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { const high = createProvider({ id: 1, name: "high", priority: 1 }); const low = createProvider({ id: 2, name: "low", priority: 10 }); const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 1 }, + apiKey: null, + } as typeof session.authState; session.setProvider(high); mocks.getCachedSystemSettings.mockResolvedValue({ discoveryEnabled: true, @@ -2156,6 +2319,941 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { } }); + test("Discovery lease conflict forces a single upstream and forbids binding writes", async () => { + const provider = createProvider({ id: 1, firstByteTimeoutStreamingMs: 100 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 7 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + racingTotalTimeoutMs: 500, + }); + mocks.acquireSessionDiscoveryLease.mockResolvedValueOnce({ + status: "conflict", + reason: "lease_held", + legacyFallbackAllowed: false, + }); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response('data: {"type":"message_stop"}\n\n', { + status: 200, + headers: { "content-type": "text/event-stream" }, + }) + ); + + const response = await ProxyForwarder.send(session); + expect(response.status).toBe(200); + expect(doForward).toHaveBeenCalledTimes(1); + expect(session.isStreamingHedgeDisabled()).toBe(true); + expect(session.isSessionBindingAllowed()).toBe(false); + expect(mocks.pickDiscoveryProviders).not.toHaveBeenCalled(); + expect(mocks.releaseSessionDiscoveryLease).not.toHaveBeenCalled(); + expect(mocks.getCachedSystemSettings).toHaveBeenCalledTimes(1); + }); + + test("foreign binding state fails closed before Discovery acquires a lease", async () => { + const provider = createProvider({ id: 1, firstByteTimeoutStreamingMs: 0 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 8 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + mocks.getCachedSystemSettings.mockResolvedValue({ discoveryEnabled: true }); + mocks.getSessionBindingSnapshot.mockResolvedValueOnce({ + status: "conflict", + reason: "legacy_owner_mismatch", + legacyFallbackAllowed: false, + }); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response('data: {"type":"message_stop"}\n\n', { + status: 200, + headers: { "content-type": "text/event-stream" }, + }) + ); + + await ProxyForwarder.send(session); + expect(doForward).toHaveBeenCalledTimes(1); + expect(session.isSessionBindingAllowed()).toBe(false); + expect(mocks.acquireSessionDiscoveryLease).not.toHaveBeenCalled(); + expect(mocks.pickDiscoveryProviders).not.toHaveBeenCalled(); + }); + + test("Sticky fallback stays held while timeout CAS and the next wave are being prepared", async () => { + vi.useFakeTimers(); + try { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const normal = createProvider({ id: 2, name: "normal", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 9 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 9, + providerId: sticky.id, + generation: "g-sticky", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 2, + discoverySlaMs: 50, + stickySlaMs: 10, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([normal]); + + let resolveClear!: (value: unknown) => void; + mocks.clearVersionedSessionProvider.mockReturnValueOnce( + new Promise((resolve) => { + resolveClear = resolve; + }) + ); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response( + new ReadableStream({ + start(controller) { + setTimeout(() => { + controller.enqueue( + new TextEncoder().encode( + 'data: {"type":"content_block_delta","delta":{"text":"sticky"}}\n\n' + ) + ); + controller.close(); + }, 15); + }, + }), + { headers: { "content-type": "text/event-stream" } } + ) + ); + doForward.mockResolvedValueOnce( + new Response( + new ReadableStream({ + start(controller) { + setTimeout(() => { + controller.enqueue( + new TextEncoder().encode( + 'data: {"type":"content_block_delta","delta":{"text":"normal"}}\n\n' + ) + ); + controller.close(); + }, 5); + }, + }), + { headers: { "content-type": "text/event-stream" } } + ) + ); + + let settledEarly = false; + const responsePromise = ProxyForwarder.send(session).then((response) => { + settledEarly = true; + return response; + }); + await vi.advanceTimersByTimeAsync(20); + expect(settledEarly).toBe(false); + + resolveClear({ + status: "ok", + legacyFallbackAllowed: false, + source: "cleared", + snapshot: { + sessionId: session.sessionId!, + keyId: 9, + providerId: null, + generation: "g-cleared", + }, + }); + await vi.advanceTimersByTimeAsync(10); + const response = await responsePromise; + expect(await response.text()).toContain('"normal"'); + expect(session.provider?.id).toBe(normal.id); + } finally { + vi.useRealTimers(); + } + }); + + test("Sticky probing does not consume a configured Discovery round", async () => { + vi.useFakeTimers(); + try { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const roundOne = createProvider({ id: 2, name: "round-one", priority: 1 }); + const roundTwo = createProvider({ id: 3, name: "round-two", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 19 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 19, + providerId: sticky.id, + generation: "g-sticky-rounds", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 2, + discoverySlaMs: 20, + stickySlaMs: 10, + racingTotalTimeoutMs: 100, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders + .mockResolvedValueOnce([roundOne]) + .mockResolvedValueOnce([roundTwo]); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementation(async (attemptSession) => { + const providerId = (attemptSession as ProxySession).provider?.id; + if (providerId !== roundTwo.id) { + return new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }); + } + return new Response( + new ReadableStream({ + start(controller) { + setTimeout(() => { + controller.enqueue( + new TextEncoder().encode( + 'data: {"type":"content_block_delta","delta":{"text":"round-two"}}\n\n' + ) + ); + controller.close(); + }, 5); + }, + }), + { headers: { "content-type": "text/event-stream" } } + ); + }); + + const responsePromise = ProxyForwarder.send(session); + await vi.advanceTimersByTimeAsync(10); + await vi.advanceTimersByTimeAsync(20); + await vi.advanceTimersByTimeAsync(5); + + const response = await responsePromise; + expect(await response.text()).toContain('"round-two"'); + expect(session.provider?.id).toBe(roundTwo.id); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(2); + expect(doForward).toHaveBeenCalledTimes(3); + } finally { + vi.useRealTimers(); + } + }); + + test("Sticky timeout still starts one normal wave when maxDiscoveryRounds is one", async () => { + vi.useFakeTimers(); + try { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const normal = createProvider({ id: 2, name: "normal", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 21 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 21, + providerId: sticky.id, + generation: "g-sticky-one-round", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 20, + stickySlaMs: 10, + racingTotalTimeoutMs: 100, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([normal]); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementation(async (attemptSession) => { + if ((attemptSession as ProxySession).provider?.id === sticky.id) { + return new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }); + } + return new Response('data: {"type":"content_block_delta","delta":{"text":"normal"}}\n\n', { + headers: { "content-type": "text/event-stream" }, + }); + }); + + const responsePromise = ProxyForwarder.send(session); + await vi.advanceTimersByTimeAsync(10); + const response = await responsePromise; + + expect(await response.text()).toContain('"normal"'); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledWith( + expect.anything(), + 1, + expect.arrayContaining([sticky.id]) + ); + expect(doForward).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); + + test("an explicit Sticky failure starts Discovery round one at full concurrency", async () => { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const normal = createProvider({ id: 2, name: "normal", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 22 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 22, + providerId: sticky.id, + generation: "g-sticky-failure", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 50, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([normal]); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockRejectedValueOnce(new Error("Sticky upstream failed")).mockResolvedValueOnce( + new Response('data: {"type":"content_block_delta","delta":{"text":"normal"}}\n\n', { + headers: { "content-type": "text/event-stream" }, + }) + ); + + const response = await ProxyForwarder.send(session); + expect(await response.text()).toContain('"normal"'); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledWith( + expect.anything(), + 2, + expect.arrayContaining([sticky.id]) + ); + expect(doForward).toHaveBeenCalledTimes(2); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledWith( + expect.objectContaining({ providerId: sticky.id, generation: "g-sticky-failure" }), + sticky.id, + 0 + ); + }); + + test("Discovery eligibility excludes WebSocket-tunneled requests", async () => { + const provider = createProvider({ id: 1, firstByteTimeoutStreamingMs: 0 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 23 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + mocks.getCachedSystemSettings.mockResolvedValue({ discoveryEnabled: true }); + mocks.isWebsocketClientRequest.mockReturnValueOnce(true); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response('data: {"type":"message_stop"}\n\n', { + status: 200, + headers: { "content-type": "text/event-stream" }, + }) + ); + + await ProxyForwarder.send(session); + expect(doForward).toHaveBeenCalledTimes(1); + expect(mocks.acquireSessionDiscoveryLease).not.toHaveBeenCalled(); + expect(mocks.pickDiscoveryProviders).not.toHaveBeenCalled(); + }); + + test("Discovery stops immediately on local database admission overload", async () => { + const provider = createProvider({ id: 1 }); + const alternative = createProvider({ id: 2 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 24 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 50, + racingTotalTimeoutMs: 200, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([alternative]); + mocks.categorizeErrorAsync.mockResolvedValueOnce(ProxyErrorCategory.LOCAL_OVERLOAD); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + const overload = new DbPoolAdmissionError("data", 32); + doForward.mockRejectedValueOnce(overload); + + await expect(ProxyForwarder.send(session)).rejects.toBe(overload); + expect(doForward).toHaveBeenCalledTimes(1); + expect(mocks.recordFailure).not.toHaveBeenCalled(); + expect(mocks.releaseSessionDiscoveryLease).toHaveBeenCalledTimes(1); + }); + + test("Discovery total deadline is not blocked by a stalled candidate selector", async () => { + vi.useFakeTimers(); + try { + const provider = createProvider({ id: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 25 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 20, + stickySlaMs: 20, + racingTotalTimeoutMs: 50, + }); + mocks.pickDiscoveryProviders.mockReturnValueOnce(new Promise(() => {})); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }) + ); + + const responsePromise = ProxyForwarder.send(session); + const observedError = responsePromise.catch((error) => error); + await vi.advanceTimersByTimeAsync(50); + expect(await observedError).toBeInstanceOf(UpstreamProxyError); + expect(mocks.releaseSessionDiscoveryLease).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + test("a fallback failure waits for the reserved wave before advancing another round", async () => { + vi.useFakeTimers(); + try { + const fallback = createProvider({ id: 1, name: "fallback", priority: 1 }); + const firstRoundLoser = createProvider({ id: 2, name: "first-round-loser", priority: 1 }); + const nextRoundWinner = createProvider({ id: 3, name: "next-round-winner", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 26 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(fallback); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 3, + discoverySlaMs: 10, + stickySlaMs: 10, + racingTotalTimeoutMs: 100, + }); + + const reservedWave = Promise.withResolvers(); + mocks.pickDiscoveryProviders + .mockResolvedValueOnce([firstRoundLoser]) + .mockReturnValueOnce(reservedWave.promise) + .mockResolvedValueOnce([ + createProvider({ id: 4, name: "unexpected-extra-round", priority: 1 }), + ]); + + const fallbackFailure = Promise.withResolvers(); + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementation(async (attemptSession) => { + const providerId = (attemptSession as ProxySession).provider?.id; + if (providerId === fallback.id) return fallbackFailure.promise; + if (providerId === nextRoundWinner.id) { + return new Response( + 'data: {"type":"content_block_delta","delta":{"text":"winner"}}\n\n', + { headers: { "content-type": "text/event-stream" } } + ); + } + return new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }); + }); + + const responsePromise = ProxyForwarder.send(session); + await vi.advanceTimersByTimeAsync(10); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(2); + + fallbackFailure.reject(new Error("fallback failed during reserved wave")); + await vi.advanceTimersByTimeAsync(0); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(2); + + reservedWave.resolve([nextRoundWinner]); + await vi.advanceTimersByTimeAsync(0); + const response = await responsePromise; + expect(await response.text()).toContain('"winner"'); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(2); + expect(doForward).toHaveBeenCalledTimes(3); + } finally { + vi.useRealTimers(); + } + }); + + test("Discovery does not immediately reselect a Provider whose launch setup failed", async () => { + const initial = createProvider({ id: 1, name: "initial" }); + const alternative = createProvider({ id: 2, name: "alternative" }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 20 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(initial); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 50, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders.mockImplementationOnce( + async (_session: ProxySession, _count: number, excludedIds: number[]) => { + expect(excludedIds).toContain(initial.id); + return [alternative]; + } + ); + + const endpointResolver = vi.spyOn( + ProxyForwarder as unknown as { + resolveStreamingHedgeEndpoint: ( + session: ProxySession, + provider: Provider + ) => Promise<{ endpointId: number | null; baseUrl: string; endpointUrl: string }>; + }, + "resolveStreamingHedgeEndpoint" + ); + endpointResolver + .mockRejectedValueOnce(new Error("initial endpoint setup failed")) + .mockResolvedValue({ + endpointId: null, + baseUrl: alternative.url, + endpointUrl: alternative.url, + }); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response('data: {"type":"content_block_delta","delta":{"text":"alternative"}}\n\n', { + headers: { "content-type": "text/event-stream" }, + }) + ); + + try { + const response = await ProxyForwarder.send(session); + expect(await response.text()).toContain('"alternative"'); + expect(doForward).toHaveBeenCalledTimes(1); + expect(session.provider?.id).toBe(alternative.id); + } finally { + endpointResolver.mockRestore(); + } + }); + + test("Discovery transfers the Provider session ref when a rectifier retries the same Provider", async () => { + const initial = createProvider({ id: 1, name: "initial", limitConcurrentSessions: 1 }); + const alternative = createProvider({ id: 2, name: "alternative", limitConcurrentSessions: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 22 }, + apiKey: null, + } as typeof session.authState; + setProviderWithSessionRef(session, initial); + withThinkingBlocks(session); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 100, + stickySlaMs: 100, + racingTotalTimeoutMs: 500, + stickyTimeoutCooldownMs: 300_000, + enableThinkingSignatureRectifier: true, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([alternative]); + + const signatureError = new UpstreamProxyError("Invalid `signature` in `thinking` block", 400, { + body: '{"error":"invalid_signature"}', + providerId: initial.id, + providerName: initial.name, + }); + let initialAttempts = 0; + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementation(async (attemptSession) => { + const runtime = attemptSession as ProxySession & AttemptRuntime; + if (runtime.provider?.id === initial.id) { + initialAttempts += 1; + if (initialAttempts === 1) throw signatureError; + + const body = runtime.request.message as { + messages: Array<{ content: Array> }>; + }; + expect(body.messages[0].content.some((block) => "signature" in block)).toBe(false); + return new Response( + 'data: {"type":"content_block_delta","delta":{"text":"rectified"}}\n\n', + { headers: { "content-type": "text/event-stream" } } + ); + } + + return new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }); + }); + + const response = await ProxyForwarder.send(session); + expect(await response.text()).toContain('"rectified"'); + + const initialAdmissionCalls = mocks.checkAndTrackProviderSession.mock.calls.filter( + ([providerId]) => providerId === initial.id + ); + const initialReleaseCalls = mocks.releaseProviderSession.mock.calls.filter( + ([providerId]) => providerId === initial.id + ); + expect(initialAttempts).toBe(2); + expect(initialAdmissionCalls).toHaveLength(0); + expect(initialReleaseCalls).toHaveLength(0); + expect(session.hasProviderSessionRef(initial.id)).toBe(true); + }); + + test("Discovery releases a transferred Provider session ref exactly once when rectifier retry setup fails", async () => { + const provider = createProvider({ id: 1, name: "initial", limitConcurrentSessions: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 23 }, + apiKey: null, + } as typeof session.authState; + setProviderWithSessionRef(session, provider); + withThinkingBlocks(session); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 100, + stickySlaMs: 100, + racingTotalTimeoutMs: 500, + stickyTimeoutCooldownMs: 300_000, + enableThinkingSignatureRectifier: true, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([]); + + const endpointResolver = vi.spyOn( + ProxyForwarder as unknown as { + resolveStreamingHedgeEndpoint: ( + session: ProxySession, + provider: Provider + ) => Promise<{ endpointId: number | null; baseUrl: string; endpointUrl: string }>; + }, + "resolveStreamingHedgeEndpoint" + ); + endpointResolver + .mockResolvedValueOnce({ + endpointId: null, + baseUrl: provider.url, + endpointUrl: provider.url, + }) + .mockRejectedValueOnce(new Error("rectifier retry endpoint setup failed")); + + const signatureError = new UpstreamProxyError("Invalid `signature` in `thinking` block", 400, { + body: '{"error":"invalid_signature"}', + providerId: provider.id, + providerName: provider.name, + }); + vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ).mockRejectedValueOnce(signatureError); + + try { + await expect(ProxyForwarder.send(session)).rejects.toBeInstanceOf(Error); + const providerAdmissionCalls = mocks.checkAndTrackProviderSession.mock.calls.filter( + ([providerId]) => providerId === provider.id + ); + const providerReleaseCalls = mocks.releaseProviderSession.mock.calls.filter( + ([providerId]) => providerId === provider.id + ); + expect(providerAdmissionCalls).toHaveLength(0); + expect(providerReleaseCalls).toHaveLength(1); + expect(session.hasProviderSessionRef(provider.id)).toBe(false); + } finally { + endpointResolver.mockRestore(); + } + }); + + test("a candidate delayed in launch setup is rolled back after another attempt wins", async () => { + vi.useFakeTimers(); + try { + const initial = createProvider({ id: 1, name: "initial" }); + const delayed = createProvider({ id: 2, name: "delayed", limitConcurrentSessions: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 10 }, + apiKey: null, + } as typeof session.authState; + setProviderWithSessionRef(session, initial); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 50, + racingTotalTimeoutMs: 200, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([delayed]); + + let resolveAdmission!: (value: unknown) => void; + mocks.checkAndTrackProviderSession.mockReturnValueOnce( + new Promise((resolve) => { + resolveAdmission = resolve; + }) + ); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockResolvedValueOnce( + new Response( + new ReadableStream({ + start(controller) { + setTimeout(() => { + controller.enqueue( + new TextEncoder().encode( + 'data: {"type":"content_block_delta","delta":{"text":"winner"}}\n\n' + ) + ); + controller.close(); + }, 5); + }, + }), + { headers: { "content-type": "text/event-stream" } } + ) + ); + + const responsePromise = ProxyForwarder.send(session); + await vi.advanceTimersByTimeAsync(10); + resolveAdmission({ allowed: true, count: 1, tracked: true, referenced: true }); + await vi.advanceTimersByTimeAsync(1); + const response = await responsePromise; + expect(await response.text()).toContain('"winner"'); + expect(doForward).toHaveBeenCalledTimes(1); + expect(mocks.releaseProviderSession).toHaveBeenCalledWith(delayed.id, session.sessionId); + } finally { + vi.useRealTimers(); + } + }); + + test("Discovery client abort preserves the captured binding and releases its lease", async () => { + const clientAbort = new AbortController(); + const provider = createProvider({ id: 1 }); + const session = createSession(clientAbort.signal); + session.authState = { + success: true, + user: null, + key: { id: 11 }, + apiKey: null, + } as typeof session.authState; + setProviderWithSessionRef(session, provider); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 100, + stickySlaMs: 100, + racingTotalTimeoutMs: 500, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([]); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementationOnce( + async (_attemptSession, _provider, _baseUrl, _audit, _count, _stream, signal) => + await new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => reject(signal.reason), { once: true }); + }) + ); + + const responsePromise = ProxyForwarder.send(session); + clientAbort.abort(new Error("client disconnected")); + await expect(responsePromise).rejects.toMatchObject({ statusCode: 499 }); + expect(mocks.clearVersionedSessionProvider).not.toHaveBeenCalled(); + expect(mocks.clearSessionProviders).not.toHaveBeenCalled(); + expect(mocks.releaseSessionDiscoveryLease).toHaveBeenCalledWith( + session.sessionId, + 11, + "lease-test" + ); + }); + + test("Discovery preserves binding state for a non-retryable client error", async () => { + const provider = createProvider({ id: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 12 }, + apiKey: null, + } as typeof session.authState; + setProviderWithSessionRef(session, provider); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 100, + stickySlaMs: 100, + racingTotalTimeoutMs: 500, + }); + mocks.pickDiscoveryProviders.mockResolvedValueOnce([]); + mocks.categorizeErrorAsync.mockResolvedValueOnce(ProxyErrorCategory.NON_RETRYABLE_CLIENT_ERROR); + const clientError = new UpstreamProxyError("invalid request", 400); + vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ).mockRejectedValueOnce(clientError); + + await expect(ProxyForwarder.send(session)).rejects.toBe(clientError); + expect(mocks.clearVersionedSessionProvider).not.toHaveBeenCalled(); + expect(mocks.clearSessionProviders).not.toHaveBeenCalled(); + expect(mocks.releaseSessionDiscoveryLease).toHaveBeenCalled(); + }); + test("removes streaming hedge client abort listener after winner response is returned", async () => { const clientAbortController = new AbortController(); const addSpy = vi.spyOn(clientAbortController.signal, "addEventListener"); diff --git a/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts b/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts index 0fa7ebafa..89d1bbb9c 100644 --- a/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts +++ b/tests/unit/proxy/proxy-forwarder-provider-session-release.test.ts @@ -22,6 +22,20 @@ describe("ProxyForwarder provider failure session release", () => { mocks.releaseProviderSession.mockClear(); }); + it("tracks baseline ownership independently for consecutive refs to one Provider", async () => { + const { ProxySession } = await import("@/app/v1/_lib/proxy/session"); + const session = Object.create(ProxySession.prototype) as ProxySession; + + session.recordProviderSessionRef(42, { retainOnSuccess: true }); + session.recordProviderSessionRef(42, { retainOnSuccess: false }); + + expect(session.shouldRetainProviderSessionRefOnSuccess(42)).toBe(true); + expect(session.consumeProviderSessionRef(42)).toBe(true); + expect(session.shouldRetainProviderSessionRefOnSuccess(42)).toBe(false); + expect(session.consumeProviderSessionRef(42)).toBe(true); + expect(session.hasProviderSessionRef(42)).toBe(false); + }); + it("标记供应商失败时仅释放本请求已获取的 provider session ref", async () => { const { ProxyForwarder } = await import("@/app/v1/_lib/proxy/forwarder"); const forwarderInternals = ProxyForwarder as unknown as { diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index ee817d348..ded21ba4b 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -104,6 +104,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), settleLeaseBudgets: vi.fn(), + releaseProviderSession: vi.fn(), }, })); @@ -114,6 +115,17 @@ vi.mock("@/lib/redis/live-chain-store", () => ({ vi.mock("@/lib/session-manager", () => ({ SessionManager: { clearSessionProvider: vi.fn(), + clearVersionedSessionProvider: vi.fn(), + compareAndSetSessionProvider: vi.fn(), + getSessionBindingSnapshot: vi.fn(), + renewSessionDiscoveryLease: vi.fn(async () => ({ + status: "renewed", + legacyFallbackAllowed: false, + })), + releaseSessionDiscoveryLease: vi.fn(async () => ({ + status: "released", + legacyFallbackAllowed: false, + })), extractCodexPromptCacheKey: vi.fn(), storeSessionResponse: vi.fn(async () => undefined), storeSessionRequestPhaseSnapshot: vi.fn(), @@ -2074,6 +2086,69 @@ describe("ProxyResponseHandler stream client abort finalization", () => { ); }); + it.each([ + { bindingIntent: "create" as const, providerId: null }, + { bindingIntent: "renew" as const, providerId: 1 }, + ])( + "preserves binding state for a client-aborted Discovery $bindingIntent stream", + async ({ bindingIntent, providerId }) => { + const controller = new AbortController(); + controller.abort(); + const session = createSession(controller.signal); + Object.assign(session, { sessionId: `session-client-abort-${bindingIntent}` }); + session.recordProviderSessionRef(1); + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValue( + "client-abort-cache-key" + ); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: bindingIntent === "create", + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + bindingIntent, + bindingSnapshot: { + sessionId: `session-client-abort-${bindingIntent}`, + keyId: 2, + providerId, + generation: `${bindingIntent}-generation`, + }, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: `session-client-abort-${bindingIntent}`, + keyId: 2, + ownerToken: `client-abort-${bindingIntent}-owner`, + ttlSeconds: 30, + }, + providerSessionRefOwned: true, + }); + + await ProxyResponseHandler.dispatch(session, createCompletedThenErroredResponsesSse()); + await drainAsyncTasks(); + + expect(SessionManager.clearVersionedSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.clearSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.updateSessionBindingSmart).not.toHaveBeenCalled(); + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledWith( + `session-client-abort-${bindingIntent}`, + 2, + `client-abort-${bindingIntent}-owner` + ); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith( + 1, + `session-client-abort-${bindingIntent}` + ); + } + ); + it("keeps a genuinely aborted upstream responses stream as 499", async () => { const controller = new AbortController(); controller.abort(); @@ -3184,6 +3259,45 @@ describe("ProxyResponseHandler stream client abort finalization", () => { ); }); + it("does not mutate non-stream bindings when the routing mode forbids it", async () => { + const controller = new AbortController(); + controller.abort(); + const session = createSession(controller.signal); + Object.assign(session, { + sessionId: "lease-conflict-non-stream", + isSessionBindingAllowed: () => false, + }); + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValue("blocked-cache-key"); + const response = new Response('{"id":"resp_lease_conflict"}', { + status: 200, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + await drainAsyncTasks(); + + expect(SessionManager.clearSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + }); + + it("does not create a non-stream Codex cache binding when binding is disabled", async () => { + const session = createSession(new AbortController().signal); + Object.assign(session, { + sessionId: "lease-conflict-non-stream-success", + isSessionBindingAllowed: () => false, + }); + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValue("blocked-cache-key"); + const response = new Response('{"id":"resp_lease_conflict"}', { + status: 200, + headers: { "content-type": "application/json" }, + }); + + await ProxyResponseHandler.dispatch(session, response); + await drainAsyncTasks(); + + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + }); + it("publishes a successful stream Codex cache binding only after durable acknowledgement", async () => { const durableAck = createDeferred(); const cacheBinding = createDeferred(); @@ -3243,6 +3357,122 @@ describe("ProxyResponseHandler stream client abort finalization", () => { } }); + it("publishes a Discovery Codex cache key only after the primary generation CAS succeeds", async () => { + const order: string[] = []; + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce( + "discovery-stream-cache-key" + ); + vi.mocked(SessionManager.compareAndSetSessionProvider).mockImplementationOnce(async () => { + order.push("primary-cas"); + return { + status: "ok", + source: "updated", + snapshot: { + sessionId: "stream-discovery-cache-binding", + keyId: 2, + providerId: 1, + generation: "discovery-updated-generation", + }, + legacyFallbackAllowed: false, + }; + }); + vi.mocked(SessionManager.updateSessionWithCodexCacheKey).mockImplementationOnce(async () => { + order.push("aux-cache-binding"); + }); + const session = createSession(new AbortController().signal); + session.sessionId = "stream-discovery-cache-binding"; + session.recordProviderSessionRef(1); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "stream-discovery-cache-binding", + keyId: 2, + providerId: null, + generation: "discovery-create-generation", + }, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: "stream-discovery-cache-binding", + keyId: 2, + ownerToken: "discovery-cache-owner", + ttlSeconds: 30, + }, + providerSessionRefOwned: true, + providerSessionRefRetainOnSuccess: true, + }); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await drainAsyncTasks(); + + expect(order).toEqual(["primary-cas", "aux-cache-binding"]); + expect(RateLimitService.releaseProviderSession).not.toHaveBeenCalled(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + }); + + it("does not publish a Discovery Codex cache key when the primary generation CAS conflicts", async () => { + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce( + "conflicted-discovery-cache-key" + ); + vi.mocked(SessionManager.compareAndSetSessionProvider).mockResolvedValueOnce({ + status: "conflict", + reason: "generation_mismatch", + legacyFallbackAllowed: false, + }); + const session = createSession(new AbortController().signal); + session.sessionId = "stream-discovery-cache-conflict"; + session.recordProviderSessionRef(1); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "stream-discovery-cache-conflict", + keyId: 2, + providerId: null, + generation: "stale-discovery-generation", + }, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: "stream-discovery-cache-conflict", + keyId: 2, + ownerToken: "conflicted-discovery-owner", + ttlSeconds: 30, + }, + providerSessionRefOwned: true, + }); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await drainAsyncTasks(); + + expect(SessionManager.compareAndSetSessionProvider).toHaveBeenCalledOnce(); + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith( + 1, + "stream-discovery-cache-conflict" + ); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + }); + it("does not publish a stream Codex cache binding for a final non-2xx outcome", async () => { vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce("stream-cache-key-2"); const session = createSession(new AbortController().signal); @@ -3276,6 +3506,33 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); }); + it("does not publish a Codex cache binding for a Discovery fallback winner", async () => { + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce( + "fallback-stream-cache-key" + ); + const session = createSession(new AbortController().signal); + session.sessionId = "stream-codex-fallback"; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 2, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + bindingIntent: "none", + }); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + await drainAsyncTasks(); + + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + }); + it("durably finalizes a Gemini non-stream passthrough body-read failure", async () => { const session = createSession(new AbortController().signal, { providerType: "gemini", diff --git a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts index af432f126..038c9378c 100644 --- a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts +++ b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts @@ -77,6 +77,11 @@ vi.mock("@/lib/session-manager", () => ({ updateSessionUsage: vi.fn(), storeSessionResponse: vi.fn(), clearSessionProvider: vi.fn(), + clearVersionedSessionProvider: vi.fn(), + compareAndSetSessionProvider: vi.fn(), + getSessionBindingSnapshot: vi.fn(), + renewSessionDiscoveryLease: vi.fn(), + releaseSessionDiscoveryLease: vi.fn(), extractCodexPromptCacheKey: vi.fn(), updateSessionBindingSmart: vi.fn(), updateSessionProvider: vi.fn(), @@ -90,6 +95,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), settleLeaseBudgets: vi.fn(), + releaseProviderSession: vi.fn(), }, })); @@ -335,6 +341,69 @@ function createSuccessStreamResponse(): Response { }); } +function createSuccessStreamResponseWithCompletion(): Response { + const sseText = + `data: ${JSON.stringify({ type: "content_block_delta", delta: { text: "ok" } })}\n\n` + + `event: message_stop\ndata: ${JSON.stringify({ type: "message_stop" })}\n\n`; + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(sseText)); + controller.close(); + }, + }); + return new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); +} + +function createMisleadingCompletionTextResponse(): Response { + const sseText = + `event: content_block_delta\ndata: ${JSON.stringify({ + type: "content_block_delta", + delta: { text: "the words message_stop and response.completed are ordinary content" }, + })}\n\n` + + `event: message_delta\ndata: ${JSON.stringify({ + type: "message_delta", + delta: { stop_reason: null }, + })}\n\n`; + return new Response(sseText, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); +} + +function createControllableSuccessStreamResponse(): { + response: Response; + complete: () => void; +} { + const encoder = new TextEncoder(); + let streamController!: ReadableStreamDefaultController; + const stream = new ReadableStream({ + start(controller) { + streamController = controller; + controller.enqueue( + encoder.encode( + `data: ${JSON.stringify({ type: "content_block_delta", delta: { text: "ok" } })}\n\n` + ) + ); + }, + }); + return { + response: new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + complete: () => { + streamController.enqueue( + encoder.encode(`data: ${JSON.stringify({ type: "message_stop" })}\n\n`) + ); + streamController.close(); + }, + }; +} + async function drainAsyncTasks(): Promise { while (asyncTasks.length > 0) { const tasks = asyncTasks.splice(0, asyncTasks.length); @@ -366,6 +435,47 @@ function setupCommonMocks() { vi.mocked(updateMessageRequestDuration).mockResolvedValue(undefined); vi.mocked(SessionManager.storeSessionResponse).mockResolvedValue(undefined); vi.mocked(SessionManager.clearSessionProvider).mockResolvedValue(undefined); + vi.mocked(SessionManager.clearVersionedSessionProvider).mockResolvedValue({ + status: "ok", + source: "cleared", + snapshot: { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "cleared", + }, + legacyFallbackAllowed: false, + }); + vi.mocked(SessionManager.compareAndSetSessionProvider).mockResolvedValue({ + status: "ok", + source: "updated", + snapshot: { + sessionId: "fake-session", + keyId: 456, + providerId: 1, + generation: "updated", + }, + legacyFallbackAllowed: false, + }); + vi.mocked(SessionManager.getSessionBindingSnapshot).mockResolvedValue({ + status: "ok", + source: "existing", + snapshot: { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "fresh", + }, + legacyFallbackAllowed: false, + }); + vi.mocked(SessionManager.renewSessionDiscoveryLease).mockResolvedValue({ + status: "renewed", + legacyFallbackAllowed: false, + }); + vi.mocked(SessionManager.releaseSessionDiscoveryLease).mockResolvedValue({ + status: "released", + legacyFallbackAllowed: false, + }); vi.mocked(SessionManager.updateSessionUsage).mockResolvedValue(undefined); vi.mocked(SessionManager.updateSessionBindingSmart).mockResolvedValue({ updated: true, @@ -383,6 +493,7 @@ function setupCommonMocks() { status: "settled", settlements: [], }); + vi.mocked(RateLimitService.releaseProviderSession).mockResolvedValue(undefined); vi.mocked(SessionTracker.refreshSession).mockResolvedValue(undefined); mockRecordFailure.mockResolvedValue(undefined); mockRecordSuccess.mockResolvedValue(undefined); @@ -428,6 +539,22 @@ describe("Endpoint circuit breaker isolation", () => { ).toBe(true); }); + it("does not clear a binding when the request routing mode forbids binding mutations", async () => { + const session = createSession(); + Object.assign(session, { isSessionBindingAllowed: () => false }); + setDeferredMeta(session, 42); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createFake200StreamResponse() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.clearSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.clearVersionedSessionProvider).not.toHaveBeenCalled(); + }); + it("高并发模式下，fake-200 流式错误仍应记录核心失败，但跳过 session 观测写入", async () => { const session = createSession(); session.setHighConcurrencyModeEnabled(true); @@ -537,4 +664,638 @@ describe("Endpoint circuit breaker isolation", () => { expect(mockRecordEndpointSuccess).not.toHaveBeenCalled(); expect(mockRecordEndpointFailure).not.toHaveBeenCalled(); }); + + it("does not clear a create binding when Discovery finishes with fake-200", async () => { + const session = createSession(); + session.recordProviderSessionRef(1); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "create-generation", + }, + providerSessionRefOwned: true, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createFake200StreamResponse() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.clearVersionedSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.clearSessionProvider).not.toHaveBeenCalled(); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledOnce(); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith(1, "fake-session"); + }); + + it("clears only the captured renew snapshot after a fake-200", async () => { + const session = createSession(); + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: 1, + generation: "renew-generation", + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "renew", + bindingSnapshot: snapshot, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createFake200StreamResponse() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.clearVersionedSessionProvider).toHaveBeenCalledWith(snapshot, 1, 0); + expect(SessionManager.clearSessionProvider).not.toHaveBeenCalled(); + }); + + it("never mutates a binding for fallback intent none", async () => { + const session = createSession(); + session.recordProviderSessionRef(1); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 2, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "none", + providerSessionRefOwned: true, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createFake200StreamResponse() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.clearVersionedSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.clearSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.updateSessionBindingSmart).not.toHaveBeenCalled(); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith(1, "fake-session"); + }); + + it("does not clear a create tombstone when the completion marker is missing", async () => { + const session = createSession(); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "incomplete-generation", + }, + requiresCompletionMarker: true, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponse() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.clearVersionedSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.clearSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + }); + + it("does not accept completion marker words embedded in ordinary SSE content", async () => { + const session = createSession(); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "misleading-content-generation", + }, + requiresCompletionMarker: true, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createMisleadingCompletionTextResponse() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + expect(updateMessageRequestDetailsDurably).toHaveBeenCalledWith( + 1, + expect.objectContaining({ + statusCode: 502, + errorMessage: "STREAM_COMPLETION_MARKER_MISSING", + }), + expect.objectContaining({ onCommitted: expect.any(Function) }) + ); + }); + + it.each([ + { + label: "OpenAI Responses", + format: "response" as const, + body: `event: response.completed\ndata: ${JSON.stringify({ + type: "response.completed", + response: { id: "resp_completed" }, + })}\n\n`, + }, + { + label: "OpenAI Chat finish reason", + format: "openai" as const, + body: `data: ${JSON.stringify({ + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + })}\n\n`, + }, + { + label: "OpenAI Chat done sentinel", + format: "openai" as const, + body: "data: [DONE]\n\n", + }, + { + label: "Gemini", + format: "gemini" as const, + body: `data: ${JSON.stringify({ + candidates: [{ finishReason: "STOP" }], + })}\n\n`, + }, + { + label: "Gemini CLI", + format: "gemini-cli" as const, + body: `data: ${JSON.stringify({ + response: { candidates: [{ finishReason: "STOP" }] }, + })}\n\n`, + }, + ])("accepts a structurally valid $label completion marker", async ({ format, body }) => { + const session = createSession(); + session.originalFormat = format; + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: `${format}-completion-generation`, + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: snapshot, + requiresCompletionMarker: true, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + new Response(body, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }) + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.compareAndSetSessionProvider).toHaveBeenCalledWith(snapshot, 1); + }); + + it("releases a create attempt ref when generation CAS loses", async () => { + const session = createSession(); + session.recordProviderSessionRef(1); + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "stale-generation", + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: snapshot, + requiresCompletionMarker: true, + providerSessionRefOwned: true, + providerSessionRefRetainOnSuccess: true, + }); + vi.mocked(SessionManager.compareAndSetSessionProvider).mockResolvedValueOnce({ + status: "conflict", + reason: "generation_mismatch", + legacyFallbackAllowed: false, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponseWithCompletion() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledOnce(); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith(1, "fake-session"); + }); + + it("retains an owned Provider ref after a renew generation CAS succeeds", async () => { + const session = createSession(); + session.recordProviderSessionRef(1); + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: 1, + generation: "renew-generation", + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "renew", + bindingSnapshot: snapshot, + requiresCompletionMarker: true, + providerSessionRefOwned: true, + providerSessionRefRetainOnSuccess: true, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponseWithCompletion() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.compareAndSetSessionProvider).toHaveBeenCalledWith(snapshot, 1); + expect(RateLimitService.releaseProviderSession).not.toHaveBeenCalled(); + }); + + it("releases an owned Provider ref after CAS success when it is not the new baseline", async () => { + const session = createSession(); + session.recordProviderSessionRef(1); + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: 1, + generation: "existing-baseline-generation", + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "renew", + bindingSnapshot: snapshot, + requiresCompletionMarker: true, + providerSessionRefOwned: true, + providerSessionRefRetainOnSuccess: false, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponseWithCompletion() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.compareAndSetSessionProvider).toHaveBeenCalledWith(snapshot, 1); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith(1, "fake-session"); + }); + + it.each([ + { + label: "lost", + leaseResult: { + status: "lost", + reason: "not_owner_or_missing", + legacyFallbackAllowed: false, + } as const, + }, + { + label: "unavailable", + leaseResult: { + status: "unavailable", + reason: "operation_failed", + capabilityState: "unavailable", + legacyFallbackAllowed: true, + } as const, + }, + ])("fails binding closed when the finalizer lease is $label", async ({ leaseResult }) => { + const session = createSession(); + session.recordProviderSessionRef(1); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "lease-guarded-generation", + }, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: "fake-session", + keyId: 456, + ownerToken: "lease-owner", + ttlSeconds: 30, + }, + providerSessionRefOwned: true, + }); + vi.mocked(SessionManager.renewSessionDiscoveryLease).mockResolvedValueOnce(leaseResult); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponseWithCompletion() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.renewSessionDiscoveryLease).toHaveBeenCalledOnce(); + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith(1, "fake-session"); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + }); + + it("renews a long-stream lease and releases it once after terminal side effects", async () => { + vi.useFakeTimers(); + try { + const order: string[] = []; + const session = createSession(); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 2, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "none", + requiresCompletionMarker: false, + discoveryLease: { + sessionId: "fake-session", + keyId: 456, + ownerToken: "lease-owner", + ttlSeconds: 2, + }, + }); + mockRecordSuccess.mockImplementationOnce(async () => { + order.push("side-effect"); + }); + vi.mocked(SessionManager.releaseSessionDiscoveryLease).mockImplementationOnce(async () => { + order.push("lease-release"); + return { status: "released", legacyFallbackAllowed: false }; + }); + const controlled = createControllableSuccessStreamResponse(); + + const clientResponse = await ProxyResponseHandler.dispatch(session, controlled.response); + const bodyPromise = clientResponse.text(); + await vi.advanceTimersByTimeAsync(0); + expect(SessionManager.renewSessionDiscoveryLease).toHaveBeenCalledTimes(1); + await vi.advanceTimersByTimeAsync(1_100); + expect(SessionManager.renewSessionDiscoveryLease).toHaveBeenCalledWith( + "fake-session", + 456, + "lease-owner", + 2 + ); + + controlled.complete(); + await bodyPromise; + await drainAsyncTasks(); + + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledWith( + "fake-session", + 456, + "lease-owner" + ); + expect(order).toEqual(["side-effect", "lease-release"]); + + const renewCalls = vi.mocked(SessionManager.renewSessionDiscoveryLease).mock.calls.length; + await vi.advanceTimersByTimeAsync(5_000); + expect(SessionManager.renewSessionDiscoveryLease).toHaveBeenCalledTimes(renewCalls); + } finally { + vi.useRealTimers(); + } + }); + + it("does not delay downstream delivery while the lease handoff renewal is pending", async () => { + const handoffRenewal = Promise.withResolvers<{ + status: "renewed"; + legacyFallbackAllowed: false; + }>(); + vi.mocked(SessionManager.renewSessionDiscoveryLease).mockReturnValueOnce( + handoffRenewal.promise + ); + const session = createSession(); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 2, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "none", + requiresCompletionMarker: false, + discoveryLease: { + sessionId: "fake-session", + keyId: 456, + ownerToken: "pending-handoff-owner", + ttlSeconds: 30, + }, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponseWithCompletion() + ); + await expect(clientResponse.text()).resolves.toContain("message_stop"); + expect(SessionManager.renewSessionDiscoveryLease).toHaveBeenCalledOnce(); + + handoffRenewal.resolve({ status: "renewed", legacyFallbackAllowed: false }); + await drainAsyncTasks(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + }); + + it("bounds a stalled lease release and still invokes compare-delete exactly once", async () => { + vi.useFakeTimers(); + try { + const session = createSession(); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 2, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "none", + requiresCompletionMarker: false, + discoveryLease: { + sessionId: "fake-session", + keyId: 456, + ownerToken: "stalled-release-owner", + ttlSeconds: 30, + }, + }); + vi.mocked(SessionManager.releaseSessionDiscoveryLease).mockImplementationOnce( + () => new Promise(() => undefined) + ); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponseWithCompletion() + ); + await clientResponse.text(); + const drainPromise = drainAsyncTasks(); + await vi.advanceTimersByTimeAsync(5_000); + await drainPromise; + + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledWith( + "fake-session", + 456, + "stalled-release-owner" + ); + } finally { + vi.useRealTimers(); + } + }); + + it("fails closed when the captured Discovery generation has expired", async () => { + const session = createSession(); + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "expired-generation", + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: snapshot, + requiresCompletionMarker: true, + providerSessionRefOwned: true, + }); + session.recordProviderSessionRef(1); + vi.mocked(SessionManager.compareAndSetSessionProvider).mockResolvedValueOnce({ + status: "conflict", + reason: "canonical_missing", + legacyFallbackAllowed: false, + }); + + const clientResponse = await ProxyResponseHandler.dispatch( + session, + createSuccessStreamResponseWithCompletion() + ); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.getSessionBindingSnapshot).not.toHaveBeenCalled(); + expect(SessionManager.compareAndSetSessionProvider).toHaveBeenCalledOnce(); + expect(SessionManager.compareAndSetSessionProvider).toHaveBeenCalledWith(snapshot, 1); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith(1, "fake-session"); + }); }); diff --git a/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts b/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts index ba7bd9b13..2524cc5ce 100644 --- a/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts +++ b/tests/unit/proxy/response-handler-gemini-stream-passthrough-timeouts.test.ts @@ -707,7 +707,7 @@ describe("ProxyResponseHandler - Gemini stream passthrough timeouts", () => { setTimeout(() => { try { res.end( - 'data: {"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":2},"finishReason":"STOP"}\n\n' + 'data: {"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":2},"candidates":[{"finishReason":"STOP"}]}\n\n' ); } catch { // ignore From 0537a6e8574d03c774adea1d7c3aee34902f9daf Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 19:35:52 -0400 Subject: [PATCH 14/23] fix(discovery): retain blocked fallback readiness --- .../v1/_lib/proxy/discovery-coordinator.ts | 14 ++++ src/app/v1/_lib/proxy/forwarder.ts | 6 +- .../unit/proxy/discovery-coordinator.test.ts | 33 ++++++++ .../proxy-forwarder-hedge-first-byte.test.ts | 80 +++++++++++++++++++ 4 files changed, 131 insertions(+), 2 deletions(-) diff --git a/src/app/v1/_lib/proxy/discovery-coordinator.ts b/src/app/v1/_lib/proxy/discovery-coordinator.ts index f1330df92..b7879a8dc 100644 --- a/src/app/v1/_lib/proxy/discovery-coordinator.ts +++ b/src/app/v1/_lib/proxy/discovery-coordinator.ts @@ -150,6 +150,20 @@ export class DiscoveryCoordinator { return this.chooseReadyNormal(); } + /** Record a ready fallback without allowing it to preempt a reserved normal wave. */ + recordReadyHeld( + id: string, + requestEpoch = this.requestEpoch, + roundEpoch = this.roundEpoch + ): boolean { + if (!this.acceptsEpoch(requestEpoch, roundEpoch) || this.isTerminal) return false; + const attempt = this.attempts.get(id); + if (!attempt?.pending || attempt.kind !== "fallback") return false; + attempt.ready = true; + this.state = "FALLBACK_READY_HELD"; + return true; + } + /** Convert a timed-out Sticky attempt into the request's fallback lane. */ demoteToFallback( id: string, diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index e31a81b17..c9943e85d 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5574,8 +5574,10 @@ export class ProxyForwarder { (fallbackPromotionBlocked || roundLaunchesInProgress > 0) ) { // The next wave has been reserved but its normal attempts may - // still be awaiting selection/endpoint setup. Keep the fallback - // ready-held until those slots are registered or exhausted. + // still be awaiting selection/endpoint setup. Persist readiness + // without promoting so the total deadline can still recover the + // buffered fallback if that setup stalls. + coordinator.recordReadyHeld(id); return; } // Record readiness even when the priority gate holds this attempt. diff --git a/tests/unit/proxy/discovery-coordinator.test.ts b/tests/unit/proxy/discovery-coordinator.test.ts index bdbb7ff20..1ce247f52 100644 --- a/tests/unit/proxy/discovery-coordinator.test.ts +++ b/tests/unit/proxy/discovery-coordinator.test.ts @@ -134,6 +134,39 @@ describe("DiscoveryCoordinator", () => { expect(coordinator.markFailed("fallback")).toEqual({ type: "none" }); }); + it("records a ready-held fallback without promoting it before the deadline", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("fallback", 1, "fallback")); + coordinator.addAttempt(attempt("normal", 1)); + + expect(coordinator.recordReadyHeld("fallback")).toBe(true); + expect(coordinator.snapshot.find((item) => item.id === "fallback")).toMatchObject({ + kind: "fallback", + ready: true, + pending: true, + }); + expect(coordinator.state).toBe("FALLBACK_READY_HELD"); + expect(coordinator.onDeadline()).toEqual({ + type: "promote_fallback", + attemptId: "fallback", + }); + }); + + it("rejects ready-held writes for stale, non-fallback, or inactive attempts", () => { + const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); + coordinator.addAttempt(attempt("normal", 1)); + coordinator.addAttempt(attempt("fallback", 1, "fallback")); + const staleEpoch = coordinator.epochs; + + expect(coordinator.recordReadyHeld("normal")).toBe(false); + coordinator.beginRound(); + expect( + coordinator.recordReadyHeld("fallback", staleEpoch.requestEpoch, staleEpoch.roundEpoch) + ).toBe(false); + coordinator.markFailed("fallback"); + expect(coordinator.recordReadyHeld("fallback")).toBe(false); + }); + it("keeps coordinator kind in sync when a running Sticky becomes fallback", () => { const coordinator = new DiscoveryCoordinator({ concurrency: 2, maxRounds: 2 }); coordinator.addAttempt(attempt("sticky", 1)); diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index bd2cafab4..f732a52ac 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -2509,6 +2509,86 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { } }); + test("a ready-held Sticky fallback survives a stalled next-wave selector until the total deadline", async () => { + vi.useFakeTimers(); + try { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 27 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 27, + providerId: sticky.id, + generation: "g-sticky-deadline", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 20, + stickySlaMs: 10, + racingTotalTimeoutMs: 50, + stickyTimeoutCooldownMs: 300_000, + }); + + const stalledSelector = Promise.withResolvers(); + mocks.pickDiscoveryProviders.mockReturnValueOnce(stalledSelector.promise); + vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ).mockResolvedValueOnce( + new Response( + new ReadableStream({ + start(controller) { + setTimeout(() => { + controller.enqueue( + new TextEncoder().encode( + 'data: {"type":"content_block_delta","delta":{"text":"sticky-fallback"}}\n\n' + ) + ); + controller.close(); + }, 15); + }, + }), + { headers: { "content-type": "text/event-stream" } } + ) + ); + + let settledEarly = false; + const responsePromise = ProxyForwarder.send(session).then((response) => { + settledEarly = true; + return response; + }); + + await vi.advanceTimersByTimeAsync(10); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(1); + await vi.advanceTimersByTimeAsync(5); + expect(settledEarly).toBe(false); + + await vi.advanceTimersByTimeAsync(35); + const response = await responsePromise; + expect(await response.text()).toContain('"sticky-fallback"'); + expect(session.provider?.id).toBe(sticky.id); + + stalledSelector.resolve([]); + await vi.advanceTimersByTimeAsync(0); + } finally { + vi.useRealTimers(); + } + }); + test("Sticky probing does not consume a configured Discovery round", async () => { vi.useFakeTimers(); try { From 12887e497070adbb4fcf088c8ccdfe9fe525b08d Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 20:00:33 -0400 Subject: [PATCH 15/23] fix(discovery): reserve sticky timeout wave --- src/app/v1/_lib/proxy/forwarder.ts | 24 +++- .../proxy-forwarder-hedge-first-byte.test.ts | 110 ++++++++++++++++++ 2 files changed, 130 insertions(+), 4 deletions(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index c9943e85d..1f1489c00 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5041,6 +5041,7 @@ export class ProxyForwarder { let roundLaunchesInProgress = 0; const roundLaunchIdleWaiters = new Set<() => void>(); let fallbackPromotionBlocked = false; + let stickyTimeoutWaveReservation: { fallbackAttemptId: string } | null = null; const hasSticky = session.shouldReuseProvider() && !!session.sessionId && @@ -5781,7 +5782,14 @@ export class ProxyForwarder { failureAction.type === "launch" || failureAction.type === "terminal_failure"; if (actionOwnsNextStep) { - await executeCoordinatorAction(failureAction); + if ( + failureAction.type === "launch" && + stickyTimeoutWaveReservation?.fallbackAttemptId === id + ) { + await launchReservedStickyTimeoutWave(failureAction.slots); + } else { + await executeCoordinatorAction(failureAction); + } } if (!actionOwnsNextStep && !committed && !settled) { const replacement = await chooseCandidate(); @@ -5868,6 +5876,13 @@ export class ProxyForwarder { } }; + const launchReservedStickyTimeoutWave = async (slots: number): Promise => { + if (!stickyTimeoutWaveReservation) return; + stickyTimeoutWaveReservation = null; + if (settled || committed) return; + await launchNextRound(slots, true); + }; + executeCoordinatorAction = async (action, terminalCancellationKind) => { if (settled || committed) return; if (action.type === "cancel" || action.type === "launch") { @@ -5965,17 +5980,18 @@ export class ProxyForwarder { sticky.kind = "fallback"; discoveryMetrics.fallbackPromoted(sticky.id, sticky.provider.id, 0); fallbackPromotionBlocked = true; + stickyTimeoutWaveReservation = { fallbackAttemptId: sticky.id }; if (bindingSnapshot && bindingSnapshot.providerId === initialProvider.id) { void clearCapturedStickyBinding( Math.ceil((settings.stickyTimeoutCooldownMs ?? 300_000) / 1000) ).finally(() => { - void launchNextRound(Math.max(0, concurrency - 1), true).catch((error) => - logger.warn("[Discovery] Sticky round launch failed", { error }) + void launchReservedStickyTimeoutWave(Math.max(0, concurrency - 1)).catch( + (error) => logger.warn("[Discovery] Sticky round launch failed", { error }) ); }); return; } - void launchNextRound(Math.max(0, concurrency - 1), true).catch((error) => + void launchReservedStickyTimeoutWave(Math.max(0, concurrency - 1)).catch((error) => logger.warn("[Discovery] Sticky round launch failed", { error }) ); } diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index f732a52ac..7bbcc1856 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -2509,6 +2509,116 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { } }); + test("Sticky timeout and fallback failure consume a single replacement-wave reservation", async () => { + vi.useFakeTimers(); + try { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const normalOne = createProvider({ id: 2, name: "normal-one", priority: 1 }); + const normalTwo = createProvider({ id: 3, name: "normal-two", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 28 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 28, + providerId: sticky.id, + generation: "g-sticky-single-wave", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 2, + discoverySlaMs: 50, + stickySlaMs: 10, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + }); + + const clearBinding = Promise.withResolvers(); + mocks.clearVersionedSessionProvider.mockReturnValueOnce(clearBinding.promise); + const replacementWave = Promise.withResolvers(); + mocks.pickDiscoveryProviders.mockReturnValueOnce(replacementWave.promise); + + const stickyAttempt = Promise.withResolvers(); + const normalOneAttempt = Promise.withResolvers(); + const normalTwoAttempt = Promise.withResolvers(); + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementation(async (attemptSession) => { + switch ((attemptSession as ProxySession).provider?.id) { + case sticky.id: + return stickyAttempt.promise; + case normalOne.id: + return normalOneAttempt.promise; + case normalTwo.id: + return normalTwoAttempt.promise; + default: + throw new Error("unexpected Provider"); + } + }); + + const responsePromise = ProxyForwarder.send(session); + await vi.advanceTimersByTimeAsync(10); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledTimes(1); + expect(mocks.pickDiscoveryProviders).not.toHaveBeenCalled(); + + stickyAttempt.reject(new Error("Sticky fallback failed while binding clear was pending")); + await vi.advanceTimersByTimeAsync(0); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(1); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledWith( + expect.anything(), + 2, + expect.arrayContaining([sticky.id]) + ); + + clearBinding.resolve({ + status: "ok", + legacyFallbackAllowed: false, + source: "cleared", + snapshot: { + sessionId: session.sessionId!, + keyId: 28, + providerId: null, + generation: "g-cleared-single-wave", + }, + }); + await vi.advanceTimersByTimeAsync(0); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(1); + + replacementWave.resolve([normalOne, normalTwo]); + await vi.advanceTimersByTimeAsync(0); + expect(doForward).toHaveBeenCalledTimes(3); + + normalOneAttempt.resolve( + new Response('data: {"type":"content_block_delta","delta":{"text":"normal-one"}}\n\n', { + headers: { "content-type": "text/event-stream" }, + }) + ); + await vi.advanceTimersByTimeAsync(0); + const response = await responsePromise; + expect(await response.text()).toContain('"normal-one"'); + expect(session.provider?.id).toBe(normalOne.id); + + normalTwoAttempt.resolve(new Response(null)); + await vi.advanceTimersByTimeAsync(0); + } finally { + vi.useRealTimers(); + } + }); + test("a ready-held Sticky fallback survives a stalled next-wave selector until the total deadline", async () => { vi.useFakeTimers(); try { From 08337d8feac1df1e857ffa3fa096e9da218a6684 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 20:05:32 -0400 Subject: [PATCH 16/23] fix(discovery): stop parsing after validity errors --- src/app/v1/_lib/proxy/discovery-validity.ts | 9 ++++++++- tests/unit/proxy/discovery-validity.test.ts | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts index 3325295a7..059bde293 100644 --- a/src/app/v1/_lib/proxy/discovery-validity.ts +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -155,6 +155,7 @@ export class DiscoveryValidityParser { constructor(readonly protocol: DiscoveryProtocol) {} push(chunk: Uint8Array | string): DiscoveryValidity { + if (this._error) return this.result; this.bytesSeen += typeof chunk === "string" ? new TextEncoder().encode(chunk).byteLength : chunk.byteLength; if (!this._ready && this.bytesSeen > DISCOVERY_PREFIX_MAX_BYTES) { @@ -172,7 +173,13 @@ export class DiscoveryValidityParser { if (this.buffered.includes("\n")) { const lines = this.buffered.split(/\r?\n/); this.buffered = lines.pop() ?? ""; - for (const line of lines) this.consumeLine(line); + for (const line of lines) { + this.consumeLine(line); + if (this._error) { + this.buffered = ""; + return this.result; + } + } } // Some providers return one raw JSON object without an SSE newline. Parse diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index 9c5ca8b39..5093298dd 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -139,6 +139,24 @@ describe("discovery validity", () => { } expect(result).toMatchObject({ ready: false, error: true, limitExceeded: true }); }); + + it("stops parsing current and future events after the event limit", () => { + const parser = new DiscoveryValidityParser("anthropic"); + const metadataEvents = Array.from( + { length: DISCOVERY_EVENT_MAX_COUNT + 1 }, + () => 'data: {"type":"ping"}\n' + ).join(""); + + const limited = parser.push(`${metadataEvents}data: {"type":"message_stop"}\n`); + expect(limited).toMatchObject({ + ready: false, + terminal: false, + error: true, + limitExceeded: true, + }); + + expect(parser.push('data: {"type":"message_stop"}\n')).toEqual(limited); + }); }); function parserForOpenAIChatToolCall(): DiscoveryValidityParser { From ef6b833f12d490fa77046ffbfec20d6e5fc16b89 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 20:43:27 -0400 Subject: [PATCH 17/23] fix(discovery): close final lifecycle gaps --- src/app/v1/_lib/proxy/forwarder.ts | 66 ++++--- src/app/v1/_lib/proxy/response-handler.ts | 55 +++++- .../integration/proxy-hedge-lifecycle.test.ts | 12 ++ .../proxy-forwarder-hedge-first-byte.test.ts | 172 ++++++++++++++++++ ...esponse-handler-client-abort-drain.test.ts | 65 +++++++ ...handler-endpoint-circuit-isolation.test.ts | 123 ++++++++++++- 6 files changed, 462 insertions(+), 31 deletions(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 1f1489c00..ffaf62d68 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5167,7 +5167,7 @@ export class ProxyForwarder { cancelLosers(null, options.cancellationKind ?? "discovery_loser"); if (!options.preserveBinding && bindingWriteAllowed && session.isSessionBindingAllowed()) { if (bindingSnapshot.providerId != null) { - void SessionManager.clearVersionedSessionProvider( + await SessionManager.clearVersionedSessionProvider( bindingSnapshot, bindingSnapshot.providerId, 0 @@ -5362,7 +5362,7 @@ export class ProxyForwarder { retainOnSuccess: boolean; }; } - ): Promise => { + ): Promise => { const transferredProviderSessionRef = options?.providerSessionRefTransfer?.owned === true; let providerSessionRefTracked = transferredProviderSessionRef; let providerSessionRefRetainOnSuccess = @@ -5376,7 +5376,7 @@ export class ProxyForwarder { }; if (settled || committed || launched.has(provider.id)) { rollbackLaunch(); - return; + return false; } launched.add(provider.id); if (!transferredProviderSessionRef && provider.id === initialProvider.id) { @@ -5404,7 +5404,7 @@ export class ProxyForwarder { } if (settled || committed) { rollbackLaunch(); - return; + return false; } let endpoint: Awaited>; try { @@ -5415,7 +5415,7 @@ export class ProxyForwarder { } if (settled || committed) { rollbackLaunch(); - return; + return false; } let attemptSession: ProxySession; try { @@ -5431,7 +5431,7 @@ export class ProxyForwarder { } if (settled || committed) { rollbackLaunch(); - return; + return false; } const controller = new AbortController(); const id = `${provider.id}:${sequence + 1}`; @@ -5505,7 +5505,7 @@ export class ProxyForwarder { }); if (!registered || settled || committed) { rollbackLaunch(); - return; + return false; } attempts.set(id, attempt); discoveryMetrics.attemptStarted({ @@ -5821,6 +5821,7 @@ export class ProxyForwarder { error, }); }); + return true; }; const launchNextRound = async (slots: number, coordinatorAlreadyAdvanced = false) => { @@ -5835,21 +5836,42 @@ export class ProxyForwarder { currentRound = nextRound.round; } if (currentRound > maxRounds || slots <= 0) return; - const candidates = await ProxyProviderResolver.pickDiscoveryProviders( - session, - slots, - Array.from(launched) - ); - if (candidates.length === 0) { - noMoreCandidates = true; - } - for (const candidate of candidates) { - try { - await launch(candidate, "normal"); - } catch (error) { - lastError = error instanceof Error ? error : new Error(String(error)); - // Launch setup failures do not establish pool exhaustion. - noMoreCandidates = false; + let remainingSlots = slots; + while (remainingSlots > 0 && !settled && !committed) { + const exclusionCountBeforeSelection = launched.size; + const candidates = await ProxyProviderResolver.pickDiscoveryProviders( + session, + remainingSlots, + Array.from(launched) + ); + if (candidates.length === 0) { + noMoreCandidates = true; + break; + } + + noMoreCandidates = false; + let registeredInBatch = 0; + for (const candidate of candidates) { + try { + if (await launch(candidate, "normal")) { + registeredInBatch += 1; + remainingSlots -= 1; + } + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + // This Provider is already excluded by launch(). Continue filling + // the same round from the remaining candidate pool. + noMoreCandidates = false; + } + if (remainingSlots <= 0 || settled || committed) break; + } + + // A selector that ignores exclusions must not create an unbounded + // setup loop. Normal selectors always advance `launched`, including + // attempts that fail before transport registration. + if (registeredInBatch === 0 && launched.size === exclusionCountBeforeSelection) { + noMoreCandidates = true; + break; } } const hasPendingAttempt = Array.from(attempts.values()).some((attempt) => attempt.pending); diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index a9ca3a7de..064483a37 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -114,7 +114,8 @@ function isSessionBindingMutationAllowed(session: ProxySession): boolean { } function startDiscoveryLeaseLifecycle(session: ProxySession): DiscoveryLeaseLifecycle { - const lease = peekDeferredStreamingFinalization(session)?.discoveryLease; + const deferred = peekDeferredStreamingFinalization(session); + const lease = deferred?.discoveryLease; if (!lease) { return { active: false, @@ -127,6 +128,12 @@ function startDiscoveryLeaseLifecycle(session: ProxySession): DiscoveryLeaseLife let renewalInFlight: Promise | null = null; let releasePromise: Promise | null = null; let ownershipState: "unknown" | "owned" | "lost" = "unknown"; + const bindingSnapshot = + (deferred?.bindingIntent === "create" || deferred?.bindingIntent === "renew") && + deferred.bindingSnapshot?.sessionId === lease.sessionId && + deferred.bindingSnapshot.keyId === lease.keyId + ? deferred.bindingSnapshot + : null; const stopRenewal = () => { if (renewalTimer) { @@ -157,6 +164,25 @@ function startDiscoveryLeaseLifecycle(session: ProxySession): DiscoveryLeaseLife }); return false; } + + if (bindingSnapshot) { + const touched = await SessionManager.touchVersionedSessionBinding(bindingSnapshot); + if ( + touched.status !== "ok" || + touched.snapshot.generation !== bindingSnapshot.generation || + touched.snapshot.providerId !== bindingSnapshot.providerId + ) { + ownershipState = "lost"; + stopRenewal(); + logger.warn("[ResponseHandler] Discovery binding heartbeat stopped", { + sessionId: lease.sessionId, + keyId: lease.keyId, + status: touched.status, + reason: "reason" in touched ? touched.reason : "snapshot_mismatch", + }); + return false; + } + } ownershipState = "owned"; return true; })() @@ -181,7 +207,14 @@ function startDiscoveryLeaseLifecycle(session: ProxySession): DiscoveryLeaseLife // the downstream response. Renew immediately so an expired/lost token is // observed before any terminal Session binding mutation is attempted. const handoffRenewal = renew(); - const renewalIntervalMs = Math.max(250, Math.floor((lease.ttlSeconds * 1000) / 3)); + const leaseRenewalIntervalMs = Math.floor((lease.ttlSeconds * 1000) / 3); + const bindingRefreshIntervalMs = bindingSnapshot + ? SessionManager.getVersionedSessionBindingRefreshIntervalMs() + : Number.POSITIVE_INFINITY; + const renewalIntervalMs = Math.max( + 250, + Math.min(leaseRenewalIntervalMs, bindingRefreshIntervalMs) + ); renewalTimer = setInterval(() => { void renew(); }, renewalIntervalMs); @@ -1468,6 +1501,10 @@ function finalizeDeferredStreamingFinalizationIfNeeded( primaryDiscoveryBindingSettled = true; resolvePrimaryDiscoveryBinding?.(updated); }; + const finalizeFailedDiscoveryBinding = async () => { + settlePrimaryDiscoveryBinding(false); + await finalizeProviderSessionRef(); + }; const confirmAuxiliarySessionBinding = async () => allowAuxiliarySessionBinding && (await primaryDiscoveryBinding); @@ -1569,12 +1606,14 @@ function finalizeDeferredStreamingFinalizationIfNeeded( // - 不在这里更新熔断/绑定（meta 缺失意味着 Forwarder 没有启用延迟结算；provider 缺失意味着无法归因）。 if (!meta || !provider) { const commitSideEffects = - shouldClearSessionBindingOnFailure || meta?.providerSessionRefOwned === true + shouldClearSessionBindingOnFailure || + meta?.providerSessionRefOwned === true || + hasDiscoveryBindingIntent ? async () => { try { if (shouldClearSessionBindingOnFailure) await clearSessionBinding(); } finally { - await finalizeProviderSessionRef(); + await finalizeFailedDiscoveryBinding(); } } : undefined; @@ -1651,7 +1690,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( // so only the Provider circuit is updated here. } } finally { - await finalizeProviderSessionRef(); + await finalizeFailedDiscoveryBinding(); } }; @@ -1695,7 +1734,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( } } } finally { - await finalizeProviderSessionRef(); + await finalizeFailedDiscoveryBinding(); } }; @@ -1766,7 +1805,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( } } } finally { - await finalizeProviderSessionRef(); + await finalizeFailedDiscoveryBinding(); } }; @@ -1830,7 +1869,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( } } } finally { - await finalizeProviderSessionRef(); + await finalizeFailedDiscoveryBinding(); } }; diff --git a/tests/integration/proxy-hedge-lifecycle.test.ts b/tests/integration/proxy-hedge-lifecycle.test.ts index 0f2b6b3b0..7707cabb6 100644 --- a/tests/integration/proxy-hedge-lifecycle.test.ts +++ b/tests/integration/proxy-hedge-lifecycle.test.ts @@ -8,6 +8,7 @@ import { ProxyResponseHandler } from "@/app/v1/_lib/proxy/response-handler"; import { type MessageContext, ProxySession } from "@/app/v1/_lib/proxy/session"; import { DbPoolAdmissionError } from "@/drizzle/admitted-client"; import { getGlobalAgentPool, resetGlobalAgentPool } from "@/lib/proxy-agent"; +import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import type { Key } from "@/types/key"; import type { Provider } from "@/types/provider"; import type { User } from "@/types/user"; @@ -125,6 +126,17 @@ vi.mock("@/lib/session-manager", async (importOriginal) => { static override async renewSessionDiscoveryLease() { return state.renewDiscoveryLease(); } + static override getVersionedSessionBindingRefreshIntervalMs() { + return 100_000; + } + static override async touchVersionedSessionBinding(snapshot: SessionBindingSnapshot) { + return { + status: "ok" as const, + source: "touched" as const, + snapshot, + legacyFallbackAllowed: false as const, + }; + } static override async releaseSessionDiscoveryLease() { return state.releaseDiscoveryLease(); } diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index 7bbcc1856..bec72c795 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -2848,6 +2848,101 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { } }); + test("Sticky Discovery refills a slot immediately after candidate setup fails", async () => { + vi.useFakeTimers(); + let endpointResolver: ReturnType | null = null; + try { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const setupFailure = createProvider({ id: 2, name: "setup-failure", priority: 1 }); + const replacement = createProvider({ id: 3, name: "replacement", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 30 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 30, + providerId: sticky.id, + generation: "setup-refill-generation", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 10, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders + .mockResolvedValueOnce([setupFailure]) + .mockResolvedValueOnce([replacement]); + + endpointResolver = vi.spyOn( + ProxyForwarder as unknown as { + resolveStreamingHedgeEndpoint: ( + session: ProxySession, + provider: Provider + ) => Promise<{ endpointId: number | null; baseUrl: string; endpointUrl: string }>; + }, + "resolveStreamingHedgeEndpoint" + ); + endpointResolver.mockImplementation(async (_attemptSession, provider) => { + if (provider.id === setupFailure.id) throw new Error("candidate endpoint setup failed"); + return { + endpointId: null, + baseUrl: provider.url, + endpointUrl: provider.url, + }; + }); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementation(async (attemptSession) => { + if ((attemptSession as ProxySession).provider?.id === sticky.id) { + return new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }); + } + return new Response( + 'data: {"type":"content_block_delta","delta":{"text":"replacement"}}\n\n', + { headers: { "content-type": "text/event-stream" } } + ); + }); + + const responsePromise = ProxyForwarder.send(session); + await vi.advanceTimersByTimeAsync(10); + await vi.advanceTimersByTimeAsync(0); + const response = await responsePromise; + + expect(await response.text()).toContain('"replacement"'); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(2); + expect(mocks.pickDiscoveryProviders).toHaveBeenNthCalledWith( + 2, + expect.anything(), + 1, + expect.arrayContaining([sticky.id, setupFailure.id]) + ); + expect(doForward).toHaveBeenCalledTimes(2); + expect(session.provider?.id).toBe(replacement.id); + } finally { + endpointResolver?.mockRestore(); + vi.useRealTimers(); + } + }); + test("an explicit Sticky failure starts Discovery round one at full concurrency", async () => { const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); const normal = createProvider({ id: 2, name: "normal", priority: 1 }); @@ -2976,6 +3071,83 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { expect(mocks.releaseSessionDiscoveryLease).toHaveBeenCalledTimes(1); }); + test("Discovery clears terminal binding state before releasing its lease", async () => { + const provider = createProvider({ id: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 29 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(provider); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 29, + providerId: provider.id, + generation: "terminal-clear-generation", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 100, + stickySlaMs: 100, + racingTotalTimeoutMs: 500, + }); + const order: string[] = []; + const clear = Promise.withResolvers<{ + status: "ok"; + legacyFallbackAllowed: false; + source: "cleared"; + snapshot: { + sessionId: string; + keyId: number; + providerId: null; + generation: string; + }; + }>(); + mocks.clearVersionedSessionProvider.mockImplementationOnce(async () => { + order.push("clear-start"); + const result = await clear.promise; + order.push("clear-end"); + return result; + }); + mocks.releaseSessionDiscoveryLease.mockImplementationOnce(async () => { + order.push("lease-release"); + return { status: "released", legacyFallbackAllowed: false }; + }); + vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ).mockRejectedValueOnce(new Error("terminal upstream failure")); + + const observed = ProxyForwarder.send(session).catch((error) => error); + for (let index = 0; index < 10 && order.length === 0; index++) { + await new Promise((resolve) => setImmediate(resolve)); + } + expect(order).toEqual(["clear-start"]); + expect(mocks.releaseSessionDiscoveryLease).not.toHaveBeenCalled(); + + clear.resolve({ + status: "ok", + legacyFallbackAllowed: false, + source: "cleared", + snapshot: { + sessionId: session.sessionId!, + keyId: 29, + providerId: null, + generation: "terminal-cleared-generation", + }, + }); + expect(await observed).toBeInstanceOf(Error); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledOnce(); + expect(mocks.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + expect(order).toEqual(["clear-start", "clear-end", "lease-release"]); + }); + test("Discovery total deadline is not blocked by a stalled candidate selector", async () => { vi.useFakeTimers(); try { diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index ded21ba4b..d52cb370e 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -10,6 +10,7 @@ import { AsyncTaskManager, shutdownAllAsyncTasks } from "@/lib/async-task-manage import { recordFailure } from "@/lib/circuit-breaker"; import { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; import { RateLimitService } from "@/lib/rate-limit"; +import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import { SessionManager } from "@/lib/session-manager"; import { updateMessageRequestCostWithBreakdown, @@ -118,6 +119,7 @@ vi.mock("@/lib/session-manager", () => ({ clearVersionedSessionProvider: vi.fn(), compareAndSetSessionProvider: vi.fn(), getSessionBindingSnapshot: vi.fn(), + getVersionedSessionBindingRefreshIntervalMs: vi.fn(() => 100_000), renewSessionDiscoveryLease: vi.fn(async () => ({ status: "renewed", legacyFallbackAllowed: false, @@ -126,6 +128,12 @@ vi.mock("@/lib/session-manager", () => ({ status: "released", legacyFallbackAllowed: false, })), + touchVersionedSessionBinding: vi.fn(async (snapshot: SessionBindingSnapshot) => ({ + status: "ok", + source: "touched", + snapshot, + legacyFallbackAllowed: false, + })), extractCodexPromptCacheKey: vi.fn(), storeSessionResponse: vi.fn(async () => undefined), storeSessionRequestPhaseSnapshot: vi.fn(), @@ -3473,6 +3481,63 @@ describe("ProxyResponseHandler stream client abort finalization", () => { expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); }); + it("settles a failed Discovery binding before rejecting its auxiliary Codex cache binding", async () => { + vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce( + "incomplete-discovery-cache-key" + ); + const session = createSession(new AbortController().signal); + session.sessionId = "stream-discovery-cache-incomplete"; + session.recordProviderSessionRef(1); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "stream-discovery-cache-incomplete", + keyId: 2, + providerId: null, + generation: "incomplete-discovery-generation", + }, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: "stream-discovery-cache-incomplete", + keyId: 2, + ownerToken: "incomplete-discovery-owner", + ttlSeconds: 30, + }, + providerSessionRefOwned: true, + }); + const incomplete = new Response( + `event: response.output_text.done\ndata: ${JSON.stringify({ + type: "response.output_text.done", + text: "partial", + })}\n\n`, + { status: 200, headers: { "content-type": "text/event-stream" } } + ); + + const downstream = await ProxyResponseHandler.dispatch(session, incomplete); + await downstream.text(); + for (let index = 0; index < 10 && !getRegisteredTask("post-terminal-side-effects"); index++) { + await new Promise((resolve) => setImmediate(resolve)); + } + const sideEffects = getRegisteredTask("post-terminal-side-effects"); + expect(sideEffects).toBeDefined(); + await expectTaskToResolveWithoutWaiting(sideEffects as Promise); + await drainAsyncTasks(); + + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.updateSessionWithCodexCacheKey).not.toHaveBeenCalled(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + }); + it("does not publish a stream Codex cache binding for a final non-2xx outcome", async () => { vi.mocked(SessionManager.extractCodexPromptCacheKey).mockReturnValueOnce("stream-cache-key-2"); const session = createSession(new AbortController().signal); diff --git a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts index 038c9378c..c6c664583 100644 --- a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts +++ b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts @@ -80,8 +80,10 @@ vi.mock("@/lib/session-manager", () => ({ clearVersionedSessionProvider: vi.fn(), compareAndSetSessionProvider: vi.fn(), getSessionBindingSnapshot: vi.fn(), + getVersionedSessionBindingRefreshIntervalMs: vi.fn(), renewSessionDiscoveryLease: vi.fn(), releaseSessionDiscoveryLease: vi.fn(), + touchVersionedSessionBinding: vi.fn(), extractCodexPromptCacheKey: vi.fn(), updateSessionBindingSmart: vi.fn(), updateSessionProvider: vi.fn(), @@ -397,7 +399,7 @@ function createControllableSuccessStreamResponse(): { }), complete: () => { streamController.enqueue( - encoder.encode(`data: ${JSON.stringify({ type: "message_stop" })}\n\n`) + encoder.encode(`event: message_stop\ndata: ${JSON.stringify({ type: "message_stop" })}\n\n`) ); streamController.close(); }, @@ -468,10 +470,17 @@ function setupCommonMocks() { }, legacyFallbackAllowed: false, }); + vi.mocked(SessionManager.getVersionedSessionBindingRefreshIntervalMs).mockReturnValue(100_000); vi.mocked(SessionManager.renewSessionDiscoveryLease).mockResolvedValue({ status: "renewed", legacyFallbackAllowed: false, }); + vi.mocked(SessionManager.touchVersionedSessionBinding).mockImplementation(async (binding) => ({ + status: "ok", + source: "touched", + snapshot: binding, + legacyFallbackAllowed: false, + })); vi.mocked(SessionManager.releaseSessionDiscoveryLease).mockResolvedValue({ status: "released", legacyFallbackAllowed: false, @@ -1165,6 +1174,118 @@ describe("Endpoint circuit breaker isolation", () => { } }); + it("touches the captured binding while a Discovery winner remains open", async () => { + vi.useFakeTimers(); + try { + const session = createSession(); + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: null, + generation: "long-stream-generation", + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: snapshot, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: "fake-session", + keyId: 456, + ownerToken: "long-stream-owner", + ttlSeconds: 3_600, + }, + }); + vi.mocked(SessionManager.getVersionedSessionBindingRefreshIntervalMs).mockReturnValue(1_000); + const controlled = createControllableSuccessStreamResponse(); + + const clientResponse = await ProxyResponseHandler.dispatch(session, controlled.response); + await vi.advanceTimersByTimeAsync(3_000); + expect(SessionManager.touchVersionedSessionBinding).toHaveBeenCalledTimes(4); + expect(SessionManager.touchVersionedSessionBinding).toHaveBeenLastCalledWith(snapshot); + + controlled.complete(); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.compareAndSetSessionProvider).toHaveBeenCalledWith(snapshot, 1); + expect(SessionManager.getSessionBindingSnapshot).not.toHaveBeenCalled(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + } finally { + vi.useRealTimers(); + } + }); + + it("revokes Sticky writes when a binding heartbeat loses its generation", async () => { + vi.useFakeTimers(); + try { + const session = createSession(); + const snapshot = { + sessionId: "fake-session", + keyId: 456, + providerId: 1, + generation: "generation-before-conflict", + } as const; + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "test-provider", + providerPriority: 10, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.com", + upstreamStatusCode: 200, + bindingIntent: "renew", + bindingSnapshot: snapshot, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: "fake-session", + keyId: 456, + ownerToken: "conflicted-stream-owner", + ttlSeconds: 3_600, + }, + }); + vi.mocked(SessionManager.getVersionedSessionBindingRefreshIntervalMs).mockReturnValue(1_000); + vi.mocked(SessionManager.touchVersionedSessionBinding) + .mockResolvedValueOnce({ + status: "ok", + source: "touched", + snapshot, + legacyFallbackAllowed: false, + }) + .mockResolvedValueOnce({ + status: "conflict", + reason: "generation_mismatch", + legacyFallbackAllowed: false, + }); + const controlled = createControllableSuccessStreamResponse(); + + const clientResponse = await ProxyResponseHandler.dispatch(session, controlled.response); + await vi.advanceTimersByTimeAsync(1_000); + controlled.complete(); + await clientResponse.text(); + await drainAsyncTasks(); + + expect(SessionManager.touchVersionedSessionBinding).toHaveBeenCalledTimes(2); + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + expect(SessionManager.getSessionBindingSnapshot).not.toHaveBeenCalled(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + } finally { + vi.useRealTimers(); + } + }); + it("does not delay downstream delivery while the lease handoff renewal is pending", async () => { const handoffRenewal = Promise.withResolvers<{ status: "renewed"; From d915e086596d9594da110f9b2476294b52635edc Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 21:01:41 -0400 Subject: [PATCH 18/23] fix(discovery): refill setup failures within current round --- src/app/v1/_lib/proxy/forwarder.ts | 146 +++++++++--------- .../proxy-forwarder-hedge-first-byte.test.ts | 82 ++++++++++ 2 files changed, 152 insertions(+), 76 deletions(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index ffaf62d68..3e88aaee0 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5284,19 +5284,6 @@ export class ProxyForwarder { }); }; - const chooseCandidate = async (): Promise => { - const candidates = await ProxyProviderResolver.pickDiscoveryProviders( - session, - 1, - Array.from(launched) - ); - if (!candidates[0]) { - noMoreCandidates = true; - return null; - } - return candidates[0]; - }; - const clearCapturedStickyBinding = async (cooldownTtlSeconds: number): Promise => { if ( !bindingWriteAllowed || @@ -5792,19 +5779,7 @@ export class ProxyForwarder { } } if (!actionOwnsNextStep && !committed && !settled) { - const replacement = await chooseCandidate(); - if (replacement) { - try { - await launch(replacement, "normal"); - } catch (launchError) { - lastError = - launchError instanceof Error ? launchError : new Error(String(launchError)); - // A single launch failure does not prove that the remaining - // candidate pool is exhausted. The current round boundary can - // still advance or retry selection from the remaining pool. - noMoreCandidates = false; - } - } + await refillCurrentRoundSlots(1); } if ( Array.from(attempts.values()).every((candidate) => !candidate.pending) && @@ -5824,6 +5799,73 @@ export class ProxyForwarder { return true; }; + const fillDiscoverySlots = async (slots: number): Promise => { + let remainingSlots = slots; + while (remainingSlots > 0 && !settled && !committed) { + const exclusionCountBeforeSelection = launched.size; + const candidates = await ProxyProviderResolver.pickDiscoveryProviders( + session, + remainingSlots, + Array.from(launched) + ); + if (candidates.length === 0) { + noMoreCandidates = true; + break; + } + + noMoreCandidates = false; + let registeredInBatch = 0; + for (const candidate of candidates) { + try { + if (await launch(candidate, "normal")) { + registeredInBatch += 1; + remainingSlots -= 1; + } + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + // launch() excludes setup failures before throwing, so keep filling + // this round from the remaining candidate pool. + noMoreCandidates = false; + } + if (remainingSlots <= 0 || settled || committed) break; + } + + // A selector that ignores exclusions must not create an unbounded setup + // loop. Normal selectors advance `launched` even before transport setup. + if (registeredInBatch === 0 && launched.size === exclusionCountBeforeSelection) { + noMoreCandidates = true; + break; + } + } + }; + + const finishRoundLaunchBatch = async (): Promise => { + roundLaunchesInProgress = Math.max(0, roundLaunchesInProgress - 1); + if (roundLaunchesInProgress !== 0) return; + + notifyRoundLaunchIdle(); + fallbackPromotionBlocked = false; + const readyFallback = Array.from(attempts.values()).find( + (attempt) => attempt.pending && attempt.ready && attempt.kind === "fallback" + ); + if (readyFallback && !committed && !settled) { + const action = coordinator.markReady(readyFallback.id); + if (action.type === "promote_fallback") await commit(readyFallback); + } + }; + + async function refillCurrentRoundSlots(slots: number): Promise { + if (slots <= 0 || settled || committed) return; + roundLaunchesInProgress += 1; + try { + // Deliberately preserve the existing round timer. An explicit failure + // releases capacity but must not grant the replacement a fresh SLA. + await fillDiscoverySlots(slots); + } finally { + await finishRoundLaunchBatch(); + } + } + const launchNextRound = async (slots: number, coordinatorAlreadyAdvanced = false) => { if (settled || committed) return; roundLaunchesInProgress += 1; @@ -5836,44 +5878,7 @@ export class ProxyForwarder { currentRound = nextRound.round; } if (currentRound > maxRounds || slots <= 0) return; - let remainingSlots = slots; - while (remainingSlots > 0 && !settled && !committed) { - const exclusionCountBeforeSelection = launched.size; - const candidates = await ProxyProviderResolver.pickDiscoveryProviders( - session, - remainingSlots, - Array.from(launched) - ); - if (candidates.length === 0) { - noMoreCandidates = true; - break; - } - - noMoreCandidates = false; - let registeredInBatch = 0; - for (const candidate of candidates) { - try { - if (await launch(candidate, "normal")) { - registeredInBatch += 1; - remainingSlots -= 1; - } - } catch (error) { - lastError = error instanceof Error ? error : new Error(String(error)); - // This Provider is already excluded by launch(). Continue filling - // the same round from the remaining candidate pool. - noMoreCandidates = false; - } - if (remainingSlots <= 0 || settled || committed) break; - } - - // A selector that ignores exclusions must not create an unbounded - // setup loop. Normal selectors always advance `launched`, including - // attempts that fail before transport registration. - if (registeredInBatch === 0 && launched.size === exclusionCountBeforeSelection) { - noMoreCandidates = true; - break; - } - } + await fillDiscoverySlots(slots); const hasPendingAttempt = Array.from(attempts.values()).some((attempt) => attempt.pending); if (!hasPendingAttempt) { await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); @@ -5883,18 +5888,7 @@ export class ProxyForwarder { scheduleRoundBoundary(discoverySlaMs); } } finally { - roundLaunchesInProgress = Math.max(0, roundLaunchesInProgress - 1); - if (roundLaunchesInProgress === 0) { - notifyRoundLaunchIdle(); - fallbackPromotionBlocked = false; - const readyFallback = Array.from(attempts.values()).find( - (attempt) => attempt.pending && attempt.ready && attempt.kind === "fallback" - ); - if (readyFallback && !committed && !settled) { - const action = coordinator.markReady(readyFallback.id); - if (action.type === "promote_fallback") await commit(readyFallback); - } - } + await finishRoundLaunchBatch(); } }; diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index bec72c795..9d8cd16ab 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -2943,6 +2943,88 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { } }); + test("Discovery keeps refilling the current round when an error replacement fails setup", async () => { + const initialFailure = createProvider({ id: 1, name: "initial-failure", priority: 1 }); + const pending = createProvider({ id: 2, name: "pending", priority: 1 }); + const setupFailure = createProvider({ id: 3, name: "setup-failure", priority: 1 }); + const healthy = createProvider({ id: 4, name: "healthy", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 31 }, + apiKey: null, + } as typeof session.authState; + session.setProvider(initialFailure); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 50, + racingTotalTimeoutMs: 200, + stickyTimeoutCooldownMs: 300_000, + }); + mocks.pickDiscoveryProviders + .mockResolvedValueOnce([pending]) + .mockResolvedValueOnce([setupFailure]) + .mockResolvedValueOnce([healthy]); + + const endpointResolver = vi.spyOn( + ProxyForwarder as unknown as { + resolveStreamingHedgeEndpoint: ( + session: ProxySession, + provider: Provider + ) => Promise<{ endpointId: number | null; baseUrl: string; endpointUrl: string }>; + }, + "resolveStreamingHedgeEndpoint" + ); + endpointResolver.mockImplementation(async (_attemptSession, provider) => { + if (provider.id === setupFailure.id) throw new Error("replacement setup failed"); + return { + endpointId: null, + baseUrl: provider.url, + endpointUrl: provider.url, + }; + }); + + const doForward = vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ); + doForward.mockImplementation(async (attemptSession) => { + const providerId = (attemptSession as ProxySession).provider?.id; + if (providerId === initialFailure.id) throw new Error("initial Provider failed"); + if (providerId === healthy.id) { + return new Response('data: {"type":"content_block_delta","delta":{"text":"healthy"}}\n\n', { + headers: { "content-type": "text/event-stream" }, + }); + } + return new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }); + }); + + try { + const response = await ProxyForwarder.send(session); + + expect(await response.text()).toContain('"healthy"'); + expect(mocks.pickDiscoveryProviders).toHaveBeenCalledTimes(3); + expect(mocks.pickDiscoveryProviders).toHaveBeenNthCalledWith( + 3, + expect.anything(), + 1, + expect.arrayContaining([initialFailure.id, pending.id, setupFailure.id]) + ); + expect(doForward).toHaveBeenCalledTimes(3); + expect(session.provider?.id).toBe(healthy.id); + } finally { + endpointResolver.mockRestore(); + } + }); + test("an explicit Sticky failure starts Discovery round one at full concurrency", async () => { const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); const normal = createProvider({ id: 2, name: "normal", priority: 1 }); From 74bb97eaba64684791174c9edf93146099542362 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 21:57:33 -0400 Subject: [PATCH 19/23] fix(discovery): fence readiness and sticky cooldown --- src/app/v1/_lib/proxy/discovery-validity.ts | 39 +++- src/app/v1/_lib/proxy/forwarder.ts | 14 +- tests/unit/proxy/discovery-validity.test.ts | 43 ++++ .../proxy-forwarder-hedge-first-byte.test.ts | 210 ++++++++++++++++++ 4 files changed, 304 insertions(+), 2 deletions(-) diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts index 059bde293..1f8931d7d 100644 --- a/src/app/v1/_lib/proxy/discovery-validity.ts +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -49,6 +49,43 @@ function hasAnthropicContentBlock(value: unknown): boolean { return block.type === "text" ? hasContent(block.text) : true; } +function hasOpenAIResponsesOutputItem(value: unknown): boolean { + if (!value || typeof value !== "object" || Array.isArray(value)) return false; + const item = value as Record; + if (typeof item.type !== "string") return false; + + switch (item.type) { + case "message": + return hasContent(item.content); + case "reasoning": + return hasContent(item.summary) || hasContent(item.content); + case "function_call": + case "mcp_call": + return hasContent(item.name) || hasContent(item.arguments); + case "custom_tool_call": + return hasContent(item.name) || hasContent(item.input); + case "computer_call": + case "web_search_call": + case "file_search_call": + case "code_interpreter_call": + case "local_shell_call": + case "shell_call": + case "apply_patch_call": + return [ + item.action, + item.arguments, + item.input, + item.queries, + item.query, + item.code, + item.command, + item.operation, + ].some(hasContent); + default: + return false; + } +} + function classifyJson(value: unknown, protocol: DiscoveryProtocol): DiscoveryValidity { if (!value || typeof value !== "object") return { ready: false, terminal: false, error: true }; const object = value as Record; @@ -78,7 +115,7 @@ function classifyJson(value: unknown, protocol: DiscoveryProtocol): DiscoveryVal ready: (object.type === "response.output_text.delta" && hasContent(object.delta)) || (object.type === "response.function_call_arguments.delta" && hasContent(object.delta)) || - (object.type === "response.output_item.added" && hasContent(object.item)), + (object.type === "response.output_item.added" && hasOpenAIResponsesOutputItem(object.item)), terminal: false, error: false, }; diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 7f73b8a0b..8016b0502 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5063,6 +5063,7 @@ export class ProxyForwarder { const roundLaunchIdleWaiters = new Set<() => void>(); let fallbackPromotionBlocked = false; let stickyTimeoutWaveReservation: { fallbackAttemptId: string } | null = null; + let stickyTimeoutCooldownPromise: Promise | null = null; const hasSticky = session.shouldReuseProvider() && !!session.sessionId && @@ -5186,6 +5187,10 @@ export class ProxyForwarder { if (roundTimer) clearTimeout(roundTimer); if (stickyTimer) clearTimeout(stickyTimer); cancelLosers(null, options.cancellationKind ?? "discovery_loser"); + // Sticky timeout owns its cooldown mutation. A terminal deadline/error or + // client abort may race that Redis CAS, but must not issue a second + // zero-cooldown clear against the same captured generation. + if (stickyTimeoutCooldownPromise) await stickyTimeoutCooldownPromise; if (!options.preserveBinding && bindingWriteAllowed && session.isSessionBindingAllowed()) { if (bindingSnapshot.providerId != null) { await SessionManager.clearVersionedSessionProvider( @@ -5338,6 +5343,13 @@ export class ProxyForwarder { } }; + const ensureStickyTimeoutCooldown = (cooldownTtlSeconds: number): Promise => { + if (!stickyTimeoutCooldownPromise) { + stickyTimeoutCooldownPromise = clearCapturedStickyBinding(cooldownTtlSeconds); + } + return stickyTimeoutCooldownPromise; + }; + const scheduleRoundBoundary = (delayMs: number) => { clearRoundTimer(); const epoch = coordinator.epochs; @@ -6019,7 +6031,7 @@ export class ProxyForwarder { fallbackPromotionBlocked = true; stickyTimeoutWaveReservation = { fallbackAttemptId: sticky.id }; if (bindingSnapshot && bindingSnapshot.providerId === initialProvider.id) { - void clearCapturedStickyBinding( + void ensureStickyTimeoutCooldown( Math.ceil((settings.stickyTimeoutCooldownMs ?? 300_000) / 1000) ).finally(() => { void launchReservedStickyTimeoutWave(Math.max(0, concurrency - 1)).catch( diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index 5093298dd..ed2690d6c 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -63,6 +63,49 @@ describe("discovery validity", () => { ).toBe(true); }); + it("holds Responses output-item metadata until a text delta is deliverable", () => { + const parser = new DiscoveryValidityParser("openai-responses"); + + expect( + parser.push( + 'data: {"type":"response.output_item.added","item":{"id":"msg_1","type":"message","status":"in_progress","content":[]}}\n\n' + ) + ).toEqual({ ready: false, terminal: false, error: false }); + expect(parser.push('data: {"type":"response.output_text.delta","delta":"hello"}\n\n')).toEqual({ + ready: true, + terminal: false, + error: false, + }); + }); + + it("does not let Responses output-item metadata mask a later error", () => { + const parser = new DiscoveryValidityParser("openai-responses"); + + expect( + parser.push( + 'data: {"type":"response.output_item.added","item":{"id":"msg_1","type":"message","status":"in_progress"}}\n\n' + ) + ).toMatchObject({ ready: false, error: false }); + expect( + parser.push('data: {"type":"response.failed","error":{"message":"upstream failed"}}\n\n') + ).toEqual({ ready: false, terminal: true, error: true }); + }); + + it("accepts only an explicit non-empty Responses tool payload", () => { + expect( + classifyDiscoveryChunk( + 'data: {"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","status":"in_progress"}}\n', + "openai-responses" + ).ready + ).toBe(false); + expect( + classifyDiscoveryChunk( + 'data: {"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","name":"lookup","arguments":"{}"}}\n', + "openai-responses" + ).ready + ).toBe(true); + }); + it("consumes split SSE lines incrementally without waiting for the full stream", () => { const parser = new DiscoveryValidityParser("openai-chat"); expect(parser.push('data: {"choices":[{"delta":{"content":"hel')).toEqual({ diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index 225e382db..9f37194bc 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -2640,6 +2640,216 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { } }); + test("Sticky timeout cooldown completes once before a racing total deadline settles", async () => { + vi.useFakeTimers(); + try { + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const session = createSession(); + session.authState = { + success: true, + user: null, + key: { id: 32 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 32, + providerId: sticky.id, + generation: "g-sticky-cooldown-deadline", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 20, + stickySlaMs: 10, + racingTotalTimeoutMs: 30, + stickyTimeoutCooldownMs: 300_000, + }); + + const cooldownClear = Promise.withResolvers<{ + status: "ok"; + legacyFallbackAllowed: false; + source: "cleared"; + snapshot: { + sessionId: string; + keyId: number; + providerId: null; + generation: string; + }; + }>(); + const order: string[] = []; + mocks.clearVersionedSessionProvider.mockImplementationOnce( + async (_snapshot, _providerId, cooldownTtlSeconds) => { + order.push(`cooldown-start:${cooldownTtlSeconds}`); + const result = await cooldownClear.promise; + order.push("cooldown-end"); + return result; + } + ); + mocks.releaseSessionDiscoveryLease.mockImplementationOnce(async () => { + order.push("lease-release"); + return { status: "released", legacyFallbackAllowed: false }; + }); + vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ).mockResolvedValueOnce( + new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }) + ); + + let requestSettled = false; + const observed = ProxyForwarder.send(session).then( + (response) => { + requestSettled = true; + return response; + }, + (error) => { + requestSettled = true; + return error; + } + ); + + await vi.advanceTimersByTimeAsync(10); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledWith( + expect.objectContaining({ generation: "g-sticky-cooldown-deadline" }), + sticky.id, + 300 + ); + expect(order).toEqual(["cooldown-start:300"]); + + await vi.advanceTimersByTimeAsync(20); + expect(requestSettled).toBe(false); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledOnce(); + expect(mocks.releaseSessionDiscoveryLease).not.toHaveBeenCalled(); + + cooldownClear.resolve({ + status: "ok", + legacyFallbackAllowed: false, + source: "cleared", + snapshot: { + sessionId: session.sessionId!, + keyId: 32, + providerId: null, + generation: "g-sticky-cooldown-applied", + }, + }); + await vi.advanceTimersByTimeAsync(0); + + expect(await observed).toBeInstanceOf(UpstreamProxyError); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledOnce(); + expect(mocks.pickDiscoveryProviders).not.toHaveBeenCalled(); + expect(order).toEqual(["cooldown-start:300", "cooldown-end", "lease-release"]); + } finally { + vi.useRealTimers(); + } + }); + + test("client abort preserves an already-reserved Sticky timeout cooldown", async () => { + vi.useFakeTimers(); + try { + const clientAbort = new AbortController(); + const sticky = createProvider({ id: 1, name: "sticky", priority: 1 }); + const session = createSession(clientAbort.signal); + session.authState = { + success: true, + user: null, + key: { id: 33 }, + apiKey: null, + } as typeof session.authState; + session.request.message.messages = [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ]; + session.setProvider(sticky); + session.setSessionBindingSnapshot({ + sessionId: session.sessionId!, + keyId: 33, + providerId: sticky.id, + generation: "g-sticky-cooldown-abort", + }); + mocks.getCachedSystemSettings.mockResolvedValue({ + discoveryEnabled: true, + discoveryConcurrency: 2, + maxDiscoveryRounds: 1, + discoverySlaMs: 50, + stickySlaMs: 10, + racingTotalTimeoutMs: 100, + stickyTimeoutCooldownMs: 300_000, + }); + + const cooldownClear = Promise.withResolvers<{ + status: "ok"; + legacyFallbackAllowed: false; + source: "cleared"; + snapshot: { + sessionId: string; + keyId: number; + providerId: null; + generation: string; + }; + }>(); + mocks.clearVersionedSessionProvider.mockReturnValueOnce(cooldownClear.promise); + vi.spyOn( + ProxyForwarder as unknown as { + doForward: (...args: unknown[]) => Promise; + }, + "doForward" + ).mockResolvedValueOnce( + new Response(new ReadableStream(), { + headers: { "content-type": "text/event-stream" }, + }) + ); + + let requestSettled = false; + const observed = ProxyForwarder.send(session).catch((error) => { + requestSettled = true; + return error; + }); + await vi.advanceTimersByTimeAsync(10); + clientAbort.abort(new Error("client disconnected")); + await vi.advanceTimersByTimeAsync(0); + + expect(requestSettled).toBe(false); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledOnce(); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledWith( + expect.objectContaining({ generation: "g-sticky-cooldown-abort" }), + sticky.id, + 300 + ); + + cooldownClear.resolve({ + status: "ok", + legacyFallbackAllowed: false, + source: "cleared", + snapshot: { + sessionId: session.sessionId!, + keyId: 33, + providerId: null, + generation: "g-sticky-cooldown-abort-applied", + }, + }); + await vi.advanceTimersByTimeAsync(0); + + const error = await observed; + expect(error).toBeInstanceOf(UpstreamProxyError); + expect((error as UpstreamProxyError).statusCode).toBe(499); + expect(mocks.clearVersionedSessionProvider).toHaveBeenCalledOnce(); + expect(mocks.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + } finally { + vi.useRealTimers(); + } + }); + test("a ready-held Sticky fallback survives a stalled next-wave selector until the total deadline", async () => { vi.useFakeTimers(); try { From 1d23d6cbc195cbfb9230dc7a079ccb555c462443 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 22:33:31 -0400 Subject: [PATCH 20/23] fix(discovery): release non-SSE winner resources --- src/app/v1/_lib/proxy/response-handler.ts | 23 +++++++ ...esponse-handler-client-abort-drain.test.ts | 66 ++++++++++++++++++- 2 files changed, 88 insertions(+), 1 deletion(-) diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index f7608e16a..94f41e043 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -283,6 +283,25 @@ async function releaseOwnedProviderSessionRef( } } +async function finalizeNonStreamDiscoveryResources( + session: ProxySession, + lifecycle: DiscoveryLeaseLifecycle +): Promise { + const deferred = peekDeferredStreamingFinalization(session); + if (!deferred?.discoveryLease) return; + + // Non-SSE responses do not run the protocol completion finalizer, so they + // cannot safely create or renew Sticky. Consume the metadata and release + // both request-scoped resources after body processing instead of waiting + // for the Redis lease TTL. + const meta = consumeDeferredStreamingFinalization(session); + try { + await releaseOwnedProviderSessionRef(session, meta, false); + } finally { + await lifecycle.release(); + } +} + function startHedgeBindingHeartbeat(session: ProxySession): void { const deferred = peekDeferredStreamingFinalization(session); const authorityPromise = deferred?.hedgeBindingAuthorityPromise; @@ -2185,9 +2204,11 @@ export class ProxyResponseHandler { ): Promise { const messageContext = session.messageContext; const provider = session.provider; + const discoveryLeaseLifecycle = startDiscoveryLeaseLifecycle(session); if (!provider) { discardBeforeResponseBodySnapshot(session); releaseSessionAgent(session); + void finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); return response; } @@ -2458,6 +2479,7 @@ export class ProxyResponseHandler { } } finally { cleanupTaskAbortBinding(); + await finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); releaseSessionAgent(session); } }; @@ -3076,6 +3098,7 @@ export class ProxyResponseHandler { } finally { cleanupTaskAbortBinding(); cleanupClientAbortListener(); + await finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); releaseSessionAgent(session); } }; diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index d52cb370e..17a125f67 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -5,7 +5,10 @@ import { ProxyResponseHandler, } from "@/app/v1/_lib/proxy/response-handler"; import { ProxySession } from "@/app/v1/_lib/proxy/session"; -import { setDeferredStreamingFinalization } from "@/app/v1/_lib/proxy/stream-finalization"; +import { + peekDeferredStreamingFinalization, + setDeferredStreamingFinalization, +} from "@/app/v1/_lib/proxy/stream-finalization"; import { AsyncTaskManager, shutdownAllAsyncTasks } from "@/lib/async-task-manager"; import { recordFailure } from "@/lib/circuit-breaker"; import { emitProxyLangfuseTrace } from "@/lib/langfuse/emit-proxy-trace"; @@ -3187,6 +3190,67 @@ describe("ProxyResponseHandler stream client abort finalization", () => { ); }); + it("releases Discovery resources for a non-SSE Gemini winner", async () => { + const session = createSession(new AbortController().signal, { + providerType: "gemini", + originalFormat: "gemini", + endpoint: "/v1beta/models/gemini-2.0-flash:streamGenerateContent", + model: "gemini-2.0-flash", + }); + session.sessionId = "non-sse-gemini-discovery"; + session.recordProviderSessionRef(1); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "gemini-discovery", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 2, + isFirstAttempt: false, + isFailoverSuccess: true, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + bindingIntent: "create", + bindingSnapshot: { + sessionId: "non-sse-gemini-discovery", + keyId: 2, + providerId: null, + generation: "non-sse-generation", + }, + requiresCompletionMarker: true, + discoveryLease: { + sessionId: "non-sse-gemini-discovery", + keyId: 2, + ownerToken: "non-sse-owner", + ttlSeconds: 30, + }, + providerSessionRefOwned: true, + providerSessionRefRetainOnSuccess: true, + }); + const response = new Response( + '{"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]}}]}}', + { status: 200, headers: { "content-type": "application/json" } } + ); + + const returned = await ProxyResponseHandler.dispatch(session, response); + expect(returned).toBe(response); + await drainAsyncTasks(); + + expect(SessionManager.renewSessionDiscoveryLease).toHaveBeenCalled(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledOnce(); + expect(SessionManager.releaseSessionDiscoveryLease).toHaveBeenCalledWith( + "non-sse-gemini-discovery", + 2, + "non-sse-owner" + ); + expect(RateLimitService.releaseProviderSession).toHaveBeenCalledWith( + 1, + "non-sse-gemini-discovery" + ); + expect(SessionManager.compareAndSetSessionProvider).not.toHaveBeenCalled(); + expect(peekDeferredStreamingFinalization(session)).toBeNull(); + }); + it("persists one durable 502 before Provider circuit mutation on non-stream response timeout", async () => { const durableAck = createDeferred(); vi.mocked(updateMessageRequestDetailsDurably).mockImplementationOnce( From 4c1f16628167071deffd66c6b7de4429bf3d127e Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 22:59:50 -0400 Subject: [PATCH 21/23] fix(discovery): refresh cleared binding authority --- src/app/v1/_lib/proxy/discovery-validity.ts | 8 +++-- src/app/v1/_lib/proxy/provider-selector.ts | 19 ++++++---- tests/unit/proxy/discovery-validity.test.ts | 18 ++++++++++ ...er-selector-model-mismatch-binding.test.ts | 36 +++++++++++++++++++ 4 files changed, 72 insertions(+), 9 deletions(-) diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts index 1f8931d7d..b03f19907 100644 --- a/src/app/v1/_lib/proxy/discovery-validity.ts +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -155,9 +155,7 @@ export function classifyDiscoveryChunk( protocol: DiscoveryProtocol ): DiscoveryValidity { const text = typeof chunk === "string" ? chunk : new TextDecoder().decode(chunk); - if (!text.trim() || text.trim().startsWith(":")) - return { ready: false, terminal: false, error: false }; - if (text.includes("[DONE]")) return { ready: false, terminal: true, error: false }; + if (!text.trim()) return { ready: false, terminal: false, error: false }; const lines = text.split(/\r?\n/); let sawTerminal = false; @@ -166,6 +164,10 @@ export function classifyDiscoveryChunk( for (const line of lines) { const candidate = line.startsWith("data:") ? line.slice(5).trim() : line.trim(); if (!candidate || candidate.startsWith(":")) continue; + if (candidate === "[DONE]") { + sawTerminal = true; + continue; + } try { const result = classifyJson(JSON.parse(candidate), protocol); sawTerminal ||= result.terminal; diff --git a/src/app/v1/_lib/proxy/provider-selector.ts b/src/app/v1/_lib/proxy/provider-selector.ts index e66e18272..a8df56e09 100644 --- a/src/app/v1/_lib/proxy/provider-selector.ts +++ b/src/app/v1/_lib/proxy/provider-selector.ts @@ -507,7 +507,14 @@ export class ProxyProviderResolver { // Read the binding once and retain its generation for Discovery timeout // cleanup/finalization. Re-reading here would allow an older request to // clear a newer binding (ABA). + const sessionId = session.sessionId; const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + const clearRejectedProviderBinding = async (providerId: number): Promise => { + await SessionManager.clearSessionProvider(sessionId, providerId, keyId); + // A clear attempt can advance or race the canonical generation. Force + // Discovery to read authoritative state instead of this old snapshot. + if (keyId != null) session.setSessionBindingSnapshot(null); + }; let providerId: number | null = null; if (keyId != null) { const binding = await SessionManager.getSessionBindingSnapshot(session.sessionId, keyId); @@ -534,7 +541,7 @@ export class ProxyProviderResolver { sessionId: session.sessionId, providerId, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -544,7 +551,7 @@ export class ProxyProviderResolver { providerId: provider.id, providerName: provider.name, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -558,7 +565,7 @@ export class ProxyProviderResolver { activeTimeEnd: provider.activeTimeEnd, timezone: systemTimezone, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -599,7 +606,7 @@ export class ProxyProviderResolver { providerType: provider.providerType, originalFormat: session.originalFormat, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -618,7 +625,7 @@ export class ProxyProviderResolver { // 清除过时绑定，避免 SET NX 死锁 // 当 session 内请求模型发生变化时，旧绑定已无意义， // 清除后新的成功请求可通过 SET NX 重新绑定匹配的 provider - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); logger.info("ProviderSelector: Cleared stale provider binding (model mismatch)", { sessionId: session.sessionId, staleProviderId: provider.id, @@ -674,7 +681,7 @@ export class ProxyProviderResolver { ], }, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index ed2690d6c..6c323b853 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -27,6 +27,24 @@ describe("discovery validity", () => { expect(classifyDiscoveryChunk("data: [DONE]\n", "openai-chat").terminal).toBe(true); }); + it("keeps stateless readiness when content and DONE share a chunk", () => { + expect( + classifyDiscoveryChunk( + 'data: {"choices":[{"delta":{"content":"done"}}]}\n\ndata: [DONE]\n\n', + "openai-chat" + ) + ).toEqual({ ready: true, terminal: true, error: false }); + }); + + it("keeps stateless readiness when a comment precedes content in one chunk", () => { + expect( + classifyDiscoveryChunk( + ': keepalive\ndata: {"choices":[{"delta":{"content":"ready"}}]}\n\n', + "openai-chat" + ) + ).toEqual({ ready: true, terminal: false, error: false }); + }); + it("rejects errors even when a later chunk contains content", () => { const parser = new DiscoveryValidityParser("openai-responses"); expect(parser.push('{"type":"response.failed","error":{"message":"no"}}').error).toBe(true); diff --git a/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts b/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts index 5efbe753a..76c495136 100644 --- a/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts +++ b/tests/unit/proxy/provider-selector-model-mismatch-binding.test.ts @@ -16,6 +16,7 @@ vi.mock("@/lib/vendor-type-circuit-breaker", () => vendorTypeCircuitMocks); const sessionManagerMocks = vi.hoisted(() => ({ SessionManager: { + getSessionBindingSnapshot: vi.fn(), getSessionProvider: vi.fn(async () => null as number | null), clearSessionProvider: vi.fn(async () => undefined), }, @@ -147,6 +148,41 @@ describe("findReusable - model mismatch clears stale binding", () => { ); }); + test("should invalidate a cleared versioned snapshot before Discovery", async () => { + const { ProxyProviderResolver } = await import("@/app/v1/_lib/proxy/provider-selector"); + const snapshot = { + sessionId: "versioned-model-mismatch", + keyId: 456, + providerId: 78, + generation: "stale-generation", + }; + sessionManagerMocks.SessionManager.getSessionBindingSnapshot.mockResolvedValueOnce({ + status: "ok", + snapshot, + }); + providerRepositoryMocks.findProviderById.mockResolvedValueOnce(createHaikuOnlyProvider()); + const setSessionBindingSnapshot = vi.fn(); + const session = { + sessionId: snapshot.sessionId, + shouldReuseProvider: () => true, + getOriginalModel: () => "claude-opus-4-6", + authState: { key: { id: snapshot.keyId } }, + getCurrentModel: () => null, + setSessionBindingSnapshot, + } as any; + + const result = await (ProxyProviderResolver as any).findReusable(session); + + expect(result).toBeNull(); + expect(sessionManagerMocks.SessionManager.clearSessionProvider).toHaveBeenCalledWith( + snapshot.sessionId, + snapshot.providerId, + snapshot.keyId + ); + expect(setSessionBindingSnapshot).toHaveBeenNthCalledWith(1, snapshot); + expect(setSessionBindingSnapshot).toHaveBeenNthCalledWith(2, null); + }); + test("should clear stale binding when bound provider type is incompatible with request format", async () => { const { ProxyProviderResolver } = await import("@/app/v1/_lib/proxy/provider-selector"); From 1c82d521a89a4426a6c8e09d148ceb6de0c0e345 Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Mon, 20 Jul 2026 23:37:34 -0400 Subject: [PATCH 22/23] fix(discovery): parse SSE events by complete data frames --- src/app/v1/_lib/proxy/discovery-validity.ts | 118 +++++++++++------- src/app/v1/_lib/proxy/response-handler.ts | 7 +- tests/unit/proxy/discovery-validity.test.ts | 94 +++++++++++--- ...handler-endpoint-circuit-isolation.test.ts | 5 + 4 files changed, 160 insertions(+), 64 deletions(-) diff --git a/src/app/v1/_lib/proxy/discovery-validity.ts b/src/app/v1/_lib/proxy/discovery-validity.ts index b03f19907..613b056eb 100644 --- a/src/app/v1/_lib/proxy/discovery-validity.ts +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -154,35 +154,12 @@ export function classifyDiscoveryChunk( chunk: Uint8Array | string, protocol: DiscoveryProtocol ): DiscoveryValidity { - const text = typeof chunk === "string" ? chunk : new TextDecoder().decode(chunk); - if (!text.trim()) return { ready: false, terminal: false, error: false }; - - const lines = text.split(/\r?\n/); - let sawTerminal = false; - let sawError = false; - let sawReady = false; - for (const line of lines) { - const candidate = line.startsWith("data:") ? line.slice(5).trim() : line.trim(); - if (!candidate || candidate.startsWith(":")) continue; - if (candidate === "[DONE]") { - sawTerminal = true; - continue; - } - try { - const result = classifyJson(JSON.parse(candidate), protocol); - sawTerminal ||= result.terminal; - sawError ||= result.error; - sawReady ||= result.ready; - } catch { - // A raw JSON response may arrive in a single chunk. Plain text is not - // a protocol-safe winner; keep waiting for a parseable event. - } - } - return { ready: sawReady && !sawError, terminal: sawTerminal, error: sawError }; + return new DiscoveryValidityParser(protocol).push(chunk); } export class DiscoveryValidityParser { private buffered = ""; + private dataLines: string[] = []; private readonly decoder = new TextDecoder(); private _ready = false; private _terminal = false; @@ -201,21 +178,24 @@ export class DiscoveryValidityParser { this._error = true; this._limitExceeded = true; this.buffered = ""; + this.dataLines = []; return this.result; } this.buffered += typeof chunk === "string" ? chunk : this.decoder.decode(chunk, { stream: true }); - // SSE streams are line framed. Consume each completed line once instead - // of reparsing the complete prefix on every chunk (which is quadratic on - // long streams). Keep only the unfinished line for the next push. + // SSE streams are line framed and events end on a blank line. Consume + // completed lines once, while preserving all data: lines for the current + // event so multi-line payloads are joined according to the SSE spec. if (this.buffered.includes("\n")) { - const lines = this.buffered.split(/\r?\n/); + const lines = this.buffered.split("\n"); this.buffered = lines.pop() ?? ""; - for (const line of lines) { + for (const rawLine of lines) { + const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine; this.consumeLine(line); if (this._error) { this.buffered = ""; + this.dataLines = []; return this.result; } } @@ -225,15 +205,11 @@ export class DiscoveryValidityParser { // it only when the complete object is available; incomplete JSON remains // buffered and is not repeatedly scanned as a protocol event. const tail = this.buffered.trim(); - if (tail) { - const candidate = tail.startsWith("data:") ? tail.slice(5).trim() : tail; - if (candidate === "[DONE]") { - this._terminal = true; - this.buffered = ""; - } else if (candidate.startsWith("{") || candidate.startsWith("[")) { + if (this.dataLines.length === 0 && tail && !this.isSseField(tail)) { + if (tail.startsWith("{") || tail.startsWith("[")) { try { - const value = JSON.parse(candidate) as unknown; - this.consumeValue(value); + const value = JSON.parse(tail) as unknown; + this.consumeEventValue(value); this.buffered = ""; } catch { // Keep incomplete raw JSON until the next chunk completes it. @@ -245,32 +221,86 @@ export class DiscoveryValidityParser { } private consumeLine(line: string): void { - const candidate = line.startsWith("data:") ? line.slice(5).trim() : line.trim(); - if (!candidate || candidate.startsWith(":")) return; - if (candidate === "[DONE]") { + if (line === "") { + this.flushSseEvent(); + return; + } + + if (line.startsWith(":")) return; + + const colonIndex = line.indexOf(":"); + const field = colonIndex === -1 ? line : line.slice(0, colonIndex); + if (field === "data") { + let value = colonIndex === -1 ? "" : line.slice(colonIndex + 1); + if (value.startsWith(" ")) value = value.slice(1); + this.dataLines.push(value); + return; + } + + // event/id/retry and unknown SSE fields carry framing metadata only. A + // bare JSON line is supported for providers returning non-SSE JSON, but + // never while an SSE data event is pending. + if (field === "event" || field === "id" || field === "retry" || this.dataLines.length > 0) { + return; + } + const candidate = line.trim(); + if (candidate.startsWith("{") || candidate.startsWith("[")) { + try { + this.consumeEventValue(JSON.parse(candidate) as unknown); + } catch { + // Plain text and incomplete/non-JSON lines cannot establish validity. + } + } + } + + private flushSseEvent(): void { + if (this.dataLines.length === 0) return; + const candidate = this.dataLines.join("\n"); + this.dataLines = []; + if (!this.beginEvent()) return; + if (candidate.trim() === "[DONE]") { this._terminal = true; return; } try { this.consumeValue(JSON.parse(candidate) as unknown); } catch { - // Ignore comments and incomplete/non-JSON protocol lines. + // A complete but non-JSON SSE event cannot establish protocol validity. } } - private consumeValue(value: unknown): void { + private consumeEventValue(value: unknown): void { + if (!this.beginEvent()) return; + this.consumeValue(value); + } + + private beginEvent(): boolean { this.eventsSeen += 1; if (!this._ready && this.eventsSeen > DISCOVERY_EVENT_MAX_COUNT) { this._error = true; this._limitExceeded = true; - return; + return false; } + return true; + } + + private consumeValue(value: unknown): void { const result = classifyJson(value, this.protocol); this._ready ||= result.ready; this._terminal ||= result.terminal; this._error ||= result.error; } + private isSseField(line: string): boolean { + return ( + line.startsWith(":") || + line.startsWith("data:") || + line.startsWith("event:") || + line.startsWith("id:") || + line.startsWith("retry:") + ); + } + get ready(): boolean { return this._ready && !this._error; } diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index 94f41e043..45807eed1 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -1315,12 +1315,7 @@ function hasStreamCompletionMarker(text: string, format: ProxySession["originalF isRecord(event.data.response) ); case "claude": - return events.some( - (event) => - event.event === "message_stop" && - isRecord(event.data) && - event.data.type === "message_stop" - ); + return events.some((event) => isRecord(event.data) && event.data.type === "message_stop"); case "openai": return events.some( (event) => diff --git a/tests/unit/proxy/discovery-validity.test.ts b/tests/unit/proxy/discovery-validity.test.ts index 6c323b853..fc76b1b19 100644 --- a/tests/unit/proxy/discovery-validity.test.ts +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -21,10 +21,10 @@ describe("discovery validity", () => { it("accepts OpenAI Chat delta and rejects DONE", () => { expect( - classifyDiscoveryChunk('data: {"choices":[{"delta":{"content":"hi"}}]}\n', "openai-chat") + classifyDiscoveryChunk('data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', "openai-chat") .ready ).toBe(true); - expect(classifyDiscoveryChunk("data: [DONE]\n", "openai-chat").terminal).toBe(true); + expect(classifyDiscoveryChunk("data: [DONE]\n\n", "openai-chat").terminal).toBe(true); }); it("keeps stateless readiness when content and DONE share a chunk", () => { @@ -54,19 +54,19 @@ describe("discovery validity", () => { it("does not promote empty tool or content events", () => { expect( classifyDiscoveryChunk( - 'data: {"type":"content_block_start","content_block":{"type":"text","text":""}}\n', + 'data: {"type":"content_block_start","content_block":{"type":"text","text":""}}\n\n', "anthropic" ).ready ).toBe(false); expect( classifyDiscoveryChunk( - 'data: {"choices":[{"delta":{"tool_calls":[{"function":{}}]}}]}\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"function":{}}]}}]}\n\n', "openai-chat" ).ready ).toBe(false); expect( classifyDiscoveryChunk( - 'data: {"type":"response.output_text.delta","delta":" "}\n', + 'data: {"type":"response.output_text.delta","delta":" "}\n\n', "openai-responses" ).ready ).toBe(false); @@ -75,7 +75,7 @@ describe("discovery validity", () => { it("accepts a non-empty function call delta as deliverable content", () => { expect( classifyDiscoveryChunk( - 'data: {"type":"response.function_call_arguments.delta","delta":"{\\"x\\":1}"}\n', + 'data: {"type":"response.function_call_arguments.delta","delta":"{\\"x\\":1}"}\n\n', "openai-responses" ).ready ).toBe(true); @@ -112,13 +112,13 @@ describe("discovery validity", () => { it("accepts only an explicit non-empty Responses tool payload", () => { expect( classifyDiscoveryChunk( - 'data: {"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","status":"in_progress"}}\n', + 'data: {"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","status":"in_progress"}}\n\n', "openai-responses" ).ready ).toBe(false); expect( classifyDiscoveryChunk( - 'data: {"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","name":"lookup","arguments":"{}"}}\n', + 'data: {"type":"response.output_item.added","item":{"id":"fc_1","type":"function_call","name":"lookup","arguments":"{}"}}\n\n', "openai-responses" ).ready ).toBe(true); @@ -138,6 +138,53 @@ describe("discovery validity", () => { }); }); + it("joins all data lines in one SSE event before parsing", () => { + const parser = new DiscoveryValidityParser("openai-chat"); + + expect(parser.push('event: message\nid: 42\ndata: {"choices":[{"delta":\n')).toEqual({ + ready: false, + terminal: false, + error: false, + }); + expect(parser.push('data: {"content":"hello"}}]}\n\n')).toEqual({ + ready: true, + terminal: false, + error: false, + }); + }); + + it("ignores comments and SSE metadata instead of parsing them as payloads", () => { + const parser = new DiscoveryValidityParser("anthropic"); + + expect(parser.push(": keepalive\nevent: message\nid: 7\nretry: 1000\n\n")).toEqual({ + ready: false, + terminal: false, + error: false, + }); + expect( + parser.push('data: {"type":"content_block_delta",\ndata: "delta":{"text":"hi"}}\n\n') + ).toMatchObject({ ready: true, error: false }); + }); + + it("does not parse raw JSON while an SSE data event is pending", () => { + const parser = new DiscoveryValidityParser("openai-chat"); + + expect(parser.push('data: {"choices":[{"delta":\n')).toEqual({ + ready: false, + terminal: false, + error: false, + }); + expect(parser.push('{"content":"wrongly standalone"}}]}')).toEqual({ + ready: false, + terminal: false, + error: false, + }); + expect(parser.push('\ndata: {"content":"ready"}}]}\n\n')).toMatchObject({ + ready: true, + error: false, + }); + }); + it("keeps ready when content and the terminal marker arrive in one read", () => { const parser = new DiscoveryValidityParser("openai-chat"); @@ -174,13 +221,13 @@ describe("discovery validity", () => { it("accepts Anthropic tool-use starts and partial JSON deltas", () => { expect( classifyDiscoveryChunk( - 'data: {"type":"content_block_start","content_block":{"type":"tool_use","id":"tu_1","name":"search","input":{}}}\n', + 'data: {"type":"content_block_start","content_block":{"type":"tool_use","id":"tu_1","name":"search","input":{}}}\n\n', "anthropic" ).ready ).toBe(true); expect( classifyDiscoveryChunk( - 'data: {"type":"content_block_delta","delta":{"type":"input_json_delta","partial_json":"{\\"q\\":1}"}}\n', + 'data: {"type":"content_block_delta","delta":{"type":"input_json_delta","partial_json":"{\\"q\\":1}"}}\n\n', "anthropic" ).ready ).toBe(true); @@ -196,19 +243,38 @@ describe("discovery validity", () => { const parser = new DiscoveryValidityParser("anthropic"); let result = parser.push(""); for (let index = 0; index <= DISCOVERY_EVENT_MAX_COUNT; index += 1) { - result = parser.push('data: {"type":"ping"}\n'); + result = parser.push('data: {"type":"ping"}\n\n'); } expect(result).toMatchObject({ ready: false, error: true, limitExceeded: true }); }); + it("counts a multi-line data payload as one complete SSE event", () => { + const parser = new DiscoveryValidityParser("anthropic"); + const metadataEvents = Array.from( + { length: DISCOVERY_EVENT_MAX_COUNT - 1 }, + () => 'data: {"type":"ping"}\n\n' + ).join(""); + + expect(parser.push(`${metadataEvents}data: {\ndata: "type":"ping"}\n\n`)).toEqual({ + ready: false, + terminal: false, + error: false, + }); + expect(parser.push('data: {"type":"ping"}\n\n')).toMatchObject({ + ready: false, + error: true, + limitExceeded: true, + }); + }); + it("stops parsing current and future events after the event limit", () => { const parser = new DiscoveryValidityParser("anthropic"); const metadataEvents = Array.from( { length: DISCOVERY_EVENT_MAX_COUNT + 1 }, - () => 'data: {"type":"ping"}\n' + () => 'data: {"type":"ping"}\n\n' ).join(""); - const limited = parser.push(`${metadataEvents}data: {"type":"message_stop"}\n`); + const limited = parser.push(`${metadataEvents}data: {"type":"message_stop"}\n\n`); expect(limited).toMatchObject({ ready: false, terminal: false, @@ -216,7 +282,7 @@ describe("discovery validity", () => { limitExceeded: true, }); - expect(parser.push('data: {"type":"message_stop"}\n')).toEqual(limited); + expect(parser.push('data: {"type":"message_stop"}\n\n')).toEqual(limited); }); }); diff --git a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts index 0e1eafa29..d7bbdd36a 100644 --- a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts +++ b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts @@ -866,6 +866,11 @@ describe("Endpoint circuit breaker isolation", () => { response: { id: "resp_completed" }, })}\n\n`, }, + { + label: "Anthropic data-only", + format: "claude" as const, + body: `data: ${JSON.stringify({ type: "message_stop" })}\n\n`, + }, { label: "OpenAI Chat finish reason", format: "openai" as const, From 3dd24e7d637a29ddd1a7e2408109944494096d3f Mon Sep 17 00:00:00 2001 From: Brisbanehuang Date: Tue, 21 Jul 2026 00:49:40 -0400 Subject: [PATCH 23/23] fix(discovery): validate fallback stream completion --- src/app/v1/_lib/proxy/forwarder.ts | 4 +++- tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/app/v1/_lib/proxy/forwarder.ts b/src/app/v1/_lib/proxy/forwarder.ts index 8016b0502..36a03c136 100644 --- a/src/app/v1/_lib/proxy/forwarder.ts +++ b/src/app/v1/_lib/proxy/forwarder.ts @@ -5279,7 +5279,9 @@ export class ProxyForwarder { ? "create" : "renew", bindingSnapshot, - requiresCompletionMarker: attempt.kind !== "fallback", + // Fallbacks cannot create Sticky, but an incomplete fallback stream must + // still be classified as failed rather than as a successful truncated 200. + requiresCompletionMarker: true, discoveryLease: lease, providerSessionRefOwned: attempt.providerSessionRefOwned, providerSessionRefRetainOnSuccess: attempt.providerSessionRefRetainOnSuccess, diff --git a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts index 9f37194bc..bf8f53777 100644 --- a/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts +++ b/tests/unit/proxy/proxy-forwarder-hedge-first-byte.test.ts @@ -2920,8 +2920,11 @@ describe("ProxyForwarder - first-byte hedge scheduling", () => { await vi.advanceTimersByTimeAsync(35); const response = await responsePromise; + const deferred = peekDeferredStreamingFinalization(session); expect(await response.text()).toContain('"sticky-fallback"'); expect(session.provider?.id).toBe(sticky.id); + expect(deferred?.bindingIntent).toBe("none"); + expect(deferred?.requiresCompletionMarker).toBe(true); stalledSelector.resolve([]); await vi.advanceTimersByTimeAsync(0);