diff --git a/.env.example b/.env.example index b6c6d72b8..177141c29 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 new file mode 100644 index 000000000..b7879a8dc --- /dev/null +++ b/src/app/v1/_lib/proxy/discovery-coordinator.ts @@ -0,0 +1,347 @@ +/** + * 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[]; promoteAttemptId?: string } + | { + type: "launch"; + slots: number; + cancelAttemptIds?: string[]; + promoteAttemptId?: string; + } + | { 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 }; + } + + 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 }; + } + + 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); + } + + /** 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 === "FALLBACK_ACTIVE" || + 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; + 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(); + } + + /** 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, + 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, + 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.pending = false; + attempt.ready = false; + const readyAction = this.chooseReadyNormal(); + if (readyAction.type !== "none") return readyAction; + 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 winner = readyNormal[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) { + const cancelAttemptIds = pendingNormal.map((attempt) => attempt.id); + for (const attempt of pendingNormal) attempt.pending = false; + if (this.round < this.maxRounds) { + this.beginRound(); + this.state = "DISCOVERY_RACING"; + return { + type: "launch", + slots: Math.max(0, this.concurrency - 1), + cancelAttemptIds, + }; + } + return { type: "cancel", attemptIds: cancelAttemptIds }; + } + 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 (this.round < this.maxRounds) { + this.beginRound(); + this.state = "DISCOVERY_RACING"; + return { + type: "launch", + slots: Math.max(0, this.concurrency - 1), + cancelAttemptIds: losers, + promoteAttemptId: fallback.id, + }; + } + return { type: "cancel", attemptIds: losers, promoteAttemptId: fallback.id }; + } + + 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 + ); + 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 = attempt.kind === "fallback" ? "FALLBACK_ACTIVE" : "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(); + + // 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); + } + return { type: "none" }; + } + + private finishOrLaunch(): DiscoveryAction { + if (this.round >= this.maxRounds) { + this.state = "TERMINAL_FAILED"; + return { type: "terminal_failure" }; + } + this.beginRound(); + this.state = "DISCOVERY_RACING"; + 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 new file mode 100644 index 000000000..613b056eb --- /dev/null +++ b/src/app/v1/_lib/proxy/discovery-validity.ts @@ -0,0 +1,326 @@ +export type DiscoveryProtocol = + | "anthropic" + | "openai-chat" + | "openai-responses" + | "gemini" + | "unknown"; + +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; + 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", + "function", + "arguments", + "partial_json", + "id", + "name", + "input", + "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 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; + 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" && hasOpenAIResponsesOutputItem(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" && hasAnthropicContentBlock(object.content_block)) || + hasContent(object.content), + terminal: false, + error: false, + }; +} + +export function classifyDiscoveryChunk( + chunk: Uint8Array | string, + protocol: DiscoveryProtocol +): DiscoveryValidity { + 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; + private _error = false; + private _limitExceeded = false; + private bytesSeen = 0; + private eventsSeen = 0; + + 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) { + 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 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("\n"); + this.buffered = lines.pop() ?? ""; + 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; + } + } + } + + // 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 (this.dataLines.length === 0 && tail && !this.isSseField(tail)) { + if (tail.startsWith("{") || tail.startsWith("[")) { + try { + const value = JSON.parse(tail) as unknown; + this.consumeEventValue(value); + this.buffered = ""; + } catch { + // Keep incomplete raw JSON until the next chunk completes it. + } + } + } + + return this.result; + } + + private consumeLine(line: string): void { + 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 { + // A complete but non-JSON SSE event cannot establish protocol validity. + } + } + + 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 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; + } + get terminal(): boolean { + return this._terminal; + } + 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 b93c9c7bf..a3c6f23ba 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"; @@ -63,6 +69,8 @@ import { buildProxyUrl } from "../url"; import { rectifyBillingHeader } from "./billing-header-rectifier"; import { bindClientAbortListener } from "./client-abort-listener"; import { deriveClientSafeUpstreamErrorMessage } from "./client-error-message"; +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"; import { @@ -193,6 +201,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; @@ -202,6 +225,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; @@ -1210,6 +1268,22 @@ export class ProxyForwarder { throw new Error("代理上下文缺少供应商或鉴权信息"); } + 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; + } + if (ProxyForwarder.shouldUseStreamingHedge(session)) { const hedgePromise = ProxyForwarder.sendStreamingWithHedge(session); void hedgePromise.catch(() => undefined); @@ -1485,6 +1559,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", { @@ -1674,7 +1749,7 @@ export class ProxyForwarder { } // ⭐ 成功后绑定 session 到供应商（智能绑定策略） - if (session.sessionId) { + if (session.sessionId && session.isSessionBindingAllowed()) { // 使用智能绑定策略（故障转移优先 + 稳定性优化） const result = await SessionManager.updateSessionBindingSmart( session.sessionId, @@ -2411,7 +2486,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"); @@ -3048,12 +3124,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 = { @@ -3263,6 +3343,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; @@ -3817,6 +3906,9 @@ export class ProxyForwarder { private static shouldUseStreamingHedge(session: ProxySession): boolean { const endpointPolicy = ProxyForwarder.getEndpointPolicy(session); + if (session.isStreamingHedgeDisabled()) { + return false; + } return ( (endpointPolicy?.allowRetry ?? true) && (endpointPolicy?.allowProviderSwitch ?? true) && @@ -3825,6 +3917,117 @@ export class ProxyForwarder { ); } + private static async prepareStreamingDiscovery( + session: ProxySession, + settings: SystemSettings, + requestStartedAt: number + ): Promise { + if (settings.discoveryEnabled !== true) { + return null; + } + const endpointPolicy = ProxyForwarder.getEndpointPolicy(session); + const protocol = ProxyForwarder.discoveryProtocol(session); + const message = session.request.message as Record; + 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 { + 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; @@ -4578,7 +4781,7 @@ export class ProxyForwarder { // here as well would perform a duplicate binding read/CAS before the // stream has passed its final validation. let hedgeBindingAuthorityPromise: Promise | undefined; - if (session.sessionId && isActualHedgeWin) { + if (session.sessionId && isActualHedgeWin && session.isSessionBindingAllowed()) { hedgeBindingAuthorityPromise = (async () => { const bindingResult = await SessionManager.updateSessionBindingSmart( session.sessionId!, @@ -4643,6 +4846,7 @@ export class ProxyForwarder { upstreamStatusCode: attempt.response.status, isHedgeWinner: isActualHedgeWin, billHedgeLosers, + bindingIntent: session.isSessionBindingAllowed() ? undefined : "none", hedgeBindingAuthorityPromise, }); @@ -4688,7 +4892,9 @@ export class ProxyForwarder { } if (checkResult.referenced) { - session.recordProviderSessionRef(provider.id); + session.recordProviderSessionRef(provider.id, { + retainOnSuccess: checkResult.tracked, + }); } } @@ -4804,6 +5010,1111 @@ 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, + prepared: PreparedStreamingDiscovery + ): Promise { + const initialProvider = session.provider; + if (!initialProvider) throw new Error("代理上下文缺少供应商"); + 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); + // 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 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 & { + id: string; + kind: "normal" | "fallback"; + controller: AbortController; + parser: DiscoveryValidityParser; + chunks: Uint8Array[]; + pending: boolean; + ready: boolean; + round: number; + readerTransferred: boolean; + readerCancelled: boolean; + providerSessionRefOwned: boolean; + providerSessionRefRetainOnSuccess: boolean; + providerSessionRefReleased: boolean; + cancellationKind: DiscoveryCancellationKind | null; + } + >(); + 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 roundLaunchesInProgress = 0; + 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 && + 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, + 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 = (attempt: (typeof winner & { id: string }) | null) => { + if (!attempt?.providerSessionRefOwned || attempt.providerSessionRefReleased) return; + attempt.providerSessionRefReleased = true; + ProxyForwarder.releaseProviderSessionRef(session, attempt.provider.id); + }; + + 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) return; + attempt.pending = false; + if (cancellationKind && !attempt.cancellationKind) { + attempt.cancellationKind = cancellationKind; + } + if (!attempt.controller.signal.aborted) { + try { + attempt.controller.abort( + cancellationKind + ? new DiscoveryCancellationError(cancellationKind) + : new Error("discovery_attempt_failed") + ); + } catch { + /* abort is best effort */ + } + } + 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); + attempt.chunks.length = 0; + discoveryMetrics.attemptFinished(attempt.id, { + providerId: attempt.provider.id, + outcome: cancellationKind ? "cancelled" : "failed", + cancellationKind, + }); + }; + + 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, cancellationKind); + } + }; + + 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(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( + 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), + }); + }); + } + } + 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.ready || + !attempt.response || + !attempt.reader + ) + return; + 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); + 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); + + 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" || !bindingWriteAllowed || !session.isSessionBindingAllowed() + ? "none" + : bindingSnapshot?.providerId == null + ? "create" + : "renew", + bindingSnapshot, + // 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, + }); + leaseTransferred = true; + 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 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 ensureStickyTimeoutCooldown = (cooldownTtlSeconds: number): Promise => { + if (!stickyTimeoutCooldownPromise) { + stickyTimeoutCooldownPromise = clearCapturedStickyBinding(cooldownTtlSeconds); + } + return stickyTimeoutCooldownPromise; + }; + + 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 false; + } + launched.add(provider.id); + 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, + session.sessionId, + limit + ); + if (!check.allowed) { + throw new ProxyError(check.reason || "Provider concurrent limit reached", 503); + } + if (check.referenced) { + session.recordProviderSessionRef(provider.id, { + retainOnSuccess: check.tracked, + }); + providerSessionRefTracked = true; + providerSessionRefRetainOnSuccess = check.tracked; + } + } + if (settled || committed) { + rollbackLaunch(); + return false; + } + let endpoint: Awaited>; + try { + endpoint = await ProxyForwarder.resolveStreamingHedgeEndpoint(session, provider); + } catch (error) { + rollbackLaunch(); + throw error; + } + if (settled || committed) { + rollbackLaunch(); + return false; + } + 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 false; + } + 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, + readerTransferred: false, + readerCancelled: false, + providerSessionRefOwned: providerSessionRefTracked, + providerSessionRefRetainOnSuccess, + providerSessionRefReleased: false, + cancellationKind: null, + provider, + session: attemptSession, + baseUrl: endpoint.baseUrl, + endpointAudit: { endpointId: endpoint.endpointId, endpointUrl: endpoint.endpointUrl }, + modelRedirect: undefined, + responseController: null, + clearResponseTimeout: null, + firstByteTimeoutMs: 0, + sequence: ++sequence, + requestAttemptCount: options?.requestAttemptCount ?? 1, + reactiveRectifierRetryState: options?.retryState ?? { + 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; + readerTransferred: boolean; + readerCancelled: boolean; + providerSessionRefOwned: boolean; + providerSessionRefRetainOnSuccess: boolean; + providerSessionRefReleased: boolean; + cancellationKind: DiscoveryCancellationKind | null; + }; + const registered = coordinator.addAttempt({ + id, + providerId: provider.id, + priority: ProxyProviderResolver.resolveEffectivePriorityForSession(provider, session), + kind, + ready: false, + pending: true, + round: currentRound, + launchOrder: attempt.sequence, + }); + if (!registered || settled || committed) { + rollbackLaunch(); + return false; + } + 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, + { ...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 (!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 + // 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); + // 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. 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. + // 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); + // 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. Stop + // reading so later chunks are not consumed before promotion. + if (action.type === "none") return; + return; + } + }) + .catch(async (error) => { + if (committed || settled || !attempt.pending) return; + 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, + }); + if ( + !(lastError instanceof DiscoveryValidityLimitError) && + lastErrorCategory === ErrorCategory.PROVIDER_ERROR && + !(lastError instanceof ProxyError && lastError.statusCode === 404) + ) { + await recordFailure(provider.id, lastError).catch(() => undefined); + } + 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" || + failureAction.type === "launch" || + failureAction.type === "terminal_failure"; + if (actionOwnsNextStep) { + if ( + failureAction.type === "launch" && + stickyTimeoutWaveReservation?.fallbackAttemptId === id + ) { + await launchReservedStickyTimeoutWave(failureAction.slots); + } else { + await executeCoordinatorAction(failureAction); + } + } + if (!actionOwnsNextStep && !committed && !settled) { + await refillCurrentRoundSlots(1); + } + if ( + Array.from(attempts.values()).every((candidate) => !candidate.pending) && + noMoreCandidates + ) { + await settleFailure( + ProxyForwarder.resolveHedgeTerminalError(lastError, lastErrorCategory) + ); + } + }) + .catch((error) => { + logger.warn("[Discovery] Attempt completion handler failed", { + providerId: provider.id, + error, + }); + }); + 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; + clearRoundTimer(); + try { + if (coordinatorAlreadyAdvanced) { + currentRound = coordinator.round; + } else { + const nextRound = coordinator.beginRound(); + currentRound = nextRound.round; + } + if (currentRound > maxRounds || slots <= 0) return; + await fillDiscoverySlots(slots); + const hasPendingAttempt = Array.from(attempts.values()).some((attempt) => attempt.pending); + if (!hasPendingAttempt) { + await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError)); + return; + } + if (!committed && !settled) { + scheduleRoundBoundary(discoverySlaMs); + } + } finally { + await finishRoundLaunchBatch(); + } + }; + + 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") { + 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_sla_timeout"); + } + if (action.promoteAttemptId) { + const fallback = attempts.get(action.promoteAttemptId); + 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); + if (attempt) await commit(attempt); + return; + } + if (action.type === "launch") { + await launchNextRound(action.slots, true); + return; + } + if (action.type === "terminal_failure") { + await settleFailure(ProxyForwarder.buildAllProvidersUnavailableError(lastError), { + cancellationKind: terminalCancellationKind, + }); + return; + } + }; + + const cleanupAbort = bindClientAbortListener(session.clientAbortSignal, () => { + if (settled || committed) return; + 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(), "request_deadline").catch((error) => + logger.warn("[Discovery] Deadline action failed", { error }) + ); + }, + Math.max(0, racingDeadlineAt - Date.now()) + ); + + 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)); + } + if (settled || committed) return; + + if (hasSticky) { + 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 + ); + 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; + stickyTimeoutWaveReservation = { fallbackAttemptId: sticky.id }; + if (bindingSnapshot && bindingSnapshot.providerId === initialProvider.id) { + void ensureStickyTimeoutCooldown( + Math.ceil((settings.stickyTimeoutCooldownMs ?? 300_000) / 1000) + ).finally(() => { + void launchReservedStickyTimeoutWave(Math.max(0, concurrency - 1)).catch( + (error) => logger.warn("[Discovery] Sticky round launch failed", { error }) + ); + }); + return; + } + void launchReservedStickyTimeoutWave(Math.max(0, concurrency - 1)).catch((error) => + logger.warn("[Discovery] Sticky round launch failed", { error }) + ); + } + }, + Math.min(stickySlaMs, Math.max(0, racingDeadlineAt - Date.now())) + ); + } + } 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; + } finally { + cleanupAbort(); + 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), + }); + }); + } + } + } + private static async resolveStreamingHedgeEndpoint( session: ProxySession, provider: Provider @@ -5034,7 +6345,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); } @@ -5043,7 +6354,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); } @@ -5063,15 +6374,23 @@ 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; - } - - void RateLimitService.releaseProviderSession(providerId, session.sessionId); + if (!providerSessionRefConsumer?.call(session, providerId)) return false; + 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; } private static buildAllProvidersUnavailableError(finalError?: Error | null): ProxyError { diff --git a/src/app/v1/_lib/proxy/provider-selector.ts b/src/app/v1/_lib/proxy/provider-selector.ts index e81d32593..6d8b5aa72 100644 --- a/src/app/v1/_lib/proxy/provider-selector.ts +++ b/src/app/v1/_lib/proxy/provider-selector.ts @@ -320,7 +320,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", { @@ -483,6 +485,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） */ @@ -491,9 +530,29 @@ 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 sessionId = session.sessionId; const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; - const providerId = await SessionManager.getSessionProvider(session.sessionId, keyId); + 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); + 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, @@ -508,7 +567,7 @@ export class ProxyProviderResolver { sessionId: session.sessionId, providerId, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -518,7 +577,7 @@ export class ProxyProviderResolver { providerId: provider.id, providerName: provider.name, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -532,7 +591,7 @@ export class ProxyProviderResolver { activeTimeEnd: provider.activeTimeEnd, timezone: systemTimezone, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -573,7 +632,7 @@ export class ProxyProviderResolver { providerType: provider.providerType, originalFormat: session.originalFormat, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -592,7 +651,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, @@ -648,7 +707,7 @@ export class ProxyProviderResolver { ], }, }); - await SessionManager.clearSessionProvider(session.sessionId, providerId, keyId); + await clearRejectedProviderBinding(providerId); return null; } @@ -1155,6 +1214,13 @@ export class ProxyProviderResolver { return provider.priority ?? 0; } + static resolveEffectivePriorityForSession(provider: Provider, session: ProxySession): number { + return ProxyProviderResolver.resolveEffectivePriority( + provider, + getEffectiveProviderGroup(session) + ); + } + /** * 优先级分层：只选择最高优先级的供应商（支持分组优先级覆盖） */ diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index d5a522a66..45807eed1 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"; @@ -59,6 +61,7 @@ import type { ProxySession } from "./session"; import { consumeDeferredStreamingFinalization, type DeferredStreamingBindingHeartbeat, + type DeferredStreamingFinalization, peekDeferredStreamingFinalization, } from "./stream-finalization"; @@ -96,8 +99,209 @@ 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 deferred = peekDeferredStreamingFinalization(session); + const lease = deferred?.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 bindingSnapshot = + (deferred?.bindingIntent === "create" || deferred?.bindingIntent === "renew") && + deferred.bindingSnapshot?.sessionId === lease.sessionId && + deferred.bindingSnapshot.keyId === lease.keyId + ? deferred.bindingSnapshot + : null; + + 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; + } + + 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; + })() + .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 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); + 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), + }); + } +} + +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; @@ -1065,8 +1269,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 + ); +} /** * 判断流式响应文本中是否存在“与格式匹配的终止完成标记”，用以区分 @@ -1075,16 +1302,33 @@ 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) => 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( @@ -1208,6 +1452,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; }; /** @@ -1233,13 +1483,14 @@ 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 hedgeAuthority = meta?.isHedgeWinner ? await meta.hedgeBindingAuthorityPromise : undefined; @@ -1262,11 +1513,126 @@ function finalizeDeferredStreamingFinalizationIfNeeded( return; } const keyId = session.authState?.key?.id ?? session.messageContext?.key?.id ?? null; + 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.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 finalizeFailedDiscoveryBinding = async () => { + settlePrimaryDiscoveryBinding(false); + await finalizeProviderSessionRef(); + }; + const confirmAuxiliarySessionBinding = async () => + allowAuxiliarySessionBinding && (await primaryDiscoveryBinding); // 仅在“上游 HTTP=200 且流自然结束”时做“假 200”检测： // - 非 200：HTTP 已经表明失败（无需额外启发式） @@ -1277,6 +1643,10 @@ function finalizeDeferredStreamingFinalizationIfNeeded( const detected = shouldDetectFake200 ? detectUpstreamErrorFromSseOrJsonText(allContent) : ({ isError: false } as const); + const completionMarkerMissing = + meta?.requiresCompletionMarker === true && + streamEndedNormally && + !hasStreamCompletionMarker(allContent, session.originalFormat); let clientAbortGateUsage: FinalizeDeferredStreamingResult["clientAbortGateUsage"]; const clientAbortCompleteSuccess = (() => { if (!clientAborted || upstreamStatusCode < 200 || upstreamStatusCode >= 300) { @@ -1295,7 +1665,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; } @@ -1321,6 +1691,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; @@ -1351,6 +1724,7 @@ function finalizeDeferredStreamingFinalizationIfNeeded( const shouldClearSessionBindingOnFailure = ((clientAborted || !streamEndedNormally) && !clientAbortCompleteSuccess) || detected.isError || + completionMarkerMissing || (upstreamStatusCode >= 400 && errorMessage !== null); if (shouldClearSessionBindingOnFailure) { meta?.hedgeBindingHeartbeat?.stop(); @@ -1361,6 +1735,18 @@ function finalizeDeferredStreamingFinalizationIfNeeded( // - 不在这里更新熔断/绑定（meta 缺失意味着 Forwarder 没有启用延迟结算；provider 缺失意味着无法归因）。 if (!meta || !provider) { meta?.hedgeBindingHeartbeat?.stop(); + const commitSideEffects = + shouldClearSessionBindingOnFailure || + meta?.providerSessionRefOwned === true || + hasDiscoveryBindingIntent + ? async () => { + try { + if (shouldClearSessionBindingOnFailure) await clearSessionBinding(); + } finally { + await finalizeFailedDiscoveryBinding(); + } + } + : undefined; return { effectiveStatusCode, errorMessage, @@ -1368,7 +1754,10 @@ function finalizeDeferredStreamingFinalizationIfNeeded( isHedgeWinner, billHedgeLosers, clientAbortGateUsage, - commitSideEffects: shouldClearSessionBindingOnFailure ? clearSessionBinding : undefined, + commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1412,22 +1801,70 @@ 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. } + } finally { + await finalizeFailedDiscoveryBinding(); + } + }; - // Stream aborts are key-level errors. The endpoint delivered HTTP 200, - // so only the Provider circuit is updated here. + return { + effectiveStatusCode, + errorMessage, + providerIdForPersistence, + isHedgeWinner, + billHedgeLosers, + clientAbortGateUsage, + commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, + }; + } + + 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 () => { + 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 finalizeFailedDiscoveryBinding(); } }; @@ -1439,6 +1876,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1475,23 +1915,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 finalizeFailedDiscoveryBinding(); } }; @@ -1503,6 +1947,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1533,22 +1980,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 finalizeFailedDiscoveryBinding(); } }; @@ -1560,6 +2011,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1583,84 +2037,107 @@ function finalizeDeferredStreamingFinalizationIfNeeded( // complete() is idempotent and permanently stops after any authority conflict. const hedgeBindingCompletion = meta.hedgeBindingHeartbeat?.complete(); const commitSideEffects = async () => { + let primaryDiscoveryBindingUpdated = false; await hedgeBindingCompletion; + 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, + }); + } + } - 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, + 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, - }); - } + // 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, + meta.providerPriority, + meta.isFirstAttempt, + meta.isFailoverSuccess, + keyId + ); - // 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 - ); + primaryDiscoveryBindingUpdated = isDiscoveryBinding && result.updated; + retainProviderSessionRef = + primaryDiscoveryBindingUpdated && meta.providerSessionRefRetainOnSuccess === true; - 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 (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 } - ); - }); + 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 { @@ -1671,6 +2148,9 @@ function finalizeDeferredStreamingFinalizationIfNeeded( billHedgeLosers, clientAbortGateUsage, commitSideEffects, + finalizeAttemptResources: finalizeProviderSessionRef, + allowAuxiliarySessionBinding, + confirmAuxiliarySessionBinding, }; } @@ -1719,9 +2199,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; } @@ -1926,7 +2408,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; @@ -1992,6 +2474,7 @@ export class ProxyResponseHandler { } } finally { cleanupTaskAbortBinding(); + await finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); releaseSessionAgent(session); } }; @@ -2101,7 +2584,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: @@ -2235,7 +2720,8 @@ export class ProxyResponseHandler { statusCode >= 200 && statusCode < 300 && session.sessionId && - provider.id + provider.id && + isSessionBindingMutationAllowed(session) ) { try { const responseData = JSON.parse(responseText) as Record; @@ -2607,6 +3093,7 @@ export class ProxyResponseHandler { } finally { cleanupTaskAbortBinding(); cleanupClientAbortListener(); + await finalizeNonStreamDiscoveryResources(session, discoveryLeaseLifecycle); releaseSessionAgent(session); } }; @@ -2651,10 +3138,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; } @@ -2754,16 +3247,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(); + } }, }); }; @@ -2959,9 +3462,11 @@ export class ProxyResponseHandler { statusCode, streamEndedNormally, clientAborted, + discoveryLeaseLifecycle, abortReason ); latestCommitSideEffects = finalized.commitSideEffects; + latestFinalizeAttemptResources = finalized.finalizeAttemptResources; const finalizedUsage = await finalizeRequestStats( session, allContent, @@ -3029,9 +3534,11 @@ export class ProxyResponseHandler { statusCode, false, clientAborted, + discoveryLeaseLifecycle, abortReason ); latestCommitSideEffects = finalized.commitSideEffects; + latestFinalizeAttemptResources = finalized.finalizeAttemptResources; await finalizeRequestStats( session, @@ -3072,6 +3579,12 @@ export class ProxyResponseHandler { } } finally { releaseTransportResources(); + if (!commitSideEffectsScheduled) { + void (async () => { + await latestFinalizeAttemptResources?.(); + await discoveryLeaseLifecycle.release(); + })(); + } } }; @@ -3356,15 +3869,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; @@ -3415,11 +3943,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; @@ -3516,7 +4046,8 @@ export class ProxyResponseHandler { effectiveStatusCode >= 200 && effectiveStatusCode < 300 && session.sessionId && - provider.id + provider.id && + finalized.allowAuxiliarySessionBinding ) { try { const sseEvents = parseSSEData(allContent); @@ -3704,6 +4235,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, @@ -4126,6 +4658,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 10d6b4740..0ef7901c1 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, @@ -120,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; @@ -202,9 +208,15 @@ 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 + // second read that could race with another request's binding update. + private sessionBindingSnapshot: SessionBindingSnapshot | null = null; private constructor(init: { startTime: number; @@ -349,25 +361,58 @@ export class ProxySession { } } - recordProviderSessionRef(providerId: number): void { + setSessionBindingSnapshot(snapshot: SessionBindingSnapshot | null): void { + this.sessionBindingSnapshot = snapshot; + } + + getSessionBindingSnapshot(): SessionBindingSnapshot | null { + return this.sessionBindingSnapshot; + } + + 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 c1eafe42b..1184f2afd 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; +}; + export type DeferredStreamingBindingHeartbeat = { stop: () => void; complete: () => Promise; @@ -48,6 +55,17 @@ 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; + /** 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; /** Binding authority established by the legacy Hedge winner's first-byte write. */ hedgeBindingAuthorityPromise?: Promise; /** ResponseHandler-owned runtime lifecycle; attached when streaming starts. */ 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/config/system-settings-cache.ts b/src/lib/config/system-settings-cache.ts index cc1211c18..0ac909bef 100644 --- a/src/lib/config/system-settings-cache.ts +++ b/src/lib/config/system-settings-cache.ts @@ -195,6 +195,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/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/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/integration/proxy-hedge-lifecycle.test.ts b/tests/integration/proxy-hedge-lifecycle.test.ts index b4ea4c546..f810c642c 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"; @@ -16,10 +17,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 +78,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 +100,52 @@ 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 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(); + } + 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 +241,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 +476,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 +485,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 +498,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 new file mode 100644 index 000000000..1ce247f52 --- /dev/null +++ b/tests/unit/proxy/discovery-coordinator.test.ts @@ -0,0 +1,196 @@ +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("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)); + 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).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); + }); + + 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", + }); + expect(coordinator.markFailed("a")).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")); + expect(coordinator.markReady("a")).toEqual({ type: "promote_fallback", attemptId: "a" }); + 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)); + 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" }); + }); + + 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"], + }); + }); + + 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", + }); + }); + + 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("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)); + + 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 new file mode 100644 index 000000000..fc76b1b19 --- /dev/null +++ b/tests/unit/proxy/discovery-validity.test.ts @@ -0,0 +1,291 @@ +import { describe, expect, it } from "vitest"; +import { + DISCOVERY_EVENT_MAX_COUNT, + DISCOVERY_PREFIX_MAX_BYTES, + 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\n', "openai-chat") + .ready + ).toBe(true); + expect(classifyDiscoveryChunk("data: [DONE]\n\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); + 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\n', + "anthropic" + ).ready + ).toBe(false); + expect( + classifyDiscoveryChunk( + 'data: {"choices":[{"delta":{"tool_calls":[{"function":{}}]}}]}\n\n', + "openai-chat" + ).ready + ).toBe(false); + expect( + classifyDiscoveryChunk( + 'data: {"type":"response.output_text.delta","delta":" "}\n\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\n', + "openai-responses" + ).ready + ).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\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\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({ + ready: false, + terminal: false, + error: false, + }); + expect(parser.push('lo"}}]}\n\n')).toEqual({ + ready: true, + terminal: false, + error: false, + }); + }); + + 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"); + + expect( + parser.push('data: {"choices":[{"delta":{"content":"done"}}]}\n\ndata: [DONE]\n\n') + ).toEqual({ + ready: true, + terminal: true, + error: false, + }); + }); + + 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( + '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\n', + "anthropic" + ).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\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\n' + ).join(""); + + const limited = parser.push(`${metadataEvents}data: {"type":"message_stop"}\n\n`); + expect(limited).toMatchObject({ + ready: false, + terminal: false, + error: true, + limitExceeded: true, + }); + + expect(parser.push('data: {"type":"message_stop"}\n\n')).toEqual(limited); + }); +}); + +function parserForOpenAIChatToolCall(): DiscoveryValidityParser { + return new DiscoveryValidityParser("openai-chat"); +} 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/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"); 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 6322286c2..bd2d4ee90 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,10 @@ 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"), @@ -37,6 +41,34 @@ const mocks = vi.hoisted(() => ({ storeSessionSpecialSettings: vi.fn(async () => {}), 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", () => ({ @@ -64,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, @@ -90,6 +127,12 @@ 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, @@ -103,6 +146,8 @@ vi.mock("@/lib/session-manager", () => ({ vi.mock("@/app/v1/_lib/proxy/provider-selector", () => ({ ProxyProviderResolver: { pickRandomProviderWithExclusion: mocks.pickRandomProviderWithExclusion, + pickDiscoveryProviders: mocks.pickDiscoveryProviders, + resolveEffectivePriorityForSession: mocks.resolveEffectivePriorityForSession, }, })); @@ -124,8 +169,10 @@ 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 { peekDeferredStreamingFinalization } from "@/app/v1/_lib/proxy/stream-finalization"; +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; @@ -223,6 +270,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, @@ -334,14 +383,134 @@ 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, 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"; @@ -2254,6 +2423,1687 @@ 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.authState = { + success: true, + user: null, + key: { id: 1 }, + apiKey: null, + } as typeof session.authState; + 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("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 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("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 { + 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; + 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); + } 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("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("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 }); + 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 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 { + 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 efc8a9028..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 { @@ -68,6 +82,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 { 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..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,11 +5,15 @@ 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"; import { RateLimitService } from "@/lib/rate-limit"; +import type { SessionBindingSnapshot } from "@/lib/redis/session-binding"; import { SessionManager } from "@/lib/session-manager"; import { updateMessageRequestCostWithBreakdown, @@ -104,6 +108,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), settleLeaseBudgets: vi.fn(), + releaseProviderSession: vi.fn(), }, })); @@ -114,6 +119,24 @@ 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(), + getVersionedSessionBindingRefreshIntervalMs: vi.fn(() => 100_000), + renewSessionDiscoveryLease: vi.fn(async () => ({ + status: "renewed", + legacyFallbackAllowed: false, + })), + releaseSessionDiscoveryLease: vi.fn(async () => ({ + 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(), @@ -2074,6 +2097,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(); @@ -3104,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( @@ -3184,6 +3331,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 +3429,179 @@ 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("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); @@ -3276,6 +3635,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 9960a59c9..432881396 100644 --- a/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts +++ b/tests/unit/proxy/response-handler-endpoint-circuit-isolation.test.ts @@ -74,13 +74,17 @@ vi.mock("@/repository/message", () => ({ vi.mock("@/lib/session-manager", () => ({ SessionManager: { - clearVersionedSessionProvider: vi.fn(), updateSessionUsage: vi.fn(), storeSessionResponse: vi.fn(), clearSessionProvider: vi.fn(), - extractCodexPromptCacheKey: vi.fn(), + 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(), updateSessionWithCodexCacheKey: vi.fn(), @@ -93,6 +97,7 @@ vi.mock("@/lib/rate-limit", () => ({ trackUserDailyCost: vi.fn(), decrementLeaseBudget: vi.fn(), settleLeaseBudgets: vi.fn(), + releaseProviderSession: vi.fn(), }, })); @@ -374,18 +379,50 @@ 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 sourceController: ReadableStreamDefaultController | null = null; + let streamController!: ReadableStreamDefaultController; const stream = new ReadableStream({ start(controller) { - sourceController = controller; + streamController = controller; controller.enqueue( encoder.encode( - `event: message_delta\ndata: ${JSON.stringify({ usage: { input_tokens: 100, output_tokens: 50 } })}\n\n` + `data: ${JSON.stringify({ type: "content_block_delta", delta: { text: "ok" } })}\n\n` ) ); }, @@ -395,7 +432,12 @@ function createControllableSuccessStreamResponse(): { status: 200, headers: { "content-type": "text/event-stream" }, }), - complete: () => sourceController?.close(), + complete: () => { + streamController.enqueue( + encoder.encode(`event: message_stop\ndata: ${JSON.stringify({ type: "message_stop" })}\n\n`) + ); + streamController.close(); + }, }; } @@ -438,17 +480,47 @@ function setupCommonMocks() { sessionId: "fake-session", keyId: 456, providerId: null, - generation: "cleared-generation", + 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.getVersionedSessionBindingRefreshIntervalMs).mockReturnValue(100_000); - vi.mocked(SessionManager.touchVersionedSessionBinding).mockImplementation(async (snapshot) => ({ + vi.mocked(SessionManager.renewSessionDiscoveryLease).mockResolvedValue({ + status: "renewed", + legacyFallbackAllowed: false, + }); + vi.mocked(SessionManager.touchVersionedSessionBinding).mockImplementation(async (binding) => ({ status: "ok", source: "touched", - snapshot, + snapshot: binding, legacyFallbackAllowed: false, })); + vi.mocked(SessionManager.releaseSessionDiscoveryLease).mockResolvedValue({ + status: "released", + legacyFallbackAllowed: false, + }); vi.mocked(SessionManager.updateSessionBindingSmart).mockResolvedValue({ updated: true, reason: "test", @@ -465,6 +537,7 @@ function setupCommonMocks() { status: "settled", settlements: [], }); + vi.mocked(RateLimitService.releaseProviderSession).mockResolvedValue(undefined); vi.mocked(SessionTracker.refreshSession).mockResolvedValue(undefined); mockRecordFailure.mockResolvedValue(undefined); mockRecordSuccess.mockResolvedValue(undefined); @@ -510,6 +583,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("OpenAI Responses response.failed with HTTP 200 should be treated as provider failure", async () => { const session = createSession(); setDeferredMeta(session, 42); @@ -649,6 +738,757 @@ describe("Endpoint circuit breaker isolation", () => { 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: "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, + 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("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"; + 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"); + }); + it("keeps the captured legacy Hedge winner binding alive throughout a long stream", async () => { vi.useFakeTimers(); try { @@ -740,8 +1580,6 @@ describe("Endpoint circuit breaker isolation", () => { const session = createSession(); setDeferredMeta(session, 42, { isHedgeWinner: true, - // The first-byte write read G0, but a concurrent request established - // Provider 1 at G1 before its CAS. No authority over G1 was acquired. hedgeBindingAuthorityPromise: Promise.resolve({ snapshot: null, legacyClearAllowed: false, @@ -823,8 +1661,6 @@ describe("Endpoint circuit breaker isolation", () => { await bodyPromise; await drainAsyncTasks(); - // Completion must not retry after an administrator or concurrent request - // advances the generation. expect(SessionManager.touchVersionedSessionBinding).toHaveBeenCalledTimes(2); expect(SessionManager.updateSessionBindingSmart).not.toHaveBeenCalled(); } finally { 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